diff --git a/examples/cai/simple_test.py b/examples/cai/simple_test.py new file mode 100644 index 00000000..007a79ea --- /dev/null +++ b/examples/cai/simple_test.py @@ -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()) \ No newline at end of file diff --git a/src/cai/agents/__init__.py b/src/cai/agents/__init__.py index 66dee8eb..29ca71b6 100644 --- a/src/cai/agents/__init__.py +++ b/src/cai/agents/__init__.py @@ -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 diff --git a/src/cai/agents/one_tool.py b/src/cai/agents/one_tool.py index 8e748318..f9537603 100644 --- a/src/cai/agents/one_tool.py +++ b/src/cai/agents/one_tool.py @@ -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 ")` + - To send input to a session: + `generic_linux_command("", "", + session_id="")` + - To terminate a session: + `generic_linux_command("session", "kill ")` + + """ 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 ")` - - To send input to a session: - `generic_linux_command("", "", - session_id="")` - - To terminate a session: - `generic_linux_command("session", "kill ")` - - """, + instructions=instructions, tools=[ generic_linux_command, ], + model=OpenAIChatCompletionsModel( + model=model_name, + openai_client=openai_client, + ) ) diff --git a/src/cai/agents/thought.py b/src/cai/agents/thought.py index 51d1182e..187daaef 100644 --- a/src/cai/agents/thought.py +++ b/src/cai/agents/thought.py @@ -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 diff --git a/src/cai/cli.py b/src/cai/cli.py index d35affed..553b4918 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -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) diff --git a/tests/conftest.py b/tests/conftest.py index ba0d8822..e096797e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/testing_processor.py b/tests/testing_processor.py index a38c3956..895f348b 100644 --- a/tests/testing_processor.py +++ b/tests/testing_processor.py @@ -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"]