"""LangGraph :class:`AgentRuntime` implementation for the Gradio UI.""" from __future__ import annotations import os import sys import threading from collections.abc import Iterator from pathlib import Path from typing import Any _AGENT_REDACT_ROOT = Path(__file__).resolve().parents[1] if str(_AGENT_REDACT_ROOT) not in sys.path: sys.path.insert(0, str(_AGENT_REDACT_ROOT)) _SHARED_DIR = _AGENT_REDACT_ROOT / "shared" if str(_SHARED_DIR) not in sys.path: sys.path.insert(0, str(_SHARED_DIR)) from agent_runtime import ( AgentRuntime, AgentRuntimeError, AgentStreamEvent, ) # noqa: E402 from langchain_core.messages import AIMessage, HumanMessage, ToolMessage # noqa: E402 from redaction_langgraph.graph import build_redaction_agent # noqa: E402 from redaction_langgraph.llm_errors import ( # noqa: E402 is_context_overflow_error, is_tool_call_json_parse_error, ) from redaction_langgraph.message_context import ( # noqa: E402 get_trim_stats, reset_trim_stats, set_aggressive_trim, ) from redaction_langgraph.tools import reset_langgraph_tool_session_state # noqa: E402 from redaction_langgraph.workflow_continue import ( # noqa: E402 build_tool_call_json_retry_prompt, build_workflow_continue_prompt, clarification_requested, identical_tool_error_streak, langgraph_auto_continue_enabled, langgraph_identical_error_stop_streak, langgraph_max_continuations, redaction_workflow_incomplete, ) class LangGraphAgentRuntime(AgentRuntime): """Session-scoped LangGraph ReAct agent (curated tools, no shell).""" def __init__(self, *, session_hash: str | None = None) -> None: self._session_hash = session_hash self._graph: Any = None self._system_message: Any = None self._messages: list[Any] = [] self._running = False self._prompt_stream_depth = 0 self._abort_requested = False self._is_compacting = False self._lock = threading.Lock() self._pending_ui_notices: list[dict[str, Any]] = [] self._pending_ui_history: list[dict[str, Any]] = [] @property def orchestrator(self) -> str: return "langgraph" @property def running(self) -> bool: return self._running @property def prompt_stream_active(self) -> bool: return self._prompt_stream_depth > 0 def start(self) -> None: if self._graph is None: self._graph, self._system_message = build_redaction_agent( self._session_hash ) self._running = True def close(self) -> None: self._running = False self._graph = None self._messages = [] self._is_compacting = False def abort(self) -> None: self._abort_requested = True def new_session(self) -> None: self._messages = [] self._abort_requested = False self._is_compacting = False reset_langgraph_tool_session_state(self._session_hash) def set_model(self, provider: str, model_id: str) -> dict[str, Any]: os.environ["AGENT_DEFAULT_PROVIDER"] = provider os.environ["AGENT_DEFAULT_MODEL"] = model_id if provider == "llama-cpp": os.environ["AGENT_LLAMA_MODEL_ID"] = model_id self._graph = None self.start() return {"provider": provider, "model": model_id} def apply_backend(self, provider: str, model_id: str) -> None: self.set_model(provider, model_id) self.new_session() def get_state(self) -> dict[str, Any]: return { "isStreaming": self.prompt_stream_active, "isCompacting": self._is_compacting, "provider": os.environ.get("AGENT_DEFAULT_PROVIDER"), "model": { "provider": os.environ.get("AGENT_DEFAULT_PROVIDER"), "id": os.environ.get("AGENT_DEFAULT_MODEL") or os.environ.get("AGENT_LLAMA_MODEL_ID"), }, } def get_messages(self) -> list[dict[str, Any]]: out: list[dict[str, Any]] = [] for message in self._messages: if isinstance(message, HumanMessage): out.append({"role": "user", "content": str(message.content)}) elif isinstance(message, AIMessage): out.append({"role": "assistant", "content": str(message.content or "")}) return out def stage_ui_chat_notice(self, label: str, message: str) -> None: text = message.strip() if not text: return self._pending_ui_history.append( {"role": "user", "content": f"_**{label}:**_ {text}"} ) self._pending_ui_history.append({"role": "assistant", "content": ""}) def drain_pending_ui_history(self) -> list[dict[str, Any]]: pending = self._pending_ui_history[:] self._pending_ui_history.clear() return pending def _yield_message_updates( self, msg: Any, *, assistant_chunks: list[str], tool_names_seen: set[str], tool_outputs: list[tuple[str, str]], loop_break_requested: list[bool] | None = None, ) -> Iterator[AgentStreamEvent]: if isinstance(msg, AIMessage): text = self._stringify_content(msg.content) if text: assistant_chunks.append(text) yield AgentStreamEvent(kind="text_snapshot", text=text) for call in msg.tool_calls or []: name = str(call.get("name") or "tool") args = call.get("args") if isinstance(call.get("args"), dict) else {} yield AgentStreamEvent( kind="tool_start", tool_name=name, tool_args=args, text=name, ) elif isinstance(msg, ToolMessage): name = str(msg.name or "tool") tool_names_seen.add(name) output = str(msg.content or "") tool_outputs.append((name, output)) yield AgentStreamEvent( kind="tool_end", tool_name=name, tool_output=output, is_error=False, ) stop_streak = langgraph_identical_error_stop_streak() if identical_tool_error_streak(tool_outputs, min_streak=stop_streak): if loop_break_requested is not None: loop_break_requested[0] = True self._abort_requested = True yield AgentStreamEvent( kind="status", text=( f"Identical tool error repeated {stop_streak} times — " "breaking retry loop…" ), ) def _yield_compaction_notice_if_needed( self, *, compaction_emitted: list[bool] ) -> Iterator[AgentStreamEvent]: if compaction_emitted[0]: return stats = get_trim_stats() if stats is None or not stats.trimmed: return compaction_emitted[0] = True self._is_compacting = True yield AgentStreamEvent( kind="compaction_start", text=( "Context compaction — trimming older tool history so the prompt " "fits the model window…" ), ) yield AgentStreamEvent( kind="compaction_end", text=( f"Context compaction finished ({stats.tokens_before:,} → " f"{stats.tokens_after:,} tokens; dropped " f"{stats.messages_dropped} message(s)). The agent continues " "with a shorter history." ), meta={ "tokens_before": stats.tokens_before, "tokens_after": stats.tokens_after, "messages_dropped": stats.messages_dropped, }, ) self._is_compacting = False def _stream_graph_round( self, graph: Any, graph_messages: list[Any], stream_config: dict[str, Any], *, assistant_chunks: list[str], tool_names_seen: set[str], tool_outputs: list[tuple[str, str]], loop_break_requested: list[bool] | None = None, ) -> Iterator[AgentStreamEvent]: reset_trim_stats() compaction_emitted = [False] for event in graph.stream( {"messages": graph_messages}, stream_mode="updates", config=stream_config, ): if self._abort_requested: yield AgentStreamEvent(kind="done", text="Agent aborted.") return yield from self._yield_compaction_notice_if_needed( compaction_emitted=compaction_emitted ) for _node, update in event.items(): for msg in update.get("messages") or []: graph_messages.append(msg) yield from self._yield_message_updates( msg, assistant_chunks=assistant_chunks, tool_names_seen=tool_names_seen, tool_outputs=tool_outputs, loop_break_requested=loop_break_requested, ) if self._abort_requested: yield AgentStreamEvent( kind="done", text=( "Agent aborted (identical tool-error loop)." if loop_break_requested and loop_break_requested[0] else "Agent aborted." ), ) return yield from self._yield_compaction_notice_if_needed( compaction_emitted=compaction_emitted ) def _rebuild_graph(self, *, aggressive_compaction: bool = False) -> None: self._graph, self._system_message = build_redaction_agent( self._session_hash, aggressive_compaction=aggressive_compaction, ) def prompt_events(self, message: str) -> Iterator[AgentStreamEvent]: self._prompt_stream_depth += 1 self._abort_requested = False self._is_compacting = False try: if not self._running: self.start() if self._graph is None: raise AgentRuntimeError("LangGraph agent is not initialized.") from eval.arize_monitoring import ( agent_turn_span, arize_session_context, langgraph_trace_config, ) from redaction_langgraph.graph import graph_recursion_limit with arize_session_context(self._session_hash): with agent_turn_span( "langgraph.agent", session_hash=self._session_hash, message=message, agent_name="langgraph", ) as turn_meta: yield AgentStreamEvent( kind="status", text="LangGraph agent started…" ) graph_messages: list[Any] = [ self._system_message, *self._messages, HumanMessage(content=message), ] self._messages.append(HumanMessage(content=message)) assistant_chunks: list[str] = [] tool_names_seen: set[str] = set() tool_outputs: list[tuple[str, str]] = [] stream_config = langgraph_trace_config( self._session_hash, recursion_limit=graph_recursion_limit(), ) max_rounds = 1 + ( langgraph_max_continuations() if langgraph_auto_continue_enabled() else 0 ) for round_idx in range(max_rounds): if round_idx > 0: yield AgentStreamEvent( kind="status", text="Workflow incomplete — nudging agent to continue…", ) round_messages_start = len(graph_messages) chunks_at_round_start = len(assistant_chunks) tools_at_round_start = set(tool_names_seen) outputs_at_round_start = len(tool_outputs) overflow_retried = False tool_json_retried = False loop_break_requested = [False] user_aborted = False while True: try: self._abort_requested = False for evt in self._stream_graph_round( self._graph, graph_messages, stream_config, assistant_chunks=assistant_chunks, tool_names_seen=tool_names_seen, tool_outputs=tool_outputs, loop_break_requested=loop_break_requested, ): if evt.kind == "done": if loop_break_requested[0]: # Soft stop: keep messages, continue outer rounds. yield AgentStreamEvent( kind="status", text=evt.text or "Breaking tool-error loop…", ) break user_aborted = True yield evt break yield evt else: # stream finished without done break if user_aborted: return if loop_break_requested[0]: break break except Exception as exc: # Roll back state from the failed stream attempt. def _rollback_round() -> None: del graph_messages[round_messages_start:] del assistant_chunks[chunks_at_round_start:] tool_names_seen.clear() tool_names_seen.update(tools_at_round_start) del tool_outputs[outputs_at_round_start:] if not overflow_retried and is_context_overflow_error( exc ): overflow_retried = True _rollback_round() yield AgentStreamEvent( kind="status", text=( "Prompt exceeded model context — retrying once " "with aggressive compaction…" ), ) set_aggressive_trim(True) try: self._rebuild_graph(aggressive_compaction=True) except Exception: set_aggressive_trim(False) raise continue if ( not tool_json_retried and is_tool_call_json_parse_error(exc) ): tool_json_retried = True _rollback_round() yield AgentStreamEvent( kind="status", text=( "Tool-call JSON was truncated or invalid — " "retrying once with a compact-script nudge…" ), ) graph_messages.append( HumanMessage( content=build_tool_call_json_retry_prompt() ) ) # Keep the nudge if a later overflow rollback runs. round_messages_start = len(graph_messages) chunks_at_round_start = len(assistant_chunks) tools_at_round_start = set(tool_names_seen) outputs_at_round_start = len(tool_outputs) continue raise if overflow_retried: set_aggressive_trim(False) # Restore normal (non-aggressive) graph for later rounds. self._rebuild_graph(aggressive_compaction=False) streak = identical_tool_error_streak( tool_outputs, min_streak=langgraph_identical_error_stop_streak(), ) assistant_so_far = "\n".join(assistant_chunks) workflow_incomplete = redaction_workflow_incomplete( tool_names_seen, tool_outputs, assistant_text=assistant_so_far, ) if not workflow_incomplete and not streak: break if round_idx >= max_rounds - 1: break # Always prefer the identical-error / pending-script continue prompt. graph_messages.append( HumanMessage( content=build_workflow_continue_prompt( tool_names_seen, tool_outputs ) ) ) self._abort_requested = False if assistant_chunks: self._messages.append( AIMessage(content="\n".join(assistant_chunks)) ) assistant_final = "\n".join(assistant_chunks) awaiting_clarification = clarification_requested( tool_outputs, assistant_text=assistant_final ) workflow_incomplete = redaction_workflow_incomplete( tool_names_seen, tool_outputs, assistant_text=assistant_final, ) turn_meta["workflow_incomplete"] = workflow_incomplete turn_meta["awaiting_clarification"] = awaiting_clarification done_text = "Agent finished." if awaiting_clarification: done_text = ( "Agent paused — awaiting your clarification " "(reply in chat to continue Pass 1)." ) elif workflow_incomplete: done_text = ( "Agent finished (workflow incomplete — review_apply not run; " "use **Send** to continue or restart the task)." ) yield AgentStreamEvent( kind="done", text=done_text, meta={ "workflow_incomplete": workflow_incomplete, "awaiting_clarification": awaiting_clarification, }, ) finally: set_aggressive_trim(False) self._is_compacting = False self._prompt_stream_depth = max(0, self._prompt_stream_depth - 1) @staticmethod def _stringify_content(content: Any) -> str: if isinstance(content, str): return content if isinstance(content, list): parts: list[str] = [] for block in content: if isinstance(block, str): parts.append(block) elif isinstance(block, dict) and block.get("type") == "text": parts.append(str(block.get("text") or "")) return "".join(parts) return str(content or "")