mirror of
https://github.com/hoshikawa2/compass_backoffice.git
synced 2026-07-10 06:14:19 +00:00
first commit
This commit is contained in:
813
src/components/clients/tais_kb_client.py
Normal file
813
src/components/clients/tais_kb_client.py
Normal file
@@ -0,0 +1,813 @@
|
||||
"""
|
||||
HTTP/DB client for the TAIS knowledge base.
|
||||
|
||||
Generates the query embedding via OCI GenAI (Cohere multilingual) and runs a
|
||||
vector-similarity search against an Oracle Autonomous Database using the
|
||||
native async oracledb driver.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import warnings
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
import oci
|
||||
import oci.exceptions
|
||||
import oci.generative_ai_inference.models
|
||||
import oci.retry
|
||||
import oracledb
|
||||
|
||||
from src.agent.local_prompts.preprocess_tais_kb_query import preprocess_tais_kb_query_pt
|
||||
from src.agent.local_prompts.postprocess_tais_kb_query import postprocess_tais_kb_query_pt
|
||||
from src.core.config import settings
|
||||
from src.core.prompt_manager import get_prompt
|
||||
from src.components.clients.exceptions.tais_kb_exceptions import TaisKbClientError
|
||||
from src.providers.llm_provider import chat_llm_with_usage, classification_llm, tais_kb_llm
|
||||
from src.utils.observer import trace_tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Return CLOBs as native strings instead of LOB handles to keep the search code simple.
|
||||
oracledb.defaults.fetch_lobs = False
|
||||
|
||||
|
||||
class Product(str, Enum):
|
||||
"""Supported TAIS products."""
|
||||
MOVEL = "Móvel"
|
||||
FIBRA = "Fibra"
|
||||
|
||||
|
||||
# Segments allowed per product
|
||||
_ALLOWED_MOVEL_SEGMENTS: dict[str, list[str]] = {
|
||||
"corporativo": ["SMB", "Top Clients", "M2M", "IoT"],
|
||||
"pospago": ["Fatura", "Express"],
|
||||
"controle": ["Fatura", "Express"],
|
||||
"prepago": ["Pré-Pago"],
|
||||
"beta": ["Beta"],
|
||||
"fixogsm": ["Fixo (GSM)"],
|
||||
}
|
||||
|
||||
_ALLOWED_FIBRA_SEGMENTS: dict[str, list[str]] = {
|
||||
"bandalarga": ["Banda Larga"],
|
||||
"wttx": ["WTTX"],
|
||||
}
|
||||
|
||||
|
||||
class TaisKbClient:
|
||||
"""Async client for the TAIS knowledge base (Oracle ADB + OCI embeddings)."""
|
||||
|
||||
_embed_client: oci.generative_ai_inference.GenerativeAiInferenceClient | None = None
|
||||
|
||||
@classmethod
|
||||
def _get_embed_client(cls) -> oci.generative_ai_inference.GenerativeAiInferenceClient:
|
||||
if cls._embed_client is not None:
|
||||
return cls._embed_client
|
||||
|
||||
if not settings.TAIS_GENAI_ENDPOINT or not settings.TAIS_GENAI_COMPARTMENT_ID:
|
||||
raise TaisKbClientError(
|
||||
"TAIS GenAI not configured (TAIS_GENAI_ENDPOINT, TAIS_GENAI_COMPARTMENT_ID)."
|
||||
)
|
||||
|
||||
import os
|
||||
from agent_framework.config.settings import settings as fw_settings
|
||||
|
||||
oci_config = oci.config.from_file(
|
||||
os.path.expanduser(getattr(fw_settings, "OCI_CONFIG_FILE", "~/.oci/config")),
|
||||
getattr(fw_settings, "OCI_PROFILE", None) or settings.OCI_CONFIG_PROFILE,
|
||||
)
|
||||
if getattr(fw_settings, "OCI_REGION", None):
|
||||
oci_config["region"] = fw_settings.OCI_REGION
|
||||
|
||||
# OCI client treats `timeout` as a single int — tuple form is silently
|
||||
# reduced to its first value. Use the configured TAIS timeout directly.
|
||||
cls._embed_client = oci.generative_ai_inference.GenerativeAiInferenceClient(
|
||||
config=oci_config,
|
||||
service_endpoint=settings.TAIS_GENAI_ENDPOINT,
|
||||
retry_strategy=oci.retry.NoneRetryStrategy(),
|
||||
timeout=settings.TAIS_DB_TIMEOUT,
|
||||
)
|
||||
return cls._embed_client
|
||||
|
||||
async def _preprocess_query(self, query_text: str) -> str:
|
||||
"""
|
||||
Preprocess a query using LLM to optimize it for semantic search.
|
||||
|
||||
Transforms the query into a semantically enriched form that maximizes
|
||||
similarity for RAG and embedding-based retrieval.
|
||||
|
||||
Args:
|
||||
query_text: Original query text to preprocess
|
||||
|
||||
Returns:
|
||||
Reformulated query optimized for semantic search
|
||||
|
||||
Raises:
|
||||
TaisKbClientError: If LLM call fails
|
||||
"""
|
||||
llm = classification_llm
|
||||
|
||||
# Get prompt from Langfuse with local fallback
|
||||
prompt = get_prompt("preprocess_tais_kb_query_pt", preprocess_tais_kb_query_pt)
|
||||
|
||||
# Format the message with the prompt and query
|
||||
message = f"{prompt}\n\nTranscrição original:\n{query_text}"
|
||||
|
||||
try:
|
||||
# Call LLM to preprocess the query (com retry de JSON decode)
|
||||
llm_resp = chat_llm_with_usage(llm, message, expect_json=True)
|
||||
content = llm_resp.content
|
||||
|
||||
logger.debug(f"Query preprocessing LLM response: {content}")
|
||||
|
||||
parsed_response = llm_resp.parsed_json or {}
|
||||
reformulated_query = parsed_response.get("reformulado", query_text)
|
||||
|
||||
logger.info(f"Query preprocessed successfully. Original: {query_text[:100]}... -> Reformulated: {reformulated_query[:100]}...")
|
||||
|
||||
return reformulated_query
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Failed to parse query preprocessing response as JSON: {e}. Returning original query.")
|
||||
return query_text
|
||||
except Exception as e:
|
||||
logger.warning(f"Query preprocessing failed: {e}. Proceeding with original query.")
|
||||
return query_text
|
||||
|
||||
async def _postprocess_results(self, query_text: str, results: list[dict], reformulated_query: str | None = None) -> dict:
|
||||
"""
|
||||
Postprocess search results using LLM to synthesize an answer.
|
||||
|
||||
Sends the query and top results to the LLM using the postprocess prompt,
|
||||
which synthesizes a comprehensive answer based on the retrieved documents.
|
||||
The LLM returns a JSON response with 'conteudo' (content) and 'id_procs' (document IDs).
|
||||
|
||||
Args:
|
||||
query_text: Original query text
|
||||
results: List of search result documents with id_proc, title_proc, description_proc, content
|
||||
reformulated_query: Query after preprocessing (optional)
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
- content: Synthesized answer from the LLM (from JSON 'conteudo' field)
|
||||
- id_procs: List of document IDs used in the answer (from JSON 'id_procs' field)
|
||||
- filled_prompt: The complete prompt with all variables filled in
|
||||
- postprocessing_succeeded: True on full success, False on JSON parse error
|
||||
or LLM call failure (fail-open: caller still receives best-effort content)
|
||||
- postprocessing_failure_reason: short label of the failure (only when failed)
|
||||
- postprocessing_http_meta: url/status/response_text/latency_ms of the LLM call
|
||||
(only when failed; consumed by AGA.036 emission downstream)
|
||||
"""
|
||||
llm = tais_kb_llm
|
||||
|
||||
# Get prompt from Langfuse with local fallback
|
||||
prompt = get_prompt("postprocess_tais_kb_query_pt", postprocess_tais_kb_query_pt)
|
||||
|
||||
# Format documents for the prompt
|
||||
formatted_docs = []
|
||||
for doc in results:
|
||||
doc_text = f"""
|
||||
Documento: {doc.get('id_proc', 'N/A')}
|
||||
Título: {doc.get('title_proc', 'N/A')}
|
||||
Descrição: {doc.get('description_proc', 'N/A')}
|
||||
Conteúdo: {doc.get('content', 'N/A')}
|
||||
"""
|
||||
formatted_docs.append(doc_text)
|
||||
|
||||
docs_context = "\n".join(formatted_docs)
|
||||
pp_start = time.perf_counter()
|
||||
|
||||
try:
|
||||
# Build query context with both original and reformulated (if available)
|
||||
query_context = f"Query: {query_text}"
|
||||
|
||||
# Format the message with the prompt, query, and documents
|
||||
filled_prompt = f"{prompt}\n\nPergunta do Operador:\n{query_context}\n\n---\n\nDocumentos de Referência:\n{docs_context}"
|
||||
|
||||
# Call LLM to postprocess the results (com retry de JSON decode)
|
||||
try:
|
||||
llm_resp = chat_llm_with_usage(llm, filled_prompt, expect_json=True)
|
||||
except json.JSONDecodeError as je:
|
||||
# Esgotou as tentativas de obter JSON válido — fail-open com raw content vazio.
|
||||
logger.warning(f"Failed to parse LLM JSON response after retries: {je}.")
|
||||
return {
|
||||
"content": "",
|
||||
"id_procs": [],
|
||||
"filled_prompt": filled_prompt,
|
||||
"postprocessing_succeeded": False,
|
||||
"postprocessing_failure_reason": "JSON parse error na resposta do LLM",
|
||||
"postprocessing_http_meta": {
|
||||
"url": settings.TAIS_GENAI_ENDPOINT or "N/A",
|
||||
"status_code": 200,
|
||||
"response_text": f"JSONDecodeError after retries: {str(je)[:200]}",
|
||||
"latency_ms": int((time.perf_counter() - pp_start) * 1000),
|
||||
},
|
||||
}
|
||||
|
||||
content = llm_resp.content
|
||||
logger.debug(f"Query postprocessing LLM response: {content[:200]}...")
|
||||
|
||||
response_json = llm_resp.parsed_json or {}
|
||||
postprocessing_content = response_json.get("conteudo", "")
|
||||
postprocessing_id_procs = response_json.get("id_procs", [])
|
||||
|
||||
logger.info(f"Results postprocessed successfully. Query: {query_text[:100]}... ID procs: {postprocessing_id_procs}")
|
||||
|
||||
return {
|
||||
"content": postprocessing_content,
|
||||
"id_procs": postprocessing_id_procs if isinstance(postprocessing_id_procs, list) else [],
|
||||
"filled_prompt": filled_prompt,
|
||||
"postprocessing_succeeded": True,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Results postprocessing failed: {e}. Returning empty result.")
|
||||
return {
|
||||
"content": "",
|
||||
"id_procs": [],
|
||||
"filled_prompt": "",
|
||||
"postprocessing_succeeded": False,
|
||||
"postprocessing_failure_reason": f"Falha na chamada do LLM ({type(e).__name__})",
|
||||
"postprocessing_http_meta": {
|
||||
"url": settings.TAIS_GENAI_ENDPOINT or "N/A",
|
||||
"status_code": getattr(e, "status", "N/A"),
|
||||
"response_text": str(e)[:500],
|
||||
"latency_ms": int((time.perf_counter() - pp_start) * 1000),
|
||||
},
|
||||
}
|
||||
|
||||
def _embed_sync(self, text: str) -> tuple[list[float], int]:
|
||||
client = self._get_embed_client()
|
||||
embed_detail = oci.generative_ai_inference.models.EmbedTextDetails()
|
||||
embed_detail.serving_mode = oci.generative_ai_inference.models.OnDemandServingMode(
|
||||
model_id=settings.TAIS_GENAI_EMBED_MODEL_ID
|
||||
)
|
||||
embed_detail.inputs = [text]
|
||||
embed_detail.truncate = "NONE"
|
||||
embed_detail.compartment_id = settings.TAIS_GENAI_COMPARTMENT_ID
|
||||
embed_detail.input_type = "SEARCH_QUERY"
|
||||
|
||||
endpoint = settings.TAIS_GENAI_ENDPOINT
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
response = client.embed_text(embed_detail)
|
||||
except oci.exceptions.ServiceError as exc:
|
||||
latency_ms = int((time.perf_counter() - start) * 1000)
|
||||
raise TaisKbClientError(
|
||||
f"OCI embedding service error: {exc}",
|
||||
status_code=exc.status,
|
||||
url=endpoint,
|
||||
response_text=str(getattr(exc, "message", exc)),
|
||||
latency_ms=latency_ms,
|
||||
) from exc
|
||||
except (oci.exceptions.RequestException, oci.exceptions.ConnectTimeout) as exc:
|
||||
latency_ms = int((time.perf_counter() - start) * 1000)
|
||||
raise TaisKbClientError(
|
||||
f"OCI embedding request error: {exc}",
|
||||
url=endpoint,
|
||||
response_text=str(exc),
|
||||
latency_ms=latency_ms,
|
||||
) from exc
|
||||
|
||||
return response.data._embeddings[0], response.status
|
||||
|
||||
async def _embed(self, text: str) -> tuple[list[float], int]:
|
||||
# OCI Python SDK has no async client; isolate the blocking call.
|
||||
return await asyncio.to_thread(self._embed_sync, text)
|
||||
|
||||
@staticmethod
|
||||
def _validate_product(product: Product | None) -> None:
|
||||
"""Validate that product is a valid Product enum value."""
|
||||
if product is None:
|
||||
raise ValueError(
|
||||
f"product is required. Valid values: Product.MOVEL, Product.FIBRA"
|
||||
)
|
||||
if not isinstance(product, Product):
|
||||
valid_values = [p.name for p in Product]
|
||||
raise TypeError(
|
||||
f"product must be a Product enum value (e.g. Product.MOVEL), got {type(product).__name__}: {product!r}. "
|
||||
f"Valid values: {valid_values}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_segments(segments: list[str], product: Product) -> None:
|
||||
"""Validate segments are known and allowed for the product."""
|
||||
# Get allowed mapping for product
|
||||
if product == Product.MOVEL:
|
||||
allowed_mapping = _ALLOWED_MOVEL_SEGMENTS
|
||||
elif product == Product.FIBRA:
|
||||
allowed_mapping = _ALLOWED_FIBRA_SEGMENTS
|
||||
else:
|
||||
raise ValueError(f"Unknown product: {product}")
|
||||
|
||||
allowed_segments = set(allowed_mapping.keys())
|
||||
|
||||
for seg in segments:
|
||||
seg_lower = seg.lower()
|
||||
if seg_lower not in allowed_segments:
|
||||
raise ValueError(
|
||||
f"Unrecognized segment '{seg}' for product {product.value}. "
|
||||
f"Allowed values: {sorted(allowed_segments)}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_sub_segments_for_segments(
|
||||
segments: list[str],
|
||||
sub_segments: list[str],
|
||||
product: Product
|
||||
) -> None:
|
||||
"""Validate that sub_segments are allowed for the given segments."""
|
||||
if not sub_segments:
|
||||
return # No sub_segments to validate
|
||||
|
||||
# Get allowed mapping for product
|
||||
if product == Product.MOVEL:
|
||||
allowed_mapping = _ALLOWED_MOVEL_SEGMENTS
|
||||
elif product == Product.FIBRA:
|
||||
allowed_mapping = _ALLOWED_FIBRA_SEGMENTS
|
||||
else:
|
||||
raise ValueError(f"Unknown product: {product}")
|
||||
|
||||
# Build set of allowed sub_segments for given segments
|
||||
allowed_for_segments: set[str] = set()
|
||||
for seg in segments:
|
||||
seg_lower = seg.lower()
|
||||
if seg_lower in allowed_mapping:
|
||||
allowed_for_segments.update(allowed_mapping[seg_lower])
|
||||
|
||||
# Validate each sub_segment is allowed and is case-insensitive match
|
||||
for sub_seg in sub_segments:
|
||||
sub_seg_lower = sub_seg.lower()
|
||||
|
||||
# Check if it exists in allowed set (case-insensitive)
|
||||
found = False
|
||||
for allowed_sub in allowed_for_segments:
|
||||
if allowed_sub.lower() == sub_seg_lower:
|
||||
found = True
|
||||
break
|
||||
|
||||
if not found:
|
||||
raise ValueError(
|
||||
f"Sub-segment '{sub_seg}' is not allowed for segments {segments}. "
|
||||
f"Allowed sub-segments: {sorted(allowed_for_segments)}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _fill_sql_with_bind_params(sql: str, bind_params: dict[str, object]) -> str:
|
||||
"""Replace bind parameters in SQL with their actual values (for debugging).
|
||||
|
||||
Note: Vectors are not filled for readability.
|
||||
"""
|
||||
filled_sql = sql
|
||||
|
||||
for key, value in bind_params.items():
|
||||
if key == "query_embedding":
|
||||
# Skip vectors for readability
|
||||
filled_sql = filled_sql.replace(f":{key}", "[VECTOR]")
|
||||
elif isinstance(value, str):
|
||||
# Escape single quotes in strings
|
||||
escaped_val = value.replace("'", "''")
|
||||
filled_sql = filled_sql.replace(f":{key}", f"'{escaped_val}'")
|
||||
elif isinstance(value, (int, float)):
|
||||
filled_sql = filled_sql.replace(f":{key}", str(value))
|
||||
elif value is None:
|
||||
filled_sql = filled_sql.replace(f":{key}", "NULL")
|
||||
else:
|
||||
filled_sql = filled_sql.replace(f":{key}", f"'{str(value)}'")
|
||||
|
||||
return filled_sql
|
||||
|
||||
@trace_tool
|
||||
async def search_documents(
|
||||
self,
|
||||
query_text: str,
|
||||
product: Product | None = None,
|
||||
segments: list[str] | None = None,
|
||||
sub_segments: list[str] | None = None,
|
||||
top_k: int = 3,
|
||||
check_expiration_date: bool = True,
|
||||
fetch_limit_multiplier: int = 100,
|
||||
preprocess: bool = True,
|
||||
postprocess: bool = True,
|
||||
deduplicate: bool = False,
|
||||
telemetry_top_n: int = 20,
|
||||
) -> dict[str, Any]:
|
||||
"""Search TAIS knowledge base using vector similarity with product-based filtering.
|
||||
|
||||
Args:
|
||||
query_text: Query text to embed and search
|
||||
product: Product to search within (MOVEL or FIBRA); required
|
||||
segments: List of segments to filter by (product-specific); optional
|
||||
sub_segments: List of sub-segments to filter by; optional
|
||||
top_k: Number of unique results to return (default 3)
|
||||
check_expiration_date: Whether to exclude expired documents (default True)
|
||||
fetch_limit_multiplier: Multiplier for database fetch limit (default 100)
|
||||
preprocess: Whether to preprocess the query with OCI GenAI before searching (default True)
|
||||
postprocess: Whether to postprocess results with LLM to synthesize an answer (default True)
|
||||
deduplicate: Whether to deduplicate results by title_proc (default False)
|
||||
Returns:
|
||||
dict with 'sql' (filled SQL for debugging), 'results' (unique records), 'reformulated_query', 'postprocessing', 'postprocessing_prompt'
|
||||
|
||||
Raises:
|
||||
TaisKbClientError: If DB not configured or OCI embedding fails
|
||||
ValueError: If product is None, segment is invalid, or sub-segment validation fails
|
||||
TypeError: If product is not a Product enum
|
||||
NotImplementedError: If product is FIBRA (not yet implemented)
|
||||
"""
|
||||
if top_k < 1:
|
||||
raise ValueError("top_k must be at least 1.")
|
||||
|
||||
# Validate product parameter
|
||||
self._validate_product(product)
|
||||
|
||||
if product == Product.FIBRA:
|
||||
raise NotImplementedError("Fibra product support is not yet implemented.")
|
||||
|
||||
if not all([settings.MONGODB_DB_USER, settings.MONGODB_DB_PASSWORD, settings.TAIS_DB_DSN]):
|
||||
raise TaisKbClientError(
|
||||
"TAIS DB not configured (MONGODB_DB_USER, MONGODB_DB_PASSWORD, TAIS_DB_DSN).",
|
||||
url="N/A",
|
||||
)
|
||||
|
||||
# Normalize to mutable lists
|
||||
active_segments: list[str] = list(segments) if segments else []
|
||||
active_sub_segments: list[str] = list(sub_segments) if sub_segments else []
|
||||
|
||||
# Validate inputs
|
||||
if active_segments:
|
||||
self._validate_segments(active_segments, product)
|
||||
|
||||
if active_sub_segments:
|
||||
# If we have sub_segments, we need at least one segment to validate against
|
||||
if not active_segments:
|
||||
raise ValueError(
|
||||
"sub_segments requires at least one segment to validate against"
|
||||
)
|
||||
self._validate_sub_segments_for_segments(
|
||||
active_segments, active_sub_segments, product
|
||||
)
|
||||
|
||||
# Product-based segment injection
|
||||
if product == Product.MOVEL:
|
||||
if "todosossegmentosmovel" not in [s.lower() for s in active_segments]:
|
||||
active_segments.append("todosossegmentosmovel")
|
||||
elif product == Product.FIBRA:
|
||||
if "todosossegmentosultrafibra" not in [s.lower() for s in active_segments]:
|
||||
active_segments.append("todosossegmentosultrafibra")
|
||||
|
||||
# Track original query and reformulated query
|
||||
original_query = query_text
|
||||
reformulated_query: str | None = None
|
||||
|
||||
if preprocess:
|
||||
# Preprocess the query with OCI GenAI (e.g. for better embedding quality)
|
||||
try:
|
||||
reformulated_query = await self._preprocess_query(original_query)
|
||||
query_text = reformulated_query
|
||||
except Exception as e:
|
||||
logger.warning(f"Query preprocessing failed, proceeding with original query. Error: {e}")
|
||||
reformulated_query = None
|
||||
|
||||
embed_start = time.perf_counter()
|
||||
embedding, embed_status = await self._embed(query_text)
|
||||
embedding_str = json.dumps(embedding)
|
||||
fetch_limit = max(top_k, 1) * fetch_limit_multiplier
|
||||
|
||||
bind_params: dict[str, object] = {
|
||||
"query_embedding": embedding_str,
|
||||
"fetch_limit": fetch_limit,
|
||||
}
|
||||
|
||||
conditions: list[str] = []
|
||||
|
||||
# Build segment filter: segments come as "controle + pospago + beta"
|
||||
# We need to find if ANY of our segments match ANY of the delimited values
|
||||
if active_segments:
|
||||
segment_conditions = []
|
||||
for i, seg in enumerate(active_segments):
|
||||
key = f"seg_{i}"
|
||||
# Search for " seg " or "seg " or " seg" to handle word boundaries with "+"
|
||||
# Oracle INSTR is case-insensitive by default when using LOWER()
|
||||
segment_conditions.append(
|
||||
f"INSTR(' ' || REPLACE(LOWER(segment), ' + ', ' ') || ' ', ' ' || LOWER(:{key}) || ' ') > 0"
|
||||
)
|
||||
bind_params[key] = seg.lower()
|
||||
conditions.append(f"({' OR '.join(segment_conditions)})")
|
||||
|
||||
# Build sub_segment filter: sub_segments come as "sub1, sub2, sub3"
|
||||
# We need to find if ANY of our sub_segments match ANY of the delimited values
|
||||
if active_sub_segments:
|
||||
sub_segment_conditions = []
|
||||
for i, sub_seg in enumerate(active_sub_segments):
|
||||
key = f"sub_seg_{i}"
|
||||
# Search for ",sub," or at start/end to handle word boundaries with ","
|
||||
# Match case-insensitively
|
||||
sub_segment_conditions.append(
|
||||
f"INSTR(',' || LOWER(REPLACE(sub_segments, ' ', '')) || ',', ',' || LOWER(:{key}) || ',') > 0"
|
||||
)
|
||||
bind_params[key] = sub_seg.lower().replace(" ", "")
|
||||
conditions.append(f"({' OR '.join(sub_segment_conditions)})")
|
||||
|
||||
# ── expiration filter ─────────────────────────────────────────────
|
||||
if check_expiration_date:
|
||||
# TRUNC(SYSDATE) strips the time component from the current server
|
||||
# date, ensuring a clean date-only comparison against expiration_date.
|
||||
conditions.append("(expiration_date IS NULL OR expiration_date >= TRUNC(SYSDATE))")
|
||||
|
||||
# ── distance filter ───────────────────────────────────────────────
|
||||
# Only return results with cosine distance lower than 0.5
|
||||
conditions.append("VECTOR_DISTANCE(embedding, TO_VECTOR(:query_embedding), COSINE) < 0.5")
|
||||
|
||||
where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""
|
||||
|
||||
sql = f"""
|
||||
SELECT
|
||||
id,
|
||||
doc_name,
|
||||
id_proc,
|
||||
CAST(title_proc AS VARCHAR2(4000)) AS title_proc,
|
||||
description_proc,
|
||||
updated_proc,
|
||||
segments,
|
||||
content,
|
||||
created_at,
|
||||
updated_at,
|
||||
uuid,
|
||||
subject,
|
||||
consultant_segments,
|
||||
expiration_date,
|
||||
sub_segments,
|
||||
segment,
|
||||
VECTOR_DISTANCE(embedding, TO_VECTOR(:query_embedding), COSINE) AS distance
|
||||
FROM {settings.TAIS_TABLE_CHUNKS}
|
||||
{where_clause}
|
||||
ORDER BY distance ASC
|
||||
FETCH FIRST :fetch_limit ROWS ONLY
|
||||
"""
|
||||
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
async with oracledb.connect_async(
|
||||
user=settings.MONGODB_DB_USER,
|
||||
password=settings.MONGODB_DB_PASSWORD,
|
||||
dsn=settings.TAIS_DB_DSN,
|
||||
tcp_connect_timeout=settings.TAIS_DB_TIMEOUT,
|
||||
) as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(sql, bind_params)
|
||||
rows = await cur.fetchall()
|
||||
cols = [c[0].lower() for c in cur.description]
|
||||
except oracledb.Error as exc:
|
||||
latency_ms = int((time.perf_counter() - start) * 1000)
|
||||
raise TaisKbClientError(
|
||||
f"TAIS DB error: {exc}",
|
||||
url=settings.TAIS_GENAI_ENDPOINT or "TAIS_DB",
|
||||
response_text=str(exc),
|
||||
latency_ms=latency_ms,
|
||||
) from exc
|
||||
|
||||
latency_ms = int((time.perf_counter() - start) * 1000)
|
||||
records = [dict(zip(cols, row)) for row in rows]
|
||||
|
||||
# Apply deduplication if enabled
|
||||
if deduplicate:
|
||||
seen: set[str] = set()
|
||||
unique: list[dict] = []
|
||||
for r in records:
|
||||
title = r.get("title_proc")
|
||||
if not title or title in seen:
|
||||
continue
|
||||
seen.add(title)
|
||||
unique.append(r)
|
||||
if len(unique) >= top_k:
|
||||
break
|
||||
else:
|
||||
# No deduplication: just take the first top_k records
|
||||
unique = records[:top_k]
|
||||
|
||||
# Return with filled SQL for debugging (except vector value)
|
||||
filled_sql = self._fill_sql_with_bind_params(sql, bind_params)
|
||||
|
||||
# Postprocess results if enabled
|
||||
postprocessing_content = None
|
||||
postprocessing_id_procs = None
|
||||
postprocessing_id_procs_map = None # Map of id_proc -> title_proc
|
||||
postprocessing_prompt = None
|
||||
postprocessing_succeeded = None # None when postprocess skipped
|
||||
postprocessing_failure_reason = None
|
||||
postprocessing_http_meta = None
|
||||
if postprocess and unique:
|
||||
try:
|
||||
postprocessing_response = await self._postprocess_results(
|
||||
query_text,
|
||||
unique,
|
||||
reformulated_query=reformulated_query
|
||||
)
|
||||
postprocessing_content = postprocessing_response.get("content")
|
||||
postprocessing_id_procs = postprocessing_response.get("id_procs")
|
||||
postprocessing_prompt = postprocessing_response.get("filled_prompt")
|
||||
postprocessing_succeeded = postprocessing_response.get("postprocessing_succeeded")
|
||||
postprocessing_failure_reason = postprocessing_response.get("postprocessing_failure_reason")
|
||||
postprocessing_http_meta = postprocessing_response.get("postprocessing_http_meta")
|
||||
|
||||
# Build mapping of id_proc -> title_proc for the returned id_procs
|
||||
if postprocessing_id_procs:
|
||||
postprocessing_id_procs_map = {}
|
||||
for doc in unique:
|
||||
doc_id = doc.get("id_proc")
|
||||
if doc_id in postprocessing_id_procs:
|
||||
postprocessing_id_procs_map[doc_id] = doc.get("title_proc", "")
|
||||
|
||||
logger.debug(f"Results postprocessed successfully. Length: {len(postprocessing_content) if postprocessing_content else 0}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Results postprocessing failed: {e}. Continuing without postprocessing.")
|
||||
postprocessing_succeeded = False
|
||||
postprocessing_failure_reason = f"Exception inesperada no postprocess ({type(e).__name__})"
|
||||
postprocessing_http_meta = {
|
||||
"url": settings.TAIS_GENAI_ENDPOINT or "N/A",
|
||||
"status_code": getattr(e, "status", "N/A"),
|
||||
"response_text": str(e)[:500],
|
||||
"latency_ms": 0,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"TAIS KB search completed. raw=%d unique=%d top_k=%d product=%s segments=%s sub_segments=%s check_expiration=%s deduplicate=%s postprocessing=%s",
|
||||
len(records), len(unique), top_k, product.name if product else None,
|
||||
active_segments, active_sub_segments, check_expiration_date, deduplicate, bool(postprocessing_content),
|
||||
)
|
||||
|
||||
http_meta = {
|
||||
"url": settings.TAIS_GENAI_ENDPOINT or "N/A",
|
||||
"status_code": embed_status,
|
||||
"response_text": f"raw_records={len(records)} unique={len(unique)}",
|
||||
"latency_ms": int((time.perf_counter() - embed_start) * 1000),
|
||||
}
|
||||
|
||||
# Build a slim "retrieved" list for telemetry (AGAs that carry
|
||||
# ragRetrievedDocuments). Capped at `telemetry_top_n` to keep IC
|
||||
# payload size sane — `records` may have hundreds of candidates.
|
||||
# Already ordered by distance ASC from the SQL. Chunks/content are
|
||||
# dropped here because consumers of `retrieved_documents` only need
|
||||
# the doc identifier for observability.
|
||||
retrieved_documents = [
|
||||
{
|
||||
"documentId": r.get("id_proc"),
|
||||
"title": r.get("title_proc"),
|
||||
"distance": float(r["distance"]) if r.get("distance") is not None else None,
|
||||
}
|
||||
for r in records[:telemetry_top_n]
|
||||
]
|
||||
|
||||
return {
|
||||
"sql": filled_sql,
|
||||
"results": unique,
|
||||
"retrieved_documents": retrieved_documents,
|
||||
"reformulated_query": reformulated_query,
|
||||
"postprocessing_content": postprocessing_content,
|
||||
"postprocessing_id_procs": postprocessing_id_procs,
|
||||
"postprocessing_id_procs_map": postprocessing_id_procs_map,
|
||||
"postprocessing_prompt": postprocessing_prompt,
|
||||
"postprocessing_succeeded": postprocessing_succeeded,
|
||||
"postprocessing_failure_reason": postprocessing_failure_reason,
|
||||
"postprocessing_http_meta": postprocessing_http_meta,
|
||||
"http_meta": http_meta,
|
||||
}
|
||||
|
||||
@trace_tool
|
||||
async def search_documents_legacy(
|
||||
self,
|
||||
query_text: str,
|
||||
top_k: int,
|
||||
segment_filter: str | None = None,
|
||||
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
"""DEPRECATED: Use search_documents() instead with product and segments parameters.
|
||||
|
||||
Legacy API wrapper that maps old segment_filter to new product-based approach.
|
||||
This method maintains backward compatibility with existing code.
|
||||
|
||||
Returns a tuple (response_dict, http_meta_dict) following the same pattern as
|
||||
SiebelClient.open_service_request_with_retry and ImdbClient.get_imdb_access_data_with_retry.
|
||||
"""
|
||||
warnings.warn(
|
||||
"search_documents_legacy() is deprecated. Use search_documents(query_text, product=Product.MOVEL, ...) instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2
|
||||
)
|
||||
|
||||
# Map legacy segment_filter to new API
|
||||
# For backward compatibility, we resolve segments using a simple mapping
|
||||
segments = None
|
||||
if segment_filter:
|
||||
key = segment_filter.strip().lower()
|
||||
# Try to map old segment names to new allowed segments
|
||||
# For simplicity, if we can find a matching key in MOVEL segments, use it
|
||||
if key in _ALLOWED_MOVEL_SEGMENTS:
|
||||
segments = [key]
|
||||
|
||||
response = await self.search_documents(
|
||||
query_text=query_text,
|
||||
product=Product.MOVEL,
|
||||
segments=segments,
|
||||
top_k=top_k,
|
||||
check_expiration_date=True,
|
||||
fetch_limit_multiplier=50, # Keep old fetch limit for backward compatibility
|
||||
)
|
||||
http_meta = response.pop("http_meta", {
|
||||
"url": settings.TAIS_GENAI_ENDPOINT or "N/A",
|
||||
"status_code": "N/A",
|
||||
"response_text": "N/A",
|
||||
"latency_ms": 0,
|
||||
})
|
||||
return response, http_meta
|
||||
|
||||
@trace_tool
|
||||
async def get_content_by_id_proc(
|
||||
self,
|
||||
id_proc: str,
|
||||
return_as: str = "html"
|
||||
) -> dict[str, Any]:
|
||||
"""Fetch document content from TAIS_BASE_CHUNKS by id_proc with format conversion.
|
||||
|
||||
Args:
|
||||
id_proc: Procedure ID to search for
|
||||
return_as: Format to return (html or markdown), case insensitive
|
||||
|
||||
Returns:
|
||||
dict with 'sql' (filled SQL for debugging), 'results' (records), 'return_as' (format)
|
||||
|
||||
Raises:
|
||||
TaisKbClientError: If DB not configured
|
||||
ValueError: If id_proc is empty or return_as is invalid
|
||||
"""
|
||||
if not all([settings.MONGODB_DB_USER, settings.MONGODB_DB_PASSWORD, settings.TAIS_DB_DSN]):
|
||||
raise TaisKbClientError(
|
||||
"TAIS DB not configured (MONGODB_DB_USER, MONGODB_DB_PASSWORD, TAIS_DB_DSN)."
|
||||
)
|
||||
|
||||
if not id_proc or not id_proc.strip():
|
||||
raise ValueError("id_proc cannot be empty.")
|
||||
|
||||
# Normalize return_as to lowercase
|
||||
return_as_lower = return_as.strip().lower()
|
||||
if return_as_lower not in ("html", "markdown"):
|
||||
raise ValueError(f"return_as must be 'html' or 'markdown', got '{return_as}'")
|
||||
|
||||
# Build SELECT clause based on return_as
|
||||
content_column = "content_html" if return_as_lower == "html" else "content_markdown"
|
||||
|
||||
sql = f"""
|
||||
SELECT
|
||||
id,
|
||||
doc_name,
|
||||
id_proc,
|
||||
title_proc,
|
||||
description_proc,
|
||||
updated_proc,
|
||||
segments,
|
||||
{content_column} AS {content_column},
|
||||
created_at,
|
||||
updated_at,
|
||||
uuid,
|
||||
subject,
|
||||
consultant_segments,
|
||||
expiration_date,
|
||||
sub_segments,
|
||||
segment
|
||||
FROM {settings.TAIS_TABLE_FILES}
|
||||
WHERE id_proc = :id_proc
|
||||
"""
|
||||
|
||||
bind_params: dict[str, object] = {
|
||||
"id_proc": id_proc.strip(),
|
||||
}
|
||||
|
||||
try:
|
||||
async with oracledb.connect_async(
|
||||
user=settings.MONGODB_DB_USER,
|
||||
password=settings.MONGODB_DB_PASSWORD,
|
||||
dsn=settings.TAIS_DB_DSN,
|
||||
tcp_connect_timeout=settings.TAIS_DB_TIMEOUT,
|
||||
) as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(sql, bind_params)
|
||||
rows = await cur.fetchall()
|
||||
cols = [c[0].lower() for c in cur.description]
|
||||
except oracledb.Error as exc:
|
||||
raise TaisKbClientError(f"TAIS DB error: {exc}") from exc
|
||||
|
||||
records = [dict(zip(cols, row)) for row in rows]
|
||||
|
||||
# Fill the SQL with bind params for debugging (except for LOB fields)
|
||||
filled_sql = self._fill_sql_with_bind_params(sql, bind_params)
|
||||
|
||||
logger.info(
|
||||
"TAIS KB get_content_by_id_proc completed. id_proc=%s return_as=%s results=%d",
|
||||
id_proc, return_as_lower, len(records),
|
||||
)
|
||||
return {
|
||||
"sql": filled_sql,
|
||||
"results": records,
|
||||
"return_as": return_as_lower
|
||||
}
|
||||
Reference in New Issue
Block a user