from __future__ import annotations from collections.abc import Mapping import logging import time from typing import Any from uuid import uuid4 from agente_contas_tim.agent.llm_gateway import LLMCapabilityGateway from agente_contas_tim.factory import CommandFactory from agente_contas_tim.observability import get_session_id from agente_contas_tim.workflows.actions.discovery import ensure_actions_loaded from agente_contas_tim.workflows.actions.registry import ( DEFAULT_ACTION_REGISTRY, ActionRegistry, WorkflowRuntimeContext, ) from agente_contas_tim.workflows.compiler import build_initial_state, compile_workflow from agente_contas_tim.workflows.execution_store import ( ExecutionStore, PostgresExecutionStore, ) from agente_contas_tim.workflows.exceptions import ( WorkflowConfigurationError, WorkflowExecutionStateError, WorkflowInputError, ) from agente_contas_tim.workflows.repositories.base import WorkflowRepository from agente_contas_tim.workflows.runtime_types import WorkflowRunResponse logger = logging.getLogger(__name__) class WorkflowService: def __init__( self, repository: WorkflowRepository, factory: CommandFactory, *, llm_gateway: LLMCapabilityGateway | None = None, postgres_dsn: str | None = None, action_registry: ActionRegistry | None = None, execution_store: ExecutionStore | None = None, checkpointer: Any | None = None, checkpointer_manager: Any | None = None, ) -> None: self._repository = repository self._factory = factory self._runtime = WorkflowRuntimeContext( factory=factory, llm_gateway=llm_gateway, workflow_runner=lambda workflow_name, input_payload, execution_id=None, version=None: self.run( workflow_name, input_payload, execution_id=execution_id, version=version, ), ) self._registry = action_registry or DEFAULT_ACTION_REGISTRY self._compiled_cache: dict[tuple[str, int], Any] = {} self._checkpointer_manager = checkpointer_manager if execution_store is not None: self._execution_store = execution_store else: if not postgres_dsn: raise ValueError( "postgres_dsn e obrigatorio quando execution_store nao e informado" ) self._execution_store = PostgresExecutionStore(postgres_dsn) if checkpointer is not None: self._checkpointer = checkpointer else: if not postgres_dsn: raise ValueError( "postgres_dsn e obrigatorio quando checkpointer nao e informado" ) ( self._checkpointer, self._checkpointer_manager, ) = self._create_checkpointer(postgres_dsn) ensure_actions_loaded("agente_contas_tim.workflows.actions") def run( self, workflow_name: str, input_payload: Mapping[str, Any], *, execution_id: str | None = None, version: int | None = None, ) -> WorkflowRunResponse: started_at = time.monotonic() input_keys = list(input_payload.keys()) if isinstance(input_payload, Mapping) else [] logger.info( "workflow.run.start name=%s version=%s execution_id=%s input_keys=%s", workflow_name, version, execution_id, input_keys, ) if execution_id is None: result = self._start_execution( workflow_name=workflow_name, input_payload=dict(input_payload), version=version, ) else: result = self._resume_execution( execution_id=execution_id, workflow_name=workflow_name, input_payload=dict(input_payload), version=version, ) elapsed_ms = round((time.monotonic() - started_at) * 1000, 2) logger.info( "workflow.run.end name=%s execution_id=%s status=%s elapsed_ms=%s", workflow_name, result.execution_id, result.status, elapsed_ms, ) return result def _start_execution( self, *, workflow_name: str, input_payload: dict[str, Any], version: int | None, ) -> WorkflowRunResponse: definition = ( self._repository.get_version(workflow_name, version) if version is not None else self._repository.get_active(workflow_name) ) graph = self._get_graph(definition) execution_id = str(uuid4()) session_id = self._current_session_id() message_id = self._extract_message_id(input_payload) self._execution_store.create( execution_id=execution_id, workflow_name=definition.name, workflow_version=definition.version, session_id=session_id, started_by_message_id=message_id, ) initial_state = build_initial_state(definition, input_payload) config = self._config(execution_id) try: graph.invoke(initial_state, config=config) except Exception as exc: self._execution_store.mark_status( execution_id, status="FAILED", current_node=None, resume_from=None, expected_input_key=None, ) logger.exception("Falha inesperada ao iniciar workflow: %s", exc) return WorkflowRunResponse( execution_id=execution_id, status="FAILED", error="Erro inesperado ao executar workflow", metadata={ "workflow": definition.name, "version": definition.version, "error_code": "WORKFLOW_INTERNAL_ERROR", }, ) return self._build_response( execution_id, definition.name, definition.version, graph ) def _resume_execution( self, *, execution_id: str, workflow_name: str, input_payload: dict[str, Any], version: int | None, ) -> WorkflowRunResponse: record = self._execution_store.claim_resume( execution_id=execution_id, workflow_name=workflow_name, workflow_version=version, session_id=self._current_session_id(), message_id=self._extract_message_id(input_payload), ) definition = self._repository.get_version( record.workflow_name, record.workflow_version, ) graph = self._get_graph(definition) config = self._config(execution_id) checkpoint_summary = self._checkpoint_debug_summary(config) logger.info( "workflow.resume.checkpoint execution_id=%s workflow=%s " "version=%s checkpoint=%s record_status=%s current_node=%s " "resume_from=%s expected_input_key=%s", execution_id, definition.name, definition.version, checkpoint_summary, record.status, record.current_node, record.resume_from, record.expected_input_key, ) try: from langgraph.types import Command except ModuleNotFoundError as exc: raise ModuleNotFoundError( "langgraph nao esta instalado. " "Adicione a dependencia para usar workflows." ) from exc try: graph.invoke(Command(resume=input_payload), config=config) except WorkflowInputError: self._execution_store.mark_status( execution_id, status="WAITING_INPUT", current_node=record.current_node, resume_from=record.resume_from, expected_input_key=record.expected_input_key, ) raise except Exception as exc: self._execution_store.mark_status( execution_id, status="FAILED", current_node=None, resume_from=None, expected_input_key=None, ) logger.info( "workflow.resume.failed.summary execution_id=%s workflow=%s " "version=%s error_type=%s error=%s checkpoint=%s " "record_status=%s current_node=%s resume_from=%s " "expected_input_key=%s", execution_id, definition.name, definition.version, type(exc).__name__, str(exc), checkpoint_summary, record.status, record.current_node, record.resume_from, record.expected_input_key, ) logger.error( "workflow.resume.failed execution_id=%s workflow=%s version=%s " "error_type=%s error=%s checkpoint=%s record_status=%s " "current_node=%s resume_from=%s expected_input_key=%s", execution_id, definition.name, definition.version, type(exc).__name__, str(exc), checkpoint_summary, record.status, record.current_node, record.resume_from, record.expected_input_key, exc_info=True, ) return WorkflowRunResponse( execution_id=execution_id, status="FAILED", error="Erro inesperado ao executar workflow", metadata={ "workflow": definition.name, "version": definition.version, "error_code": "WORKFLOW_INTERNAL_ERROR", }, ) return self._build_response( execution_id, definition.name, definition.version, graph ) def _checkpoint_debug_summary(self, config: dict[str, Any]) -> dict[str, Any]: get_tuple = getattr(self._checkpointer, "get_tuple", None) if not callable(get_tuple): return {"available": False, "reason": "checkpointer_without_get_tuple"} try: checkpoint_tuple = get_tuple(config) except Exception as exc: logger.error( "workflow.resume.checkpoint_read_failed thread_id=%s " "error_type=%s error=%s", config.get("configurable", {}).get("thread_id"), type(exc).__name__, str(exc), exc_info=True, ) return { "available": False, "read_error_type": type(exc).__name__, "read_error": str(exc), } if checkpoint_tuple is None: return {"available": False, "reason": "checkpoint_not_found"} checkpoint = dict(getattr(checkpoint_tuple, "checkpoint", {}) or {}) channel_values = dict(checkpoint.get("channel_values") or {}) pending_writes = list(getattr(checkpoint_tuple, "pending_writes", []) or []) config_values = dict(getattr(checkpoint_tuple, "config", {}) or {}) checkpoint_config = dict(config_values.get("configurable", {}) or {}) return { "available": True, "checkpoint_id": checkpoint_config.get("checkpoint_id"), "channel_keys": sorted(str(key) for key in channel_values.keys()), "pending_writes_count": len(pending_writes), "pending_write_channels": sorted( {str(write[1]) for write in pending_writes if len(write) >= 2} ), } def _build_response( self, execution_id: str, workflow_name: str, workflow_version: int, graph: Any, ) -> WorkflowRunResponse: snapshot = graph.get_state(self._config(execution_id)) values = dict(getattr(snapshot, "values", {}) or {}) status = str(values.get("status") or "") if not status: logger.warning( "workflow.response.status_missing execution_id=%s workflow=%s " "version=%s snapshot=%s", execution_id, workflow_name, workflow_version, self._snapshot_debug_summary(values), ) # Salvaguarda: alguns checkpointers podem nao propagar a channel # `status` ate o snapshot final apos FINISH_NODE. Inferir pelo # par final_data/final_error setado no action_node antes da # transicao para o no terminal. if values.get("final_error"): status = "FAILED" elif values.get("final_data") is not None: status = "COMPLETED" else: status = "FAILED" if status == "WAITING_INPUT": pending = values.get("pending_interrupt") if not isinstance(pending, dict): raise WorkflowConfigurationError( f"Workflow {workflow_name!r} pausou sem pending_interrupt" ) self._execution_store.mark_status( execution_id, status="WAITING_INPUT", current_node=str(pending.get("node_id", "")) or None, resume_from=str(pending.get("resume_from", "")) or None, expected_input_key=str(pending.get("expected_input_key", "")) or None, ) return WorkflowRunResponse( execution_id=execution_id, status="WAITING_INPUT", data=pending.get("payload"), metadata={ "workflow": workflow_name, "version": workflow_version, "paused_at": pending.get("node_id"), "resume_from": pending.get("resume_from"), "expected_input_key": pending.get("expected_input_key"), "allowed_values": pending.get("allowed_values", []), "normalize": pending.get("normalize"), }, ) if status == "COMPLETED": self._execution_store.mark_status( execution_id, status="COMPLETED", current_node=None, resume_from=None, expected_input_key=None, ) metadata = dict(values.get("final_metadata", {}) or {}) metadata.setdefault("workflow", workflow_name) metadata.setdefault("version", workflow_version) return WorkflowRunResponse( execution_id=execution_id, status="COMPLETED", data=values.get("final_data"), metadata=metadata, ) if status == "FAILED": logger.error( "workflow.response.failed execution_id=%s workflow=%s " "version=%s snapshot=%s", execution_id, workflow_name, workflow_version, self._snapshot_debug_summary(values), ) self._execution_store.mark_status( execution_id, status="FAILED", current_node=None, resume_from=None, expected_input_key=None, ) metadata = dict(values.get("final_metadata", {}) or {}) metadata.setdefault("workflow", workflow_name) metadata.setdefault("version", workflow_version) return WorkflowRunResponse( execution_id=execution_id, status="FAILED", data=None, error=str(values.get("final_error") or "Erro na execucao do workflow"), metadata=metadata, ) raise WorkflowExecutionStateError( f"Estado final invalido para workflow: {status!r}" ) @staticmethod def _snapshot_debug_summary(values: dict[str, Any]) -> dict[str, Any]: trace = values.get("trace") trace_tail = trace[-3:] if isinstance(trace, list) else [] final_metadata = values.get("final_metadata") return { "status": values.get("status"), "current_node": values.get("current_node"), "last_node": values.get("last_node"), "final_error": values.get("final_error"), "final_metadata": ( final_metadata if isinstance(final_metadata, dict) else {} ), "final_data_type": type(values.get("final_data")).__name__, "pending_interrupt_present": isinstance( values.get("pending_interrupt"), dict, ), "trace_tail": trace_tail, "channel_keys": sorted(str(key) for key in values.keys()), } def _get_graph(self, definition: Any) -> Any: key = (definition.name, definition.version) graph = self._compiled_cache.get(key) if graph is not None: return graph graph = compile_workflow( definition, action_registry=self._registry, runtime=self._runtime, checkpointer=self._checkpointer, ) self._compiled_cache[key] = graph logger.info( "Workflow compilado com LangGraph: name=%s version=%s", definition.name, definition.version, ) return graph @staticmethod def _config(execution_id: str) -> dict[str, Any]: return {"configurable": {"thread_id": execution_id}} @staticmethod def _current_session_id() -> str | None: session_id = str(get_session_id() or "").strip() return session_id if session_id and session_id != "-" else None @staticmethod def _extract_message_id(input_payload: Mapping[str, Any]) -> str | None: for key in ("message_id", "messageId"): value = str(input_payload.get(key, "") or "").strip() if value: return value return None @staticmethod def _create_checkpointer(postgres_dsn: str) -> tuple[Any, Any | None]: try: from langgraph.checkpoint.postgres import PostgresSaver except ModuleNotFoundError as exc: raise ModuleNotFoundError( "langgraph.checkpoint.postgres nao esta instalado. " "Adicione as dependencias para usar workflows com PostgreSQL." ) from exc manager = PostgresSaver.from_conn_string(postgres_dsn) if hasattr(manager, "__enter__") and hasattr(manager, "__exit__"): saver = manager.__enter__() setup = getattr(saver, "setup", None) if callable(setup): setup() return saver, manager setup = getattr(manager, "setup", None) if callable(setup): setup() return manager, None def close(self) -> None: close_store = getattr(self._execution_store, "close", None) if callable(close_store): close_store() manager = getattr(self, "_checkpointer_manager", None) if manager is not None and hasattr(manager, "__exit__"): manager.__exit__(None, None, None) return close_checkpointer = getattr(self._checkpointer, "close", None) if callable(close_checkpointer): close_checkpointer()