Source code for aidputils.agents.tools.service.kb_service

from src.volume_service_python_client.volume_service import KnowledgebaseClient

from aidputils.agents.auth.util import auth_utils
import logging
import time
from aidputils.agents.toolkit.service_metrics_util import get_service_metrics

logger = logging.getLogger(__name__)




class KBService:
    def __init__(self, conf, **context_vars):
        self.conf = conf
        self.dh_user_principal = auth_utils.get_auth_context().dh_user_principal
        self.context_vars = context_vars
        self.client = self._create_client()
        _init_kbservice_metrics(self)

    def _create_client(self):
        client = KnowledgebaseClient(self.conf, **self.context_vars)
        if not self.dh_user_principal:
            raise ValueError("Required authentication information not found in the context")
        if self.dh_user_principal:
            client.base_client.session.headers.update({"dh-user-principal": self.dh_user_principal})
        return client



    def fetch_knowledge_base(self, catalog, schema, kb_name):
        return self.client.get_knowledge_base(catalog, schema, kb_name)




    def fetch_vector_credentials(self, catalog, schema, kb_name):
        start_time = time.time()
        try:
            creds = self.client.get_vector_store_credentials(catalog, schema, kb_name)
            elapsed_ms = (time.time() - start_time) * 1000.0
            self._record_cred_fetch_success(elapsed_ms)
            return creds
        except Exception:
            elapsed_ms = (time.time() - start_time) * 1000.0
            self._record_cred_fetch_failure(elapsed_ms)
            raise


    # --- Metrics helpers for vector credential fetch ---


    def _record_cred_fetch_success(self, elapsed_ms: float):
        try:
            if getattr(self, "_service_metrics", None) and self._cred_fetch_success_counter:
                self._service_metrics.increment_counter(self._cred_fetch_success_counter, 1)
            if getattr(self, "_service_metrics", None) and self._cred_fetch_latency_hist:
                self._service_metrics.record_histogram(self._cred_fetch_latency_hist, elapsed_ms)
        except Exception:
            logger.exception("Failed to record ragtool.connection.cred.fetch.success/latency_ms")

    def _record_cred_fetch_failure(self, elapsed_ms: float):
        try:
            if getattr(self, "_service_metrics", None) and self._cred_fetch_failure_counter:
                self._service_metrics.increment_counter(self._cred_fetch_failure_counter, 1)
            if getattr(self, "_service_metrics", None) and self._cred_fetch_latency_hist:
                self._service_metrics.record_histogram(self._cred_fetch_latency_hist, elapsed_ms)
        except Exception:
            logger.exception("Failed to record ragtool.connection.cred.fetch.failure/latency_ms")
        except Exception:
            logger.exception("Failed to record ragtool.connection.cred.fetch.failure/latency_ms")



def _init_kbservice_metrics(self):
    try:
        self._service_metrics = get_service_metrics()
        self._cred_fetch_success_counter = self._service_metrics.create_counter(
            "ragtool.connection.cred.fetch.success",
            description="Count of successful vector credential fetch operations for RAG tool",
        )
        self._cred_fetch_failure_counter = self._service_metrics.create_counter(
            "ragtool.connection.cred.fetch.failure",
            description="Count of failed vector credential fetch operations for RAG tool",
        )
        self._cred_fetch_latency_hist = self._service_metrics.create_histogram(
            "ragtool.connection.cred.fetch.latency_ms",
            description="Latency (ms) of vector credential fetch operations for RAG tool",
        )
    except Exception:
        logger.exception("Failed to initialize RAGTool vector credential metrics")
        self._service_metrics = None
        self._cred_fetch_success_counter = None
        self._cred_fetch_failure_counter = None
        self._cred_fetch_latency_hist = None