add types

Author: xzrderek

eval_protocol/__init__.py

@@ -62,6 +62,13 @@
 except ImportError:
     LangSmithAdapter = None
 
+# Remote server types
+from .types.remote_rollout_processor import (
+    InitRequest,
+    RolloutMetadata,
+    StatusResponse,
+    create_langfuse_config_tags,
+)
 
 warnings.filterwarnings("default", category=DeprecationWarning, module="eval_protocol")
 
@@ -110,6 +117,11 @@
     # Submodules
     "rewards",
     "mcp",
+    # Remote server types
+    "InitRequest",
+    "RolloutMetadata",
+    "StatusResponse",
+    "create_langfuse_config_tags",
 ]
 
 from . import _version

eval_protocol/pytest/remote_rollout_processor.py

@@ -28,7 +28,6 @@ class RemoteRolloutProcessor(RolloutProcessor):
           "run_id": str | null,
           "row_id": str | null
         },
-        "num_turns": int
       }
       Returns: {"ok": true}
 
@@ -40,15 +39,13 @@ def __init__(
         self,
         *,
         remote_base_url: Optional[str] = None,
-        num_turns: int = 2,
         poll_interval: float = 1.0,
         timeout_seconds: float = 120.0,
         output_data_loader: Callable[[str], DynamicDataLoader],
     ):
         # Prefer constructor-provided configuration. These can be overridden via
         # config.kwargs at call time for backward compatibility.
         self._remote_base_url = remote_base_url
-        self._num_turns = num_turns
         self._poll_interval = poll_interval
         self._timeout_seconds = timeout_seconds
         self._output_data_loader = output_data_loader
@@ -58,15 +55,13 @@ def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) ->
 
         # Start with constructor values
         remote_base_url: Optional[str] = self._remote_base_url
-        num_turns: int = self._num_turns
         poll_interval: float = self._poll_interval
         timeout_seconds: float = self._timeout_seconds
 
         # Backward compatibility: allow overrides via config.kwargs
         if config.kwargs:
             if remote_base_url is None:
                 remote_base_url = config.kwargs.get("remote_base_url", remote_base_url)
-            num_turns = int(config.kwargs.get("num_turns", num_turns))
             poll_interval = float(config.kwargs.get("poll_interval", poll_interval))
             timeout_seconds = float(config.kwargs.get("timeout_seconds", timeout_seconds))
 
@@ -121,7 +116,6 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
                 "messages": clean_messages,
                 "tools": row.tools,
                 "metadata": meta,
-                "num_turns": num_turns,
             }
 
             # Fire-and-poll

eval_protocol/types/remote_rollout_processor.py

@@ -0,0 +1,45 @@
+"""
+Request and response models for remote rollout processor servers.
+"""
+
+from typing import Any, Dict, List, Optional
+from pydantic import BaseModel, Field
+from eval_protocol.models import Message
+
+
+class RolloutMetadata(BaseModel):
+    """Metadata for rollout execution."""
+
+    invocation_id: str
+    experiment_id: str
+    rollout_id: str
+    run_id: str
+    row_id: str
+
+
+class InitRequest(BaseModel):
+    """Request model for POST /init endpoint."""
+
+    rollout_id: str
+    model: str
+    messages: List[Message] = Field(min_length=1)
+    tools: Optional[List[Dict[str, Any]]] = None
+    metadata: RolloutMetadata
+
+
+class StatusResponse(BaseModel):
+    """Response model for GET /status endpoint."""
+
+    terminated: bool
+
+
+def create_langfuse_config_tags(init_request: InitRequest) -> List[str]:
+    """Create Langfuse tags from InitRequest metadata."""
+    metadata = init_request.metadata
+    return [
+        f"invocation_id:{metadata.invocation_id}",
+        f"experiment_id:{metadata.experiment_id}",
+        f"rollout_id:{metadata.rollout_id}",
+        f"run_id:{metadata.run_id}",
+        f"row_id:{metadata.row_id}",
+    ]

tests/chinook/langfuse/remote_server.py

@@ -1,37 +1,32 @@
 import os
 import threading
-from typing import Any, Dict
+from typing import Any, Dict, List
 
 import uvicorn
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
 from langfuse.openai import openai  # pyright: ignore[reportPrivateImportUsage]
 
