"""LangChain integration — execute step, result type, and trace converter.""" from __future__ import annotations import logging from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Tuple from ..core.context import ACEStepContext from ..implementations.prompts import wrap_skillbook_for_external_agent logger = logging.getLogger(__name__) try: from langchain_core.runnables import Runnable LANGCHAIN_AVAILABLE = True except ImportError: LANGCHAIN_AVAILABLE = False Runnable = None # type: ignore try: from langchain.agents import AgentExecutor AGENT_EXECUTOR_AVAILABLE = True except ImportError: AGENT_EXECUTOR_AVAILABLE = False AgentExecutor = None # type: ignore try: from langgraph.graph.state import CompiledStateGraph LANGGRAPH_AVAILABLE = True except ImportError: LANGGRAPH_AVAILABLE = False CompiledStateGraph = None # type: ignore # --------------------------------------------------------------------------- # Input / Output types # --------------------------------------------------------------------------- @dataclass class LangChainResult: """Output from a LangChain Runnable execution. This is the integration-specific result — not yet in ACE trace format. Use ``LangChainToTrace`` to convert to a standardised trace dict. ``result_type`` indicates the source variant: - ``"simple"`` — basic chain (prompt | llm) - ``"agent"`` — AgentExecutor with intermediate_steps - ``"langgraph"`` — LangGraph CompiledStateGraph with messages - ``"error"`` — execution failed """ task: str output: str = "" result_type: str = "simple" success: bool = True error: Optional[str] = None intermediate_steps: List[Tuple[Any, Any]] = field(default_factory=list) messages: List[Any] = field(default_factory=list) raw_result: Any = None # --------------------------------------------------------------------------- # Execute step # --------------------------------------------------------------------------- class LangChainExecuteStep: """INJECT skillbook context and EXECUTE a LangChain Runnable. Reads input from ``ctx.sample`` (string, dict, or message list), writes a ``LangChainResult`` to ``ctx.trace``. Handles three Runnable variants automatically: - Simple chains (prompt | llm) - AgentExecutor (intermediate_steps tracing) - LangGraph CompiledStateGraph (message-based I/O) """ requires = frozenset({"sample", "skillbook"}) provides = frozenset({"trace"}) def __init__( self, runnable: Any, output_parser: Optional[Callable[[Any], str]] = None, ) -> None: if not LANGCHAIN_AVAILABLE: raise ImportError( "LangChain is not installed. Install with: " "pip install langchain-core" ) self.runnable = runnable self.output_parser = output_parser or self._default_output_parser def __call__(self, ctx: ACEStepContext) -> ACEStepContext: task = self._get_task_str(ctx.sample) # -- INJECT -- enhanced_input = self._inject_context(ctx.sample, ctx.skillbook) # -- EXECUTE -- is_agent = self._is_agent_executor() is_langgraph = self._is_langgraph() original_setting = False if is_agent: original_setting = getattr( self.runnable, "return_intermediate_steps", False ) self.runnable.return_intermediate_steps = True try: raw = self.runnable.invoke(enhanced_input) except Exception as exc: if is_agent: self.runnable.return_intermediate_steps = original_setting result = LangChainResult( task=task, output=f"Failed: {exc}", result_type="error", success=False, error=str(exc), ) return ctx.replace(trace=result) finally: if is_agent: self.runnable.return_intermediate_steps = original_setting # -- BUILD RESULT -- if is_agent and isinstance(raw, dict) and "intermediate_steps" in raw: result = self._build_agent_result(task, raw) elif is_langgraph and isinstance(raw, dict) and "messages" in raw: result = self._build_langgraph_result(task, raw) else: result = self._build_simple_result(task, raw) return ctx.replace(trace=result) # ------------------------------------------------------------------ # Injection # ------------------------------------------------------------------ @staticmethod def _inject_context(original_input: Any, skillbook: Any) -> Any: if skillbook is None: return original_input context = wrap_skillbook_for_external_agent(skillbook) if not context: return original_input if isinstance(original_input, str): return f"{original_input}\n\n{context}" if isinstance(original_input, dict) and "messages" in original_input: messages = original_input["messages"] if messages and hasattr(messages[0], "content"): enhanced = list(messages) first = enhanced[0] enhanced[0] = type(first)(content=f"{context}\n\n{first.content}") return { "messages": enhanced, **{k: v for k, v in original_input.items() if k != "messages"}, } return original_input if isinstance(original_input, dict) and "input" in original_input: enhanced_dict = original_input.copy() enhanced_dict["input"] = f"{original_input['input']}\n\n{context}" return enhanced_dict if isinstance(original_input, dict): enhanced_dict = original_input.copy() enhanced_dict["skillbook_context"] = context return enhanced_dict return original_input # ------------------------------------------------------------------ # Result builders # ------------------------------------------------------------------ def _build_simple_result(self, task: str, raw: Any) -> LangChainResult: return LangChainResult( task=task, output=self.output_parser(raw), result_type="simple", raw_result=raw, ) def _build_agent_result(self, task: str, raw: Dict[str, Any]) -> LangChainResult: output = raw.get("output", "") steps = raw.get("intermediate_steps", []) intermediate: List[Tuple[Any, Any]] = [] for step_tuple in steps: if len(step_tuple) == 2: intermediate.append(tuple(step_tuple)) # type: ignore[arg-type] return LangChainResult( task=task, output=str(output), result_type="agent", intermediate_steps=intermediate, raw_result=raw, ) def _build_langgraph_result( self, task: str, raw: Dict[str, Any] ) -> LangChainResult: messages = raw.get("messages", []) output = self._extract_langgraph_output(raw) intermediate = self._extract_langgraph_steps(raw) return LangChainResult( task=task, output=output, result_type="langgraph", intermediate_steps=intermediate, messages=list(messages), raw_result=raw, ) # ------------------------------------------------------------------ # LangGraph helpers # ------------------------------------------------------------------ @staticmethod def _extract_langgraph_output(result: Dict[str, Any]) -> str: for msg in reversed(result.get("messages", [])): if hasattr(msg, "content") and msg.content: msg_type = getattr(msg, "type", msg.__class__.__name__.lower()) if msg_type != "tool": return str(msg.content) return "" @staticmethod def _extract_langgraph_steps( result: Dict[str, Any], ) -> List[Tuple[Any, Any]]: intermediate: List[Tuple[Any, Any]] = [] for msg in result.get("messages", []): msg_type = getattr(msg, "type", msg.__class__.__name__.lower()) content = getattr(msg, "content", str(msg)) if msg_type == "ai": for tc in getattr(msg, "tool_calls", []): intermediate.append((tc, None)) elif msg_type == "tool": for i in range(len(intermediate) - 1, -1, -1): if intermediate[i][1] is None: intermediate[i] = (intermediate[i][0], content) break return intermediate # ------------------------------------------------------------------ # Utilities # ------------------------------------------------------------------ def _is_agent_executor(self) -> bool: if not AGENT_EXECUTOR_AVAILABLE or AgentExecutor is None: return False return isinstance(self.runnable, AgentExecutor) def _is_langgraph(self) -> bool: if not LANGGRAPH_AVAILABLE or CompiledStateGraph is None: return False return isinstance(self.runnable, CompiledStateGraph) @staticmethod def _get_task_str(original_input: Any) -> str: if isinstance(original_input, str): return original_input if isinstance(original_input, dict): if "messages" in original_input: messages = original_input["messages"] if messages and hasattr(messages[0], "content"): return str(messages[0].content) return ( original_input.get("input") or original_input.get("question") or original_input.get("query") or str(original_input) ) return str(original_input) @staticmethod def _default_output_parser(result: Any) -> str: if isinstance(result, str): return result if hasattr(result, "content"): return str(result.content) if isinstance(result, dict): for key in ("output", "answer", "result", "text"): if key in result: return str(result[key]) return str(result) return str(result) # --------------------------------------------------------------------------- # Convert step — LangChainResult → standardised trace dict # --------------------------------------------------------------------------- class LangChainToTrace: """Convert a ``LangChainResult`` on ``ctx.trace`` to the standardised trace dict that the learning tail (``ReflectStep``) expects. """ requires = frozenset({"trace"}) provides = frozenset({"trace"}) def __call__(self, ctx: ACEStepContext) -> ACEStepContext: r: LangChainResult = ctx.trace # type: ignore[assignment] reasoning = self._build_reasoning(r) if r.success: feedback = self._build_feedback(r) else: feedback = f"Chain execution failed. Error: {r.error}" trace: dict = { "question": r.task, "reasoning": reasoning, "answer": r.output, "skill_ids": [], "feedback": feedback, "ground_truth": None, } return ctx.replace(trace=trace) # ------------------------------------------------------------------ # Reasoning formatters per result type # ------------------------------------------------------------------ @staticmethod def _build_reasoning(r: LangChainResult) -> str: if r.result_type == "error": return ( f"Question/Task: {r.task}\n\n" f"Execution Result: FAILED\nError: {r.error}" ) if r.result_type == "agent": parts = [f"Question/Task: {r.task}", ""] parts.append( f"=== AGENT EXECUTION TRACE ({len(r.intermediate_steps)} steps) ===" ) for i, (action, observation) in enumerate(r.intermediate_steps, 1): parts.append(f"\n--- Step {i} ---") if hasattr(action, "log") and action.log: parts.append(f"Thought: {action.log}") if hasattr(action, "tool"): parts.append(f"Action: {action.tool}") parts.append(f"Action Input: {str(action.tool_input)[:300]}") parts.append(f"Observation: {str(observation)[:300]}") parts.append("\n=== END TRACE ===") parts.append(f"\nFinal Answer: {r.output}") return "\n".join(parts) if r.result_type == "langgraph": msg_parts: list[str] = [] for msg in r.messages: msg_type = getattr(msg, "type", msg.__class__.__name__.lower()) content = getattr(msg, "content", str(msg)) if msg_type == "human": msg_parts.append(f"Human: {str(content)[:300]}") elif msg_type == "ai": if content: msg_parts.append(f"Assistant: {str(content)[:300]}") for tc in getattr(msg, "tool_calls", []): name = ( tc.get("name", "unknown") if isinstance(tc, dict) else getattr(tc, "name", "unknown") ) msg_parts.append(f" Tool Call: {name}") elif msg_type == "tool": msg_parts.append(f"Tool Result: {str(content)[:300]}") trace_str = "\n".join(msg_parts) return ( f"Question/Task: {r.task}\n\n" f"=== LANGGRAPH EXECUTION TRACE ({len(r.messages)} messages) ===\n" f"{trace_str}\n" f"=== END TRACE ===\n\n" f"Final Answer: {r.output}" ) # simple return ( f"Question/Task: {r.task}\n\n" f"Chain Output: {r.output}\n\n" f"Note: External LangChain chain execution." ) @staticmethod def _build_feedback(r: LangChainResult) -> str: if r.result_type == "agent": return f"Agent completed task in {len(r.intermediate_steps)} steps" if r.result_type == "langgraph": return f"LangGraph agent completed in {len(r.messages)} messages" return f"External chain completed for task: {r.task[:200]}" __all__ = [ "LangChainExecuteStep", "LangChainResult", "LangChainToTrace", ]