Source code for aidputils.agents.tools.rag

import oci
import json

from aidputils.agents.auth.util import auth_utils
from aidputils.agents.tools.base_tool import BaseTool
from aidputils.agents.tools.ragtool.helpers.vector_search import execute_vector_search
from aidputils.agents.tools.ragtool.helpers.wallet_utils import create_wallet_from_credentials
from aidputils.agents.tools.service.kb_service import KBService
from aidputils.agents.tools.service.genai_service import GenAIService
from aidputils.agents.toolkit.metrics_util import MetricUtility
from aidputils.agents.auth.signer.custom_remote_signer import CustomRemoteSigner
from aidputils.agents.toolkit import chat_context

import logging

from aidputils.agents.tools.utils import SystemUtils

logger = logging.getLogger(__name__)




@BaseTool.register
class RAGTool(BaseTool):

    metricsUtil = MetricUtility("RAG_TOOL")
    success_counter = metricsUtil.create_counter("RAG_TOOL_SUCCESS_COUNTER")
    failure_counter = metricsUtil.create_counter("RAG_TOOL_FAILURE_COUNTER")

    @classmethod
    def _invoke_tool(cls, conf, runtime_params, **context_vars):
        is_success = False
        try:
            tool_name = chat_context.tool_context_var.get()
            cls.logger().info(f"context variable for tool name {tool_name}")
            catalog = conf['catalog']
            schema = conf['schema']
            kb_name = conf['knowledgeBase']

            query = runtime_params['query']
            cls.logger(tool_name="agent-app-rag").info(f"Invoking rag tool for {query}")

            top_k = conf.get('top_k', conf.get('topK'))
            if top_k is None:
                # If not found in conf, try runtime_params with both spellings
                top_k = runtime_params.get('top_k', runtime_params.get('topK'))

            if not catalog or not schema or not kb_name:
                cls.logger().error("Catalog / schema / kb is not configured")
                return {
                    "result": {
                        "content": [
                            {
                                "type": "text",
                                "text": "RAG Tool is configured incorrectly"
                            }
                        ],
                        "isError": True
                    }
                }

            if not query or top_k is None:
                return {
                    "result": {
                        "content": [
                            {
                                "type": "text",
                                "text": "Missing parameters. One of query or top_k is not provided"
                            }
                        ],
                        "isError": True
                    }
                }

            # 1. Fetch Knowledge Base and Credentials
            datalakeId = context_vars.get("datalake_id") or SystemUtils.get_env_var("DATALAKE_ID")
            lakeproxy_endpoint = auth_utils.get_lakeproxy_endpoint(datalakeId)
            context_vars.setdefault("service_endpoint", lakeproxy_endpoint)
            context_vars.setdefault("signer", oci.auth.signers.get_resource_principals_signer())

            cls.logger().info("Starting tool invocation for knowledge base: %s, hitting endpoint %s", kb_name,
                              lakeproxy_endpoint)

            cls.logger().info("Initializing KBService")
            kb_service = KBService(conf, **context_vars)
            cls.logger().debug(f"KBService initialized: {kb_service}")

            cls.logger().info("Fetching knowledge base: catalog=%s, schema=%s, kb_name=%s", catalog, schema, kb_name)
            kb = kb_service.fetch_knowledge_base(catalog, schema, kb_name)
            cls.logger().debug(f"KB fetched: {kb}")

            cls.logger().info("Fetching vector search credentials")
            credentials = kb_service.fetch_vector_credentials(catalog, schema, kb_name)
            cls.logger().debug(f"credentials fetched: {credentials}")

            # 2. Prepare Wallet
            db_user = credentials.data.user
            db_password = credentials.data.password
            tns_alias = credentials.data.tns_alias
            wallet_password = credentials.data.wallet_password

            wallet_path = create_wallet_from_credentials(credentials)

            cls.logger().info("Finding Oracle Cloud region by wallet")
            genai_service_endpoint = conf.get("llm", {}).get("endpoint")
            cls.logger().debug(
                f"service_endpoint: {genai_service_endpoint}")

            embedding_model_name = kb.data.embedding_model_name
            vector_embedded_table = kb.data.vector_table
            embedding_model_type = kb.data.embedding_model_source_type

            # 4. Instantiate GenAI service
            cls.logger().info("Instantiating GenAIService for endpoint: %s",
                              genai_service_endpoint)
            gen_ai_signer = CustomRemoteSigner()
            genai_service = GenAIService(conf, signer=gen_ai_signer, service_endpoint=genai_service_endpoint)

            # 5. Run Vector Search
            cls.logger().info("Executing vector search for query: '%s'", query)
            result_list = execute_vector_search(
                db_user=db_user,
                db_password=db_password,
                dsn=tns_alias,
                wallet_password = wallet_password,
                wallet_path=wallet_path,
                model_type=embedding_model_type,
                model_name=embedding_model_name,
                vector_table=vector_embedded_table,
                query=query,
                top_k=top_k,
                genai_service=genai_service
            )
            cls.logger().info("Vector search returned %d results", len(result_list) if result_list else 0)

            # 6. Generate LLM Response
            cls.logger().info("Generating LLM response")
            llm_result = genai_service.generate_response(conf, runtime_params=runtime_params, result_list=result_list)

            cls.logger().info("Tool invocation completed successfully")
            cls.logger(tool_name="agent-app-rag").info(f"Response from rag tool {llm_result}")
            is_success = True
            return llm_result
        except Exception as e:
            cls.logger().exception("Exception during tool invocation", exc_info=True)
            is_success = False
            return {"error": f"An error occurred during RAG Tool invocation: {str(e)}"}
        finally:
            if is_success:
                cls.metricsUtil.increment_counter(cls.success_counter)
            else:
                cls.metricsUtil.increment_counter(cls.failure_counter)

    @classmethod
    def _format_output_as_mcp(cls, response: dict) -> dict:
        if "error" in response:
            return {
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": response["error"]
                        }
                    ],
                    "isError": True
                }
            }
        # Prepare the MCP structured output
        structured_result = {
            "answer": response.get("answer", ""),
            "retrieved_chunks": response.get("retrieved_chunks", [])
        }
        text_output = json.dumps(structured_result, ensure_ascii=False, separators=(",", ":"))

        return {
            "result": {
                "content": [
                    {
                        "type": "text",
                        "text": text_output
                    }
                ],
                "structuredContent": structured_result
            }
        }