adk-typescript/dist/tools/retrieval/VertexAiRagRetrieval.js

TypeScript port of Google's Agent Development Kit (ADK)

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.VertexAiRagRetrieval = void 0;
const BaseRetrievalTool_1 = require("./BaseRetrievalTool");
/**
 * Retrieval tool that uses Vertex AI RAG to retrieve information
 */
class VertexAiRagRetrieval extends BaseRetrievalTool_1.BaseRetrievalTool {
    /**
     * Create a new Vertex AI RAG retrieval tool
     *
     * @param options Options for the Vertex AI RAG retrieval tool
     */
    constructor(options) {
        super(options);
        this.vertexRagStore = {
            ragCorpora: options.ragCorpora,
            ragResources: options.ragResources,
            similarityTopK: options.similarityTopK,
            vectorDistanceThreshold: options.vectorDistanceThreshold
        };
    }
    /**
     * Process an LLM request to add RAG capabilities
     *
     * @param options The options for processing
     * @returns Nothing
     */
    async processLlmRequest(options) {
        const { toolContext, llmRequest } = options;
        // Use Gemini built-in Vertex AI RAG tool for Gemini 2 models
        if (llmRequest.model && llmRequest.model.startsWith('gemini-2')) {
            llmRequest.config = llmRequest.config || {};
            llmRequest.config.tools = llmRequest.config.tools || [];
            llmRequest.config.tools.push({
                retrieval: {
                    vertexRagStore: this.vertexRagStore
                }
            });
        }
        else {
            // For other models, add the function declaration to the tools
            // This would be implemented in the BaseRetrievalTool class
            // and would be called here with await super.processLlmRequest(options);
        }
    }
    /**
     * Retrieve information from Vertex AI RAG
     *
     * @param query The query to retrieve information for
     * @param maxResults Maximum number of results to return
     * @param context The context for the retrieval
     * @returns The retrieved information
     */
    async retrieve(query, maxResults, context) {
        try {
            console.log(`Retrieving from Vertex AI RAG: ${query}`);
            // This is a placeholder implementation
            // In a real implementation, this would call the Vertex AI RAG API
            const response = await this.retrievalQuery({
                text: query,
                ragResources: this.vertexRagStore.ragResources,
                ragCorpora: this.vertexRagStore.ragCorpora,
                similarityTopK: this.vertexRagStore.similarityTopK,
                vectorDistanceThreshold: this.vertexRagStore.vectorDistanceThreshold
            });
            if (!response.contexts || !response.contexts.contexts || response.contexts.contexts.length === 0) {
                return `No matching result found with the config: ${JSON.stringify(this.vertexRagStore)}`;
            }
            return response.contexts.contexts.map((context) => context.text);
        }
        catch (error) {
            console.error("Error retrieving from Vertex AI RAG:", error);
            return `Error retrieving information: ${error instanceof Error ? error.message : String(error)}`;
        }
    }
    /**
     * Placeholder implementation of the Vertex AI RAG API
     *
     * @param options Options for the retrieval query
     * @returns A mock response
     */
    async retrievalQuery(options) {
        // This is a placeholder implementation
        // In a real implementation, this would call the Vertex AI RAG API
        console.log(`RAG query: ${options.text}`);
        console.log(`RAG config: ${JSON.stringify({
            ragResources: options.ragResources,
            ragCorpora: options.ragCorpora,
            similarityTopK: options.similarityTopK,
            vectorDistanceThreshold: options.vectorDistanceThreshold
        })}`);
        // Return a mock response
        return {
            contexts: {
                contexts: [
                    {
                        text: `This is a placeholder response for "${options.text}". In a real implementation, this would retrieve information from Vertex AI RAG.`
                    }
                ]
            }
        };
    }
}
exports.VertexAiRagRetrieval = VertexAiRagRetrieval;