first commit

legacy_reference/workflows/actions/rag/actions.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import logging
+
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import copy_context
+from typing import Any
+from agente_contas_tim.workflows.actions.registry import (
+    WorkflowRuntimeContext,
+    workflow_action,
+)
+from agente_contas_tim.workflows.runtime_types import ActionResult
+from agente_contas_tim.workflows.actions.common.helpers import (
+    _result_failed_or_missing_data,
+    _runtime_llm_callbacks,
+    _runtime_llm_metadata,
+    _to_dict,
+)
+from agente_contas_tim.workflows.actions.rag.helpers import (
+    _build_rag_answer_item,
+    _extract_rag_documents,
+    _is_selected_rag_document,
+    _normalize_rag_queries,
+    _rag_document_text,
+    _rag_document_title,
+)
+
+_RAG_FALLBACK_MESSAGE = (
+    "Não foi possível buscar essa informação no momento. "
+    "Por favor, aguarde na linha para que possamos te ajudar de outra forma."
+)
+
+logger = logging.getLogger("agente_contas_tim.workflows.actions.tim_actions")
+
+
+@workflow_action("buscar_informacao_rag")
+def buscar_informacao_rag(
+    state: dict[str, Any],
+    params: dict[str, Any],
+    runtime: WorkflowRuntimeContext,
+) -> ActionResult:
+    queries = _normalize_rag_queries(params)
+    if not queries:
+        return ActionResult.fail("Informe uma pergunta para buscar na base.")
+
+    raw_top_k = params.get("top_k")
+    resolved_top_k: int | None = None
+    if raw_top_k is not None and str(raw_top_k).strip():
+        try:
+            resolved_top_k = int(raw_top_k)
+        except (TypeError, ValueError):
+            return ActionResult.fail("top_k invalido: informe um numero inteiro")
+
+    segment = str(params.get("segment", "")).strip()
+    results: list[dict[str, Any]] = []
+    documents: list[dict[str, Any]] = []
+    answer_parts: list[str] = []
+    retrieved_titles: list[str] = []
+    selected_titles: list[str] = []
+    metadata: dict[str, Any] = {}
+
+    def _execute(query: str) -> Any:
+        return runtime.factory.create_rag_search(
+            query=query,
+            top_k=resolved_top_k,
+            segment=segment,
+        ).execute()
+
+    if len(queries) == 1:
+        raw_results = [_execute(queries[0])]
+    else:
+        with ThreadPoolExecutor(
+            max_workers=min(len(queries), 6),
+            thread_name_prefix="rag-search",
+        ) as executor:
+            futures = [
+                executor.submit(copy_context().run, _execute, query)
+                for query in queries
+            ]
+            raw_results = [future.result() for future in futures]
+
+    for query, result in zip(queries, raw_results):
+        if _result_failed_or_missing_data(result, state=state):
+            return ActionResult.fail(
+                result.error or "Falha na busca RAG",
+                **result.metadata,
+            )
+
+        metadata.update(result.metadata)
+        payload = _to_dict(result.data)
+        query_documents = _extract_rag_documents(payload)
+        query_total = len(query_documents)
+        query_message = (
+            f"Encontrei {query_total} trecho(s) relevante(s) na base."
+            if query_total
+            else "Nao encontrei trechos relevantes para essa busca."
+        )
+        query_answer = _build_rag_answer_item(query, query_documents)
+        results.append(
+            {
+                "query": query,
+                "total": query_total,
+                "documents": query_documents,
+                "message": query_message,
+                "answer": query_answer,
+            }
+        )
+        documents.extend(query_documents)
+        answer_parts.append(query_answer)
+
+        for doc in query_documents:
+            title = _rag_document_title(doc)
+            if not title:
+                continue
+            retrieved_titles.append(title)
+            if _is_selected_rag_document(doc):
+                selected_titles.append(title)
+
+    total = len(documents)
+    message = (
+        f"Encontrei {total} trecho(s) relevante(s) na base."
+        if total
+        else "Nao encontrei trechos relevantes para essa busca."
+    )
+
+    return ActionResult.ok(
+        {
+            "success": True,
+            "query": queries[0],
+            "queries": queries,
+            "total": total,
+            "documents": documents,
+            "results": results,
+            "answer": "\n\n".join(answer_parts),
+            "message": message,
+            "ragRetrievedDocuments": "|".join(retrieved_titles),
+            "ragSelectedDocuments": "|".join(selected_titles),
+            "noMatchRag": total == 0,
+        },
+        **metadata,
+    )
+
+
+@workflow_action("reescrever_resposta_buscar_informacao")
+def reescrever_resposta_buscar_informacao(
+    state: dict[str, Any],
+    params: dict[str, Any],
+    runtime: WorkflowRuntimeContext,
+) -> ActionResult:
+    raw_queries = params.get("queries")
+    queries: list[str] = []
+    if isinstance(raw_queries, (list, tuple)):
+        for item in raw_queries:
+            text = str(item or "").strip()
+            if text:
+                queries.append(text)
+    elif isinstance(raw_queries, str) and raw_queries.strip():
+        queries.append(raw_queries.strip())
+
+    raw_documents = params.get("documents")
+    documents: list[dict[str, Any]] = []
+    if isinstance(raw_documents, (list, tuple)):
+        for item in raw_documents:
+            if isinstance(item, dict):
+                documents.append(item)
+
+    rag_context_parts: list[str] = []
+    for doc in documents:
+        text = _rag_document_text(doc)
+        if text:
+            rag_context_parts.append(text)
+    rag_context = "\n\n".join(rag_context_parts).strip()
+    if not rag_context:
+        rag_context = str(params.get("answer", "") or "").strip()
+
+    no_match = bool(params.get("noMatchRag", False))
+
+    def _payload(answer: str, *, success: bool = True) -> dict[str, Any]:
+        return {
+            "success": success,
+            "answer": answer,
+            "noMatchRag": no_match,
+        }
+
+    if runtime.llm_gateway is None:
+        return ActionResult.ok(_payload(_RAG_FALLBACK_MESSAGE))
+
+    queries_text = "; ".join(queries) if queries else ""
+    try:
+        llm_result = runtime.llm_gateway.execute(
+            capability_id="fluxo_buscar_informacao_reescrita",
+            variables={
+                "queries": queries_text,
+                "rag_context": rag_context,
+            },
+            user_text=queries_text,
+            callbacks=_runtime_llm_callbacks(runtime),
+            tags=["workflow_action"],
+            metadata=_runtime_llm_metadata(runtime),
+        )
+        answer = str(getattr(llm_result, "content", "") or "").strip()
+    except Exception:
+        logger.exception(
+            "reescrever_resposta_buscar_informacao: falha ao invocar "
+            "capability fluxo_buscar_informacao_reescrita"
+        )
+        return ActionResult.ok(_payload(_RAG_FALLBACK_MESSAGE))
+
+    if not answer:
+        return ActionResult.ok(_payload(_RAG_FALLBACK_MESSAGE))
+
+    return ActionResult.ok(_payload(answer))
+
+
+__all__ = [
+    'buscar_informacao_rag',
+    'reescrever_resposta_buscar_informacao',
+]