simplify further

Author: benjibc

eval_protocol/pytest/default_langchain_rollout_processor.py

@@ -1,110 +1,144 @@
 import asyncio
-import time
-from typing import List, Any, cast
+from typing import Any, Callable, Dict, List, Optional
 
-try:
-    from langchain_core.messages import BaseMessage as LCBaseMessage, HumanMessage  # type: ignore
-except ImportError:  # pragma: no cover - optional dependency path
-    # Minimal fallbacks to satisfy typing when langchain is not present
-    class LCBaseMessage:  # type: ignore
-        content: str
-        type: str
-
-        def __init__(self, content: str = "", msg_type: str = "assistant"):
-            self.content = content
-            self.type = msg_type
-
-    class HumanMessage(LCBaseMessage):  # type: ignore
-        def __init__(self, content: str):
-            super().__init__(content=content, msg_type="human")
-
-
-from eval_protocol.models import EvaluationRow, Message
+from eval_protocol.models import EvaluationRow, Status, Message
 from eval_protocol.pytest.rollout_processor import RolloutProcessor
-from eval_protocol.pytest.types import RolloutProcessorConfig
+from eval_protocol.pytest.types import CompletionParams, RolloutProcessorConfig
 
 
 class LangGraphRolloutProcessor(RolloutProcessor):
-    """Generic rollout processor for LangChain agents.
-
-    Accepts an async factory that returns a target to invoke. The target can be:
-    - An object with `.graph.ainvoke(payload)` (e.g., LangGraph compiled graph)
-    - An object with `.ainvoke(payload)`
-    - A callable that accepts `payload` and returns the result dict
     """
+    Generic rollout processor for LangGraph graphs.
 
-    def __init__(self, get_invoke_target):
-        self.get_invoke_target = get_invoke_target
-
-    def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig):
-        tasks: List[asyncio.Task] = []
-
-        async def _process_row(row: EvaluationRow) -> EvaluationRow:
-            start_time = time.perf_counter()
-
-            # Build LC messages from EP row (minimal: last user to HumanMessage)
-            lm_messages: List[LCBaseMessage] = []
-            if row.messages:
-                last_user = [m for m in row.messages if m.role == "user"]
-                if last_user:
-                    content = last_user[-1].content or ""
-                    if isinstance(content, list):
-                        content = "".join([getattr(p, "text", str(p)) for p in content])
-                    lm_messages.append(HumanMessage(content=str(content)))
-            if not lm_messages:
-                lm_messages = [HumanMessage(content="")]
-
-            target = await self.get_invoke_target(config)
-
-            # Resolve the appropriate async invoke function
-            if hasattr(target, "graph") and hasattr(target.graph, "ainvoke"):
+    Configure with:
+    - to_input(row): build the input payload for graph.ainvoke (default: {"messages": row.messages})
+    - apply_result(row, result): write graph outputs back onto the row (default: row.messages = result["messages"])
+    - build_graph_kwargs(cp): map completion_params to graph kwargs (default: {})
 
