From aef3f5e73c25d6357186fc9c0be5fc9514b93638 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Sun, 30 Mar 2025 09:21:25 +0200 Subject: [PATCH] Streaming now working for OLLAMA models, but not the rest. Figure it out MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- examples/cai/simple_one_tool_test_streamed.py | 61 ++++++ examples/cai/simple_test.py | 30 --- .../agents/models/openai_chatcompletions.py | 179 ++++++++++++++---- 3 files changed, 205 insertions(+), 65 deletions(-) create mode 100644 examples/cai/simple_one_tool_test_streamed.py delete mode 100644 examples/cai/simple_test.py diff --git a/examples/cai/simple_one_tool_test_streamed.py b/examples/cai/simple_one_tool_test_streamed.py new file mode 100644 index 00000000..08ea6a35 --- /dev/null +++ b/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()) \ No newline at end of file diff --git a/examples/cai/simple_test.py b/examples/cai/simple_test.py deleted file mode 100644 index 007a79ea..00000000 --- a/examples/cai/simple_test.py +++ /dev/null @@ -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()) \ No newline at end of file diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 65dd284e..7bcdd17f 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/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: - if tc_delta.index not in state.function_calls: - state.function_calls[tc_delta.index] = ResponseFunctionToolCall( + tool_calls = delta.tool_calls + elif isinstance(delta, dict) and 'tool_calls' in delta and delta['tool_calls']: + 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 - - state.function_calls[tc_delta.index].arguments += ( - tc_function.arguments if tc_function else "" - ) or "" - state.function_calls[tc_delta.index].name += ( - tc_function.name if tc_function else "" - ) or "" - state.function_calls[tc_delta.index].call_id += tc_delta.id or "" + + tc_function = None + if hasattr(tc_delta, 'function'): + tc_function = tc_delta.function + elif isinstance(tc_delta, dict) and 'function' in tc_delta: + tc_function = tc_delta['function'] + + if tc_function: + # 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, - output_tokens=usage.completion_tokens, - total_tokens=usage.total_tokens, + input_tokens=usage.prompt_tokens if usage and hasattr(usage, 'prompt_tokens') else 0, + output_tokens=usage.completion_tokens if usage and hasattr(usage, 'completion_tokens') else 0, + 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 - else None + if usage is not 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, - api_base=get_ollama_api_base(), - custom_llm_provider="openai" + **ollama_kwargs, + api_base=ollama_api_base, + 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, - api_base=get_ollama_api_base(), - custom_llm_provider="openai" + **ollama_kwargs, + api_base=ollama_api_base, + custom_llm_provider="ollama" ) return ret else: