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

CLI responding, but only to proprietary models 

Need to investigate why self-hosted models via OLLAMA fail API-wise

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
This commit is contained in:
Víctor Mayoral Vilches 2025-03-28 09:48:16 +01:00
parent f0b93aec91
commit 944128583c
7 changed files with 108 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
examples/cai/simple_test.py Normal file
View File

@ -0,0 +1,30 @@
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())

4
src/cai/agents/__init__.py
View File

@ -139,7 +139,7 @@ def get_agent_module(agent_name: str) -> str:
Returns:
The full module name where the agent
is defined (e.g., 'cai.agents.basic')
is defined (e.g., 'cai.sdk.agents.basic')
"""
# Try to import all agents from the agents folder
for _, name, _ in pkgutil.iter_modules(__path__,
@ -202,7 +202,7 @@ def get_agent_by_name(agent_name: str) -> Agent:
# # Special handling for one_tool agent
# if agent_name == "one_tool_agent":
# from cai.agents.one_tool import one_tool_agent
# from cai.sdk.agents.one_tool import one_tool_agent
# # Create handoffs between agents
# # Add a handoff from one_tool_agent to flag_discriminator

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

@ -2,46 +2,62 @@
CTF Agent with one tool
"""
import os
from cai.sdk.agents import Agent
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
from cai.tools.reconnaissance.generic_linux_command import generic_linux_command # noqa
model = os.getenv('CAI_MODEL', "qwen2.5:14b")
from openai import AsyncOpenAI
# Get model from environment or use default
model_name = os.getenv('CAI_MODEL', "qwen2.5:14b")
# Create OpenAI client for the agent
openai_client = AsyncOpenAI(
base_url = os.getenv('LITELLM_BASE_URL', 'http://localhost:4000'),
api_key=os.getenv('LITELLM_API_KEY', 'key')
)
# 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
challenge.
INSTRUCTIONS:
1. Execute the generic_linux_command tool without any
explanation.
2. Never assume the flag format - it could be any string
3. YOU MUST USE THE flag_discriminator to check the flag
when you have it in DECODED 64-bit format
## generic_linux_command usage
# 1. Regular command execution
# 2. Shell Session Management: create and manage
# interactive shell sessions for commands like netcat,
# SSH, telnet.
- To start a new session: Use `generic_linux_command` with
commands like `ssh`
- To list active sessions:
`generic_linux_command("session", "list")`
- To get output from a session:
`generic_linux_command("session", "output <session_id>")`
- To send input to a session:
`generic_linux_command("<command>", "<args>",
session_id="<session_id>")`
- To terminate a session:
`generic_linux_command("session", "kill <session_id>")`
"""
one_tool_agent = Agent(
model=model,
name="CTF agent",
description="""Agent focused on conquering security challenges using generic linux commands
Expert in cybersecurity and exploitation.""",
instructions="""You are a Cybersecurity expert Leader facing a CTF
challenge.
INSTRUCTIONS:
1. Execute the generic_linux_command tool without any
explanation.
2. Never assume the flag format - it could be any string
3. YOU MUST USE THE flag_discriminator to check the flag
when you have it in DECODED 64-bit format
## generic_linux_command usage
# 1. Regular command execution
# 2. Shell Session Management: create and manage
# interactive shell sessions for commands like netcat,
# SSH, telnet.
- To start a new session: Use `generic_linux_command` with
commands like `ssh`
- To list active sessions:
`generic_linux_command("session", "list")`
- To get output from a session:
`generic_linux_command("session", "output <session_id>")`
- To send input to a session:
`generic_linux_command("<command>", "<args>",
session_id="<session_id>")`
- To terminate a session:
`generic_linux_command("session", "kill <session_id>")`
""",
instructions=instructions,
tools=[
generic_linux_command,
],
model=OpenAIChatCompletionsModel(
model=model_name,
openai_client=openai_client,
)
)

2
src/cai/agents/thought.py
View File

@ -3,7 +3,7 @@ First prototype of a reasoner agent
using reasoner as a tool call
support meta agent may better @cai.agents.meta.reasoner_support
support meta agent may better @cai.sdk.agents.meta.reasoner_support
"""
from cai.tools.misc.reasoning import thought
from cai.sdk.agents import Agent # pylint: disable=import-error

