Streaming now working for OLLAMA models, but not the rest. Figure it out

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>

examples/cai/simple_one_tool_test_streamed.py

@@ -0,0 +1,61 @@
+"""
+A simple example to test the one_tool agent with streaming.
+
+This script demonstrates how to initialize and run the one_tool agent
+using Runner.run_streamed() with a simple hello message to verify everything
+is working correctly, with streaming output.
+"""
+
+import os
+import asyncio
+from dotenv import load_dotenv
+from openai import AsyncOpenAI
+from openai.types.responses import ResponseTextDeltaEvent
+from cai.sdk.agents import Runner, set_default_openai_client
+from cai.agents import get_agent_by_name
+from cai.util import fix_litellm_transcription_annotations, color
+
+
+# Load environment variables
+load_dotenv()
+
+# Initialize OpenAI client
+external_client = AsyncOpenAI(
+    base_url=os.getenv('LITELLM_BASE_URL', 'http://localhost:4000'),
+    api_key=os.getenv('LITELLM_API_KEY', 'key')
+)
+set_default_openai_client(external_client)
+
+async def main():
+    # Apply litellm patch to fix the __annotations__ error
+    patch_applied = fix_litellm_transcription_annotations()
+    if not patch_applied:
+        print(color("Something went wrong patching LiteLLM fix_litellm_transcription_annotations", color="red"))
+        
+    # Get the one_tool agent
+    agent = get_agent_by_name("one_tool_agent")
+    
+    print("Testing one_tool agent with a simple hello message (streaming mode)...")
+    print(f"Using model: {os.getenv('CAI_MODEL', 'default')}")
+    
+    # Stream indicator
+    print("\nAgent response (streaming):")
+    print("-" * 40)
+    print("Agent: ", end="", flush=True)
+    
+    # Run the agent with a simple test message in streaming mode
+    result = Runner.run_streamed(agent, "Explain what you can do in a few sentences. Don't use any tools.")
+    
+    # Process the streaming response
+    async for event in result.stream_events():
+        if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
+            # Print the delta with a visible marker for each token
+            # print(f"{event.data.delta}|", end="", flush=True)
+
+            print(f"{event.data.delta}", end="", flush=True)
+    
+    print("\n" + "-" * 40)
+    print("\nTest completed successfully!")
+
+if __name__ == "__main__":
+    asyncio.run(main()) 

examples/cai/simple_test.py

@@ -1,30 +0,0 @@
-from cai.sdk.agents import Agent, Runner, AsyncOpenAI, OpenAIChatCompletionsModel
-import asyncio
-
-spanish_agent = Agent(
-    name="Spanish agent",
-    instructions="You only speak Spanish.",
-    model="o3-mini", 
-)
-
-english_agent = Agent(
-    name="English agent",
-    instructions="You only speak English",
-    model=OpenAIChatCompletionsModel( 
-        model="gpt-4o",
-        openai_client=AsyncOpenAI()
-    ),
-)
-
-triage_agent = Agent(
-    name="Triage agent",
-    instructions="Handoff to the appropriate agent based on the language of the request.",
-    handoffs=[spanish_agent, english_agent],
-    model="gpt-3.5-turbo",
-)
-
-async def main():
-    result = await Runner.run(english_agent, input="Hey there, what's up?")
-    print(result.final_output)
-
-asyncio.run(main())

src/cai/sdk/agents/models/openai_chatcompletions.py

@@ -216,15 +216,43 @@ class OpenAIChatCompletionsModel(Model):
                     )
 
                 # The usage is only available in the last chunk
-                usage = chunk.usage
+                if hasattr(chunk, 'usage'):
+                    usage = chunk.usage
+                # For Ollama/LiteLLM streams that don't have usage attribute
+                else:
+                    usage = None
 
-                if not chunk.choices or not chunk.choices[0].delta:
+                # Handle different stream chunk formats
+                if hasattr(chunk, 'choices') and chunk.choices:
+                    choices = chunk.choices
+                elif hasattr(chunk, 'delta') and chunk.delta:
+                    # Some providers might return delta directly
+                    choices = [{"delta": chunk.delta}]
+                elif isinstance(chunk, dict) and 'choices' in chunk:
+                    choices = chunk['choices']
+                else:
+                    # Skip chunks that don't contain choice data
+                    continue
+                
+                if not choices or len(choices) == 0:
+                    continue
+                
+                # Get the delta content
+                delta = choices[0].get('delta', None)
+                if not delta and hasattr(choices[0], 'delta'):
+                    delta = choices[0].delta
+                
+                if not delta:
                     continue
 
-                delta = chunk.choices[0].delta
-
                 # Handle text