-
-app = FastAPI()
+from eval_protocol.types.remote_rollout_processor import (
+    InitRequest,
+    StatusResponse,
+    create_langfuse_config_tags,
+)
+from eval_protocol.models import Message
 
 
-class InitRequest(BaseModel):
-    rollout_id: str
-    model: str
-    messages: list[dict]
-    tools: list[dict] | None = None
-    metadata: dict
-    num_turns: int = 2
+app = FastAPI()
 
 
 _STATE: Dict[str, Dict[str, Any]] = {}
 
-
 ALLOWED_MESSAGE_FIELDS = {"role", "content", "tool_calls", "tool_call_id", "name"}
 
 
-def _clean_messages_for_api(messages: list[dict]) -> list[dict]:
+def _clean_messages_for_api(messages: List[Message]) -> list[dict]:
     cleaned: list[dict] = []
     for msg in messages:
-        if not isinstance(msg, dict):
-            continue
-        cm = {k: v for k, v in msg.items() if k in ALLOWED_MESSAGE_FIELDS and v is not None}
+        msg_dict = msg.model_dump()
+        cm = {k: v for k, v in msg_dict.items() if k in ALLOWED_MESSAGE_FIELDS and v is not None}
         # Some providers dislike empty content on assistant messages; keep if present
         cleaned.append(cm)
     return cleaned
@@ -42,53 +37,25 @@ def init(req: InitRequest):
     # Persist state
     _STATE[req.rollout_id] = {"terminated": False}
 
-    # Kick off worker thread that runs multi-turn chat via Langfuse OpenAI integration
+    # Kick off worker thread that does a single-turn chat via Langfuse OpenAI integration
     def _worker():
         try:
-            # Prepare tags for Langfuse filtering
-            metadata = {
-                "langfuse_tags": [
-                    f"invocation_id:{req.metadata.get('invocation_id')}",
-                    f"experiment_id:{req.metadata.get('experiment_id')}",
-                    f"rollout_id:{req.metadata.get('rollout_id')}",
-                    f"run_id:{req.metadata.get('run_id')}",
-                    f"row_id:{req.metadata.get('row_id')}",
-                ]
+            metadata = {"langfuse_tags": create_langfuse_config_tags(req)}
+
+            completion_kwargs = {
+                "model": req.model,
+                "messages": _clean_messages_for_api(req.messages),
+                "metadata": metadata,
             }
 
-            messages = req.messages
-
-            # Simulate N-1 assistant turns (single-shot or simple echo)
-            for _ in range(max(1, req.num_turns)):
-                completion_kwargs = {
-                    "model": req.model,
-                    "messages": _clean_messages_for_api(messages),
-                    "metadata": metadata,
-                }
-
-                if req.tools:
-                    completion_kwargs["tools"] = req.tools
-
-                completion = openai.chat.completions.create(**completion_kwargs)
-                assistant_message = completion.choices[0].message
-
-                # Convert to dict format for next turn
-                assistant_dict = {"role": "assistant", "content": assistant_message.content}
-                if assistant_message.tool_calls:
-                    assistant_dict["tool_calls"] = [
-                        {
-                            "id": tc.id,
-                            "type": tc.type,
-                            "function": {"name": tc.function.name, "arguments": tc.function.arguments},
-                        }
-                        for tc in assistant_message.tool_calls
-                    ]
-
-                # Append assistant for next turn
-                messages = messages + [assistant_dict]
-
-        except Exception:
+            if req.tools:
+                completion_kwargs["tools"] = req.tools
+
+            completion = openai.chat.completions.create(**completion_kwargs)
+
+        except Exception as e:
             # Best-effort; mark as done even on error to unblock polling
+            print(f"❌ Error in rollout {req.rollout_id}: {e}")
             pass
         finally:
             _STATE[req.rollout_id]["terminated"] = True
@@ -98,12 +65,12 @@ def _worker():
     return {"ok": True}
 
 
-@app.get("/status")
-def status(rollout_id: str):
+@app.get("/status", response_model=StatusResponse)
+def status(rollout_id: str) -> StatusResponse:
     st = _STATE.get(rollout_id)
     if not st:
         raise HTTPException(status_code=404, detail="unknown rollout_id")
-    return {"terminated": bool(st.get("terminated", False))}
+    return StatusResponse(terminated=bool(st.get("terminated", False)))
 
 
 def main():

tests/chinook/langfuse/test_remote_langfuse_chinook.py

@@ -1,12 +1,14 @@
+# MANUAL SERVER STARTUP REQUIRED:
+# Before running this test, start the remote server manually:
+# cd /Users/derekxu/Documents/code/python-sdk
+# python -m tests.chinook.langfuse.remote_server
+#
+# The server should be running on http://127.0.0.1:7077
+
 import os
-import multiprocessing
-import time
-from datetime import datetime, timedelta
 from typing import List
-import atexit
 
 import pytest
-import requests
 
 from eval_protocol.data_loader.dynamic_data_loader import DynamicDataLoader
 from eval_protocol.models import EvaluationRow, Message
@@ -33,7 +35,7 @@ def fetch_langfuse_traces(rollout_id: str) -> List[EvaluationRow]:
     ROLLOUT_IDS.add(rollout_id)
 
     adapter = create_langfuse_adapter()
-    return adapter.get_evaluation_rows(tags=[f"rollout_id:{rollout_id}"])
+    return adapter.get_evaluation_rows(tags=[f"rollout_id:{rollout_id}"], max_retries=5)
 
 
 def langfuse_output_data_loader(rollout_id: str) -> DynamicDataLoader:
@@ -42,51 +44,8 @@ def langfuse_output_data_loader(rollout_id: str) -> DynamicDataLoader:
     )
 
 