-                async def _invoke_graph(payload):
-                    return await target.graph.ainvoke(payload)  # type: ignore[attr-defined]
-
-                invoke_fn = _invoke_graph
-            elif hasattr(target, "ainvoke"):
-
-                async def _invoke_direct(payload):
-                    return await target.ainvoke(payload)  # type: ignore[attr-defined]
-
-                invoke_fn = _invoke_direct
-            elif callable(target):
-
-                async def _invoke_wrapper(payload):
-                    result = target(payload)
-                    if asyncio.iscoroutine(result):
-                        return await result
-                    return result
-
-                invoke_fn = _invoke_wrapper
-            else:
-                raise TypeError("Unsupported invoke target for LangGraphRolloutProcessor")
-
-            result_obj = await invoke_fn({"messages": lm_messages})
-            if isinstance(result_obj, dict):
-                result_messages: List[LCBaseMessage] = result_obj.get("messages", [])
-            else:
-                result_messages = getattr(result_obj, "messages", [])
+    Compatible with eval_protocol.pytest.evaluation_test.
+    """
 
-            def _serialize_message(msg: LCBaseMessage) -> Message:
-                try:
-                    from eval_protocol.adapters.langchain import serialize_lc_message_to_ep as _ser
-                except ImportError:
-                    content = getattr(msg, "content", "")
-                    return Message(role=getattr(msg, "type", "assistant"), content=str(content))
-                return _ser(cast(Any, msg))
+    def __init__(
+        self,
+        *,
+        graph_factory: Callable[[Dict[str, Any]], Any],
+        to_input: Optional[Callable[[EvaluationRow], Dict[str, Any]]] = None,
+        apply_result: Optional[Callable[[EvaluationRow, Any], EvaluationRow]] = None,
+        build_graph_kwargs: Optional[Callable[[CompletionParams], Dict[str, Any]]] = None,
+        input_key: str = "messages",
+        output_key: str = "messages",
+    ) -> None:
+        # Build the graph per-call using completion_params
+        self._graph_factory = graph_factory
+        self._to_input = to_input
+        self._apply_result = apply_result
+        self._build_graph_kwargs = build_graph_kwargs
+        self._input_key = input_key
+        self._output_key = output_key
+
+    def _default_to_input(self, row: EvaluationRow) -> Dict[str, Any]:
+        messages = row.messages or []
+        from eval_protocol.adapters.langchain import serialize_ep_messages_to_lc as _to_lc
+
+        return {self._input_key: _to_lc(messages)}
+
+    def _default_apply_result(self, row: EvaluationRow, result: Any) -> EvaluationRow:
+        # Expect dict with output_key → list of messages; coerce to EP messages
+        maybe_msgs = None
+        if isinstance(result, dict):
+            maybe_msgs = result.get(self._output_key)
+
+        if maybe_msgs is None:
+            return row
 
-            row.messages = [_serialize_message(m) for m in result_messages]
+        # If already EP messages, assign directly
+        if isinstance(maybe_msgs, list) and all(isinstance(m, Message) for m in maybe_msgs):
+            row.messages = maybe_msgs
+            return row
 
-            row.execution_metadata.duration_seconds = time.perf_counter() - start_time
+        # Try to convert from LangChain messages; preserve EP Message items as-is
+        try:
+            from langchain_core.messages import BaseMessage as _LCBase
+            from eval_protocol.adapters.langchain import serialize_lc_message_to_ep as _to_ep
+
+            if isinstance(maybe_msgs, list) and any(isinstance(m, _LCBase) for m in maybe_msgs):
+                converted: List[Message] = []
+                for m in maybe_msgs:
+                    if isinstance(m, Message):
+                        converted.append(m)
+                    elif isinstance(m, _LCBase):
+                        converted.append(_to_ep(m))
+                    elif isinstance(m, dict):
+                        role = m.get("role") or "assistant"
+                        content = m.get("content")
+                        converted.append(Message(role=role, content=content))
+                    else:
+                        # Best-effort for LC-like objects without importing LC types
+                        role_like = getattr(m, "type", None)
+                        content_like = getattr(m, "content", None)
+                        if content_like is not None:
+                            role_value = "assistant"
+                            if isinstance(role_like, str):
+                                rl = role_like.lower()
+                                if rl in ("human", "user"):
+                                    role_value = "user"
+                                elif rl in ("ai", "assistant"):
+                                    role_value = "assistant"
+                                elif rl in ("system",):
+                                    role_value = "system"
+                            converted.append(Message(role=role_value, content=str(content_like)))
+                        else:
+                            converted.append(Message(role="assistant", content=str(m)))
+                row.messages = converted
+                return row
+        except ImportError:
+            # If LC is not available, fall back to best-effort below
+            pass
+
+        # Generic best-effort fallback: stringify to assistant messages
+        if isinstance(maybe_msgs, list):
+            row.messages = [Message(role="assistant", content=str(m)) for m in maybe_msgs]
+        else:
+            row.messages = [Message(role="assistant", content=str(maybe_msgs))]
+        return row
+
+    def _default_build_graph_kwargs(self, _: CompletionParams) -> Dict[str, Any]:
+        # Keep generic: callers can override to map to their graph’s expected kwargs
+        return {}
+
+    def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[asyncio.Task[EvaluationRow]]:
+        tasks: List[asyncio.Task[EvaluationRow]] = []
+
+        to_input = self._to_input or self._default_to_input
+        apply_result = self._apply_result or self._default_apply_result
+        build_kwargs = self._build_graph_kwargs or self._default_build_graph_kwargs
+
+        graph_config: Optional[Dict[str, Any]] = None
+        if config.completion_params:
+            graph_config = build_kwargs(config.completion_params)
+
+        # (Re)build the graph for this call using the graph kwargs
+        graph_target = self._graph_factory(graph_config or {})
 
-            return row
+        async def _process_row(row: EvaluationRow) -> EvaluationRow:
+            try:
+                payload = to_input(row)
+                if graph_config is not None:
+                    result = await graph_target.ainvoke(payload, config=graph_config)
+                else:
+                    result = await graph_target.ainvoke(payload)
+                row = apply_result(row, result)
+                row.rollout_status = Status.rollout_finished()
+                return row
+            except (RuntimeError, ValueError, TypeError, KeyError, AttributeError, ImportError) as e:  # noqa: BLE001
+                row.rollout_status = Status.rollout_error(str(e))
+                return row
 
         for r in rows:
             tasks.append(asyncio.create_task(_process_row(r)))
 
         return tasks
 
     def cleanup(self) -> None:
+        # No-op by default
         return None

eval_protocol/pytest/langgraph_processor.py

@@ -1,3 +1 @@
 {"name":"p1","prompt":"Say hello in one sentence","gt":"hello"}
-{"name":"p2","prompt":"Introduce yourself briefly","gt":"intro"}
-{"name":"p3","prompt":"Respond with a fun fact about space","gt":"space"}

examples/langgraph/data/simple_prompts.jsonl

@@ -2,7 +2,7 @@
 
 from eval_protocol.models import EvaluationRow, EvaluateResult, Message
 from eval_protocol.pytest import evaluation_test
-from eval_protocol.pytest.langgraph_processor import LangGraphRolloutProcessor
+from eval_protocol.pytest.default_langchain_rollout_processor import LangGraphRolloutProcessor
 from eval_protocol.pytest.types import RolloutProcessorConfig as _UnusedRolloutProcessorConfig  # noqa: F401
 
 from examples.langgraph.simple_graph import build_simple_graph