-                if delta.content:
+                content = None
+                if hasattr(delta, 'content') and delta.content is not None:
+                    content = delta.content
+                elif isinstance(delta, dict) and 'content' in delta and delta['content'] is not None:
+                    content = delta['content']
+                
+                if content:
                     if not state.text_content_index_and_output:
                         # Initialize a content tracker for streaming text
                         state.text_content_index_and_output = (
@@ -263,16 +291,22 @@ class OpenAIChatCompletionsModel(Model):
                     # Emit the delta for this segment of content
                     yield ResponseTextDeltaEvent(
                         content_index=state.text_content_index_and_output[0],
-                        delta=delta.content,
+                        delta=content,
                         item_id=FAKE_RESPONSES_ID,
                         output_index=0,
                         type="response.output_text.delta",
                     )
                     # Accumulate the text into the response part
-                    state.text_content_index_and_output[1].text += delta.content
+                    state.text_content_index_and_output[1].text += content
 
                 # Handle refusals (model declines to answer)
+                refusal_content = None
                 if hasattr(delta, 'refusal') and delta.refusal:
+                    refusal_content = delta.refusal
+                elif isinstance(delta, dict) and 'refusal' in delta and delta['refusal']:
+                    refusal_content = delta['refusal']
+                
+                if refusal_content:
                     if not state.refusal_content_index_and_output:
                         # Initialize a content tracker for streaming refusal text
                         state.refusal_content_index_and_output = (
@@ -307,36 +341,66 @@ class OpenAIChatCompletionsModel(Model):
                     # Emit the delta for this segment of refusal
                     yield ResponseRefusalDeltaEvent(
                         content_index=state.refusal_content_index_and_output[0],
-                        delta=delta.refusal,
+                        delta=refusal_content,
                         item_id=FAKE_RESPONSES_ID,
                         output_index=0,
                         type="response.refusal.delta",
                     )
                     # Accumulate the refusal string in the output part
-                    state.refusal_content_index_and_output[1].refusal += delta.refusal
+                    state.refusal_content_index_and_output[1].refusal += refusal_content
 
                 # Handle tool calls
                 # Because we don't know the name of the function until the end of the stream, we'll
                 # save everything and yield events at the end
+                tool_calls = None
                 if hasattr(delta, 'tool_calls') and delta.tool_calls:
-                    for tc_delta in delta.tool_calls:
+                    tool_calls = delta.tool_calls
-                        if tc_delta.index not in state.function_calls:
+                elif isinstance(delta, dict) and 'tool_calls' in delta and delta['tool_calls']:
-                            state.function_calls[tc_delta.index] = ResponseFunctionToolCall(
+                    tool_calls = delta['tool_calls']
+                
+                if tool_calls:
+                    for tc_delta in tool_calls:
+                        tc_index = tc_delta.index if hasattr(tc_delta, 'index') else tc_delta.get('index', 0)
+                        if tc_index not in state.function_calls:
+                            state.function_calls[tc_index] = ResponseFunctionToolCall(
                                 id=FAKE_RESPONSES_ID,
                                 arguments="",
                                 name="",
                                 type="function_call",
                                 call_id="",
                             )
-                        tc_function = tc_delta.function
+                        
-
+                        tc_function = None
-                        state.function_calls[tc_delta.index].arguments += (
+                        if hasattr(tc_delta, 'function'):
-                            tc_function.arguments if tc_function else ""
+                            tc_function = tc_delta.function
-                        ) or ""
+                        elif isinstance(tc_delta, dict) and 'function' in tc_delta:
-                        state.function_calls[tc_delta.index].name += (
+                            tc_function = tc_delta['function']
-                            tc_function.name if tc_function else ""
+                            
-                        ) or ""
+                        if tc_function:
-                        state.function_calls[tc_delta.index].call_id += tc_delta.id or ""
+                            # Handle both object and dict formats
+                            args = ""
+                            if hasattr(tc_function, 'arguments'):
+                                args = tc_function.arguments or ""
+                            elif isinstance(tc_function, dict) and 'arguments' in tc_function:
+                                args = tc_function.get('arguments', "") or ""
+                                
+                            name = ""
+                            if hasattr(tc_function, 'name'):
+                                name = tc_function.name or ""
+                            elif isinstance(tc_function, dict) and 'name' in tc_function:
+                                name = tc_function.get('name', "") or ""
+                                
+                            state.function_calls[tc_index].arguments += args
+                            state.function_calls[tc_index].name += name
+                        
+                        # Handle call_id in both formats
+                        call_id = ""
+                        if hasattr(tc_delta, 'id'):
+                            call_id = tc_delta.id or ""
+                        elif isinstance(tc_delta, dict) and 'id' in tc_delta:
+                            call_id = tc_delta.get('id', "") or ""
+                            
+                        state.function_calls[tc_index].call_id += call_id
 
             function_call_starting_index = 0
             if state.text_content_index_and_output:
@@ -423,31 +487,38 @@ class OpenAIChatCompletionsModel(Model):
 
             final_response = response.model_copy()
             final_response.output = outputs
+
             final_response.usage = (
                 ResponseUsage(
-                    input_tokens=usage.prompt_tokens,
+                    input_tokens=usage.prompt_tokens if usage and hasattr(usage, 'prompt_tokens') else 0,
-                    output_tokens=usage.completion_tokens,
+                    output_tokens=usage.completion_tokens if usage and hasattr(usage, 'completion_tokens') else 0,
-                    total_tokens=usage.total_tokens,
+                    total_tokens=usage.total_tokens if usage and hasattr(usage, 'total_tokens') else 0,
                     output_tokens_details=OutputTokensDetails(
                         reasoning_tokens=usage.completion_tokens_details.reasoning_tokens
-                        if hasattr(usage, 'completion_tokens_details') 
+                        if usage and hasattr(usage, 'completion_tokens_details') 
                         and usage.completion_tokens_details
                         and hasattr(usage.completion_tokens_details, 'reasoning_tokens')
                         and usage.completion_tokens_details.reasoning_tokens
                         else 0
                     ),
                     input_tokens_details={
-                        "prompt_tokens": usage.prompt_tokens if usage.prompt_tokens else 0,
+                        "prompt_tokens": usage.prompt_tokens if usage and hasattr(usage, 'prompt_tokens') else 0,
                         "cached_tokens": usage.prompt_tokens_details.cached_tokens
-                        if hasattr(usage, 'prompt_tokens_details')
+                        if usage and hasattr(usage, 'prompt_tokens_details')
                         and usage.prompt_tokens_details
                         and hasattr(usage.prompt_tokens_details, 'cached_tokens')
                         and usage.prompt_tokens_details.cached_tokens
                         else 0
                     },
                 )
-                if usage
+                if usage is not None
-                else None
+                else ResponseUsage(
+                    input_tokens=0,
+                    output_tokens=0,
+                    total_tokens=0,
+                    output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+                    input_tokens_details={"prompt_tokens": 0, "cached_tokens": 0}
+                )
             )
 
             yield ResponseCompletedEvent(
@@ -462,6 +533,11 @@ class OpenAIChatCompletionsModel(Model):
                     "input_tokens": usage.prompt_tokens,
                     "output_tokens": usage.completion_tokens,
                 }
+            else:
+                span_generation.span_data.usage = {
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                }
 
     @overload
     async def _fetch_response(
@@ -559,7 +635,38 @@ class OpenAIChatCompletionsModel(Model):
             "extra_headers": _HEADERS,
         }
 
+        # Previously articulated as
+        # ret = await self._get_client().chat.completions.create(**kwargs)
+        # but now we're using litellm
+
         if self.is_ollama:
+            # Filter out parameters not supported by Ollama
+            ollama_supported_params = {
+                "model": kwargs["model"],
+                "messages": kwargs["messages"],
+                "temperature": kwargs["temperature"] if kwargs["temperature"] is not NOT_GIVEN else None,
+                "top_p": kwargs["top_p"] if kwargs["top_p"] is not NOT_GIVEN else None,
+                "max_tokens": kwargs["max_tokens"] if kwargs["max_tokens"] is not NOT_GIVEN else None,
+                "stream": kwargs["stream"],
+                "extra_headers": kwargs["extra_headers"]
+            }
+            
+            # Modify the messages to remove system message for Ollama
+            if ollama_supported_params["messages"] and ollama_supported_params["messages"][0].get("role") == "system":
+                # Extract the system message
+                system_content = ollama_supported_params["messages"][0].get("content", "")
+                # Remove it from the messages
+                ollama_supported_params["messages"] = ollama_supported_params["messages"][1:]
+                # If there are user messages, prepend system to first user
+                if ollama_supported_params["messages"] and ollama_supported_params["messages"][0].get("role") == "user":
+                    # Prepend the system instruction to the first user message, with a separator
+                    user_content = ollama_supported_params["messages"][0].get("content", "")
+                    if isinstance(user_content, str):
+                        ollama_supported_params["messages"][0]["content"] = f"System: {system_content}\n\nUser: {user_content}"
+
+            # Remove None values
+            ollama_kwargs = {k: v for k, v in ollama_supported_params.items() if v is not None}
+            
             if stream:
                 # For streaming with Ollama, we need to create a Response object first
                 response = Response(
@@ -590,18 +697,20 @@ class OpenAIChatCompletionsModel(Model):
                     },
                 )
                 # Get the streaming object
+                ollama_api_base = get_ollama_api_base().rstrip('/v1')  # Remove /v1 if present
                 stream_obj = await litellm.acompletion(
-                    **kwargs,
+                    **ollama_kwargs,
-                    api_base=get_ollama_api_base(),
+                    api_base=ollama_api_base,
-                    custom_llm_provider="openai"
+                    custom_llm_provider="ollama"
                 )
                 return response, stream_obj
             else:
                 # Non-streaming mode
+                ollama_api_base = get_ollama_api_base().rstrip('/v1')  # Remove /v1 if present
                 ret = litellm.completion(
-                    **kwargs,
+                    **ollama_kwargs,
-                    api_base=get_ollama_api_base(),
+                    api_base=ollama_api_base,
-                    custom_llm_provider="openai"
+                    custom_llm_provider="ollama"
                 )
                 return ret
         else: