mirror of
https://github.com/hoshikawa2/compass_backoffice.git
synced 2026-07-10 06:14:19 +00:00
814 lines
34 KiB
Python
814 lines
34 KiB
Python
"""
|
|
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
|
|
}
|