Switching Single Turn to LiteLLM

eval_protocol/pytest/default_single_turn_rollout_process.py

@@ -1,21 +1,16 @@
 import asyncio
 from typing import List
 
-from openai import AsyncOpenAI
+from litellm import acompletion
 
-from eval_protocol.auth import get_fireworks_api_base, get_fireworks_api_key
 from eval_protocol.models import EvaluationRow, Message
 from eval_protocol.pytest.types import RolloutProcessorConfig
 
 
 async def default_single_turn_rollout_processor(
     rows: List[EvaluationRow], config: RolloutProcessorConfig
 ) -> List[EvaluationRow]:
-    """Generate a single response from a Fireworks model concurrently."""
-
-    api_key = get_fireworks_api_key()
-    api_base = get_fireworks_api_base()
-    client = AsyncOpenAI(api_key=api_key, base_url=f"{api_base}/inference/v1")
+    """Generate a single response from any supported model provider using LiteLLM."""
 
     async def process_row(row: EvaluationRow) -> EvaluationRow:
         """Process a single row asynchronously."""
@@ -24,10 +19,13 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
 
         messages_payload = [{"role": m.role, "content": m.content} for m in row.messages]
 
-        create_kwargs = dict(model=config.model, messages=messages_payload, **config.input_params)
+        request_params = {"model": config.model, "messages": messages_payload, **config.input_params}
+
         if row.tools is not None:
-            create_kwargs["tools"] = row.tools
-        response = await client.chat.completions.create(**create_kwargs)
+            request_params["tools"] = row.tools
+
+        response = await acompletion(**request_params)
+
         assistant_content = response.choices[0].message.content or ""
         tool_calls = response.choices[0].message.tool_calls if response.choices[0].message.tool_calls else None
         messages = list(row.messages) + [