lambdagent/lambdagent_guard/langchain.py

"""
lambdagent_guard.langchain — LangChain guard wrapper (I11).

Usage:
    from lambdagent_guard import guard_langchain

    executor = AgentExecutor(agent=agent, tools=tools, max_iterations=20)
    guarded = guard_langchain(executor, cost_budget=5.0)
    result = guarded.invoke({"input": "Analyze this repo"})

Catches: LangChain #10997 (type crash), #2495 (infinite loop), #26019 (tool re-invoke)
"""

from __future__ import annotations
from typing import Any, Callable, Optional
from .core import GuardConfig, RuntimeMonitor, run_compile_checks
import warnings


def guard_langchain(
    executor,
    cost_budget: float = float("inf"),
    type_check: bool = True,
    loop_detection: bool = True,
    cost_alert: Optional[Callable[[float], None]] = None,
    parallel_safety: bool = False,
):
    """
    Wrap a LangChain AgentExecutor with lambdagent safety guards.

    Non-invasive: returns the same executor with hooks injected.
    User code changes: 2 lines (import + wrap).

    Args:
        executor: LangChain AgentExecutor
        cost_budget: Max cost in USD, auto-pause if exceeded
        type_check: Verify tool I/O type compatibility at wrap time
        loop_detection: Detect repeated identical states at runtime
        cost_alert: Callback(cost_usd) on each step
    """
    guard_cfg = GuardConfig(
        cost_budget=cost_budget,
        type_check=type_check,
        loop_detection=loop_detection,
        cost_alert=cost_alert,
        parallel_safety=parallel_safety,
    )

    # Phase 1: Compile-time checks
    try:
        from lambdagent.extractors import extract_config
        config = extract_config(executor, framework="langchain")
        issues = run_compile_checks(config, guard_cfg)
        for issue in issues:
            warnings.warn(f"[lambdagent-guard] {issue}")
    except Exception as e:
        warnings.warn(f"[lambdagent-guard] Compile check skipped: {e}")

    # Phase 2: Runtime hooks
    monitor = RuntimeMonitor(guard_cfg)

    # Hook into AgentExecutor's step method
    original_take_step = None
    if hasattr(executor, '_take_next_step'):
        original_take_step = executor._take_next_step
    elif hasattr(executor, 'agent') and hasattr(executor.agent, 'plan'):
        original_take_step = executor.agent.plan

    if original_take_step is not None:
        def guarded_step(*args, **kwargs):
            result = original_take_step(*args, **kwargs)
            # Extract info for monitoring
            output_str = str(result)[:500] if result else ""
            monitor.on_step({
                "output": output_str,
                "tokens": 0,  # LangChain doesn't expose tokens per step easily
            })
            return result

        if hasattr(executor, '_take_next_step'):
            executor._take_next_step = guarded_step
        elif hasattr(executor, 'agent') and hasattr(executor.agent, 'plan'):
            executor.agent.plan = guarded_step

    # Attach monitor for result extraction
    executor._lambdagent_monitor = monitor

    return executor