take out litellm

Author: xzrderek

eval_protocol/pytest/remote_rollout_processor.py

@@ -4,7 +4,7 @@
 
 import requests
 
-from eval_protocol.models import EvaluationRow
+from eval_protocol.models import EvaluationRow, Status
 from eval_protocol.data_loader.dynamic_data_loader import DynamicDataLoader
 from .rollout_processor import RolloutProcessor
 from .types import RolloutProcessorConfig
@@ -167,12 +167,15 @@ def _load_data():
 
             output_rows: List[EvaluationRow] = [row for result in results for row in result.rows]
 
-            assert len(output_rows) == 1, "Dataloader used for RemoteRolloutProcessor should have exactly one row"
-
-            langfuse_row = output_rows[0]
-            langfuse_row.input_metadata.completion_params = row.input_metadata.completion_params
-
-            return langfuse_row
+            if len(output_rows) == 0:  # Fallback to original row if no Langfuse data found
+                row.rollout_status = Status(code=Status.Code.NOT_FOUND, message="No Langfuse data found for rollout")
+                return row
+            elif len(output_rows) == 1:  # Return the Langfuse row
+                langfuse_row = output_rows[0]
+                langfuse_row.input_metadata.completion_params = row.input_metadata.completion_params
+                return langfuse_row
+            else:
+                raise ValueError("RemoteRolloutProcessor's output_data_loader should return exactly one row.")
 
         for r in rows:
             tasks.append(asyncio.create_task(_process_row(r)))

tests/chinook/langfuse/remote_server.py

@@ -5,7 +5,7 @@
 import uvicorn
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-import requests
+from langfuse.openai import openai  # pyright: ignore[reportPrivateImportUsage]
 
 
 app = FastAPI()
@@ -42,52 +42,50 @@ def init(req: InitRequest):
     # Persist state
     _STATE[req.rollout_id] = {"terminated": False}
 
-    # Kick off worker thread that runs multi-turn chat via LiteLLM proxy
+    # Kick off worker thread that runs multi-turn chat via Langfuse OpenAI integration
     def _worker():
         try:
-            base_url = os.getenv(
-                "LITELLM_BASE_URL",
-                "https://litellm-cloud-proxy-prod-644257448872.us-central1.run.app",
-            )
-            url = f"{base_url}/v1/chat/completions"
-            headers = {
-                "Authorization": f"Bearer {os.getenv('FIREWORKS_API_KEY', '')}",
-                "Content-Type": "application/json",
-            }
-
-            # Prepare metadata payload to attach for Langfuse filtering
+            # Prepare tags for Langfuse filtering
             metadata = {
-                "tags": [
+                "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')}",
-                ],
-                "invocation_id": req.metadata.get("invocation_id"),
-                "experiment_id": req.metadata.get("experiment_id"),
-                "rollout_id": req.metadata.get("rollout_id"),
-                "run_id": req.metadata.get("run_id"),
-                "row_id": req.metadata.get("row_id"),
+                ]
             }
 
             messages = req.messages
 
             # Simulate N-1 assistant turns (single-shot or simple echo)
             for _ in range(max(1, req.num_turns)):
-                payload = {
+                completion_kwargs = {
                     "model": req.model,
                     "messages": _clean_messages_for_api(messages),
                     "metadata": metadata,
                 }
+
                 if req.tools:
-                    payload["tools"] = req.tools
-                r = requests.post(url, json=payload, headers=headers, timeout=60)
-                r.raise_for_status()
-                data = r.json()
-                assistant = data.get("choices", [{}])[0].get("message", {})
+                    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]
+                messages = messages + [assistant_dict]
 
         except Exception:
             # Best-effort; mark as done even on error to unblock polling

tests/chinook/langfuse/test_remote_langfuse_chinook.py

@@ -25,11 +25,7 @@ def check_rollout_coverage():
     ROLLOUT_IDS.clear()
     yield
 
-    # Verify we've seen the expected number of rollout_ids after test is done
-    expected_rollout_count = 3
-    assert len(ROLLOUT_IDS) == expected_rollout_count, (
-        f"Expected to see {expected_rollout_count} rollout_ids, but only saw {len(ROLLOUT_IDS)}: {ROLLOUT_IDS}"
-    )
+    assert len(ROLLOUT_IDS) == 3, f"Expected to see {ROLLOUT_IDS} rollout_ids, but only saw {ROLLOUT_IDS}"
 
 
 def fetch_langfuse_traces(rollout_id: str) -> List[EvaluationRow]:
@@ -95,7 +91,7 @@ def remote_langfuse_data_generator() -> List[EvaluationRow]:
 
 
 @pytest.mark.skipif(os.environ.get("CI") == "true", reason="Only run this test locally (skipped in CI)")
-@pytest.mark.parametrize("completion_params", [{"model": "fireworks_ai/accounts/fireworks/models/kimi-k2-instruct"}])
+@pytest.mark.parametrize("completion_params", [{"model": "gpt-4o"}])
 @evaluation_test(
     data_loaders=DynamicDataLoader(
         generators=[remote_langfuse_data_generator],
@@ -115,6 +111,8 @@ async def test_remote_rollout_and_fetch_langfuse(row: EvaluationRow) -> Evaluati
     - 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 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, (
         f"Row rollout_id {row.execution_metadata.rollout_id} should be in tracked rollout_ids: {ROLLOUT_IDS}"
     )