first_contas/legacy_reference/workflows/actions/rag/actions.py

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',
]