58
src/cai/cli.py
View File

@ -99,7 +99,7 @@ Usage Examples:
import os
from dotenv import load_dotenv
from openai import AsyncOpenAI
from cai.sdk.agents import OpenAIChatCompletionsModel, Agent, Runner
from cai.sdk.agents import OpenAIChatCompletionsModel, Agent, Runner, AsyncOpenAI
from cai.sdk.agents import set_default_openai_client, set_tracing_disabled
from openai.types.responses import ResponseTextDeltaEvent
from rich.console import Console
@ -126,22 +126,23 @@ external_client = AsyncOpenAI(
set_default_openai_client(external_client)
set_tracing_disabled(True)
# # llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini')
# llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini')
# # llm_model=os.getenv('LLM_MODEL', 'claude-3-7')
# llm_model=os.getenv('LLM_MODEL', 'qwen2.5:14b')
llm_model=os.getenv('LLM_MODEL', 'qwen2.5:14b')
# # For Qwen models, we need to skip system instructions as they're not supported
# instructions = None if "qwen" in llm_model.lower() else "You are a helpful assistant"
# For Qwen models, we need to skip system instructions as they're not supported
instructions = None if "qwen" in llm_model.lower() else "You are a helpful assistant"
# agent = Agent(
# name="Assistant",
# instructions=instructions,
# model=OpenAIChatCompletionsModel(
# model=llm_model,
# openai_client=external_client,
# )
# )
agent = Agent(
name="Assistant",
instructions=instructions,
model=OpenAIChatCompletionsModel(
model=llm_model,
# openai_client=AsyncOpenAI() # original OpenAI servers
openai_client = external_client
)
)
def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=float('inf')):
"""
@ -226,35 +227,20 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
turn_count += 1
except KeyboardInterrupt:
break
except Exception as e:
import traceback
import sys
exc_type, exc_value, exc_traceback = sys.exc_info()
tb_info = traceback.extract_tb(exc_traceback)
filename, line, func, text = tb_info[-1]
console.print(f"[bold red]Error: {str(e)}[/bold red]")
console.print(f"[bold red]Traceback: {tb_info}[/bold red]")
# except Exception as e:
# import traceback
# import sys
# exc_type, exc_value, exc_traceback = sys.exc_info()
# tb_info = traceback.extract_tb(exc_traceback)
# filename, line, func, text = tb_info[-1]
# console.print(f"[bold red]Error: {str(e)}[/bold red]")
# console.print(f"[bold red]Traceback: {tb_info}[/bold red]")
def main():
# Get agent type from environment variables or use default
agent_type = os.getenv('CAI_AGENT_TYPE', "one_tool_agent")
llm_model=os.getenv('LLM_MODEL', 'qwen2.5:14b')
# llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini')
# For Qwen models, we need to skip system instructions as they're not supported
instructions = None if "qwen" in llm_model.lower() else "You are a helpful assistant"
agent = Agent(
name="Assistant",
instructions=instructions,
model=OpenAIChatCompletionsModel(
model=llm_model,
openai_client=external_client,
)
)
# Get the agent instance by name
agent = get_agent_by_name(agent_type)

10
tests/conftest.py
View File

@ -2,11 +2,11 @@ from __future__ import annotations
import pytest
from agents.models import _openai_shared
from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
from agents.models.openai_responses import OpenAIResponsesModel
from agents.tracing import set_trace_processors
from agents.tracing.setup import GLOBAL_TRACE_PROVIDER
from cai.sdk.agents.models import _openai_shared
from cai.sdk.agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
from cai.sdk.agents.models.openai_responses import OpenAIResponsesModel
from cai.sdk.agents.tracing import set_trace_processors
from cai.sdk.agents.tracing.setup import GLOBAL_TRACE_PROVIDER
from .testing_processor import SPAN_PROCESSOR_TESTING

2
tests/testing_processor.py
View File

@ -4,7 +4,7 @@ import threading
from datetime import datetime
from typing import Any, Literal
from agents.tracing import Span, Trace, TracingProcessor
from cai.sdk.agents.tracing import Span, Trace, TracingProcessor
TestSpanProcessorEvent = Literal["trace_start", "trace_end", "span_start", "span_end"]