RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Made local models work by switching into litellm python APIs 

Still some issues with stream versions which needs further
investigation

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
This commit is contained in:
Víctor Mayoral Vilches 2025-03-29 20:17:42 +01:00
parent 944128583c
commit cff20480b7
5 changed files with 179 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
examples/cai/simple_one_tool_test.py Normal file
View File

@ -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())

6
src/cai/agents/one_tool.py
View File

@ -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

6
src/cai/cli.py
View File

@ -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")

132
src/cai/sdk/agents/models/openai_chatcompletions.py
View File

@ -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(

34
src/cai/util.py
View File

@ -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