-def _start_remote_server():
-    # Starts FastAPI server defined in remote_server.py using absolute import
-    import importlib
-
-    os.environ.setdefault("REMOTE_SERVER_HOST", "127.0.0.1")
-    os.environ.setdefault("REMOTE_SERVER_PORT", "7077")
-    mod = importlib.import_module("tests.chinook.langfuse.remote_server")
-    mod.main()
-
-
-def _ensure_server_running():
-    host = os.getenv("REMOTE_SERVER_HOST", "127.0.0.1")
-    port = int(os.getenv("REMOTE_SERVER_PORT", "7077"))
-    base_url = f"http://{host}:{port}"
-
-    def _is_up() -> bool:
-        try:
-            r = requests.get(f"{base_url}/status", params={"rollout_id": "ping"}, timeout=1.0)
-            return r.status_code in (200, 404)
-        except Exception:
-            return False
-
-    if _is_up():
-        return None
-
-    # Launch in a background process
-    proc = multiprocessing.Process(target=_start_remote_server, daemon=True)
-    proc.start()
-
-    # Poll for readiness up to 10s
-    deadline = time.time() + 10
-    while time.time() < deadline:
-        if _is_up():
-            break
-        time.sleep(0.5)
-    return proc
-
-
 def remote_langfuse_data_generator() -> List[EvaluationRow]:
-    # Ensure server is running BEFORE rollouts start (evaluation_test triggers rollouts before test body)
-    _SERVER_PROC = _ensure_server_running()
-    atexit.register(lambda: (_SERVER_PROC and _SERVER_PROC.is_alive() and _SERVER_PROC.terminate()))
-
-    # Minimal single-user-turn message to trigger a response
-    row = EvaluationRow(messages=[Message(role="user", content="Hello there! Please say hi back.")])
+    row = EvaluationRow(messages=[Message(role="user", content="What is the capital of France?")])
     return [row, row, row]
 
 
@@ -98,19 +57,18 @@ def remote_langfuse_data_generator() -> List[EvaluationRow]:
     ),
     rollout_processor=RemoteRolloutProcessor(
         remote_base_url="http://127.0.0.1:7077",
-        num_turns=2,
         timeout_seconds=30,
         output_data_loader=langfuse_output_data_loader,
     ),
 )
 async def test_remote_rollout_and_fetch_langfuse(row: EvaluationRow) -> EvaluationRow:
     """
     End-to-end test:
-    - remote server started at import time
+    - REQUIRES MANUAL SERVER STARTUP: python -m tests.chinook.langfuse.remote_server
     - trigger remote rollout via RemoteRolloutProcessor (calls init/status)
     - fetch traces from Langfuse filtered by metadata via output_data_loader; FAIL if none found
     """
-    assert row.messages[0].content == "Hello there! Please say hi back.", "Row should have correct message content"
+    assert row.messages[0].content == "What is the capital of France?", "Row should have correct message content"
     assert len(row.messages) > 1, "Row should have a response. If this fails, we fellback to the original row."
 
     assert row.execution_metadata.rollout_id in ROLLOUT_IDS, (