diff --git a/examples/cai/simple_one_tool_test.py b/examples/cai/simple_one_tool_test.py new file mode 100644 index 00000000..056ad636 --- /dev/null +++ b/examples/cai/simple_one_tool_test.py @@ -0,0 +1,51 @@ +""" +A simple example to test the one_tool agent. + +This script demonstrates how to initialize and run the one_tool agent +using Runner.run() with a simple hello message to verify everything +is working correctly. +""" + +import os +import asyncio +from dotenv import load_dotenv +from openai import AsyncOpenAI +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...") + print(f"Using model: {os.getenv('CAI_MODEL', 'default')}") + + # Run the agent with a simple test message + result = await Runner.run(agent, "Hello! Can you introduce yourself and explain what you can do?") + + # Print the result + print("\nAgent response:") + print("-" * 40) + print(result.final_output) + print("-" * 40) + print("\nTest completed successfully!") + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/src/cai/agents/one_tool.py b/src/cai/agents/one_tool.py index f9537603..6d08c576 100644 --- a/src/cai/agents/one_tool.py +++ b/src/cai/agents/one_tool.py @@ -15,9 +15,11 @@ openai_client = AsyncOpenAI( api_key=os.getenv('LITELLM_API_KEY', 'key') ) +# # Check if we're using a Qwen model +# is_qwen = "qwen" in model_name.lower() + # For Qwen models, we need to skip system instructions as they're not supported -is_qwen = "qwen" in model_name.lower() -instructions = None if is_qwen else """You are a Cybersecurity expert Leader facing a CTF +instructions = """You are a Cybersecurity expert Leader facing a CTF challenge. INSTRUCTIONS: 1. Execute the generic_linux_command tool without any diff --git a/src/cai/cli.py b/src/cai/cli.py index 553b4918..4efc90f8 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -104,6 +104,7 @@ from cai.sdk.agents import set_default_openai_client, set_tracing_disabled from openai.types.responses import ResponseTextDeltaEvent from rich.console import Console import asyncio +from cai.util import fix_litellm_transcription_annotations, color # Import modules from cai.repl from cai.repl.commands import FuzzyCommandCompleter, handle_command as commands_handle_command @@ -238,6 +239,11 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= 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 agent type from environment variables or use default agent_type = os.getenv('CAI_AGENT_TYPE', "one_tool_agent") diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 3bece718..65dd284e 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -3,9 +3,13 @@ from __future__ import annotations import dataclasses import json import time +import os +import litellm + from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast, overload +from cai.util import get_ollama_api_base from openai import NOT_GIVEN, AsyncOpenAI, AsyncStream, NotGiven from openai.types import ChatModel @@ -103,6 +107,8 @@ class OpenAIChatCompletionsModel(Model): ) -> None: self.model = model self._client = openai_client + # Check if we're using OLLAMA models + self.is_ollama = os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false' def _non_null_or_not_given(self, value: Any) -> Any: return value if value is not None else NOT_GIVEN @@ -266,7 +272,7 @@ class OpenAIChatCompletionsModel(Model): state.text_content_index_and_output[1].text += delta.content # Handle refusals (model declines to answer) - if delta.refusal: + if hasattr(delta, 'refusal') and delta.refusal: if not state.refusal_content_index_and_output: # Initialize a content tracker for streaming refusal text state.refusal_content_index_and_output = ( @@ -312,7 +318,7 @@ class OpenAIChatCompletionsModel(Model): # 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 - if delta.tool_calls: + 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( @@ -424,14 +430,19 @@ class OpenAIChatCompletionsModel(Model): total_tokens=usage.total_tokens, output_tokens_details=OutputTokensDetails( reasoning_tokens=usage.completion_tokens_details.reasoning_tokens - if usage.completion_tokens_details + if 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, "cached_tokens": usage.prompt_tokens_details.cached_tokens - if usage.prompt_tokens_details and usage.prompt_tokens_details.cached_tokens + if 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 }, ) @@ -510,12 +521,14 @@ class OpenAIChatCompletionsModel(Model): ) tool_choice = _Converter.convert_tool_choice(model_settings.tool_choice) response_format = _Converter.convert_response_format(output_schema) - converted_tools = [ToolConverter.to_openai(tool) for tool in tools] if tools else [] for handoff in handoffs: converted_tools.append(ToolConverter.convert_handoff_tool(handoff)) + # if self.is_ollama: + # converted_tools = [] + if _debug.DONT_LOG_MODEL_DATA: logger.debug("Calling LLM") else: @@ -525,44 +538,77 @@ class OpenAIChatCompletionsModel(Model): f"Stream: {stream}\n" f"Tool choice: {tool_choice}\n" f"Response format: {response_format}\n" + f"Using OLLAMA: {self.is_ollama}\n" ) - ret = await self._get_client().chat.completions.create( - model=self.model, - messages=converted_messages, - tools=converted_tools or NOT_GIVEN, - temperature=self._non_null_or_not_given(model_settings.temperature), - top_p=self._non_null_or_not_given(model_settings.top_p), - frequency_penalty=self._non_null_or_not_given(model_settings.frequency_penalty), - presence_penalty=self._non_null_or_not_given(model_settings.presence_penalty), - max_tokens=self._non_null_or_not_given(model_settings.max_tokens), - tool_choice=tool_choice, - response_format=response_format, - parallel_tool_calls=parallel_tool_calls, - stream=stream, - stream_options={"include_usage": True} if stream else NOT_GIVEN, - extra_headers=_HEADERS, - ) + # Prepare kwargs for the API call + kwargs = { + "model": self.model, + "messages": converted_messages, + "tools": converted_tools or NOT_GIVEN, + "temperature": self._non_null_or_not_given(model_settings.temperature), + "top_p": self._non_null_or_not_given(model_settings.top_p), + "frequency_penalty": self._non_null_or_not_given(model_settings.frequency_penalty), + "presence_penalty": self._non_null_or_not_given(model_settings.presence_penalty), + "max_tokens": self._non_null_or_not_given(model_settings.max_tokens), + "tool_choice": tool_choice, + "response_format": response_format, + "parallel_tool_calls": parallel_tool_calls, + "stream": stream, + "stream_options": {"include_usage": True} if stream else NOT_GIVEN, + "extra_headers": _HEADERS, + } - if isinstance(ret, ChatCompletion): + if self.is_ollama: + if stream: + # For streaming with Ollama, we need to create a Response object first + response = Response( + id=FAKE_RESPONSES_ID, + created_at=time.time(), + model=self.model, + object="response", + output=[], + tool_choice="auto" if tool_choice is None else cast(Literal["auto", "required", "none"], tool_choice) + if tool_choice != NOT_GIVEN + else "auto", + top_p=model_settings.top_p, + temperature=model_settings.temperature, + tools=[], + parallel_tool_calls=parallel_tool_calls or False, + usage={ + "completion_tokens": 0, + "prompt_tokens": 0, + "total_tokens": 0, + "input_tokens": 0, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens": 0, + "output_tokens_details": { + "reasoning_tokens": 0 + } + }, + ) + # Get the streaming object + stream_obj = await litellm.acompletion( + **kwargs, + api_base=get_ollama_api_base(), + custom_llm_provider="openai" + ) + return response, stream_obj + else: + # Non-streaming mode + ret = litellm.completion( + **kwargs, + api_base=get_ollama_api_base(), + custom_llm_provider="openai" + ) + return ret + else: + # Standard LiteLLM handling + ret = litellm.completion(**kwargs) return ret - response = Response( - id=FAKE_RESPONSES_ID, - created_at=time.time(), - model=self.model, - object="response", - output=[], - tool_choice=cast(Literal["auto", "required", "none"], tool_choice) - if tool_choice != NOT_GIVEN - else "auto", - top_p=model_settings.top_p, - temperature=model_settings.temperature, - tools=[], - parallel_tool_calls=parallel_tool_calls or False, - ) - return response, ret - def _get_client(self) -> AsyncOpenAI: if self._client is None: self._client = AsyncOpenAI() @@ -575,7 +621,7 @@ class _Converter: cls, tool_choice: Literal["auto", "required", "none"] | str | None ) -> ChatCompletionToolChoiceOptionParam | NotGiven: if tool_choice is None: - return NOT_GIVEN + return None elif tool_choice == "auto": return "auto" elif tool_choice == "required": @@ -595,7 +641,7 @@ class _Converter: cls, final_output_schema: AgentOutputSchema | None ) -> ResponseFormat | NotGiven: if not final_output_schema or final_output_schema.is_plain_text(): - return NOT_GIVEN + return None return { "type": "json_schema", @@ -621,17 +667,17 @@ class _Converter: message_item.content.append( ResponseOutputText(text=message.content, type="output_text", annotations=[]) ) - if message.refusal: + if hasattr(message, 'refusal') and message.refusal: message_item.content.append( ResponseOutputRefusal(refusal=message.refusal, type="refusal") ) - if message.audio: + if hasattr(message, 'audio') and message.audio: raise AgentsException("Audio is not currently supported") if message_item.content: items.append(message_item) - if message.tool_calls: + if hasattr(message, 'tool_calls') and message.tool_calls: for tool_call in message.tool_calls: items.append( ResponseFunctionToolCall( diff --git a/src/cai/util.py b/src/cai/util.py index 8224ee2c..d986b6dc 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -8,12 +8,11 @@ import pathlib from rich.console import Console from rich.tree import Tree from mako.template import Template # pylint: disable=import-error +from wasabi import color -def get_ollama_api_base() -> str: - """ - Get the Ollama API base URL from the environment variable. - """ - return os.getenv("OLLAMA_API_BASE", "http://host.docker.internal:8000/v1") +def get_ollama_api_base(): + """Get the Ollama API base URL from environment variable or default to localhost:8000.""" + return os.environ.get("OLLAMA_API_BASE", "http://localhost:8000/v1") def load_prompt_template(template_path): """ @@ -129,3 +128,28 @@ def visualize_agent_graph(start_agent): # Start recursive traversal from root agent add_agent_node(start_agent) console.print(tree) + +def fix_litellm_transcription_annotations(): + """ + Apply a monkey patch to fix the TranscriptionCreateParams.__annotations__ issue in LiteLLM. + + This is a temporary fix until the issue is fixed in the LiteLLM library itself. + """ + try: + import litellm.litellm_core_utils.model_param_helper as model_param_helper + + # Override the problematic method to avoid the error + original_get_transcription_kwargs = model_param_helper.ModelParamHelper._get_litellm_supported_transcription_kwargs + + def safe_get_transcription_kwargs(): + """A safer version that doesn't rely on __annotations__.""" + return set(["file", "model", "language", "prompt", "response_format", + "temperature", "api_base", "api_key", "api_version", + "timeout", "custom_llm_provider"]) + + # Apply the monkey patch + model_param_helper.ModelParamHelper._get_litellm_supported_transcription_kwargs = safe_get_transcription_kwargs + return True + except (ImportError, AttributeError): + # If the import fails or the attribute doesn't exist, the patch couldn't be applied + return False