diff --git a/.gitignore b/.gitignore index 2ab5a819..37d20996 100644 --- a/.gitignore +++ b/.gitignore @@ -142,3 +142,7 @@ cython_debug/ # PyPI configuration file .pypirc + +# CAI files +.cai/ +.vscode/ diff --git a/README.md b/README.md index bbd4a5a0..0b008cb2 100644 --- a/README.md +++ b/README.md @@ -17,11 +17,11 @@ Notably, our SDK [is compatible](https://openai.github.io/openai-agents-python/m ## Get started -1. Set up your Python environment +1. Set up your Python environment (note that within Cursor this is big issue :warning:, so probably set it up in another folder) ``` -python -m venv env -source env/bin/activate +python3 -m venv cai +source cai/bin/activate ``` 2. Install Agents SDK diff --git a/docs/ref/agent.md b/docs/ref/agent.md index 3af92687..49b6546c 100644 --- a/docs/ref/agent.md +++ b/docs/ref/agent.md @@ -1,3 +1,3 @@ # `Agents` -::: cai.agents.agent +::: cai.sdk.agents.agent diff --git a/docs/ref/agent_output.md b/docs/ref/agent_output.md index f56c2a9d..21734da6 100644 --- a/docs/ref/agent_output.md +++ b/docs/ref/agent_output.md @@ -1,3 +1,3 @@ # `Agent output` -::: cai.agents.agent_output +::: cai.sdk.agents.agent_output diff --git a/docs/ref/exceptions.md b/docs/ref/exceptions.md index c11bd248..232015c2 100644 --- a/docs/ref/exceptions.md +++ b/docs/ref/exceptions.md @@ -1,3 +1,3 @@ # `Exceptions` -::: cai.agents.exceptions +::: cai.sdk.agents.exceptions diff --git a/docs/ref/extensions/handoff_filters.md b/docs/ref/extensions/handoff_filters.md index cc11e149..244655bb 100644 --- a/docs/ref/extensions/handoff_filters.md +++ b/docs/ref/extensions/handoff_filters.md @@ -1,3 +1,3 @@ # `Handoff filters` -::: cai.agents.extensions.handoff_filters +::: cai.sdk.agents.extensions.handoff_filters diff --git a/docs/ref/extensions/handoff_prompt.md b/docs/ref/extensions/handoff_prompt.md index ab823ab6..3dcf939c 100644 --- a/docs/ref/extensions/handoff_prompt.md +++ b/docs/ref/extensions/handoff_prompt.md @@ -1,6 +1,6 @@ # `Handoff prompt` -::: cai.agents.extensions.handoff_prompt +::: cai.sdk.agents.extensions.handoff_prompt options: members: diff --git a/docs/ref/function_schema.md b/docs/ref/function_schema.md index 18a521ca..f8ae4291 100644 --- a/docs/ref/function_schema.md +++ b/docs/ref/function_schema.md @@ -1,3 +1,3 @@ # `Function schema` -::: cai.agents.function_schema +::: cai.sdk.agents.function_schema diff --git a/docs/ref/guardrail.md b/docs/ref/guardrail.md index c0153249..d7c6107f 100644 --- a/docs/ref/guardrail.md +++ b/docs/ref/guardrail.md @@ -1,3 +1,3 @@ # `Guardrails` -::: cai.agents.guardrail +::: cai.sdk.agents.guardrail diff --git a/docs/ref/handoffs.md b/docs/ref/handoffs.md index acd6fc2e..097e3685 100644 --- a/docs/ref/handoffs.md +++ b/docs/ref/handoffs.md @@ -1,3 +1,3 @@ # `Handoffs` -::: cai.agents.handoffs +::: cai.sdk.agents.handoffs diff --git a/docs/ref/index.md b/docs/ref/index.md index 5f5e0b6b..db03c5ff 100644 --- a/docs/ref/index.md +++ b/docs/ref/index.md @@ -1,6 +1,6 @@ # Agents module -::: cai.agents +::: cai.sdk.agents options: members: diff --git a/docs/ref/items.md b/docs/ref/items.md index b6e8c06b..9b67bb88 100644 --- a/docs/ref/items.md +++ b/docs/ref/items.md @@ -1,3 +1,3 @@ # `Items` -::: cai.agents.items +::: cai.sdk.agents.items diff --git a/docs/ref/lifecycle.md b/docs/ref/lifecycle.md index 55ec4692..02057193 100644 --- a/docs/ref/lifecycle.md +++ b/docs/ref/lifecycle.md @@ -1,6 +1,6 @@ # `Lifecycle` -::: cai.agents.lifecycle +::: cai.sdk.agents.lifecycle options: show_source: false diff --git a/docs/ref/model_settings.md b/docs/ref/model_settings.md index a009d327..0334a366 100644 --- a/docs/ref/model_settings.md +++ b/docs/ref/model_settings.md @@ -1,3 +1,3 @@ # `Model settings` -::: cai.agents.model_settings +::: cai.sdk.agents.model_settings diff --git a/docs/ref/models/interface.md b/docs/ref/models/interface.md index 8d08472c..627330ca 100644 --- a/docs/ref/models/interface.md +++ b/docs/ref/models/interface.md @@ -1,3 +1,3 @@ # `Model interface` -::: cai.agents.models.interface +::: cai.sdk.agents.models.interface diff --git a/docs/ref/models/openai_chatcompletions.md b/docs/ref/models/openai_chatcompletions.md index 6ae90ae7..e4364ce8 100644 --- a/docs/ref/models/openai_chatcompletions.md +++ b/docs/ref/models/openai_chatcompletions.md @@ -1,3 +1,3 @@ # `OpenAI Chat Completions model` -::: cai.agents.models.openai_chatcompletions +::: cai.sdk.agents.models.openai_chatcompletions diff --git a/docs/ref/models/openai_responses.md b/docs/ref/models/openai_responses.md index a19564bd..3d208a0e 100644 --- a/docs/ref/models/openai_responses.md +++ b/docs/ref/models/openai_responses.md @@ -1,3 +1,3 @@ # `OpenAI Responses model` -::: cai.agents.models.openai_responses +::: cai.sdk.agents.models.openai_responses diff --git a/docs/ref/result.md b/docs/ref/result.md index 7fdcb657..e38fdec2 100644 --- a/docs/ref/result.md +++ b/docs/ref/result.md @@ -1,3 +1,3 @@ # `Results` -::: cai.agents.result +::: cai.sdk.agents.result diff --git a/docs/ref/run.md b/docs/ref/run.md index 1533808f..21bd197c 100644 --- a/docs/ref/run.md +++ b/docs/ref/run.md @@ -1,6 +1,6 @@ # `Runner` -::: cai.agents.run +::: cai.sdk.agents.run options: members: diff --git a/docs/ref/run_context.md b/docs/ref/run_context.md index 8a224087..a8ca4cb4 100644 --- a/docs/ref/run_context.md +++ b/docs/ref/run_context.md @@ -1,3 +1,3 @@ # `Run context` -::: cai.agents.run_context +::: cai.sdk.agents.run_context diff --git a/docs/ref/stream_events.md b/docs/ref/stream_events.md index 5f8fe9e0..4a9de0e8 100644 --- a/docs/ref/stream_events.md +++ b/docs/ref/stream_events.md @@ -1,3 +1,3 @@ # `Streaming events` -::: cai.agents.stream_events +::: cai.sdk.agents.stream_events diff --git a/docs/ref/tool.md b/docs/ref/tool.md index 5b2ce120..adf931d2 100644 --- a/docs/ref/tool.md +++ b/docs/ref/tool.md @@ -1,3 +1,3 @@ # `Tools` -::: cai.agents.tool +::: cai.sdk.agents.tool diff --git a/docs/ref/tracing/create.md b/docs/ref/tracing/create.md index 5a885bfc..c21c0ebd 100644 --- a/docs/ref/tracing/create.md +++ b/docs/ref/tracing/create.md @@ -1,3 +1,3 @@ # `Creating traces/spans` -::: cai.agents.tracing.create +::: cai.sdk.agents.tracing.create diff --git a/docs/ref/tracing/index.md b/docs/ref/tracing/index.md index ad9ae885..e94d2724 100644 --- a/docs/ref/tracing/index.md +++ b/docs/ref/tracing/index.md @@ -1,3 +1,3 @@ # Tracing module -::: cai.agents.tracing +::: cai.sdk.agents.tracing diff --git a/docs/ref/tracing/processor_interface.md b/docs/ref/tracing/processor_interface.md index 8b469f2f..18e5fd35 100644 --- a/docs/ref/tracing/processor_interface.md +++ b/docs/ref/tracing/processor_interface.md @@ -1,3 +1,3 @@ # `Processor interface` -::: cai.agents.tracing.processor_interface +::: cai.sdk.agents.tracing.processor_interface diff --git a/docs/ref/tracing/processors.md b/docs/ref/tracing/processors.md index 984f0753..26ec51b2 100644 --- a/docs/ref/tracing/processors.md +++ b/docs/ref/tracing/processors.md @@ -1,3 +1,3 @@ # `Processors` -::: cai.agents.tracing.processors +::: cai.sdk.agents.tracing.processors diff --git a/docs/ref/tracing/scope.md b/docs/ref/tracing/scope.md index 1543e12f..40fc5c79 100644 --- a/docs/ref/tracing/scope.md +++ b/docs/ref/tracing/scope.md @@ -1,3 +1,3 @@ # `Scope` -::: cai.agents.tracing.scope +::: cai.sdk.agents.tracing.scope diff --git a/docs/ref/tracing/setup.md b/docs/ref/tracing/setup.md index 69fd88a3..371ac539 100644 --- a/docs/ref/tracing/setup.md +++ b/docs/ref/tracing/setup.md @@ -1,3 +1,3 @@ # `Setup` -::: cai.agents.tracing.setup +::: cai.sdk.agents.tracing.setup diff --git a/docs/ref/tracing/span_data.md b/docs/ref/tracing/span_data.md index d9aa1635..5f9138dd 100644 --- a/docs/ref/tracing/span_data.md +++ b/docs/ref/tracing/span_data.md @@ -1,3 +1,3 @@ # `Span data` -::: cai.agents.tracing.span_data +::: cai.sdk.agents.tracing.span_data diff --git a/docs/ref/tracing/spans.md b/docs/ref/tracing/spans.md index 37643fc4..3b004e74 100644 --- a/docs/ref/tracing/spans.md +++ b/docs/ref/tracing/spans.md @@ -1,6 +1,6 @@ # `Spans` -::: cai.agents.tracing.spans +::: cai.sdk.agents.tracing.spans options: members: diff --git a/docs/ref/tracing/traces.md b/docs/ref/tracing/traces.md index cc377a92..792f00c3 100644 --- a/docs/ref/tracing/traces.md +++ b/docs/ref/tracing/traces.md @@ -1,3 +1,3 @@ # `Traces` -::: cai.agents.tracing.traces +::: cai.sdk.agents.tracing.traces diff --git a/docs/ref/tracing/util.md b/docs/ref/tracing/util.md index ab2198ba..23acc564 100644 --- a/docs/ref/tracing/util.md +++ b/docs/ref/tracing/util.md @@ -1,3 +1,3 @@ # `Util` -::: cai.agents.tracing.util +::: cai.sdk.agents.tracing.util diff --git a/docs/ref/usage.md b/docs/ref/usage.md index a3501242..dae3b702 100644 --- a/docs/ref/usage.md +++ b/docs/ref/usage.md @@ -1,3 +1,3 @@ # `Usage` -::: cai.agents.usage +::: cai.sdk.agents.usage diff --git a/docs/ref/voice/events.md b/docs/ref/voice/events.md index 4a67939c..ed9a7485 100644 --- a/docs/ref/voice/events.md +++ b/docs/ref/voice/events.md @@ -1,3 +1,3 @@ # `Events` -::: cai.agents.voice.events +::: cai.sdk.agents.voice.events diff --git a/docs/ref/voice/exceptions.md b/docs/ref/voice/exceptions.md index b9b7c189..6dc182ed 100644 --- a/docs/ref/voice/exceptions.md +++ b/docs/ref/voice/exceptions.md @@ -1,3 +1,3 @@ # `Exceptions` -::: cai.agents.voice.exceptions +::: cai.sdk.agents.voice.exceptions diff --git a/docs/ref/voice/input.md b/docs/ref/voice/input.md index 77d704c7..b5f04dde 100644 --- a/docs/ref/voice/input.md +++ b/docs/ref/voice/input.md @@ -1,3 +1,3 @@ # `Input` -::: cai.agents.voice.input +::: cai.sdk.agents.voice.input diff --git a/docs/ref/voice/model.md b/docs/ref/voice/model.md index 34fe7ddf..bbd7942c 100644 --- a/docs/ref/voice/model.md +++ b/docs/ref/voice/model.md @@ -1,3 +1,3 @@ # `Model` -::: cai.agents.voice.model +::: cai.sdk.agents.voice.model diff --git a/docs/ref/voice/models/openai_provider.md b/docs/ref/voice/models/openai_provider.md index bc800817..dd136e50 100644 --- a/docs/ref/voice/models/openai_provider.md +++ b/docs/ref/voice/models/openai_provider.md @@ -1,3 +1,3 @@ # `OpenAIVoiceModelProvider` -::: cai.agents.voice.models.openai_model_provider +::: cai.sdk.agents.voice.models.openai_model_provider diff --git a/docs/ref/voice/models/openai_stt.md b/docs/ref/voice/models/openai_stt.md index 37ae36d2..b6acaa27 100644 --- a/docs/ref/voice/models/openai_stt.md +++ b/docs/ref/voice/models/openai_stt.md @@ -1,3 +1,3 @@ # `OpenAI STT` -::: cai.agents.voice.models.openai_stt +::: cai.sdk.agents.voice.models.openai_stt diff --git a/docs/ref/voice/models/openai_tts.md b/docs/ref/voice/models/openai_tts.md index ba21c4fa..15268cc8 100644 --- a/docs/ref/voice/models/openai_tts.md +++ b/docs/ref/voice/models/openai_tts.md @@ -1,3 +1,3 @@ # `OpenAI TTS` -::: cai.agents.voice.models.openai_tts +::: cai.sdk.agents.voice.models.openai_tts diff --git a/docs/ref/voice/pipeline.md b/docs/ref/voice/pipeline.md index ba82e7d9..a78a9544 100644 --- a/docs/ref/voice/pipeline.md +++ b/docs/ref/voice/pipeline.md @@ -1,3 +1,3 @@ # `Pipeline` -::: cai.agents.voice.pipeline +::: cai.sdk.agents.voice.pipeline diff --git a/docs/ref/voice/pipeline_config.md b/docs/ref/voice/pipeline_config.md index 88f1f769..fbe88cca 100644 --- a/docs/ref/voice/pipeline_config.md +++ b/docs/ref/voice/pipeline_config.md @@ -1,3 +1,3 @@ # `Pipeline Config` -::: cai.agents.voice.pipeline_config +::: cai.sdk.agents.voice.pipeline_config diff --git a/docs/ref/voice/result.md b/docs/ref/voice/result.md index 9d1ba599..ee89ce0c 100644 --- a/docs/ref/voice/result.md +++ b/docs/ref/voice/result.md @@ -1,3 +1,3 @@ # `Result` -::: cai.agents.voice.result +::: cai.sdk.agents.voice.result diff --git a/docs/ref/voice/utils.md b/docs/ref/voice/utils.md index 4dc8a426..4680ea30 100644 --- a/docs/ref/voice/utils.md +++ b/docs/ref/voice/utils.md @@ -1,3 +1,3 @@ # `Utils` -::: cai.agents.voice.utils +::: cai.sdk.agents.voice.utils diff --git a/docs/ref/voice/workflow.md b/docs/ref/voice/workflow.md index 2f2241b9..5bba5cca 100644 --- a/docs/ref/voice/workflow.md +++ b/docs/ref/voice/workflow.md @@ -1,3 +1,3 @@ # `Workflow` -::: cai.agents.voice.workflow +::: cai.sdk.agents.voice.workflow diff --git a/examples/model_providers/README.md b/examples/model_providers/README.md index e5fd2dac..310b2f79 100644 --- a/examples/model_providers/README.md +++ b/examples/model_providers/README.md @@ -43,3 +43,8 @@ curl -s http://localhost:4000/v1/chat/completions -H "Content-Type: application/ # gpt-4o curl -s http://localhost:4000/v1/chat/completions -H "Content-Type: application/json" -d '{"model": "gpt-4o", "messages": [{"role": "user", "content": "Say hi"}], "max_tokens": 10}' | jq ``` + +When using virtual keys: +```bash +curl -s http://localhost:4000/v1/chat/completions -H "Content-Type: application/json" -H "Authorization: Bearer sk-pNCn8ZA0SCtWMpkZNUWe5g" -d '{"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "Say hi"}], "max_tokens": 10}' | jq +``` \ No newline at end of file diff --git a/examples/model_providers/litellm.py b/examples/model_providers/litellm.py index d3242708..5742bb47 100644 --- a/examples/model_providers/litellm.py +++ b/examples/model_providers/litellm.py @@ -1,8 +1,8 @@ import os from dotenv import load_dotenv from openai import AsyncOpenAI -from cai.agents import OpenAIChatCompletionsModel,Agent,Runner -from cai.agents import set_default_openai_client, set_tracing_disabled +from cai.sdk.agents import OpenAIChatCompletionsModel,Agent,Runner +from cai.sdk.agents import set_default_openai_client, set_tracing_disabled from openai.types.responses import ResponseTextDeltaEvent # Load environment variables from .env file diff --git a/examples/model_providers/litellm_config.yaml b/examples/model_providers/litellm_config.yaml index 8c029636..ab5b9da9 100644 --- a/examples/model_providers/litellm_config.yaml +++ b/examples/model_providers/litellm_config.yaml @@ -4,6 +4,11 @@ model_list: model: gpt-4o api_key: "os.environ/OPENAI_API_KEY" api_base: https://api.openai.com/v1 + - model_name: gpt-4o-mini + litellm_params: + model: gpt-4o-mini + api_key: "os.environ/OPENAI_API_KEY" + api_base: https://api.openai.com/v1 - model_name: claude-3-7 litellm_params: model: claude-3-7-sonnet-20250219 diff --git a/mkdocs.yml b/mkdocs.yml index f2f10c0f..f1b6e865 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -144,4 +144,4 @@ extra_css: - stylesheets/extra.css watch: - - "src/cai/agents" + - "src/cai/sdk/agents" diff --git a/pyproject.toml b/pyproject.toml index d9360154..2bb8f013 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,11 @@ dependencies = [ "requests>=2.0, <3", "types-requests>=2.0, <3", "openinference-instrumentation-openai>=0.1.22", + "wasabi>=1.1.3", + "rich>=13.9.4", + "prompt_toolkit>=3.0.39", + "dotenv>=0.9.9", + "litellm>=1.63.7" ] classifiers = [ "Typing :: Typed", @@ -106,7 +111,7 @@ module = "sounddevice.*" ignore_missing_imports = true [tool.coverage.run] -source = ["tests", "src/cai/agents"] +source = ["tests", "src/cai/sdk/agents"] [tool.coverage.report] show_missing = true diff --git a/src/cai/__init__.py b/src/cai/__init__.py index e69de29b..5e89eb6c 100644 --- a/src/cai/__init__.py +++ b/src/cai/__init__.py @@ -0,0 +1,46 @@ +""" +A library to build Bug Bounty-level grade Cybersecurity AIs (CAIs). +""" + +def is_pentestperf_available(): + """ + Check if pentestperf is available + """ + try: + from pentestperf.ctf import CTF # pylint: disable=import-error,import-outside-toplevel,unused-import # noqa: E501,F401 + except ImportError: + return False + return True + + +def is_caiextensions_report_available(): + """ + Check if caiextensions report is available + """ + try: + from caiextensions.report.common import get_base_instructions # pylint: disable=import-error,import-outside-toplevel,unused-import # noqa: E501,F401 + except ImportError: + return False + return True + + +def is_caiextensions_memory_available(): + """ + Check if caiextensions memory is available + """ + try: + from caiextensions.memory import is_memory_installed # pylint: disable=import-error,import-outside-toplevel,unused-import # noqa: E501,F401 + except ImportError: + return False + return True + + +def is_caiextensions_platform_available(): + """ + Check if caiextensions-platform is available + """ + try: + from caiextensions.platform.base import platform_manager # pylint: disable=import-error,import-outside-toplevel,unused-import # noqa: E501,F401 + except ImportError: + return False + return True diff --git a/src/cai/cli.py b/src/cai/cli.py new file mode 100644 index 00000000..b59e9f08 --- /dev/null +++ b/src/cai/cli.py @@ -0,0 +1,131 @@ +import os +from dotenv import load_dotenv +from openai import AsyncOpenAI +from cai.sdk.agents import OpenAIChatCompletionsModel, Agent, Runner +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 + +# Import modules from cai.repl +from cai.repl.commands import FuzzyCommandCompleter, handle_command as commands_handle_command +from cai.repl.ui.keybindings import create_key_bindings +from cai.repl.ui.logging import setup_session_logging +from cai.repl.ui.banner import display_banner +from cai.repl.ui.prompt import get_user_input +from cai.repl.ui.toolbar import get_toolbar_with_refresh + +# Load environment variables from .env file +load_dotenv() + +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) +set_tracing_disabled(True) + +# 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') + + +# 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, + ) +) + +def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=float('inf')): + """ + Run a simple interactive CLI loop for CAI. + + Args: + starting_agent: The initial agent to use for the conversation + context_variables: Optional dictionary of context variables to initialize the session + stream: Boolean flag to enable/disable streaming responses (default: False) + max_turns: Maximum number of interaction turns before terminating (default: infinity) + + Returns: + None + """ + agent = starting_agent + turn_count = 0 + + console = Console() + + # Initialize command completer and key bindings + command_completer = FuzzyCommandCompleter() + current_text = [''] + kb = create_key_bindings(current_text) + + # Setup session logging + history_file = setup_session_logging() + + # Display banner + display_banner(console) + + while turn_count < max_turns: + try: + # Get user input with command completion and history + user_input = get_user_input( + command_completer, + kb, + history_file, + get_toolbar_with_refresh, + current_text + ) + + # Handle special commands + if user_input.startswith('/') or user_input.startswith('$'): + parts = user_input.strip().split() + command = parts[0] + args = parts[1:] if len(parts) > 1 else None + + # Process the command with the handler + if commands_handle_command(command, args): + continue # Command was handled, continue to next iteration + + # If command wasn't recognized, show error + console.print(f"[red]Unknown command: {command}[/red]") + continue + + # Process the conversation with the agent + if stream: + # Use streamed response + print("Agent: ", end="", flush=True) + + async def process_streamed_response(): + try: + result = Runner.run_streamed(agent, user_input) + async for event in result.stream_events(): + if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent): + print(event.data.delta, end="", flush=True) + print() # Add a newline at the end + return result + except Exception as e: + print() # Add a newline after any partial output + print(f"\n[Error occurred during streaming: {str(e)}]") + return None + + asyncio.run(process_streamed_response()) + else: + # Use non-streamed response + console.print("[dim]Thinking...[/dim]") + response = asyncio.run(Runner.run(agent, user_input)) + console.print(f"Agent: {response.final_output}") + turn_count += 1 + except KeyboardInterrupt: + break + except Exception as e: + console.print(f"[bold red]Error: {str(e)}[/bold red]") + + +if __name__ == "__main__": + run_cai_cli(agent, stream=True) \ No newline at end of file diff --git a/src/cai/agents/extensions/__init__.py b/src/cai/repl/__init__.py similarity index 100% rename from src/cai/agents/extensions/__init__.py rename to src/cai/repl/__init__.py diff --git a/src/cai/repl/commands/__init__.py b/src/cai/repl/commands/__init__.py new file mode 100644 index 00000000..95bb4a27 --- /dev/null +++ b/src/cai/repl/commands/__init__.py @@ -0,0 +1,92 @@ +""" +Commands module for CAI REPL. +This module exports all commands available +in the CAI REPL. +""" +from typing import ( + Dict, + List, +) + +from cai.repl.commands.completer import ( + FuzzyCommandCompleter +) + +# Import base command structure +from cai.repl.commands.base import ( + Command, + COMMANDS, + COMMAND_ALIASES, + register_command, + get_command, + handle_command +) + +# Import all command modules +# These imports will register the commands with the registry +from cai.repl.commands import ( # pylint: disable=import-error,unused-import,line-too-long,redefined-builtin # noqa: E501,F401 + memory, + help, + graph, + exit, + shell, + env, + platform, + kill, + model, + turns, + agent, + history, + config +) + +# Define helper functions + + +def get_command_descriptions() -> Dict[str, str]: + """Get descriptions for all commands. + + Returns: + A dictionary mapping command names to descriptions + """ + return {cmd.name: cmd.description for cmd in COMMANDS.values()} + + +def get_subcommand_descriptions() -> Dict[str, str]: + """Get descriptions for all subcommands. + + Returns: + A dictionary mapping command paths to descriptions + """ + descriptions = {} + for cmd in COMMANDS.values(): + for subcmd in cmd.get_subcommands(): + key = f"{cmd.name} {subcmd}" + descriptions[key] = cmd.get_subcommand_description(subcmd) + return descriptions + + +def get_all_commands() -> Dict[str, List[str]]: + """Get all commands and their subcommands. + + Returns: + A dictionary mapping command names to lists of subcommand names + """ + return {cmd.name: cmd.get_subcommands() for cmd in COMMANDS.values()} + + +# Import the command completer after defining the helper functions + +# Export command registry +__all__ = [ + 'Command', + 'COMMANDS', + 'COMMAND_ALIASES', + 'register_command', + 'get_command', + 'handle_command', + 'get_command_descriptions', + 'get_subcommand_descriptions', + 'get_all_commands', + 'FuzzyCommandCompleter' +] diff --git a/src/cai/repl/commands/agent.py b/src/cai/repl/commands/agent.py new file mode 100644 index 00000000..5369f82c --- /dev/null +++ b/src/cai/repl/commands/agent.py @@ -0,0 +1,351 @@ +# """ +# Agent "command" for CAI CLI abstraction + +# Provides commands for managing and switching between agents. +# """ + +# # Standard library imports +# import inspect +# import os +# import sys + +# from typing import List, Optional + +# # Third-party imports +# from rich.console import Console # pylint: disable=import-error +# from rich.markdown import Markdown # pylint: disable=import-error +# from rich.table import Table # pylint: disable=import-error + +# # Local imports +# from cai.agents import get_available_agents, get_agent_module +# from cai.repl.commands.base import Command, register_command +# from cai.types import Agent +# from cai.util import visualize_agent_graph + +# console = Console() + + +# class AgentCommand(Command): +# """Command for managing and switching between agents.""" + +# def __init__(self): +# """Initialize the agent command.""" +# # Initialize with basic parameters +# super().__init__( +# name="/agent", +# description="Manage and switch between agents", +# aliases=["/a"] +# ) + +# # Add subcommands manually +# self._subcommands = { +# "list": "List available agents", +# "select": "Select an agent by name or number", +# "info": "Show information about an agent", +# "multi": "Enable multi-agent mode" +# } + +# def _get_model_display(self, agent_name: str, agent: Agent) -> str: +# """Get the display string for an agent's model. + +# Args: +# agent_name: Name of the agent +# agent: Agent instance + +# Returns: +# String to display for the agent's model +# """ +# # For code agent, always show the model +# if agent_name == "code": +# return agent.model + +# # For other agents, check if CTF_MODEL is set +# ctf_model = os.getenv('CTF_MODEL') +# if ctf_model and agent.model == ctf_model: +# # Don't show default model for CTF_MODEL in table +# # but show "Default CTF Model" in info +# return "" + +# # Show the model from environment variable if available +# env_var_name = f"CAI_{agent_name.upper()}_MODEL" +# model_env = os.getenv(env_var_name) +# if model_env: +# return model_env + +# return agent.model + +# def _get_model_display_for_info( +# self, agent_name: str, agent: Agent) -> str: +# """Get the display string for an agent's model in the info view. + +# Args: +# agent_name: Name of the agent +# agent: Agent instance + +# Returns: +# String to display for the agent's model in the info view +# """ +# # For code agent, always show the model +# if agent_name == "code": +# return agent.model + +# # For other agents, check if CTF_MODEL is set +# ctf_model = os.getenv('CTF_MODEL') +# if ctf_model and agent.model == ctf_model: +# # Show "Default CTF Model" in info +# return "Default CTF Model" + +# # Show the model from environment variable if available +# env_var_name = f"CAI_{agent_name.upper()}_MODEL" +# model_env = os.getenv(env_var_name) +# if model_env: +# return model_env + +# return agent.model + +# def get_subcommands(self) -> List[str]: +# """Get list of subcommand names. + +# Returns: +# List of subcommand names +# """ +# return list(self._subcommands.keys()) + +# def get_subcommand_description(self, subcommand: str) -> str: +# """Get description for a subcommand. + +# Args: +# subcommand: Name of the subcommand + +# Returns: +# Description of the subcommand +# """ +# return self._subcommands.get(subcommand, "") + +# def handle(self, args: Optional[List[str]] = None) -> bool: +# """Handle the agent command. + +# Args: +# args: Optional list of command arguments + +# Returns: +# True if the command was handled successfully, False otherwise +# """ +# if not args: +# return self.handle_list(args) + +# subcommand = args[0] +# if subcommand in self._subcommands: +# handler = getattr(self, f"handle_{subcommand}", None) +# if handler: +# return handler(args[1:] if len(args) > 1 else None) + +# # If not a subcommand, try to select an agent by name +# return self.handle_select(args) + +# def handle_list(self, args: Optional[List[str]] = None) -> bool: # pylint: disable=unused-argument # noqa: E501 +# """Handle /agent list command. + +# Args: +# args: Optional list of command arguments (not used) + +# Returns: +# True if the command was handled successfully +# """ +# table = Table(title="Available Agents") +# table.add_column("#", style="dim") +# table.add_column("Name", style="cyan") +# table.add_column("Module", style="magenta") +# table.add_column("Description", style="green") +# table.add_column("Pattern", style="blue") +# table.add_column("Model", style="yellow") + +# # Scan all agents from the agents folder +# agents_to_display = get_available_agents() + +# # Display all agents +# for i, (name, agent) in enumerate(agents_to_display.items(), 1): +# description = agent.description +# if not description and hasattr(agent, 'instructions'): +# if callable(agent.instructions): +# description = agent.instructions(context_variables={}) +# else: +# description = agent.instructions +# # Clean up description - remove newlines and strip spaces +# if isinstance(description, str): +# description = " ".join(description.split()) +# if len(description) > 50: +# description = description[:47] + "..." + +# # Get the module name for the agent +# module_name = get_agent_module(name) + +# # Get the pattern if it exists +# pattern = getattr(agent, 'pattern', '') +# if pattern: +# pattern = pattern.capitalize() + +# # Handle model display based on agent type +# model_display = self._get_model_display(name, agent) +# table.add_row( +# str(i), +# name, +# module_name, +# description, +# pattern, +# model_display +# ) + +# console.print(table) +# return True + +# def handle_select(self, args: Optional[List[str]] = None) -> bool: # pylint: disable=too-many-branches,line-too-long # noqa: E501 +# """Handle /agent select command. + +# Args: +# args: Optional list of command arguments + +# Returns: +# True if the command was handled successfully, False otherwise +# """ +# if not args: +# console.print("[red]Error: No agent specified[/red]") +# console.print("Usage: /agent select ") +# return False + +# agent_id = args[0] + +# # Get the list of available agents +# agents_to_display = get_available_agents() + +# # Check if agent_id is a number +# if agent_id.isdigit(): +# index = int(agent_id) +# if 1 <= index <= len(agents_to_display): +# agent_name = list(agents_to_display.keys())[index - 1] +# else: +# console.print( +# f"[red]Error: Invalid agent number: {agent_id}[/red]") +# return False +# else: +# # Treat as agent name +# agent_name = agent_id +# if agent_name not in agents_to_display: +# console.print(f"[red]Error: Unknown agent: {agent_name}[/red]") +# return False + +# # Get the agent +# agent = agents_to_display[agent_name] + +# # Set the agent as the current agent in the REPL +# # We need to avoid circular imports, so we'll use a different approach +# # to access the client and current_agent variables + +# # Import the module dynamically to avoid circular imports +# if 'cai.repl.repl' in sys.modules: +# repl_module = sys.modules['cai.repl.repl'] + +# # Check if client is initialized +# if hasattr(repl_module, 'client') and repl_module.client: +# # Update the active_agent in the client +# repl_module.client.active_agent = agent + +# # Update the global current_agent variable if it exists +# if hasattr(repl_module, 'current_agent'): +# repl_module.current_agent = agent + +# # Update the global agent variable if it exists +# if hasattr(repl_module, 'agent'): +# repl_module.agent = agent + +# # Also update the agent variable in the run_demo_loop +# # function's frame if possible +# try: +# for frame_info in inspect.stack(): +# frame = frame_info.frame +# if ('run_demo_loop' in frame.f_code.co_name and +# 'agent' in frame.f_locals): +# frame.f_locals['agent'] = agent +# break +# except Exception: # pylint: disable=broad-except # nosec +# # If this fails, we still have the global current_agent as +# # a fallback +# pass + +# console.print( +# f"[green]Switched to agent: {agent_name}[/green]") +# visualize_agent_graph(agent) +# return True +# console.print("[red]Error: CAI client not initialized[/red]") +# return False +# console.print("[red]Error: REPL module not initialized[/red]") +# return False + +# def handle_info(self, args: Optional[List[str]] = None) -> bool: +# """Handle /agent info command. + +# Args: +# args: Optional list of command arguments + +# Returns: +# True if the command was handled successfully, False otherwise +# """ +# if not args: +# console.print("[red]Error: No agent specified[/red]") +# console.print("Usage: /agent info ") +# return False + +# agent_id = args[0] + +# # Get the list of available agents +# agents_to_display = get_available_agents() + +# # Check if agent_id is a number +# if agent_id.isdigit(): +# index = int(agent_id) +# if 1 <= index <= len(agents_to_display): +# agent_name = list(agents_to_display.keys())[index - 1] +# else: +# console.print( +# f"[red]Error: Invalid agent number: {agent_id}[/red]") +# return False +# else: +# # Treat as agent name +# agent_name = agent_id +# if agent_name not in agents_to_display: +# console.print(f"[red]Error: Unknown agent: {agent_name}[/red]") +# return False + +# # Get the agent +# agent = agents_to_display[agent_name] + +# # Display agent information +# instructions = agent.instructions +# if callable(instructions): +# instructions = instructions() + +# # Handle model display based on agent type +# model_display = self._get_model_display_for_info(agent_name, agent) + +# # Create a markdown table for agent details +# markdown_content = f""" +# # Agent: {agent_name} + +# | Property | Value | +# |----------|-------| +# | Name | {agent.name} | +# | Model | {model_display} | +# | Functions | {len(agent.functions)} | +# | Parallel Tool Calls | {'Yes' if agent.parallel_tool_calls else 'No'} | + +# ## Instructions + +# {instructions} +# """ + +# console.print(Markdown(markdown_content)) +# return True + + +# # Register the command +# register_command(AgentCommand()) diff --git a/src/cai/repl/commands/base.py b/src/cai/repl/commands/base.py new file mode 100644 index 00000000..4f540f85 --- /dev/null +++ b/src/cai/repl/commands/base.py @@ -0,0 +1,159 @@ +""" +Base module for CAI REPL commands. +This module provides the base structure for all commands in the CAI REPL. +""" +from typing import ( + List, + Optional, + Dict, + Any, + Callable +) +from rich.console import Console # pylint: disable=import-error + +console = Console() + + +class Command: + """Base class for all commands.""" + + def __init__(self, name: str, description: str, aliases: List[str] = None): + """Initialize a command. + + Args: + name: The name of the command (e.g. "/memory") + description: A short description of the command + aliases: Optional list of command aliases + """ + self.name = name + self.description = description + self.aliases = aliases or [] + self.subcommands: Dict[str, Dict[str, Any]] = {} + + def add_subcommand(self, name: str, description: str, handler: Callable): + """Add a subcommand to this command. + + Args: + name: The name of the subcommand (e.g. "list") + description: A short description of the subcommand + handler: The function to call when the subcommand is invoked + """ + self.subcommands[name] = { + "description": description, + "handler": handler + } + + def get_subcommands(self) -> List[str]: + """Get a list of all subcommand names. + + Returns: + A list of subcommand names + """ + return list(self.subcommands.keys()) + + def get_subcommand_description(self, subcommand: str) -> str: + """Get the description of a subcommand. + + Args: + subcommand: The name of the subcommand + + Returns: + The description of the subcommand + """ + return self.subcommands.get(subcommand, {}).get("description", "") + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + if not args: + return self.handle_no_args() + + subcommand = args[0] + if subcommand in self.subcommands: + handler = self.subcommands[subcommand]["handler"] + return handler(args[1:] if len(args) > 1 else None) + + return self.handle_unknown_subcommand(subcommand) + + def handle_no_args(self) -> bool: + """Handle the command when no arguments are provided. + + Returns: + True if the command was handled successfully, False otherwise + """ + console.print( + f"[yellow]{ + self.name} command requires a subcommand: { + ', '.join( + self.get_subcommands())}[/yellow]") + return False + + def handle_unknown_subcommand(self, subcommand: str) -> bool: + """Handle an unknown subcommand. + + Args: + subcommand: The unknown subcommand + + Returns: + True if the command was handled successfully, False otherwise + """ + console.print( + f"[red]Unknown { + self.name} subcommand: {subcommand}[/red]") + return False + + +# Registry for all commands +COMMANDS: Dict[str, Command] = {} +COMMAND_ALIASES: Dict[str, str] = {} + + +def register_command(command: Command) -> None: + """Register a command in the global registry. + + Args: + command: The command to register + """ + COMMANDS[command.name] = command + + # Register aliases + for alias in command.aliases: + COMMAND_ALIASES[alias] = command.name + + +def get_command(name: str) -> Optional[Command]: + """Get a command by name or alias. + + Args: + name: The name or alias of the command + + Returns: + The command if found, None otherwise + """ + # Check if it's an alias + name = COMMAND_ALIASES.get(name, name) + + return COMMANDS.get(name) + + +def handle_command(command: str, args: Optional[List[str]] = None) -> bool: + """Handle a command. + + Args: + command: The command name or alias + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + cmd = get_command(command) + if cmd: + return cmd.handle(args) + + return False diff --git a/src/cai/repl/commands/completer.py b/src/cai/repl/commands/completer.py new file mode 100644 index 00000000..cadf97a5 --- /dev/null +++ b/src/cai/repl/commands/completer.py @@ -0,0 +1,526 @@ +""" +Command completer for CAI REPL. +This module provides a fuzzy command completer with autocompletion menu and +command shadowing. +""" +# Standard library imports +import datetime +import threading +import time +from functools import lru_cache +from typing import ( + List, + Optional, + Dict, + Any +) + +# Third-party imports +import requests # pylint: disable=import-error,unused-import,line-too-long # noqa: E501 +from prompt_toolkit.completion import ( # pylint: disable=import-error + Completer, + Completion +) +from prompt_toolkit.formatted_text import HTML # pylint: disable=import-error +from prompt_toolkit.styles import Style # pylint: disable=import-error +from rich.console import Console # pylint: disable=import-error + +from cai.util import get_ollama_api_base +from cai.repl.commands.base import ( + COMMANDS, + COMMAND_ALIASES +) + +console = Console() + +# Global cache for command descriptions and subcommands +COMMAND_DESCRIPTIONS_CACHE = None +SUBCOMMAND_DESCRIPTIONS_CACHE = None +ALL_COMMANDS_CACHE = None + + +class FuzzyCommandCompleter(Completer): + """Command completer with fuzzy matching for the REPL. + + This advanced completer provides intelligent suggestions for commands, + subcommands, and arguments based on what the user is typing. + It supports fuzzy matching to find commands even with typos. + + Features: + - Fuzzy matching for commands and subcommands + - Autocompletion menu with descriptions + - Command shadowing (showing hints for previously used commands) + - Model completion for the /model command + """ + + # Class-level cache for models + _cached_models = [] + _cached_model_numbers = {} + _last_model_fetch = datetime.datetime.now() - datetime.timedelta(minutes=10) + _fetch_lock = threading.Lock() + + def __init__(self): + """Initialize the command completer with cached model information.""" + super().__init__() + self.command_history = {} # Store command usage frequency + + # Fetch models in background thread to avoid blocking + threading.Thread( + target=self._background_fetch_models, + daemon=True + ).start() + + # Styling for the completion menu + self.completion_style = Style.from_dict({ + 'completion-menu': 'bg:#2b2b2b #ffffff', + 'completion-menu.completion': 'bg:#2b2b2b #ffffff', + 'completion-menu.completion.current': 'bg:#004b6b #ffffff', + 'scrollbar.background': 'bg:#2b2b2b', + 'scrollbar.button': 'bg:#004b6b', + }) + + def _background_fetch_models(self): + """Fetch models in background to avoid blocking the UI.""" + try: + self.fetch_ollama_models() + except Exception: # pylint: disable=broad-except + pass + + def fetch_ollama_models(self): # pylint: disable=too-many-branches,too-many-statements,inconsistent-return-statements,line-too-long # noqa: E501 + """Fetch available models from Ollama if it's running.""" + # Only fetch every 60 seconds to avoid excessive API calls + now = datetime.datetime.now() + + # Use a lock to prevent multiple threads from fetching simultaneously + with self._fetch_lock: + if (now - self._last_model_fetch).total_seconds() < 60: + return + + self._last_model_fetch = now + ollama_models = [] + + try: + # Get Ollama models with a short timeout to prevent hanging + api_base = get_ollama_api_base() + response = requests.get( + f"{api_base.replace('/v1', '')}/api/tags", timeout=0.5) + + if response.status_code == 200: + data = response.json() + if 'models' in data: + models = data['models'] + else: + # Fallback for older Ollama versions + models = data.get('items', []) + + ollama_models = [(model.get('name', ''), []) for model in models] + except Exception: # pylint: disable=broad-except + # Silently fail if Ollama is not available + pass + + # Standard models always available + standard_models = [ + # Claude 3.7 models + "claude-3-7-sonnet-20250219", + + # Claude 3.5 models + "claude-3-5-sonnet-20240620", + "claude-3-5-20241122", + + # Claude 3 models + "claude-3-opus-20240229", + "claude-3-sonnet-20240229", + "claude-3-haiku-20240307", + + # OpenAI O-series models + "o1", + "o1-mini", + "o3-mini", + + # OpenAI GPT models + "gpt-4o", + "gpt-4-turbo", + "gpt-3.5-turbo", + + # DeepSeek models + "deepseek-v3", + "deepseek-r1" + ] + + # Combine standard models with Ollama models + self._cached_models = standard_models + ollama_models + + # Create number mappings for models (1-based indexing) + self._cached_model_numbers = {} + for i, model in enumerate(self._cached_models, 1): + self._cached_model_numbers[str(i)] = model + + def record_command_usage(self, command: str): + """Record command usage for command shadowing. + + Args: + command: The command that was used + """ + if command.startswith('/'): + # Extract the main command + parts = command.split() + main_command = parts[0] + + # Update usage count + if main_command in self.command_history: + self.command_history[main_command] += 1 + else: + self.command_history[main_command] = 1 + + @lru_cache(maxsize=1) + def get_command_descriptions(self): + """Get descriptions for all commands. + + Returns: + A dictionary mapping command names to descriptions + """ + global COMMAND_DESCRIPTIONS_CACHE + if COMMAND_DESCRIPTIONS_CACHE is None: + COMMAND_DESCRIPTIONS_CACHE = {cmd.name: cmd.description for cmd in COMMANDS.values()} + return COMMAND_DESCRIPTIONS_CACHE + + @lru_cache(maxsize=1) + def get_subcommand_descriptions(self): + """Get descriptions for all subcommands. + + Returns: + A dictionary mapping command paths to descriptions + """ + global SUBCOMMAND_DESCRIPTIONS_CACHE + if SUBCOMMAND_DESCRIPTIONS_CACHE is None: + descriptions = {} + for cmd in COMMANDS.values(): + for subcmd in cmd.get_subcommands(): + key = f"{cmd.name} {subcmd}" + descriptions[key] = cmd.get_subcommand_description(subcmd) + SUBCOMMAND_DESCRIPTIONS_CACHE = descriptions + return SUBCOMMAND_DESCRIPTIONS_CACHE + + @lru_cache(maxsize=1) + def get_all_commands(self): + """Get all commands and their subcommands. + + Returns: + A dictionary mapping command names to lists of subcommand names + """ + global ALL_COMMANDS_CACHE + if ALL_COMMANDS_CACHE is None: + ALL_COMMANDS_CACHE = {cmd.name: cmd.get_subcommands() for cmd in COMMANDS.values()} + return ALL_COMMANDS_CACHE + + # Cache for command suggestions to avoid recalculating + _command_suggestions_cache = {} + _command_suggestions_last_update = 0 + _command_suggestions_update_interval = 1.0 # Update every second + + def get_command_suggestions(self, current_word: str) -> List[Completion]: + """Get command suggestions with fuzzy matching. + + Args: + current_word: The current word being typed + + Returns: + A list of completions for commands + """ + # Check cache first + current_time = time.time() + cache_key = current_word + + if (cache_key in self._command_suggestions_cache and + current_time - self._command_suggestions_last_update < + self._command_suggestions_update_interval): + return self._command_suggestions_cache[cache_key] + + suggestions = [] + + # Get command descriptions + command_descriptions = self.get_command_descriptions() + + # Sort commands by usage frequency (for command shadowing) + sorted_commands = sorted( + command_descriptions.items(), + key=lambda x: self.command_history.get(x[0], 0), + reverse=True + ) + + # Add command completions + for cmd, description in sorted_commands: + # Exact prefix match + if cmd.startswith(current_word): + suggestions.append(Completion( + cmd, + start_position=-len(current_word), + display=HTML( + f"{cmd:<15} " + f"{description}"), + style="fg:ansicyan bold" + )) + # Fuzzy match (contains the substring) + elif current_word in cmd and not cmd.startswith(current_word): + suggestions.append(Completion( + cmd, + start_position=-len(current_word), + display=HTML( + f"{cmd:<15} {description}"), + style="fg:ansicyan" + )) + + # Add alias completions + for alias, cmd in sorted(COMMAND_ALIASES.items()): + cmd_description = command_descriptions.get(cmd, "") + if alias.startswith(current_word): + suggestions.append(Completion( + alias, + start_position=-len(current_word), + display=HTML( + f"{alias:<15} " + f"{cmd} - {cmd_description}"), + style="fg:ansigreen bold" + )) + elif current_word in alias and not alias.startswith(current_word): + suggestions.append(Completion( + alias, + start_position=-len(current_word), + display=HTML( + f"{alias:<15} " + f"{cmd} - {cmd_description}"), + style="fg:ansigreen" + )) + + # Update cache + self._command_suggestions_cache[cache_key] = suggestions + self._command_suggestions_last_update = current_time + + return suggestions + + # Cache for command shadow + _command_shadow_cache = {} + _command_shadow_last_update = 0 + _command_shadow_update_interval = 0.2 # Update every 200ms + + @lru_cache(maxsize=100) + def _get_command_shadow_cached(self, text: str) -> Optional[str]: + """Cached version of command shadow lookup.""" + if not text or not text.startswith('/'): + return None + + # Find commands that start with the current input + matching_commands = [] + for cmd, count in self.command_history.items(): + if cmd.startswith(text) and cmd != text: + matching_commands.append((cmd, count)) + + # Sort by usage count (descending) + matching_commands.sort(key=lambda x: x[1], reverse=True) + + # Return the most frequently used command + if matching_commands: + return matching_commands[0][0] + + return None + + def get_command_shadow(self, text: str) -> Optional[str]: + """Get a command shadow suggestion based on command history. + + This method returns a suggestion for command shadowing based on + the current input and command usage history. + + Args: + text: The current input text + + Returns: + A suggested command completion or None if no suggestion + """ + # Check cache first + current_time = time.time() + + if (text in self._command_shadow_cache and + current_time - self._command_shadow_last_update < + self._command_shadow_update_interval): + return self._command_shadow_cache[text] + + # Get shadow from cached function + result = self._get_command_shadow_cached(text) + + # Update cache + self._command_shadow_cache[text] = result + self._command_shadow_last_update = current_time + + return result + + # Cache for subcommand suggestions + _subcommand_suggestions_cache = {} + _subcommand_suggestions_last_update = 0 + _subcommand_suggestions_update_interval = 1.0 # Update every second + + def get_subcommand_suggestions( + self, cmd: str, current_word: str) -> List[Completion]: + """Get subcommand suggestions with fuzzy matching. + + Args: + cmd: The main command + current_word: The current word being typed + + Returns: + A list of completions for subcommands + """ + # Check cache first + current_time = time.time() + cache_key = f"{cmd}:{current_word}" + + if (cache_key in self._subcommand_suggestions_cache and + current_time - self._subcommand_suggestions_last_update < + self._subcommand_suggestions_update_interval): + return self._subcommand_suggestions_cache[cache_key] + + suggestions = [] + + # If using an alias, get the real command + cmd = COMMAND_ALIASES.get(cmd, cmd) + + all_commands = self.get_all_commands() + subcommand_descriptions = self.get_subcommand_descriptions() + + if cmd in all_commands: + for subcmd in sorted(all_commands[cmd]): + # Get description for this subcommand if available + subcmd_description = subcommand_descriptions.get( + f"{cmd} {subcmd}", "") + + # Exact prefix match + if subcmd.startswith(current_word): + suggestions.append(Completion( + subcmd, + start_position=-len(current_word), + display=HTML( + f"{subcmd:<15} " + f"{subcmd_description}"), + style="fg:ansiyellow bold" + )) + # Fuzzy match + elif (current_word in subcmd and + not subcmd.startswith(current_word)): + suggestions.append(Completion( + subcmd, + start_position=-len(current_word), + display=HTML( + f"{subcmd:<15} " + f"{subcmd_description}"), + style="fg:ansiyellow" + )) + + # Update cache + self._subcommand_suggestions_cache[cache_key] = suggestions + self._subcommand_suggestions_last_update = current_time + + return suggestions + + def get_model_suggestions(self, current_word: str) -> List[Completion]: + """Get model suggestions for the /model command. + + Args: + current_word: The current word being typed + + Returns: + A list of completions for models + """ + suggestions = [] + + # First try to complete model numbers + for num, model_name in self._cached_model_numbers.items(): + if num.startswith(current_word): + suggestions.append(Completion( + num, + start_position=-len(current_word), + display=HTML( + f"{num:<3} " + f"{model_name}"), + style="fg:ansiwhite bold" + )) + + # Then try to complete model names + for model in self._cached_models: + model_name = model[0] if isinstance(model, tuple) else model + if model_name.startswith(current_word): + suggestions.append(Completion( + model_name, + start_position=-len(current_word), + display=HTML( + f"{model_name}"), + style="fg:ansimagenta bold" + )) + elif (current_word.lower() in model_name.lower() and + not model_name.startswith(current_word)): + suggestions.append(Completion( + model_name, + start_position=-len(current_word), + display=HTML(f"{model_name}"), + style="fg:ansimagenta" + )) + + return suggestions + + # pylint: disable=unused-argument + def get_completions(self, document, complete_event): + """Get completions for the current document + with fuzzy matching support. + + Args: + document: The document to complete + complete_event: The completion event + + Returns: + A generator of completions + """ + text = document.text_before_cursor.strip() + words = text.split() + + # Refresh Ollama models periodically + self.fetch_ollama_models() + + if not text: + # Show all main commands with descriptions + command_descriptions = self.get_command_descriptions() + + # Sort commands by usage frequency (for command shadowing) + sorted_commands = sorted( + command_descriptions.items(), + key=lambda x: self.command_history.get(x[0], 0), + reverse=True + ) + + for cmd, description in sorted_commands: + yield Completion( + cmd, + start_position=0, + display=HTML( + f"{cmd:<15} " + f"{description}"), + style="fg:ansicyan bold" + ) + return + + if text.startswith('/'): + current_word = words[-1] + + # Main command completion (first word) + if len(words) == 1: + # Get command suggestions + yield from self.get_command_suggestions(current_word) + + # Subcommand completion (second word) + elif len(words) == 2: + cmd = words[0] + + # Special handling for model command + if cmd in ["/model", "/mod"]: + yield from self.get_model_suggestions(current_word) + else: + # Get subcommand suggestions + yield from self.get_subcommand_suggestions( + cmd, current_word) diff --git a/src/cai/repl/commands/config.py b/src/cai/repl/commands/config.py new file mode 100644 index 00000000..81184e6d --- /dev/null +++ b/src/cai/repl/commands/config.py @@ -0,0 +1,317 @@ +""" +Config command for CAI via environmental variables. +""" +# Standard library imports +import os +from typing import List, Optional + +# Third party imports +from rich.console import Console # pylint: disable=import-error +from rich.table import Table # pylint: disable=import-error + +# Local imports +from cai.repl.commands.base import Command, register_command + +console = Console() + +# Define environment variables with descriptions and default values +ENV_VARS = { + # CTF variables + 1: { + "name": "CTF_NAME", + "description": "Name of the CTF challenge to run", + "default": None + }, + 2: { + "name": "CTF_CHALLENGE", + "description": "Specific challenge name within the CTF to test", + "default": None + }, + 3: { + "name": "CTF_SUBNET", + "description": "Network subnet for the CTF container", + "default": "192.168.2.0/24" + }, + 4: { + "name": "CTF_IP", + "description": "IP address for the CTF container", + "default": "192.168.2.100" + }, + 5: { + "name": "CTF_INSIDE", + "description": "Whether to conquer the CTF from within container", + "default": "true" + }, + # CAI variables + 6: { + "name": "CAI_MODEL", + "description": "Model to use for agents", + "default": "qwen2.5:14b" + }, + 7: { + "name": "CAI_DEBUG", + "description": "Set debug output level (0: Only tool outputs, 1: Verbose debug output, 2: CLI debug output)", # noqa: E501 # pylint: disable=line-too-long + "default": "1" + }, + 8: { + "name": "CAI_BRIEF", + "description": "Enable/disable brief output mode", + "default": "false" + }, + 9: { + "name": "CAI_MAX_TURNS", + "description": "Maximum number of turns for agent interactions", + "default": "inf" + }, + 10: { + "name": "CAI_TRACING", + "description": "Enable/disable OpenTelemetry tracing", + "default": "true" + }, + 11: { + "name": "CAI_AGENT_TYPE", + "description": "Specify the agents to use (boot2root, one_tool...)", # noqa: E501 # pylint: disable=line-too-long + "default": "one_tool" + }, + 12: { + "name": "CAI_STATE", + "description": "Enable/disable stateful mode", + "default": "false" + }, + 13: { + "name": "CAI_MEMORY", + "description": "Enable/disable memory mode (episodic, semantic, all)", + "default": "false" + }, + 14: { + "name": "CAI_MEMORY_ONLINE", + "description": "Enable/disable online memory mode", + "default": "false" + }, + 15: { + "name": "CAI_MEMORY_OFFLINE", + "description": "Enable/disable offline memory", + "default": "false" + }, + 16: { + "name": "CAI_ENV_CONTEXT", + "description": "Add dirs and current env to llm context", + "default": "true" + }, + 17: { + "name": "CAI_MEMORY_ONLINE_INTERVAL", + "description": "Number of turns between online memory updates", + "default": "5" + }, + 18: { + "name": "CAI_PRICE_LIMIT", + "description": "Price limit for the conversation in dollars", + "default": "1" + }, + 19: { + "name": "CAI_REPORT", + "description": "Enable/disable reporter mode (ctf, nis2, pentesting)", + "default": "ctf" + }, + 20: { + "name": "CAI_SUPPORT_MODEL", + "description": "Model to use for the support agent", + "default": "o3-mini" + }, + 21: { + "name": "CAI_SUPPORT_INTERVAL", + "description": "Number of turns between support agent executions", + "default": "5" + }, +} + + +def get_env_var_value(var_name: str) -> str: + """Get the current value of an environment variable. + + Args: + var_name: The name of the environment variable + + Returns: + The current value or the default value if not set + """ + for var_info in ENV_VARS.values(): + if var_info["name"] == var_name: + return os.environ.get(var_name, var_info["default"] or "Not set") + return "Unknown variable" + + +def set_env_var(var_name: str, value: str) -> bool: + """Set an environment variable. + + Args: + var_name: The name of the environment variable + value: The value to set + + Returns: + True if successful, False otherwise + """ + os.environ[var_name] = value + return True + + +class ConfigCommand(Command): + """Command for displaying and configuring environment variables.""" + + def __init__(self): + """Initialize the config command.""" + super().__init__( + name="/config", + description=( + "Display and configure environment variables" + ), + aliases=["/cfg"] + ) + + # Add subcommands + self.add_subcommand( + "list", + "List all environment variables and their values", + self.handle_list + ) + self.add_subcommand( + "set", + "Set an environment variable by its number", + self.handle_set + ) + self.add_subcommand( + "get", + "Get the value of an environment variable by its number", + self.handle_get + ) + + def handle_no_args(self) -> bool: + """Handle the command when no arguments are provided. + + Returns: + True if the command was handled successfully, False otherwise + """ + return self.handle_list(None) + + def handle_list(self, _: Optional[List[str]] = None) -> bool: + """List all environment variables and their values. + + Args: + _: Ignored arguments + + Returns: + True if successful + """ + table = Table( + title="Environment Variables", + show_header=True, + header_style="bold yellow" + ) + table.add_column("#", style="dim") + table.add_column("Variable", style="yellow") + table.add_column("Value", style="green") + table.add_column("Default", style="blue") + table.add_column("Description") + + for num, var_info in ENV_VARS.items(): + var_name = var_info["name"] + current_value = get_env_var_value(var_name) + default_value = var_info["default"] or "Not set" + + table.add_row( + str(num), + var_name, + current_value, + default_value, + var_info["description"] + ) + + console.print(table) + console.print( + "\nUsage: /config set to configure a variable" + ) + return True + + def handle_get(self, args: Optional[List[str]] = None) -> bool: + """Get the value of an environment variable by its number. + + Args: + args: Command arguments [var_number] + + Returns: + True if successful, False otherwise + """ + if not args or len(args) < 1: + console.print( + "[yellow]Usage: /config get [/yellow]" + ) + return False + + try: + var_num = int(args[0]) + if var_num not in ENV_VARS: + console.print( + f"[red]Error: Variable number {var_num} not found[/red]" + ) + return False + + var_info = ENV_VARS[var_num] + var_name = var_info["name"] + current_value = get_env_var_value(var_name) + + console.print( + f"[yellow]{var_name}[/yellow]: " + f"[green]{current_value}[/green] " + f"(Default: [blue]{var_info['default'] or 'Not set'}[/blue])" + ) + return True + except ValueError: + console.print( + "[red]Error: Variable number must be an integer[/red]" + ) + return False + + def handle_set(self, args: Optional[List[str]] = None) -> bool: + """Set an environment variable by its number. + + Args: + args: Command arguments [var_number, value] + + Returns: + True if successful, False otherwise + """ + if not args or len(args) < 2: + console.print( + "[yellow]Usage: /config set [/yellow]" + ) + return False + + try: + var_num = int(args[0]) + if var_num not in ENV_VARS: + console.print( + f"[red]Error: Variable number {var_num} not found[/red]" + ) + return False + + value = args[1] + var_info = ENV_VARS[var_num] + var_name = var_info["name"] + + old_value = get_env_var_value(var_name) + set_env_var(var_name, value) + + console.print( + f"[green]Set {var_name} to '{value}' " + f"(was: '{old_value}')[/green]" + ) + return True + except ValueError: + console.print( + "[red]Error: Variable number must be an integer[/red]" + ) + return False + + +# Register the command +register_command(ConfigCommand()) diff --git a/src/cai/repl/commands/env.py b/src/cai/repl/commands/env.py new file mode 100644 index 00000000..91445d9a --- /dev/null +++ b/src/cai/repl/commands/env.py @@ -0,0 +1,82 @@ +""" +Environment command for CAI REPL. +This module provides commands for displaying environment variables. +""" +import os +from typing import ( + List, + Optional +) +from rich.console import Console # pylint: disable=import-error +from rich.table import Table # pylint: disable=import-error + +from cai.repl.commands.base import Command, register_command + +console = Console() + + +class EnvCommand(Command): + """Command for displaying environment variables.""" + + def __init__(self): + """Initialize the env command.""" + super().__init__( + name="/env", + description="Display environment variables and their values", + aliases=["/e"] + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the env command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + return self.handle_env_command() + + def handle_env_command(self) -> bool: + """Display environment variables starting with CAI or CTF. + + Returns: + bool: True if the command was executed successfully + """ + # Get all environment variables + env_vars = { + k: v for k, v in os.environ.items() if k.startswith( + ('CAI_', 'CTF_'))} + + if not env_vars: + console.print( + "[yellow]No CAI_ or CTF_ environment variables found[/yellow]") + return True + + # Create a table to display the variables + table = Table( + title="Environment Variables", + show_header=True, + header_style="bold magenta") + table.add_column("Variable", style="cyan") + table.add_column("Value", style="green") + + # Add rows to the table with masked values for sensitive data + for key, value in sorted(env_vars.items()): + # Mask sensitive values (API keys, tokens, etc.) + if any(sensitive in key.lower() + for sensitive in ['key', 'token', 'secret', 'password']): + # Show first half of the value, mask the rest + half_length = len(value) // 2 + masked_value = value[:half_length] + \ + '*' * (len(value) - half_length) + table.add_row(key, masked_value) + else: + table.add_row(key, value) + + console.print(table) + return True + + +# Register the command +register_command(EnvCommand()) diff --git a/src/cai/repl/commands/exit.py b/src/cai/repl/commands/exit.py new file mode 100644 index 00000000..2a17e149 --- /dev/null +++ b/src/cai/repl/commands/exit.py @@ -0,0 +1,35 @@ +""" +Exit command for CAI REPL. +This module provides the command to exit the REPL. +""" +import sys +from typing import List, Optional + +from cai.repl.commands.base import Command, register_command + + +class ExitCommand(Command): + """Command for exiting the REPL.""" + + def __init__(self): + """Initialize the exit command.""" + super().__init__( + name="/exit", + description="Exit the CAI REPL", + aliases=["/q", "/quit"] + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the exit command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + sys.exit(0) + + +# Register the command +register_command(ExitCommand()) diff --git a/src/cai/repl/commands/graph.py b/src/cai/repl/commands/graph.py new file mode 100644 index 00000000..e3ccbc93 --- /dev/null +++ b/src/cai/repl/commands/graph.py @@ -0,0 +1,61 @@ +""" +Graph command for CAI REPL. +This module provides commands for visualizing the agent interaction graph. +""" +from typing import ( + List, + Optional +) +from rich.console import Console # pylint: disable=import-error + +from cai.repl.commands.base import Command, register_command + +console = Console() + + +class GraphCommand(Command): + """Command for visualizing the agent interaction graph.""" + + def __init__(self): + """Initialize the graph command.""" + super().__init__( + name="/graph", + description="Visualize the agent interaction graph", + aliases=["/g"] + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the graph command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + return self.handle_graph_show() + + def handle_graph_show(self) -> bool: + """Handle /graph show command""" + from cai.repl.repl import client # pylint: disable=import-error + + # Import here to avoid circular imports + + if not client or not client._graph: # pylint: disable=protected-access + console.print("[yellow]No conversation graph available.[/yellow]") + return True + + try: + console.print("\n[bold]Conversation Graph:[/bold]") + console.print("------------------") + console.print( + client._graph.ascii()) # pylint: disable=protected-access + console.print() + return True + except Exception as e: # pylint: disable=broad-except + console.print(f"[red]Error displaying graph: {e}[/red]") + return False + + +# Register the command +register_command(GraphCommand()) diff --git a/src/cai/repl/commands/help.py b/src/cai/repl/commands/help.py new file mode 100644 index 00000000..4df6e498 --- /dev/null +++ b/src/cai/repl/commands/help.py @@ -0,0 +1,901 @@ +""" +Help command for CAI REPL. +This module provides commands for displaying help information. +""" +from typing import List, Optional +try: + from rich.console import Console + from rich.table import Table + from rich.panel import Panel + from rich.text import Text +except ImportError as exc: + raise ImportError( + "The 'rich' package is required. Please install it with: " + "pip install rich" + ) from exc + +from cai.repl.commands.base import ( + Command, + register_command, + COMMANDS, + COMMAND_ALIASES +) + +try: + from cai import is_caiextensions_platform_available + from caiextensions.platform.base.platform_manager import PlatformManager + HAS_PLATFORM_EXTENSIONS = True +except ImportError: + HAS_PLATFORM_EXTENSIONS = False + +console = Console() + + +def create_styled_table( + title: str, + headers: List[tuple[str, str]], + header_style: str = "bold white" +) -> Table: + """Create a styled table with consistent formatting. + + Args: + title: The table title + headers: List of (header_name, style) tuples + header_style: Style for the header row + + Returns: + A configured Table instance + """ + table = Table( + title=title, + show_header=True, + header_style=header_style + ) + for header, style in headers: + table.add_column(header, style=style) + return table + + +def create_notes_panel( + notes: List[str], + title: str = "Notes", + border_style: str = "yellow" +) -> Panel: + """Create a notes panel with consistent formatting. + + Args: + notes: List of note strings + title: Panel title + border_style: Style for the panel border + + Returns: + A configured Panel instance + """ + notes_text = Text.from_markup( + "\n".join(f"• {note}" for note in notes) + ) + return Panel( + notes_text, + title=title, + border_style=border_style + ) + + +class HelpCommand(Command): + """Command for displaying help information.""" + + def __init__(self): + """Initialize the help command.""" + super().__init__( + name="/help", + description=( + "Display help information about commands " + "and features" + ), + aliases=["/h"] + ) + + # Add subcommands + self.add_subcommand( + "memory", + "Display help for memory commands", + self.handle_memory + ) + self.add_subcommand( + "agents", + "Display help for agent commands", + self.handle_agents + ) + self.add_subcommand( + "graph", + "Display help for graph commands", + self.handle_graph + ) + self.add_subcommand( + "platform", + "Display help for platform commands", + self.handle_platform + ) + self.add_subcommand( + "shell", + "Display help for shell commands", + self.handle_shell + ) + self.add_subcommand( + "env", + "Display help for environment commands", + self.handle_env + ) + self.add_subcommand( + "aliases", + "Display command aliases", + self.handle_aliases + ) + self.add_subcommand( + "model", + "Display help for model commands", + self.handle_model + ) + self.add_subcommand( + "turns", + "Display help for turns commands", + self.handle_turns + ) + self.add_subcommand( + "config", + "Display help for config commands", + self.handle_config + ) + + def handle_memory(self, _: Optional[List[str]] = None) -> bool: + """Show help for memory commands.""" + # Get the memory command and show its help + memory_cmd = next((cmd for cmd in COMMANDS.values() + if cmd.name == "/memory"), None) + if memory_cmd and hasattr(memory_cmd, 'show_help'): + memory_cmd.show_help() + return True + + # Fallback if memory command not found or doesn't have show_help + self.handle_help_memory() + return True + + def handle_agents(self, _: Optional[List[str]] = None) -> bool: + """Show help for agent-related features.""" + console.print(Panel( + "Agents are autonomous AI assistants that can perform specific " + "tasks.\n\n" + "[bold]Available Commands:[/bold]\n" + "• [yellow]/agent list[/yellow] - List all available agents\n" + "• [yellow]/agent use [/yellow] - Switch to a specific agent\n" + "• [yellow]/agent info [/yellow] - Show details about an " + "agent\n\n" + "[bold]Examples:[/bold]\n" + "• [green]/agent use boot2root_agent[/green] - Switch to the CLI " + "security testing agent\n" + "• [green]/agent use dns_smtp_agent[/green] - Switch to the " + "DNS/SMTP reconnaissance agent", + title="Agent Commands", + border_style="blue" + )) + return True + + def handle_graph(self, _: Optional[List[str]] = None) -> bool: + """Show help for graph visualization.""" + console.print(Panel( + "Graph visualization helps you understand the relationships " + "between different pieces of information in your session.\n\n" + "[bold]Available Commands:[/bold]\n" + "• [yellow]/graph show[/yellow] - Display the current memory " + "graph\n" + "• [yellow]/graph export [/yellow] - Export graph to a " + "file\n" + "• [yellow]/graph focus [/yellow] - Focus on a specific " + "node\n\n" + "[bold]Examples:[/bold]\n" + "• [green]/graph show[/green] - Display the current memory graph\n" + "• [green]/graph export session_graph.png[/green] - Save graph " + "as PNG", + title="Graph Visualization Commands", + border_style="blue" + )) + return True + + def handle_platform(self, _: Optional[List[str]] = None) -> bool: + """Show help for platform-specific features.""" + platform_cmd = next( + (cmd for cmd in COMMANDS.values() if cmd.name == "/platform"), + None + ) + + if platform_cmd and hasattr(platform_cmd, 'show_help'): + platform_cmd.show_help() + return True + + console.print(Panel( + "Platform commands provide access to platform-specific " + "features.\n\n" + "[bold]Available Commands:[/bold]\n" + "• [yellow]/platform list[/yellow] - List available platforms\n" + "• [yellow]/platform [/yellow] - Run " + "platform-specific command\n\n" + "[bold]Examples:[/bold]\n" + "• [green]/platform list[/green] - Show all available platforms\n" + "• [green]/p list[/green] - Shorthand for platform list", + title="Platform Commands", + border_style="blue" + )) + return True + + def handle_shell(self, _: Optional[List[str]] = None) -> bool: + """Show help for shell command execution.""" + console.print(Panel( + "Shell commands allow you to execute system commands directly.\n\n" + "[bold]Available Commands:[/bold]\n" + "• [yellow]/shell [/yellow] - Execute a shell command\n" + "• [yellow]/![/yellow] - Shorthand for /shell\n\n" + "[bold]Session Management:[/bold]\n" + "• [yellow]/shell session list[/yellow] - List active sessions\n" + "• [yellow]/shell session output [/yellow] - Get output from " + "a session\n" + "• [yellow]/shell session kill [/yellow] - Terminate a " + "session\n\n" + "[bold]Examples:[/bold]\n" + "• [green]/shell ls -la[/green] - List files in current " + "directory\n" + "• [green]/! pwd[/green] - Show current working directory", + title="Shell Commands", + border_style="blue" + )) + return True + + def handle_env(self, _: Optional[List[str]] = None) -> bool: + """Show help for environment variables.""" + console.print(Panel( + "Environment variables control CAI's behavior.\n\n" + "[bold]Key Variables:[/bold]\n" + "• [yellow]CAI_MODEL[/yellow] - Default AI model (e.g., " + "'claude-3-7-sonnet-20250219')\n" + "• [yellow]CAI_MEMORY_DIR[/yellow] - Directory for storing memory " + "collections\n" + "• [yellow]OPENAI_API_KEY[/yellow] - API key for OpenAI models\n" + "• [yellow]ANTHROPIC_API_KEY[/yellow] - API key for Anthropic " + "models\n\n" + "[bold]Available Commands:[/bold]\n" + "• [yellow]/env list[/yellow] - Show all environment variables\n" + "• [yellow]/env set [/yellow] - Set an environment " + "variable\n" + "• [yellow]/env get [/yellow] - Get the value of an " + "environment variable", + title="Environment Variables", + border_style="blue" + )) + return True + + def handle_aliases(self, _: Optional[List[str]] = None) -> bool: + """Show all command aliases.""" + return self.handle_help_aliases() + + def handle_model(self, _: Optional[List[str]] = None) -> bool: + """Show help for model selection.""" + return self.handle_help_model() + + def handle_turns(self, _: Optional[List[str]] = None) -> bool: + """Show help for managing turns.""" + return self.handle_help_turns() + + def handle_config(self, _: Optional[List[str]] = None) -> bool: + """Display help for config commands. + + Args: + _: Ignored arguments + + Returns: + True if successful + """ + return self.handle_help_config() + + def handle_no_args(self) -> bool: + """Handle the command when no arguments are provided.""" + return self.handle_help() + + def _print_command_table( + self, + title: str, + commands: List[tuple[str, str, str]], + header_style: str = "bold yellow", + command_style: str = "yellow" + ) -> None: + """Print a table of commands with consistent formatting.""" + table = create_styled_table( + title, + [ + ("Command", command_style), + ("Alias", "green"), + ("Description", "white") + ], + header_style + ) + + for cmd, alias, desc in commands: + table.add_row(cmd, alias, desc) + + console.print(table) + + def handle_help(self) -> bool: + """Display general help information. + + Returns: + True if successful + """ + console.print( + Panel( + Text.from_markup( + "Welcome to the CAI help system. " + "This system provides information about " + "available commands and features." + ), + title="CAI Help", + border_style="yellow" + ) + ) + + # Memory Commands + memory_commands = [ + ("/memory list", "/m list", + "List all available memory collections"), + ("/memory load ", "/m load ", + "Load a memory collection"), + ("/memory delete ", "/m delete ", + "Delete a memory collection"), + ("/memory create ", "/m create ", + "Create a new memory collection") + ] + self._print_command_table("Memory Commands", memory_commands) + + # Collection types info + collection_info = Text() + collection_info.append("\nCollection Types:\n", style="bold") + collection_info.append("• CTF_NAME", style="yellow") + collection_info.append( + " - Episodic memory for a specific CTF (e.g. ", + style="white" + ) + collection_info.append("baby_first", style="bold white") + collection_info.append(")\n", style="white") + collection_info.append("• _all_", style="yellow") + collection_info.append( + " - Semantic memory across all CTFs", + style="white" + ) + console.print(collection_info) + + # Graph Commands + graph_commands = [ + ("/graph", "/g", + "Show the graph of the current memory collection") + ] + self._print_command_table( + "Graph Commands", + graph_commands, + "bold blue", + "blue" + ) + + # Shell Commands + shell_commands = [ + ("/shell ", "/s ", + "Execute a shell command (can be interrupted with CTRL+C)") + ] + self._print_command_table( + "Shell Commands", + shell_commands, + "bold green", + "green" + ) + + # Config Commands + config_commands = [ + ("/config", "/cfg", + "List all environment variables and their values"), + ("/config list", "/cfg list", + "List all environment variables and their values"), + ("/config get ", "/cfg get ", + "Get the value of a specific environment variable"), + ("/config set ", "/cfg set ", + "Set the value of a specific environment variable") + ] + self._print_command_table( + "Config Commands", + config_commands, + "bold magenta", + "magenta" + ) + + # Environment Commands + env_commands = [ + ("/env", "/e", + "Display environment variables (CAI_* and CTF_*)") + ] + self._print_command_table( + "Environment Commands", + env_commands, + "bold cyan", + "cyan" + ) + + # Model Commands + model_commands = [ + ("/model", "/mod", + "Display current model and list available models"), + ("/model ", "/mod ", + "Change the model to ") + ] + self._print_command_table( + "Model Commands", + model_commands, + "bold magenta", + "magenta" + ) + + # Turns Commands + turns_commands = [ + ("/turns", "/t", "Display current maximum number of turns"), + ("/turns ", "/t ", + "Change the maximum number of turns") + ] + self._print_command_table( + "Turns Commands", + turns_commands, + "bold magenta", + "magenta" + ) + + # Platform Commands + self.handle_help_platform_manager() + + # Tips section + tips = Panel( + Text.from_markup( + "Tips:\n" + "• Use [bold]Tab[/bold] for command completion\n" + "• Use [bold]↑/↓[/bold] to navigate command history\n" + "• Use [bold]Ctrl+L[/bold] to clear the screen\n" + "• Most commands have shorter aliases (e.g. [bold]/h[/bold] " + "instead of [bold]/help[/bold])" + ), + title="Helpful Tips", + border_style="cyan" + ) + console.print(tips) + + return True + + def handle_help_aliases(self) -> bool: + """Show all command aliases in a well-formatted table.""" + # Create a styled header + console.print( + Panel( + "Command Aliases Reference", + border_style="magenta", + title="Aliases" + ) + ) + + # Create a table for aliases + alias_table = create_styled_table( + "Command Aliases", + [ + ("Alias", "green"), + ("Command", "yellow"), + ("Description", "white") + ], + "bold magenta" + ) + + # Add rows for each alias + for alias, command in sorted(COMMAND_ALIASES.items()): + cmd = COMMANDS.get(command) + description = cmd.description if cmd else "" + alias_table.add_row(alias, command, description) + + console.print(alias_table) + + # Add tips + tips = [ + "Aliases can be used anywhere the full command would be used", + ( + "Example: [green]/m list[/green] instead of " + "[yellow]/memory list[/yellow]" + ) + ] + console.print("\n") + console.print(create_notes_panel(tips, "Tips", "cyan")) + + return True + + def handle_help_memory(self) -> bool: + """Show help for memory commands with rich formatting.""" + # Create a styled header + header = Text("Memory Command Help", style="bold yellow") + console.print(Panel(header, border_style="yellow")) + + # Usage table + usage_table = create_styled_table( + "Usage", + [("Command", "yellow"), ("Description", "white")] + ) + + usage_table.add_row( + "/memory list", + "Display all available memory collections" + ) + usage_table.add_row( + "/memory load ", + "Set the active memory collection" + ) + usage_table.add_row( + "/memory delete ", + "Delete a memory collection" + ) + usage_table.add_row( + "/memory create ", + "Create a new memory collection" + ) + usage_table.add_row("/m", "Alias for /memory") + + console.print(usage_table) + + # Examples table + examples_table = create_styled_table( + "Examples", + [("Example", "cyan"), ("Description", "white")], + "bold cyan" + ) + + examples = [ + ("/memory list", "List all available collections"), + ("/memory load _all_", "Load the semantic memory collection"), + ("/memory load my_ctf", "Load the episodic memory for 'my_ctf'"), + ( + "/memory create new_collection", + "Create a new collection named 'new_collection'" + ), + ( + "/memory delete old_collection", + "Delete the collection named 'old_collection'" + ) + ] + + for example, desc in examples: + examples_table.add_row(example, desc) + + console.print(examples_table) + + # Collection types table + types_table = create_styled_table( + "Collection Types", + [("Type", "green"), ("Description", "white")], + "bold green" + ) + + types = [ + ("_all_", "Semantic memory across all CTFs"), + ("", "Episodic memory for a specific CTF"), + ("", "Custom memory collection") + ] + + for type_name, desc in types: + types_table.add_row(type_name, desc) + + console.print(types_table) + + # Notes panel + notes = [ + "Memory collections are stored in the Qdrant vector database", + "The active collection is stored in the CAI_MEMORY_COLLECTION " + "env var", + "Episodic memory is used for specific CTFs or tasks", + "Semantic memory (_all_) is used across all CTFs", + "Memory is used to provide context to the agent" + ] + + console.print(create_notes_panel(notes)) + + return True + + def handle_help_model(self) -> bool: + """Show help for model command with rich formatting.""" + # Create a styled header + header = Text("Model Command Help", style="bold magenta") + console.print(Panel(header, border_style="magenta")) + + # Usage table + usage_table = create_styled_table( + "Usage", + [("Command", "magenta"), ("Description", "white")] + ) + + usage_commands = [ + ("/model", "Display current model and list available models"), + ("/model ", "Change the model to "), + ( + "/model ", + "Change the model using its number from the list" + ), + ("/mod", "Alias for /model") + ] + + for cmd, desc in usage_commands: + usage_table.add_row(cmd, desc) + + console.print(usage_table) + + # Examples table + examples_table = create_styled_table( + "Examples", + [("Example", "cyan"), ("Description", "white")], + "bold cyan" + ) + + examples = [ + ( + "/model 1", + "Switch to the first model in the list (Claude 3.7 Sonnet)" + ), + ( + "/model claude-3-7-sonnet-20250219", + "Switch to Claude 3.7 Sonnet model" + ), + ( + "/model o1", + "Switch to OpenAI's O1 model (good for math)" + ), + ( + "/model gpt-4o", + "Switch to OpenAI's GPT-4o model" + ) + ] + + for example, desc in examples: + examples_table.add_row(example, desc) + + console.print(examples_table) + + # Model categories table + categories_table = create_styled_table( + "Model Categories", + [("Category", "green"), ("Description", "white")], + "bold green" + ) + + categories = [ + ( + "Claude 3.7", + "Best models for complex reasoning and creative tasks" + ), + ( + "Claude 3.5", + "Excellent balance of performance and efficiency" + ), + ( + "Claude 3", + "Range of models from powerful (Opus) to fast (Haiku)" + ), + ( + "OpenAI O-series", + "Specialized models with strong mathematical capabilities" + ), + ( + "OpenAI GPT-4", + "Powerful general-purpose models" + ), + ( + "Ollama", + "Local models running on your machine or Docker container" + ) + ] + + for category, desc in categories: + categories_table.add_row(category, desc) + + console.print(categories_table) + + # Notes panel + notes = [ + "The model change takes effect on the next agent interaction", + "The model is stored in the CAI_MODEL environment variable", + "Some models may require specific API keys to be set", + "OpenAI models require OPENAI_API_KEY to be set", + "Anthropic models require ANTHROPIC_API_KEY to be set", + "Ollama models require Ollama to be running locally", + ( + "Ollama is configured to run on " + "host.docker.internal:8000" + ) + ] + + console.print(create_notes_panel(notes)) + + return True + + def handle_help_turns(self) -> bool: + """Show help for turns command with rich formatting.""" + # Create a styled header + header = Text("Turns Command Help", style="bold magenta") + console.print(Panel(header, border_style="magenta")) + + # Usage table + usage_table = create_styled_table( + "Usage", + [("Command", "magenta"), ("Description", "white")] + ) + + usage_commands = [ + ("/turns", "Display current maximum number of turns"), + ("/turns ", "Change the maximum number of turns"), + ("/turns inf", "Set unlimited turns"), + ("/t", "Alias for /turns") + ] + + for cmd, desc in usage_commands: + usage_table.add_row(cmd, desc) + + console.print(usage_table) + + # Examples table + examples_table = create_styled_table( + "Examples", + [("Example", "cyan"), ("Description", "white")], + "bold cyan" + ) + + examples = [ + ("/turns", "Show current maximum turns"), + ("/turns 10", "Set maximum turns to 10"), + ("/turns inf", "Set unlimited turns"), + ("/t 5", "Set maximum turns to 5 (using alias)") + ] + + for example, desc in examples: + examples_table.add_row(example, desc) + + console.print(examples_table) + + # Notes panel + notes = [ + ( + "The maximum turns limit controls how many responses the " + "agent will give" + ), + "Setting turns to 'inf' allows unlimited responses", + ( + "The turns count is stored in the CAI_MAX_TURNS " + "environment variable" + ), + "Each agent response counts as one turn" + ] + + console.print(create_notes_panel(notes)) + + return True + + def handle_help_platform_manager(self) -> bool: + """Show help for platform manager commands.""" + if HAS_PLATFORM_EXTENSIONS and is_caiextensions_platform_available(): + try: + from caiextensions.platform.base import platform_manager + platforms = platform_manager.list_platforms() + + if not platforms: + console.print( + "[yellow]No platforms registered.[/yellow]" + ) + return True + + platform_table = create_styled_table( + "Available Platforms", + [ + ("Platform", "magenta"), + ("Description", "white") + ], + "bold magenta" + ) + + for platform_name in platforms: + platform = platform_manager.get_platform(platform_name) + description = getattr( + platform, 'description', platform_name.capitalize()) + platform_table.add_row( + platform_name, + description + ) + + console.print(platform_table) + + # Add platform command examples + examples = [] + for platform_name in platforms: + platform = platform_manager.get_platform(platform_name) + commands = platform.get_commands() + if commands: + examples.append( + f"[green]/platform {platform_name} { + commands[0]}[/green] - Example {platform_name} command") + + if examples: + console.print(Panel( + "\n".join(examples), + title="Platform Command Examples", + border_style="blue" + )) + + return True + except (ImportError, Exception) as e: + console.print( + f"[yellow]Error loading platforms: {e}[/yellow]" + ) + return True + + console.print( + "[yellow]No platform extensions available.[/yellow]" + ) + return True + + def handle_help_config(self) -> bool: + """Display help for config commands. + + Returns: + True if successful + """ + console.print( + Panel( + Text.from_markup( + "The [bold yellow]/config[/bold yellow] command allows you" + "to view and configure environment variables that control" + "the behavior of CAI." + ), + title="Config Commands", + border_style="yellow" + ) + ) + + # Create table for subcommands + table = create_styled_table( + "Available Subcommands", + [("Command", "yellow"), ("Description", "white")] + ) + + table.add_row( + "/config", + "List all environment variables and their current values" + ) + table.add_row( + "/config list", + "List all environment variables and their current values" + ) + table.add_row( + "/config get ", + "Get the value of a specific environment variable by its number" + ) + table.add_row( + "/config set ", + "Set the value of a specific environment variable by its number" + ) + + console.print(table) + + # Create notes panel + notes = [ + "Environment variables control various aspects of CAI behavior.", + "Changes environment variables only affect the current session.", + "Use the [yellow]/config list[/yellow] command to see options.", + "Each variable is assigned a number for easy reference." + ] + console.print(create_notes_panel(notes)) + + return True + + +# Register the command +register_command(HelpCommand()) diff --git a/src/cai/repl/commands/history.py b/src/cai/repl/commands/history.py new file mode 100644 index 00000000..3a7df9c9 --- /dev/null +++ b/src/cai/repl/commands/history.py @@ -0,0 +1,84 @@ +""" +History command for CAI REPL. +This module provides commands for displaying conversation history. +""" +from rich.console import Console # pylint: disable=import-error +from rich.table import Table # pylint: disable=import-error + +from cai.repl.commands.base import Command, register_command + +console = Console() + + +class HistoryCommand(Command): + """Command for displaying conversation history.""" + + def __init__(self): + """Initialize the history command.""" + super().__init__( + name="/history", + description="Display the conversation history", + aliases=["/h"] + ) + + def handle_no_args(self) -> bool: + """Handle the command when no arguments are provided. + + Returns: + True if the command was handled successfully, False otherwise + """ + # Access messages directly from repl.py's global scope + try: + from cai.repl.repl import messages # pylint: disable=import-outside-toplevel # noqa: E501 + except ImportError: + console.print( + "[red]Error: Could not access conversation history[/red]") + return False + + if not messages: + console.print("[yellow]No conversation history available[/yellow]") + return True + + # Create a table for the history + table = Table( + title="Conversation History", + show_header=True, + header_style="bold yellow" + ) + table.add_column("#", style="dim") + table.add_column("Role", style="cyan") + table.add_column("Content", style="green") + + # Add messages to the table + for idx, msg in enumerate(messages, 1): + role = msg.get("role", "unknown") + content = msg.get("content", "") + + # Truncate long content for better display + if len(content) > 100: + content = content[:97] + "..." + + # Color the role based on type + if role == "user": + role_style = "cyan" + elif role == "assistant": + role_style = "yellow" + else: + role_style = "red" + + # Add a newline between each role for better readability + if idx > 1: + table.add_row("", "", "") + + table.add_row( + str(idx), + f"[{role_style}]{role}[/{role_style}]", + content + ) + + console.print(table) + return True + + +# Register the command +register_command(HistoryCommand()) diff --git a/src/cai/repl/commands/kill.py b/src/cai/repl/commands/kill.py new file mode 100644 index 00000000..aba42cac --- /dev/null +++ b/src/cai/repl/commands/kill.py @@ -0,0 +1,81 @@ +""" +Kill command for CAI REPL. +This module provides commands for terminating active processes or sessions. +""" +import os +import signal +from typing import ( + List, + Optional +) +from rich.console import Console # pylint: disable=import-error + +from cai.repl.commands.base import Command, register_command + +console = Console() + + +class KillCommand(Command): + """Command for terminating active processes or sessions.""" + + def __init__(self): + """Initialize the kill command.""" + super().__init__( + name="/kill", + description="Terminate active processes or sessions", + aliases=["/k"] + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the kill command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + return self.handle_kill_command(args) + + def handle_kill_command(self, args: List[str]) -> bool: + """Kill a background process by PID. + + Args: + args: List containing the PID to kill + + Returns: + bool: True if the process was killed successfully + """ + if not args: + console.print("[red]Error: No PID specified[/red]") + return False + + try: + pid = int(args[0]) + + # Try to kill the process group + try: + os.killpg(pid, signal.SIGTERM) + console.print(f"[green]Process group {pid} terminated[/green]") + except BaseException: # pylint: disable=broad-exception-caught + # If killing the process group fails, try killing just the + # process + os.kill(pid, signal.SIGTERM) + console.print(f"[green]Process {pid} terminated[/green]") + + return True + except ValueError: + console.print("[red]Error: Invalid PID format[/red]") + return False + except ProcessLookupError: + console.print( + f"[yellow]No process with PID { + args[0]} found[/yellow]") + return False + except Exception as e: # pylint: disable=broad-exception-caught + console.print(f"[red]Error killing process: {str(e)}[/red]") + return False + + +# Register the command +register_command(KillCommand()) diff --git a/src/cai/repl/commands/memory.py b/src/cai/repl/commands/memory.py new file mode 100644 index 00000000..12614d92 --- /dev/null +++ b/src/cai/repl/commands/memory.py @@ -0,0 +1,4 @@ +""" +Memory command for CAI REPL. +This module provides commands for managing memory collections. +""" diff --git a/src/cai/repl/commands/model.py b/src/cai/repl/commands/model.py new file mode 100644 index 00000000..31cd8e9e --- /dev/null +++ b/src/cai/repl/commands/model.py @@ -0,0 +1,711 @@ +""" +Model command for CAI REPL. +This module provides commands for viewing and changing the current LLM model. +""" +import os +import datetime +# Standard library imports +from typing import List, Optional # Dict and Any removed as unused + +# Third-party imports +import requests # pylint: disable=import-error +from rich.console import Console # pylint: disable=import-error +from rich.table import Table # pylint: disable=import-error +from rich.panel import Panel # pylint: disable=import-error +from cai.util import get_ollama_api_base +from cai.repl.commands.base import Command, register_command + +console = Console() + +LITELLM_URL = ( + "https://raw.githubusercontent.com/BerriAI/litellm/main/" + "model_prices_and_context_window.json" +) + + +class ModelCommand(Command): + """Command for viewing and changing the current LLM model.""" + + def __init__(self): + """Initialize the model command.""" + super().__init__( + name="/model", + description="View or change the current LLM model", + aliases=["/mod"] + ) + + # Cache for model information + self.cached_models = [] + # Map of numbers to model names + self.cached_model_numbers = {} + self.last_model_fetch = ( + datetime.datetime.now() - datetime.timedelta(minutes=10) + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the model command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + return self.handle_model_command(args) + + # pylint: disable=too-many-locals,too-many-branches,too-many-statements + def handle_model_command(self, args: List[str]) -> bool: + """Change the model used by CAI. + + Args: + args: List containing the model name to use or a number to select + from the list + + Returns: + bool: True if the model was changed successfully + """ + # Define model categories and their models for easy reference + # pylint: disable=invalid-name + MODEL_CATEGORIES = { + "Claude 3.7": [ + { + "name": "claude-3-7-sonnet-20250219", + "description": ( + "Best model for complex reasoning and creative tasks" + ) + } + ], + "Claude 3.5": [ + { + "name": "claude-3-5-sonnet-20240620", + "description": ( + "Excellent balance of performance and efficiency" + ) + }, + { + "name": "claude-3-5-sonnet-20241022", + "description": ( + "Latest Claude 3.5 model with improved capabilities" + ) + } + ], + "Claude 3": [ + { + "name": "claude-3-opus-20240229", + "description": "Powerful Claude 3 model for complex tasks" + }, + { + "name": "claude-3-sonnet-20240229", + "description": "Balanced performance and speed" + }, + { + "name": "claude-3-haiku-20240307", + "description": "Fast and efficient model" + } + ], + "OpenAI O-series": [ + { + "name": "o1", + "description": ( + "Excellent for mathematical reasoning and " + "problem-solving" + ) + }, + { + "name": "o1-mini", + "description": ( + "Smaller O1 model with good math capabilities" + ) + }, + { + "name": "o3-mini", + "description": "Latest mini model in the O-series" + }, + { + "name": "gpt-4o", + "description": ( + "Latest GPT-4 model with improved capabilities" + ) + }, + { + "name": "gpt-4o-audio-preview", + "description": "GPT-4o with audio capabilities" + }, + { + "name": "gpt-4o-audio-preview-2024-12-17", + "description": "Updated GPT-4o with audio capabilities" + }, + { + "name": "gpt-4o-audio-preview-2024-10-01", + "description": "Previous GPT-4o with audio capabilities" + } + ], + "OpenAI GPT-4": [ + {"name": "gpt-4", "description": "Original GPT-4 model"}, + {"name": "gpt-4-turbo", + "description": "Fast and powerful GPT-4 model"} + ], + "OpenAI GPT-4.5": [ + { + "name": "gpt-4.5-preview", + "description": ( + "Latest non reasoning openai model with " + "improved capabilities" + ) + }, + { + "name": "gpt-4.5-preview-2025-02-27", + "description": "Specific version of GPT-4.5 preview" + } + ], + "OpenAI GPT-3.5": [ + { + "name": "gpt-3.5-turbo", + "description": "Fast and cost-effective model" + } + ], + "DeepSeek": [ + { + "name": "deepseek-v3", + "description": "DeepSeek's latest general-purpose model" + }, + { + "name": "deepseek-r1", + "description": "DeepSeek's specialized reasoning model" + } + ] + } + + # Fetch model pricing data from LiteLLM GitHub repository + model_pricing_data = {} + try: + response = requests.get(LITELLM_URL, timeout=2) + if response.status_code == 200: + model_pricing_data = response.json() + + # Add DeepSeek models with their pricing if not in data + deepseek_v3_path = "deepseek/deepseek-v3" + deepseek_r1_path = "deepseek/deepseek-r1" + + if (deepseek_v3_path in model_pricing_data and + "deepseek-v3" not in model_pricing_data): + model_pricing_data["deepseek-v3"] = ( + model_pricing_data[deepseek_v3_path] + ) + if (deepseek_r1_path in model_pricing_data and + "deepseek-r1" not in model_pricing_data): + model_pricing_data["deepseek-r1"] = ( + model_pricing_data[deepseek_r1_path] + ) + except Exception: # pylint: disable=broad-except + console.print( + "[yellow]Warning: Could not fetch model pricing data[/yellow]" + ) + + # Create a flat list of all models for numeric selection + # pylint: disable=invalid-name + ALL_MODELS = [] + for category, models in MODEL_CATEGORIES.items(): + for model in models: + # Get pricing info if available + pricing_info = model_pricing_data.get(model["name"], {}) + input_cost = pricing_info.get("input_cost_per_token", None) + output_cost = pricing_info.get("output_cost_per_token", None) + + # Convert to dollars per million tokens if values exist + input_cost_per_million = None + output_cost_per_million = None + + if input_cost is not None: + input_cost_per_million = input_cost * 1000000 + if output_cost is not None: + output_cost_per_million = output_cost * 1000000 + + ALL_MODELS.append({ + "name": model["name"], + "provider": ( + "Anthropic" if "claude" in model["name"] + else "DeepSeek" if "deepseek" in model["name"] + else "OpenAI" + ), + "category": category, + "description": model["description"], + "input_cost": input_cost_per_million, + "output_cost": output_cost_per_million + }) + + # Update cached models + self.cached_models = [model["name"] for model in ALL_MODELS] + self.cached_model_numbers = { + str(i): model["name"] + for i, model in enumerate(ALL_MODELS, 1) + } + + if not args: # pylint: disable=too-many-nested-blocks + # Display current model + model_info = os.getenv("CAI_MODEL", "Unknown") + console.print( + Panel( + f"Current model: [bold green]{model_info}[/bold green]", + border_style="green", + title="Active Model" + ) + ) + + # Show available models in a table + model_table = Table( + title="Available Models", + show_header=True, + header_style="bold yellow") + model_table.add_column("#", style="bold white", justify="right") + model_table.add_column("Model", style="cyan") + model_table.add_column("Provider", style="magenta") + model_table.add_column("Category", style="blue") + model_table.add_column( + "Input Cost ($/M)", + style="green", + justify="right") + model_table.add_column( + "Output Cost ($/M)", + style="red", + justify="right") + model_table.add_column("Description", style="white") + + # Add all predefined models with numbers + for i, model in enumerate(ALL_MODELS, 1): + # Format pricing info as dollars per million tokens + input_cost_str = ( + f"${model['input_cost']:.2f}" + if model['input_cost'] is not None else "Unknown" + ) + output_cost_str = ( + f"${model['output_cost']:.2f}" + if model['output_cost'] is not None else "Unknown" + ) + + model_table.add_row( + str(i), + model["name"], + model["provider"], + model["category"], + input_cost_str, + output_cost_str, + model["description"] + ) + + # Ollama models (if available) + # pylint: disable=too-many-nested-blocks + try: + # Get Ollama models with a short timeout to prevent hanging + api_base = get_ollama_api_base() + ollama_base = api_base.replace('/v1', '') + response = requests.get( + f"{ollama_base}/api/tags", + timeout=1 + ) + + if response.status_code == 200: + data = response.json() + ollama_models = [] + + if 'models' in data: + ollama_models = data['models'] + else: + # Fallback for older Ollama versions + ollama_models = data.get('items', []) + + # Add Ollama models to the table with continuing numbers + start_index = len(ALL_MODELS) + 1 + for i, model in enumerate(ollama_models, start_index): + model_name = model.get('name', '') + model_size = model.get('size', 0) + # Convert size to human-readable format + size_str = "" + if model_size: + size_mb = model_size / (1024 * 1024) + if model_size < 1024 * 1024 * 1024: + size_str = f"{size_mb:.1f} MB" + else: + size_gb = size_mb / 1024 + size_str = f"{size_gb:.1f} GB" + + # Ollama models are free to use locally + model_description = "Local model" + if size_str: + model_description += f" ({size_str})" + + model_table.add_row( + str(i), + model_name, + "Ollama", + "Local", + "Free", + "Free", + model_description + ) + + # Add to cached models for numeric selection + self.cached_models.append(model_name) + self.cached_model_numbers[str(i)] = model_name + except Exception: # pylint: disable=broad-except + # Add a note about Ollama if we couldn't fetch models + start_index = len(ALL_MODELS) + 1 + model_table.add_row( + str(start_index), + "llama3", + "Ollama", + "Local", + "Free", + "Free", + "Local Llama 3 model (if installed)") + model_table.add_row(str(start_index + 1), + "mistral", + "Ollama", + "Local", + "Free", + "Free", + "Local Mistral model (if installed)") + model_table.add_row(str(start_index + 2), + "...", + "Ollama", + "Local", + "Free", + "Free", + "Other local models (if installed)") + + console.print(model_table) + + # Usage instructions + console.print("\n[cyan]Usage:[/cyan]") + console.print( + " [bold]/model [/bold] - Select by name (e.g. " + "[bold]/model claude-3-7-sonnet-20250219[/bold])" + ) + console.print( + " [bold]/model [/bold] - Select by number (e.g. " + "[bold]/model 1[/bold] for first model in list)" + ) + console.print( + " [bold]/model-show[/bold] - Show all available " + "models from LiteLLM repository" + ) + return True + + model_arg = args[0] + + # Check if the argument is a number for model selection + if model_arg.isdigit(): + model_index = int(model_arg) - 1 # Convert to 0-based index + if 0 <= model_index < len(self.cached_models): + model_name = self.cached_models[model_index] + else: + # If the number is out of range, we use the number + # directly as the model name + model_name = model_arg + else: + model_name = model_arg + + # Set the model in environment variable + os.environ["CAI_MODEL"] = model_name + + # Display model change notification + change_message = ( + f"Model changed to: [bold green]{model_name}[/bold green]\n" + "[yellow]Note: This will take effect on the next agent " + "interaction[/yellow]" + ) + console.print( + Panel( + change_message, + border_style="green", + title="Model Changed" + ) + ) + return True + + +class ModelShowCommand(Command): + """Command for showing all available models from LiteLLM repository.""" + + def __init__(self): + """Initialize the model-show command.""" + super().__init__( + name="/model-show", + description="Show all available models from LiteLLM repository", + aliases=["/mod-show"] + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: # pylint: disable=too-many-locals,too-many-branches,too-many-statements,line-too-long # noqa: E501 + """Handle the model-show command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + # Check if we should only show supported models + show_only_supported = False + search_term = None + + if args: + if "supported" in args: + show_only_supported = True + # Remove 'supported' from args to handle search term + args = [arg for arg in args if arg != "supported"] + + if args: # If there are still args left, use as search term + search_term = args[0].lower() + + # Fetch model pricing data from LiteLLM GitHub repository + try: + with console.status( + "[bold blue]Fetching model data...[/bold blue]" + ): + response = requests.get(LITELLM_URL, timeout=5) + + if response.status_code != 200: + error_msg = ( + f"[red]Error fetching model data: " + f"HTTP {response.status_code}[/red]" + ) + console.print(error_msg) + return True + + model_data = response.json() + + # Create a table to display the models + title = "All Available Models" + if show_only_supported: + title = "Supported Models (with Function Calling)" + if search_term: + title += f" - Search: '{search_term}'" + + model_table = Table( + title=title, + show_header=True, + header_style="bold yellow" + ) + model_table.add_column("#", style="bold white", justify="right") + model_table.add_column("Model", style="cyan") + model_table.add_column("Provider", style="magenta") + model_table.add_column("Max Tokens", style="blue", justify="right") + model_table.add_column( + "Input Cost ($/M)", + style="green", + justify="right") + model_table.add_column( + "Output Cost ($/M)", + style="red", + justify="right") + model_table.add_column("Features", style="white") + + # Count models for summary + total_models = 0 + displayed_models = 0 + model_index = 1 + + # Process and display models + for model_name, model_info in sorted(model_data.items()): + total_models += 1 + + # Skip if showing only supported models and no function calling + supports_functions = model_info.get( + "supports_function_calling", + False + ) + if show_only_supported and not supports_functions: + continue + + # Skip if search term provided and not in model name + if search_term and search_term not in model_name.lower(): + continue + + displayed_models += 1 + + # Extract provider from litellm_provider if available + provider = model_info.get("litellm_provider", "Unknown") + if provider == "text-completion-openai": + provider = "OpenAI" + elif provider == "openai": + provider = "OpenAI" + elif "/" in model_name: + # Extract provider from model name + provider = model_name.split("/")[0].capitalize() + + # Get max tokens + max_tokens = model_info.get("max_tokens", "N/A") + + # Get pricing info + input_cost = model_info.get("input_cost_per_token", 0) + output_cost = model_info.get("output_cost_per_token", 0) + + # Convert to dollars per million tokens + input_cost_per_million = ( + input_cost * 1000000 if input_cost else 0 + ) + output_cost_per_million = ( + output_cost * 1000000 if output_cost else 0 + ) + + # Format pricing info + if input_cost_per_million: + input_cost_str = f"${input_cost_per_million:.4f}" + else: + input_cost_str = "Free" + + if output_cost_per_million: + output_cost_str = f"${output_cost_per_million:.4f}" + else: + output_cost_str = "Free" + + # Get features + features = [] + if model_info.get("supports_vision"): + features.append("Vision") + if model_info.get("supports_function_calling"): + features.append("Function calling") + if model_info.get("supports_parallel_function_calling"): + features.append("Parallel functions") + if (model_info.get("supports_audio_input") or + model_info.get("supports_audio_output")): + features.append("Audio") + if model_info.get("mode") == "embedding": + features.append("Embeddings") + if model_info.get("mode") == "image_generation": + features.append("Image generation") + + features_str = ( + ", ".join(features) if features else "Text generation" + ) + + # Add row to table + model_table.add_row( + str(model_index), + model_name, + provider, + str(max_tokens), + input_cost_str, + output_cost_str, + features_str + ) + + model_index += 1 + + # Now add Ollama models if available + try: + # Get Ollama models with a short timeout + api_base = get_ollama_api_base() + api_tags = f"{api_base.replace('/v1', '')}/api/tags" + ollama_response = requests.get(api_tags, timeout=1) + + if ollama_response.status_code == 200: + ollama_data = ollama_response.json() + ollama_models = [] + + if 'models' in ollama_data: + ollama_models = ollama_data['models'] + else: + # Fallback for older Ollama versions + ollama_models = ollama_data.get('items', []) + + # Add Ollama models to the table + for model in ollama_models: + model_name = model.get('name', '') + + # Skip if search term provided and not in model name + if (search_term and + search_term not in model_name.lower()): + continue + + total_models += 1 + displayed_models += 1 + + model_size = model.get('size', 0) + # Convert size to human-readable format + size_str = "" + if model_size: + size_mb = model_size / (1024 * 1024) + if model_size < 1024 * 1024 * 1024: + size_str = f"{size_mb:.1f} MB" + else: + size_gb = size_mb / 1024 + size_str = f"{size_gb:.1f} GB" + + # Add row to table + model_description = "Local model" + if size_str: + model_description += f" ({size_str})" + + model_table.add_row( + str(model_index), + model_name, + "Ollama", + "Varies", + "Free", + "Free", + model_description + ) + + model_index += 1 + except Exception: # pylint: disable=broad-except + # Silently fail if Ollama is not available + # This is acceptable as Ollama is optional and we don't want to + # disrupt the user experience if it's not running + console.print( + "[dim]Ollama models not available[/dim]", + style="dim") + + # Display the table + console.print(model_table) + + # Display summary + displayed_str = str(displayed_models) + total_str = str(total_models) + summary_text = ( + f"\n[cyan]Showing {displayed_str} of {total_str} models" + ) + if show_only_supported: + summary_text += " with function calling support" + if search_term: + summary_text += f" matching '{search_term}'" + summary_text += "[/cyan]" + console.print(summary_text) + + # Usage instructions + console.print("\n[cyan]Usage:[/cyan]") + console.print( + " [bold]/model-show[/bold] - Show all " + "available models") + console.print( + " [bold]/model-show supported[/bold] - Show only " + "models with function calling") + console.print( + " [bold]/model-show [/bold] - Filter " + "models by search term") + console.print( + " [bold]/model-show supported [/bold] - Filter " + "supported models by search term") + console.print( + " [bold]/model [/bold] - Select a " + "model to use") + console.print( + " [bold]/model [/bold] - Select a " + "model by its number") + + # Data source attribution + data_source = ( + "https://github.com/BerriAI/litellm/blob/main/" + "model_prices_and_context_window.json" + ) + console.print(f"\n[dim]Data source: {data_source}[/dim]") + + except Exception as e: # pylint: disable=broad-except + console.print(f"[red]Error fetching model data: {str(e)}[/red]") + + return True + + +# Register the commands +register_command(ModelCommand()) +register_command(ModelShowCommand()) diff --git a/src/cai/repl/commands/platform.py b/src/cai/repl/commands/platform.py new file mode 100644 index 00000000..22de5c58 --- /dev/null +++ b/src/cai/repl/commands/platform.py @@ -0,0 +1,223 @@ +""" +Platform command for CAI REPL. +This module provides commands for interacting with platform-specific features. +""" +from typing import ( + List, + Optional +) +from rich.console import Console # pylint: disable=import-error +from rich.panel import Panel # pylint: disable=import-error + +from cai import is_caiextensions_platform_available +from cai.repl.commands.base import Command, register_command + +console = Console() + + +class PlatformCommand(Command): + """Command for interacting with platform-specific features.""" + + def __init__(self): + """Initialize the platform command.""" + super().__init__( + name="/platform", + description="Interact with platform-specific features", + aliases=["/p"] + ) + + # Add subcommands dynamically based on available platforms + if is_caiextensions_platform_available(): + from caiextensions.platform.base import platform_manager # pylint: disable=import-error,import-outside-toplevel,unused-import,line-too-long,no-name-in-module # noqa: E501 + + # Add list subcommand + self.add_subcommand( + "list", + "List available platforms", + self.handle_list) + + # Add VPN status command + self.add_subcommand( + "vpn-status", + "Check the status of the VPN connection", + self.handle_vpn_status) + + # Add keep-vpn command + self.add_subcommand( + "keep-vpn", + "Keep VPN connection active even when interrupted", + self.handle_keep_vpn) + + # Add platform-specific subcommands + platforms = platform_manager.list_platforms() + for platform in platforms: + platform_cmds = platform_manager.get_platform( + platform).get_commands() + for cmd in platform_cmds: + # Add platform-specific commands as subcommands + self.add_subcommand( + f"{platform}:{cmd}", + f"Execute {cmd} command on {platform} platform", + lambda args, p=platform, c=cmd: + self.handle_platform_command([p, c] + (args or [])) + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the platform command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + if not is_caiextensions_platform_available(): + console.print("[red]Platform extensions are not available[/red]") + return False + + return self.handle_platform_command(args) + + def handle_list(self, args: Optional[List[str]] = None) -> bool: # pylint: disable=unused-argument # noqa: E501 + """Handle /platform list command.""" + if not is_caiextensions_platform_available(): + console.print("[red]Platform extensions are not available[/red]") + return False + + from caiextensions.platform.base import platform_manager # pylint: disable=import-error,import-outside-toplevel,unused-import,line-too-long,no-name-in-module # noqa: E501 + platforms = platform_manager.list_platforms() + + console.print(Panel( + "\n".join(f"[green]{p}[/green]" for p in platforms), + title="Available Platforms", + border_style="blue" + )) + return True + + def handle_platform_command( + self, args: Optional[List[str]] = None) -> bool: + """Handle platform specific commands.""" + if not is_caiextensions_platform_available(): + console.print("[red]Platform extensions are not available[/red]") + return False + + from caiextensions.platform.base import platform_manager # pylint: disable=import-error,import-outside-toplevel,unused-import,line-too-long,no-name-in-module # noqa: E501 + + if not args: + # Show available platforms + platforms = platform_manager.list_platforms() + console.print(Panel( + "\n".join(f"[green]{p}[/green]" for p in platforms), + title="Available Platforms", + border_style="blue" + )) + return True + + platform_name = args[0].lower() + platform = platform_manager.get_platform(platform_name) + + if not platform: + console.print(f"[red]Unknown platform: {platform_name}[/red]") + return False + + if len(args) == 1: + # Show platform help + console.print(Panel( + platform.get_help(), + title=f"{platform_name.upper()} Help", + border_style="blue" + )) + return True + + # Pass the command to the platform (without the platform name) + platform.handle_command(args[1:]) + return True + + def handle_vpn_status( + self, args: Optional[List[str]] = None) -> bool: # pylint: disable=unused-argument # noqa: E501 + """ + Check the status of the VPN connection. + + Args: + args: Optional list of command arguments (not used) + + Returns: + True if the command was handled successfully, False otherwise + """ + if not is_caiextensions_platform_available(): + console.print("[red]Platform extensions are not available[/red]") + return False + + try: + from caiextensions.platform.htb.cli import ( # pylint: disable=import-error,import-outside-toplevel,line-too-long # noqa: E501 + is_vpn_connected, get_vpn_ip, vpn_active + ) + # Check VPN connection status + if is_vpn_connected(): + status = "[green]Connected[/green]" + else: + status = "[red]Disconnected[/red]" + + # Check if VPN is set to persistent mode + if vpn_active: + persistent = "[green]Yes[/green]" + else: + persistent = "[red]No[/red]" + ip = get_vpn_ip() + + console.print(Panel( + f"Status: {status}\n" + f"Persistent: {persistent}\n" + f"IP Address: {ip}", + title="VPN Status", + border_style="blue" + )) + return True + except ImportError: + console.print("[red]HTB platform module not available[/red]") + return False + + def handle_keep_vpn( + self, args: Optional[List[str]] = None) -> bool: # pylint: disable=unused-argument # noqa: E501 + """ + Set the VPN to remain active even when the program is interrupted. + + Args: + args: Optional list of command arguments (not used) + + Returns: + True if the command was handled successfully, False otherwise + """ + if not is_caiextensions_platform_available(): + console.print("[red]Platform extensions are not available[/red]") + return False + + try: + from caiextensions.platform.htb.cli import ( # pylint: disable=import-error,import-outside-toplevel,line-too-long # noqa: E501 + is_vpn_connected + ) + if not is_vpn_connected(): + console.print("[red]No active VPN connection found[/red]") + console.print( + "[yellow]Connect to VPN first using " + "/platform htb:connect[/yellow]" + ) + return False + + # Set the VPN to persistent mode + import caiextensions.platform.htb.cli as htb_cli # pylint: disable=import-error,import-outside-toplevel,line-too-long # noqa: E501 + htb_cli.vpn_active = True + + console.print( + "[green]VPN connection set to persistent mode[/green]") + console.print( + "[yellow]VPN will remain active even if you press Ctrl+C" + "[/yellow]") + return True + except ImportError: + console.print("[red]HTB platform module not available[/red]") + return False + + +# Register the command +if is_caiextensions_platform_available(): + register_command(PlatformCommand()) diff --git a/src/cai/repl/commands/shell.py b/src/cai/repl/commands/shell.py new file mode 100644 index 00000000..d32c0ec9 --- /dev/null +++ b/src/cai/repl/commands/shell.py @@ -0,0 +1,139 @@ +""" +Shell command for CAI REPL. +This module provides commands for executing shell commands. +""" +import os +import signal +import subprocess # nosec B404 +from typing import ( + List, + Optional +) +from rich.console import Console # pylint: disable=import-error + +from cai.repl.commands.base import Command, register_command + +console = Console() + + +class ShellCommand(Command): + """Command for executing shell commands.""" + + def __init__(self): + """Initialize the shell command.""" + super().__init__( + name="/shell", + description="Execute shell commands in the current environment", + aliases=["/s", "$"] + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the shell command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + if not args: + console.print("[red]Error: No command specified[/red]") + return False + + return self.handle_shell_command(args) + + def handle_shell_command(self, command_args: List[str]) -> bool: + """Execute a shell command that can be interrupted with CTRL+C. + + Args: + command_args: The shell command and its arguments + + Returns: + bool: True if the command was executed successfully + """ + if not command_args: + console.print("[red]Error: No command specified[/red]") + return False + + shell_command = " ".join(command_args) + console.print(f"[blue]Executing:[/blue] {shell_command}") + + # Save original signal handler + original_sigint_handler = signal.getsignal(signal.SIGINT) + + try: + # Set temporary handler for SIGINT that only affects shell command + def shell_sigint_handler(sig, frame): # pylint: disable=unused-argument + # Just allow KeyboardInterrupt to propagate + signal.signal(signal.SIGINT, original_sigint_handler) + raise KeyboardInterrupt + + signal.signal(signal.SIGINT, shell_sigint_handler) + + # Check if this is a command that should run asynchronously + async_commands = [ + 'nc', + 'netcat', + 'ncat', + 'telnet', + 'ssh', + 'python -m http.server'] + is_async = any(cmd in shell_command for cmd in async_commands) + + if is_async: + # For async commands, use os.system to allow terminal + # interaction + console.print( + "[yellow]Running in async mode " + "(Ctrl+C to return to REPL)[/yellow]") + os.system(shell_command) # nosec B605 + console.print( + "[green]Async command completed or detached[/green]") + return True + + # For regular commands, use the standard approach + process = subprocess.Popen( # nosec B602 # pylint: disable=consider-using-with # noqa: E501 + shell_command, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + bufsize=1 + ) + + # Show output in real time + for line in iter(process.stdout.readline, ''): + print(line, end='') + + # Wait for process to finish + process.wait() + + if process.returncode == 0: + console.print( + "[green]Command completed successfully[/green]") + else: + console.print( + f"[yellow]Command exited with code { + process.returncode}" + f"[/yellow]") + return True + + except KeyboardInterrupt: + # Handle CTRL+C only for this command + try: + if not is_async: + process.terminate() + console.print("\n[yellow]Command interrupted by user[/yellow]") + except Exception: # pylint: disable=broad-except # nosec + pass + return True + except Exception as e: # pylint: disable=broad-except + console.print(f"[red]Error executing command: {str(e)}[/red]") + return False + finally: + # Restore original signal handler + signal.signal(signal.SIGINT, original_sigint_handler) + + +# Register the command +register_command(ShellCommand()) diff --git a/src/cai/repl/commands/turns.py b/src/cai/repl/commands/turns.py new file mode 100644 index 00000000..3c97f50f --- /dev/null +++ b/src/cai/repl/commands/turns.py @@ -0,0 +1,96 @@ +""" +Turns command for CAI REPL. +This module provides commands for viewing and changing the maximum number +of turns. +""" +import os +from typing import ( + List, + Optional +) +from rich.console import Console # pylint: disable=import-error +from rich.panel import Panel # pylint: disable=import-error + +from cai.repl.commands.base import Command, register_command + +console = Console() + + +class TurnsCommand(Command): + """Command for viewing and changing the maximum number of turns.""" + + def __init__(self): + """Initialize the turns command.""" + super().__init__( + name="/turns", + description="View or change the maximum number of turns", + aliases=["/t"] + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the turns command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + return self.handle_turns_command(args) + + def handle_turns_command(self, args: List[str]) -> bool: + """Change the maximum number of turns for CAI. + + Args: + args: List containing the number of turns + + Returns: + bool: True if the max turns was changed successfully + """ + if not args: + # Display current max turns + max_turns_info = os.getenv("CAI_MAX_TURNS", "inf") + console.print(Panel( + f"Current maximum turns: [bold green]{ + max_turns_info}[/bold green]", + border_style="green", + title="Max Turns Setting" + )) + + # Usage instructions + console.print( + "\n[cyan]Usage:[/cyan] [bold]/turns [/bold]") + console.print("[cyan]Examples:[/cyan]") + console.print(" [bold]/turns 10[/bold] - Limit to 10 turns") + console.print(" [bold]/turns inf[/bold] - Unlimited turns") + return True + + try: + turns = args[0] + # Check if it's a number or 'inf' + if turns.lower() == 'inf': + turns = 'inf' + else: + turns = int(turns) + + # Set the max turns in environment variable + os.environ["CAI_MAX_TURNS"] = turns + + console.print(Panel( + f"Maximum turns changed to: [bold green]{turns}[/bold green]\n" + "[yellow]Note: This will take effect on the next run[/yellow]", + border_style="green", + title="Max Turns Changed" + )) + return True + except ValueError: + console.print(Panel( + "Error: Max turns must be a number or 'inf'", + border_style="red", + title="Invalid Input" + )) + return False + + +# Register the command +register_command(TurnsCommand()) diff --git a/src/cai/repl/ui/__init__.py b/src/cai/repl/ui/__init__.py new file mode 100644 index 00000000..94ef22e1 --- /dev/null +++ b/src/cai/repl/ui/__init__.py @@ -0,0 +1,3 @@ +""" +UI components for the CAI REPL. +""" diff --git a/src/cai/repl/ui/banner.py b/src/cai/repl/ui/banner.py new file mode 100644 index 00000000..453b316c --- /dev/null +++ b/src/cai/repl/ui/banner.py @@ -0,0 +1,226 @@ +""" +Module for displaying the CAI banner and welcome message. +""" +# Standard library imports +import os +import glob +import logging +from configparser import ConfigParser + +# Third-party imports +import requests # pylint: disable=import-error +from rich.console import Console # pylint: disable=import-error +from rich.panel import Panel # pylint: disable=import-error +from rich.table import Table # pylint: disable=import-error + + +def get_version(): + """Get the CAI version from setup.cfg.""" + version = "unknown" + try: + config = ConfigParser() + config.read('setup.cfg') + version = config.get('metadata', 'version') + except Exception: # pylint: disable=broad-except + logging.warning("Could not read version from setup.cfg") + return version + + +def get_supported_models_count(): + """Get the count of supported models (with function calling).""" + try: + # Fetch model data from LiteLLM repository + response = requests.get( + "https://raw.githubusercontent.com/BerriAI/litellm/main/" + "model_prices_and_context_window.json", + timeout=2 + ) + + if response.status_code == 200: + model_data = response.json() + + # Count models with function calling support + function_calling_models = sum( + 1 for model_info in model_data.values() + if model_info.get("supports_function_calling", False) + ) + + # Try to get Ollama models count + try: + ollama_api_base = os.getenv( + "OLLAMA_API_BASE", + "http://host.docker.internal:8000/v1" + ) + ollama_response = requests.get( + f"{ollama_api_base.replace('/v1', '')}/api/tags", + timeout=1 + ) + + if ollama_response.status_code == 200: + ollama_data = ollama_response.json() + ollama_models = len( + ollama_data.get( + 'models', ollama_data.get('items', []) + ) + ) + return function_calling_models + ollama_models + except Exception: # pylint: disable=broad-except + logging.debug("Could not fetch Ollama models") + # Continue without Ollama models + + return function_calling_models + except Exception: # pylint: disable=broad-except + logging.warning("Could not fetch model data from LiteLLM") + + # Default count if we can't fetch the data + return "many" + + +def count_tools(): + """Count the number of tools in the CAI framework.""" + try: + # Count Python files in the tools directory + tool_files = glob.glob("cai/tools/**/*.py", recursive=True) + # Exclude __init__.py and other non-tool files + tool_files = [ + f for f in tool_files + if not f.endswith("__init__.py") and not f.endswith("__pycache__") + ] + return len(tool_files) + except Exception: # pylint: disable=broad-except + logging.warning("Could not count tools") + return "50+" + + +def count_agents(): + """Count the number of agents in the CAI framework.""" + try: + # Count Python files in the agents directory + agent_files = glob.glob("cai/agents/**/*.py", recursive=True) + # Exclude __init__.py and other non-agent files + agent_files = [ + f for f in agent_files + if not f.endswith("__init__.py") and not f.endswith("__pycache__") + ] + return len(agent_files) + except Exception: # pylint: disable=broad-except + logging.warning("Could not count agents") + return "20+" + + +def count_ctf_memories(): + """Count the number of CTF memories in the CAI framework.""" + # This is a placeholder - adjust the actual counting logic based on your + # framework structure + return "100+" + + +def display_banner(console: Console): + """ + Display a stylized CAI banner with Alias Robotics corporate colors. + + Args: + console: Rich console for output + """ + version = get_version() + + # Original banner with Alias Robotics colors (blue and white) + # Use noqa to ignore line length for the ASCII art + banner = f""" +[bold blue] CCCCCCCCCCCCC ++++++++ ++++++++ IIIIIIIIII +[bold blue] CCC::::::::::::C ++++++++++ ++++++++++ I::::::::I +[bold blue] CC:::::::::::::::C ++++++++++ ++++++++++ I::::::::I +[bold blue] C:::::CCCCCCCC::::C +++++++++ ++ +++++++++ II::::::II +[bold blue] C:::::C CCCCCC +++++++ +++++ +++++++ I::::I +[bold blue] C:::::C +++++ +++++++ +++++ I::::I +[bold blue] C:::::C ++++ ++++ I::::I +[bold blue] C:::::C ++ ++ I::::I +[bold blue] C:::::C + +++++++++++++++ + I::::I +[bold blue] C:::::C +++++++++++++++++++ I::::I +[bold blue] C:::::C +++++++++++++++++ I::::I +[bold blue] C:::::C CCCCCC +++++++++++++++ I::::I +[bold blue] C:::::CCCCCCCC::::C +++++++++++++ II::::::II +[bold blue] CC:::::::::::::::C +++++++++ I::::::::I +[bold blue] CCC::::::::::::C +++++ I::::::::I +[bold blue] CCCCCCCCCCCCC ++ IIIIIIIIII + +[bold blue] Cybersecurity AI (CAI), v{version}[/bold blue] +[white] Bug bounty-ready AI[/white] + """ + + console.print(banner) + + # # Create a table showcasing CAI framework capabilities + # # + # # reconsider in the future if necessary + # display_framework_capabilities(console) + + +def display_framework_capabilities(console: Console): + """ + Display a table showcasing CAI framework capabilities in Metasploit style. + + Args: + console: Rich console for output + """ + # Create the main table + table = Table( + title="", + box=None, + show_header=False, + show_edge=False, + padding=(0, 2) + ) + + table.add_column("Category", style="bold cyan") + table.add_column("Count", style="bold yellow") + table.add_column("Description", style="white") + + # Add rows for different capabilities + table.add_row( + "AI Models", + str(get_supported_models_count()), + "Supported AI models including GPT-4, Claude, Llama" + ) + + # table.add_row( + # "Tools", + # str(count_tools()), + # "Cybersecurity tools for reconnaissance and scanning" + # ) + + table.add_row( + "Agents", + str(count_agents()), + "Specialized AI agents for different cybersecurity tasks" + ) + + # Add the table to a panel for better visual separation + capabilities_panel = Panel( + table, + title="[bold blue]CAI Features[/bold blue]", + border_style="blue", + padding=(1, 2) + ) + + console.print(capabilities_panel) + + +def display_welcome_tips(console: Console): + """ + Display welcome message with tips for using the REPL. + + Args: + console: Rich console for output + """ + console.print(Panel( + "[white]• Use arrow keys ↑↓ to navigate command history[/white]\n" + "[white]• Press Tab for command completion[/white]\n" + "[white]• Type /help for available commands[/white]\n" + "[white]• Type /help aliases for command shortcuts[/white]\n" + "[white]• Press Ctrl+L to clear the screen[/white]\n" + "[white]• Press Esc+Enter to add a new line (multiline input)[/white]\n" + "[white]• Press Ctrl+C to exit[/white]", + title="Quick Tips", + border_style="blue" + )) diff --git a/src/cai/repl/ui/keybindings.py b/src/cai/repl/ui/keybindings.py new file mode 100644 index 00000000..2a397382 --- /dev/null +++ b/src/cai/repl/ui/keybindings.py @@ -0,0 +1,86 @@ +""" +Module for CAI REPL key bindings. +""" +import os +import subprocess # nosec B404 - Required for screen clearing +# pylint: disable=import-error +from prompt_toolkit.key_binding import KeyBindings +from prompt_toolkit.auto_suggest import AutoSuggestFromHistory +from cai.repl.commands import FuzzyCommandCompleter + + +def create_key_bindings(current_text): + """ + Create key bindings for the REPL. + + Args: + current_text: Reference to the current text for command shadowing + + Returns: + KeyBindings object with configured bindings + """ + kb = KeyBindings() + + @kb.add('c-l') + def _(event): # pylint: disable=unused-argument + """Clear the screen.""" + # Replace os.system with subprocess.run to avoid shell injection + if os.name == 'nt': + # Using fixed commands with shell=False is safe + subprocess.run( + ['cls'], + shell=False, + check=False) # nosec B603 B607 + else: + # Using fixed commands with shell=False is safe + subprocess.run( + ['clear'], + shell=False, + check=False) # nosec B603 B607 + + @kb.add('tab') + def handle_tab(event): + """Handle tab key to show completions menu or complete command.""" + buffer = event.current_buffer + text = buffer.text + + # Update current text for shadow + current_text[0] = text + + # First check if we have a history suggestion + history_suggestion = None + if text: + # Get suggestion from history + auto_suggest = AutoSuggestFromHistory() + suggestion = auto_suggest.get_suggestion(buffer, buffer.document) + if suggestion and suggestion.text: + history_suggestion = text + suggestion.text + + # If we have a history suggestion, use it + if history_suggestion: + buffer.text = history_suggestion + buffer.cursor_position = len(history_suggestion) + else: + # If no history suggestion, check for command shadow from fuzzy + # completer + shadow = FuzzyCommandCompleter().get_command_shadow(text) + if shadow and shadow.startswith(text): + # Complete with the shadow + buffer.text = shadow + buffer.cursor_position = len(shadow) + # If no shadow or shadow is the same as current text + elif buffer.complete_state: + # If completion menu is already showing, select the next item + buffer.complete_next() + else: + # Otherwise, start completion + buffer.start_completion(select_first=True) + + @kb.add('escape', 'enter') + def handle_escape_enter(event): + """ + Alternative way to insert a newline using Escape followed by Enter. + """ + event.current_buffer.insert_text('\n') + + return kb diff --git a/src/cai/repl/ui/logging.py b/src/cai/repl/ui/logging.py new file mode 100644 index 00000000..f1e1a870 --- /dev/null +++ b/src/cai/repl/ui/logging.py @@ -0,0 +1,33 @@ +""" +Module for CAI REPL session logging. +""" +from pathlib import Path + + +def setup_session_logging(): + """ + Set up session logging. + + Returns: + Tuple of (history_file, session_log, log_interaction function) + """ + # Setup history file + history_dir = Path.cwd() / ".cai" + history_dir.mkdir(exist_ok=True) + history_file = history_dir / "history.txt" + + # # Setup session log file + # session_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + # session_log = history_dir / f"session_{session_id}.log" + + # # Function to log interactions + # def log_interaction(role, content): + # with open(session_log, "a", encoding="utf-8") as f: + # f.write( + # f"\n[{ + # datetime.datetime.now().strftime('%H:%M:%S')}] { + # role.upper()}:\n") + # f.write(f"{content}\n") + + # return history_file, session_log, log_interaction + return history_file diff --git a/src/cai/repl/ui/prompt.py b/src/cai/repl/ui/prompt.py new file mode 100644 index 00000000..5a220e1f --- /dev/null +++ b/src/cai/repl/ui/prompt.py @@ -0,0 +1,115 @@ +""" +Module for CAI REPL prompt functionality. +""" +import time +from functools import lru_cache +from prompt_toolkit import prompt # pylint: disable=import-error +from prompt_toolkit.history import FileHistory # pylint: disable=import-error +from prompt_toolkit.auto_suggest import AutoSuggestFromHistory # pylint: disable=import-error # noqa: E501 +from prompt_toolkit.styles import Style # pylint: disable=import-error +from prompt_toolkit.formatted_text import HTML # pylint: disable=import-error +from cai.repl.commands import FuzzyCommandCompleter + + +# Cache for command shadow to avoid recalculating it too frequently +shadow_cache = { + 'text': '', + 'result': '', + 'last_update': 0, + 'update_interval': 0.1 # Update at most every 100ms +} + + +@lru_cache(maxsize=32) +def get_command_shadow_cached(text): + """Get command shadow suggestion with caching for repeated calls.""" + return FuzzyCommandCompleter().get_command_shadow(text) + + +def get_command_shadow(text): + """Get command shadow suggestion with throttling.""" + current_time = time.time() + + # If the text hasn't changed, return the cached result + if text == shadow_cache['text']: + return shadow_cache['result'] + + # If we've updated recently, return the cached result + if (current_time - shadow_cache['last_update'] < shadow_cache['update_interval'] + and shadow_cache['result']): + return shadow_cache['result'] + + # Update the cache + shadow = get_command_shadow_cached(text) + if shadow and shadow.startswith(text): + result = shadow[len(text):] + else: + result = "" + + # Store in cache + shadow_cache['text'] = text + shadow_cache['result'] = result + shadow_cache['last_update'] = current_time + + return result + + +def create_prompt_style(): + """Create a style for the CLI.""" + return Style.from_dict({ + 'prompt': 'bold cyan', + 'completion-menu': 'bg:#2b2b2b #ffffff', + 'completion-menu.completion': 'bg:#2b2b2b #ffffff', + 'completion-menu.completion.current': 'bg:#004b6b #ffffff', + 'scrollbar.background': 'bg:#2b2b2b', + 'scrollbar.button': 'bg:#004b6b', + }) + + +def get_user_input( + command_completer, + key_bindings, + history_file, + toolbar_func, + current_text +): + """ + Get user input with all prompt features. + + Args: + command_completer: Command completer instance + key_bindings: Key bindings instance + history_file: Path to history file + toolbar_func: Function to get toolbar content + current_text: Reference to current text for command shadowing + + Returns: + User input string + """ + # Function to update current text and get command shadow + def get_rprompt(): + """Get the right prompt with command shadow.""" + shadow = get_command_shadow(current_text[0]) + if not shadow: + return None + return HTML(f'{shadow}') + + # Get user input with all features + return prompt( + [('class:prompt', 'CAI> ')], + completer=command_completer, + style=create_prompt_style(), + history=FileHistory(str(history_file)), + auto_suggest=AutoSuggestFromHistory(), + key_bindings=key_bindings, + bottom_toolbar=toolbar_func, + complete_in_thread=True, + complete_while_typing=True, # Enable real-time completion + enable_system_prompt=True, # Enable shadow prediction + mouse_support=False, # Enable mouse support for menu navigation + enable_suspend=True, # Allow suspending with Ctrl+Z + enable_open_in_editor=True, # Allow editing with Ctrl+X Ctrl+E + multiline=False, # Enable multiline input + rprompt=get_rprompt, # Missing comma here + color_depth=None, # Auto-detect color support + ) diff --git a/src/cai/repl/ui/toolbar.py b/src/cai/repl/ui/toolbar.py new file mode 100644 index 00000000..76e9d4e9 --- /dev/null +++ b/src/cai/repl/ui/toolbar.py @@ -0,0 +1,153 @@ +""" +Module for the CAI REPL toolbar functionality. +""" +import datetime +import os +import socket +import platform +import threading +import time +from functools import lru_cache +import requests # pylint: disable=import-error +from prompt_toolkit.formatted_text import HTML # pylint: disable=import-error + +# Variable to track when to refresh the toolbar +toolbar_last_refresh = [datetime.datetime.now()] + +# Cache for toolbar data +toolbar_cache = { + 'html': "", + 'last_update': datetime.datetime.now(), + 'refresh_interval': 60 # Refresh every 60 seconds +} + +# Cache for system information that rarely changes +system_info = { + 'ip_address': None, + 'os_name': None, + 'os_version': None +} + + +@lru_cache(maxsize=1) +def get_system_info(): + """Get system information that rarely changes (cached).""" + if not system_info['ip_address']: + try: + # Get local IP addresses + hostname = socket.gethostname() + system_info['ip_address'] = socket.gethostbyname(hostname) + + # Get OS information + system_info['os_name'] = platform.system() + system_info['os_version'] = platform.release() + except Exception: # pylint: disable=broad-except + system_info['ip_address'] = "unknown" + system_info['os_name'] = "unknown" + system_info['os_version'] = "unknown" + + return system_info + + +def update_toolbar_in_background(): + """Update the toolbar cache in a background thread.""" + try: + # Get system info (cached) + sys_info = get_system_info() + ip_address = sys_info['ip_address'] + os_name = sys_info['os_name'] + os_version = sys_info['os_version'] + + # Get Ollama information + ollama_status = "unavailable" + try: + # Get Ollama models with a short timeout to prevent hanging + api_base = os.getenv( + "OLLAMA_API_BASE", + "http://host.docker.internal:8000/v1") + response = requests.get( + f"{api_base.replace('/v1', '')}/api/tags", timeout=0.5) + + if response.status_code == 200: + data = response.json() + if 'models' in data: + ollama_models = len(data['models']) + else: + # Fallback for older Ollama versions + ollama_models = len(data.get('items', [])) + ollama_status = f"{ollama_models} models" + except Exception: # pylint: disable=broad-except + # Silently fail if Ollama is not available + ollama_status = "unavailable" + + # Get current time for the toolbar refresh indicator + current_time = datetime.datetime.now().strftime("%H:%M") + + # Add timezone information to show it's local time + timezone_name = datetime.datetime.now().astimezone().tzname() + current_time_with_tz = f"{current_time} {timezone_name}" + + # Update the cache + toolbar_cache['html'] = HTML( + f"IP: { + ip_address} | " + f"OS: { + os_name} {os_version} | " + f"Ollama: { + ollama_status} | " + f"Model: { + os.getenv('CAI_MODEL', 'default')} | " + f"Max Turns: { + os.getenv('CAI_MAX_TURNS', 'inf')} | " + f"{current_time_with_tz}" + ) + toolbar_cache['last_update'] = datetime.datetime.now() + except Exception: # pylint: disable=broad-except + # If there's an error, set a simple toolbar + toolbar_cache['html'] = HTML( + f"{datetime.datetime.now().strftime('%H:%M')}" + ) + + +def get_bottom_toolbar(): + """Get the bottom toolbar with system information (cached).""" + # If the toolbar is empty, initialize it + if not toolbar_cache['html']: + # Create a simple initial toolbar while the full one loads + current_time = datetime.datetime.now().strftime("%H:%M") + timezone_name = datetime.datetime.now().astimezone().tzname() + toolbar_cache['html'] = HTML( + f"Loading system information... {current_time} {timezone_name}" + ) + # Start background update + threading.Thread( + target=update_toolbar_in_background, + daemon=True + ).start() + + # Return the cached toolbar HTML + return toolbar_cache['html'] + + +def get_toolbar_with_refresh(): + """Get toolbar with refresh control (once per minute).""" + now = datetime.datetime.now() + seconds_elapsed = (now - toolbar_cache['last_update']).total_seconds() + + # Check if we need to refresh the toolbar + if seconds_elapsed >= toolbar_cache['refresh_interval']: + # Start a background thread to update the toolbar + threading.Thread( + target=update_toolbar_in_background, + daemon=True + ).start() + + # Always return the cached version immediately + return get_bottom_toolbar() + + +# Initialize the toolbar on module import +threading.Thread( + target=update_toolbar_in_background, + daemon=True +).start() diff --git a/src/cai/agents/__init__.py b/src/cai/sdk/agents/__init__.py similarity index 100% rename from src/cai/agents/__init__.py rename to src/cai/sdk/agents/__init__.py diff --git a/src/cai/agents/_config.py b/src/cai/sdk/agents/_config.py similarity index 100% rename from src/cai/agents/_config.py rename to src/cai/sdk/agents/_config.py diff --git a/src/cai/agents/_debug.py b/src/cai/sdk/agents/_debug.py similarity index 100% rename from src/cai/agents/_debug.py rename to src/cai/sdk/agents/_debug.py diff --git a/src/cai/agents/_run_impl.py b/src/cai/sdk/agents/_run_impl.py similarity index 100% rename from src/cai/agents/_run_impl.py rename to src/cai/sdk/agents/_run_impl.py diff --git a/src/cai/agents/agent.py b/src/cai/sdk/agents/agent.py similarity index 100% rename from src/cai/agents/agent.py rename to src/cai/sdk/agents/agent.py diff --git a/src/cai/agents/agent_output.py b/src/cai/sdk/agents/agent_output.py similarity index 100% rename from src/cai/agents/agent_output.py rename to src/cai/sdk/agents/agent_output.py diff --git a/src/cai/agents/computer.py b/src/cai/sdk/agents/computer.py similarity index 100% rename from src/cai/agents/computer.py rename to src/cai/sdk/agents/computer.py diff --git a/src/cai/agents/exceptions.py b/src/cai/sdk/agents/exceptions.py similarity index 100% rename from src/cai/agents/exceptions.py rename to src/cai/sdk/agents/exceptions.py diff --git a/src/cai/agents/models/__init__.py b/src/cai/sdk/agents/extensions/__init__.py similarity index 100% rename from src/cai/agents/models/__init__.py rename to src/cai/sdk/agents/extensions/__init__.py diff --git a/src/cai/agents/extensions/handoff_filters.py b/src/cai/sdk/agents/extensions/handoff_filters.py similarity index 100% rename from src/cai/agents/extensions/handoff_filters.py rename to src/cai/sdk/agents/extensions/handoff_filters.py diff --git a/src/cai/agents/extensions/handoff_prompt.py b/src/cai/sdk/agents/extensions/handoff_prompt.py similarity index 100% rename from src/cai/agents/extensions/handoff_prompt.py rename to src/cai/sdk/agents/extensions/handoff_prompt.py diff --git a/src/cai/agents/function_schema.py b/src/cai/sdk/agents/function_schema.py similarity index 100% rename from src/cai/agents/function_schema.py rename to src/cai/sdk/agents/function_schema.py diff --git a/src/cai/agents/guardrail.py b/src/cai/sdk/agents/guardrail.py similarity index 100% rename from src/cai/agents/guardrail.py rename to src/cai/sdk/agents/guardrail.py diff --git a/src/cai/agents/handoffs.py b/src/cai/sdk/agents/handoffs.py similarity index 100% rename from src/cai/agents/handoffs.py rename to src/cai/sdk/agents/handoffs.py diff --git a/src/cai/agents/items.py b/src/cai/sdk/agents/items.py similarity index 100% rename from src/cai/agents/items.py rename to src/cai/sdk/agents/items.py diff --git a/src/cai/agents/lifecycle.py b/src/cai/sdk/agents/lifecycle.py similarity index 100% rename from src/cai/agents/lifecycle.py rename to src/cai/sdk/agents/lifecycle.py diff --git a/src/cai/agents/logger.py b/src/cai/sdk/agents/logger.py similarity index 100% rename from src/cai/agents/logger.py rename to src/cai/sdk/agents/logger.py diff --git a/src/cai/agents/model_settings.py b/src/cai/sdk/agents/model_settings.py similarity index 100% rename from src/cai/agents/model_settings.py rename to src/cai/sdk/agents/model_settings.py diff --git a/src/cai/agents/util/__init__.py b/src/cai/sdk/agents/models/__init__.py similarity index 100% rename from src/cai/agents/util/__init__.py rename to src/cai/sdk/agents/models/__init__.py diff --git a/src/cai/agents/models/_openai_shared.py b/src/cai/sdk/agents/models/_openai_shared.py similarity index 100% rename from src/cai/agents/models/_openai_shared.py rename to src/cai/sdk/agents/models/_openai_shared.py diff --git a/src/cai/agents/models/fake_id.py b/src/cai/sdk/agents/models/fake_id.py similarity index 100% rename from src/cai/agents/models/fake_id.py rename to src/cai/sdk/agents/models/fake_id.py diff --git a/src/cai/agents/models/interface.py b/src/cai/sdk/agents/models/interface.py similarity index 100% rename from src/cai/agents/models/interface.py rename to src/cai/sdk/agents/models/interface.py diff --git a/src/cai/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py similarity index 100% rename from src/cai/agents/models/openai_chatcompletions.py rename to src/cai/sdk/agents/models/openai_chatcompletions.py diff --git a/src/cai/agents/models/openai_provider.py b/src/cai/sdk/agents/models/openai_provider.py similarity index 100% rename from src/cai/agents/models/openai_provider.py rename to src/cai/sdk/agents/models/openai_provider.py diff --git a/src/cai/agents/models/openai_responses.py b/src/cai/sdk/agents/models/openai_responses.py similarity index 100% rename from src/cai/agents/models/openai_responses.py rename to src/cai/sdk/agents/models/openai_responses.py diff --git a/src/cai/agents/py.typed b/src/cai/sdk/agents/py.typed similarity index 100% rename from src/cai/agents/py.typed rename to src/cai/sdk/agents/py.typed diff --git a/src/cai/agents/result.py b/src/cai/sdk/agents/result.py similarity index 100% rename from src/cai/agents/result.py rename to src/cai/sdk/agents/result.py diff --git a/src/cai/agents/run.py b/src/cai/sdk/agents/run.py similarity index 100% rename from src/cai/agents/run.py rename to src/cai/sdk/agents/run.py diff --git a/src/cai/agents/run_context.py b/src/cai/sdk/agents/run_context.py similarity index 100% rename from src/cai/agents/run_context.py rename to src/cai/sdk/agents/run_context.py diff --git a/src/cai/agents/stream_events.py b/src/cai/sdk/agents/stream_events.py similarity index 100% rename from src/cai/agents/stream_events.py rename to src/cai/sdk/agents/stream_events.py diff --git a/src/cai/agents/strict_schema.py b/src/cai/sdk/agents/strict_schema.py similarity index 100% rename from src/cai/agents/strict_schema.py rename to src/cai/sdk/agents/strict_schema.py diff --git a/src/cai/agents/tool.py b/src/cai/sdk/agents/tool.py similarity index 100% rename from src/cai/agents/tool.py rename to src/cai/sdk/agents/tool.py diff --git a/src/cai/agents/tracing/__init__.py b/src/cai/sdk/agents/tracing/__init__.py similarity index 100% rename from src/cai/agents/tracing/__init__.py rename to src/cai/sdk/agents/tracing/__init__.py diff --git a/src/cai/agents/tracing/create.py b/src/cai/sdk/agents/tracing/create.py similarity index 100% rename from src/cai/agents/tracing/create.py rename to src/cai/sdk/agents/tracing/create.py diff --git a/src/cai/agents/tracing/logger.py b/src/cai/sdk/agents/tracing/logger.py similarity index 100% rename from src/cai/agents/tracing/logger.py rename to src/cai/sdk/agents/tracing/logger.py diff --git a/src/cai/agents/tracing/processor_interface.py b/src/cai/sdk/agents/tracing/processor_interface.py similarity index 100% rename from src/cai/agents/tracing/processor_interface.py rename to src/cai/sdk/agents/tracing/processor_interface.py diff --git a/src/cai/agents/tracing/processors.py b/src/cai/sdk/agents/tracing/processors.py similarity index 100% rename from src/cai/agents/tracing/processors.py rename to src/cai/sdk/agents/tracing/processors.py diff --git a/src/cai/agents/tracing/scope.py b/src/cai/sdk/agents/tracing/scope.py similarity index 100% rename from src/cai/agents/tracing/scope.py rename to src/cai/sdk/agents/tracing/scope.py diff --git a/src/cai/agents/tracing/setup.py b/src/cai/sdk/agents/tracing/setup.py similarity index 100% rename from src/cai/agents/tracing/setup.py rename to src/cai/sdk/agents/tracing/setup.py diff --git a/src/cai/agents/tracing/span_data.py b/src/cai/sdk/agents/tracing/span_data.py similarity index 100% rename from src/cai/agents/tracing/span_data.py rename to src/cai/sdk/agents/tracing/span_data.py diff --git a/src/cai/agents/tracing/spans.py b/src/cai/sdk/agents/tracing/spans.py similarity index 100% rename from src/cai/agents/tracing/spans.py rename to src/cai/sdk/agents/tracing/spans.py diff --git a/src/cai/agents/tracing/traces.py b/src/cai/sdk/agents/tracing/traces.py similarity index 100% rename from src/cai/agents/tracing/traces.py rename to src/cai/sdk/agents/tracing/traces.py diff --git a/src/cai/agents/tracing/util.py b/src/cai/sdk/agents/tracing/util.py similarity index 100% rename from src/cai/agents/tracing/util.py rename to src/cai/sdk/agents/tracing/util.py diff --git a/src/cai/agents/usage.py b/src/cai/sdk/agents/usage.py similarity index 100% rename from src/cai/agents/usage.py rename to src/cai/sdk/agents/usage.py diff --git a/src/cai/agents/voice/models/__init__.py b/src/cai/sdk/agents/util/__init__.py similarity index 100% rename from src/cai/agents/voice/models/__init__.py rename to src/cai/sdk/agents/util/__init__.py diff --git a/src/cai/agents/util/_coro.py b/src/cai/sdk/agents/util/_coro.py similarity index 100% rename from src/cai/agents/util/_coro.py rename to src/cai/sdk/agents/util/_coro.py diff --git a/src/cai/agents/util/_error_tracing.py b/src/cai/sdk/agents/util/_error_tracing.py similarity index 100% rename from src/cai/agents/util/_error_tracing.py rename to src/cai/sdk/agents/util/_error_tracing.py diff --git a/src/cai/agents/util/_json.py b/src/cai/sdk/agents/util/_json.py similarity index 100% rename from src/cai/agents/util/_json.py rename to src/cai/sdk/agents/util/_json.py diff --git a/src/cai/agents/util/_pretty_print.py b/src/cai/sdk/agents/util/_pretty_print.py similarity index 100% rename from src/cai/agents/util/_pretty_print.py rename to src/cai/sdk/agents/util/_pretty_print.py diff --git a/src/cai/agents/util/_transforms.py b/src/cai/sdk/agents/util/_transforms.py similarity index 100% rename from src/cai/agents/util/_transforms.py rename to src/cai/sdk/agents/util/_transforms.py diff --git a/src/cai/agents/util/_types.py b/src/cai/sdk/agents/util/_types.py similarity index 100% rename from src/cai/agents/util/_types.py rename to src/cai/sdk/agents/util/_types.py diff --git a/src/cai/agents/version.py b/src/cai/sdk/agents/version.py similarity index 100% rename from src/cai/agents/version.py rename to src/cai/sdk/agents/version.py diff --git a/src/cai/agents/voice/__init__.py b/src/cai/sdk/agents/voice/__init__.py similarity index 100% rename from src/cai/agents/voice/__init__.py rename to src/cai/sdk/agents/voice/__init__.py diff --git a/src/cai/agents/voice/events.py b/src/cai/sdk/agents/voice/events.py similarity index 100% rename from src/cai/agents/voice/events.py rename to src/cai/sdk/agents/voice/events.py diff --git a/src/cai/agents/voice/exceptions.py b/src/cai/sdk/agents/voice/exceptions.py similarity index 100% rename from src/cai/agents/voice/exceptions.py rename to src/cai/sdk/agents/voice/exceptions.py diff --git a/src/cai/agents/voice/imports.py b/src/cai/sdk/agents/voice/imports.py similarity index 100% rename from src/cai/agents/voice/imports.py rename to src/cai/sdk/agents/voice/imports.py diff --git a/src/cai/agents/voice/input.py b/src/cai/sdk/agents/voice/input.py similarity index 100% rename from src/cai/agents/voice/input.py rename to src/cai/sdk/agents/voice/input.py diff --git a/src/cai/agents/voice/model.py b/src/cai/sdk/agents/voice/model.py similarity index 100% rename from src/cai/agents/voice/model.py rename to src/cai/sdk/agents/voice/model.py diff --git a/src/cai/sdk/agents/voice/models/__init__.py b/src/cai/sdk/agents/voice/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/cai/agents/voice/models/openai_model_provider.py b/src/cai/sdk/agents/voice/models/openai_model_provider.py similarity index 100% rename from src/cai/agents/voice/models/openai_model_provider.py rename to src/cai/sdk/agents/voice/models/openai_model_provider.py diff --git a/src/cai/agents/voice/models/openai_stt.py b/src/cai/sdk/agents/voice/models/openai_stt.py similarity index 100% rename from src/cai/agents/voice/models/openai_stt.py rename to src/cai/sdk/agents/voice/models/openai_stt.py diff --git a/src/cai/agents/voice/models/openai_tts.py b/src/cai/sdk/agents/voice/models/openai_tts.py similarity index 100% rename from src/cai/agents/voice/models/openai_tts.py rename to src/cai/sdk/agents/voice/models/openai_tts.py diff --git a/src/cai/agents/voice/pipeline.py b/src/cai/sdk/agents/voice/pipeline.py similarity index 100% rename from src/cai/agents/voice/pipeline.py rename to src/cai/sdk/agents/voice/pipeline.py diff --git a/src/cai/agents/voice/pipeline_config.py b/src/cai/sdk/agents/voice/pipeline_config.py similarity index 100% rename from src/cai/agents/voice/pipeline_config.py rename to src/cai/sdk/agents/voice/pipeline_config.py diff --git a/src/cai/agents/voice/result.py b/src/cai/sdk/agents/voice/result.py similarity index 100% rename from src/cai/agents/voice/result.py rename to src/cai/sdk/agents/voice/result.py diff --git a/src/cai/agents/voice/utils.py b/src/cai/sdk/agents/voice/utils.py similarity index 100% rename from src/cai/agents/voice/utils.py rename to src/cai/sdk/agents/voice/utils.py diff --git a/src/cai/agents/voice/workflow.py b/src/cai/sdk/agents/voice/workflow.py similarity index 100% rename from src/cai/agents/voice/workflow.py rename to src/cai/sdk/agents/voice/workflow.py diff --git a/src/cai/util.py b/src/cai/util.py new file mode 100644 index 00000000..06956e4a --- /dev/null +++ b/src/cai/util.py @@ -0,0 +1,9 @@ +""" +Util model for CAI +""" + +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") diff --git a/uv.lock b/uv.lock index bfb1e76f..063fd156 100644 --- a/uv.lock +++ b/uv.lock @@ -203,13 +203,18 @@ name = "cai-framework" version = "0.4.0" source = { editable = "." } dependencies = [ + { name = "dotenv" }, { name = "griffe" }, + { name = "litellm" }, { name = "openai" }, { name = "openinference-instrumentation-openai" }, + { name = "prompt-toolkit" }, { name = "pydantic" }, { name = "requests" }, + { name = "rich" }, { name = "types-requests" }, { name = "typing-extensions" }, + { name = "wasabi" }, ] [package.optional-dependencies] @@ -243,14 +248,19 @@ dev = [ [package.metadata] requires-dist = [ + { name = "dotenv", specifier = ">=0.9.9" }, { name = "griffe", specifier = ">=1.5.6,<2" }, + { name = "litellm", specifier = ">=1.63.7" }, { name = "numpy", marker = "python_full_version >= '3.10' and extra == 'voice'", specifier = ">=2.2.0,<3" }, { name = "openai", specifier = ">=1.68.2" }, { name = "openinference-instrumentation-openai", specifier = ">=0.1.22" }, + { name = "prompt-toolkit", specifier = ">=3.0.39" }, { name = "pydantic", specifier = ">=2.10,<3" }, { name = "requests", specifier = ">=2.0,<3" }, + { name = "rich", specifier = ">=13.9.4" }, { name = "types-requests", specifier = ">=2.0,<3" }, { name = "typing-extensions", specifier = ">=4.12.2,<5" }, + { name = "wasabi", specifier = ">=1.1.3" }, { name = "websockets", marker = "extra == 'voice'", specifier = ">=15.0,<16" }, ] provides-extras = ["voice"] @@ -542,6 +552,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277 }, ] +[[package]] +name = "dotenv" +version = "0.9.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dotenv" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/b7/545d2c10c1fc15e48653c91efde329a790f2eecfbbf2bd16003b5db2bab0/dotenv-0.9.9-py2.py3-none-any.whl", hash = "sha256:29cf74a087b31dafdb5a446b6d7e11cbce8ed2741540e2339c69fbef92c94ce9", size = 1892 }, +] + [[package]] name = "evdev" version = "1.9.1" @@ -1675,6 +1696,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/6d/84533aa3fcc395235d58c3412fb86013653b697d91fc53f379c83bbb0b79/prisma-0.15.0-py3-none-any.whl", hash = "sha256:de949cc94d3d91243615f22ff64490aa6e2d7cb81aabffce53d92bd3977c09a4", size = 173809 }, ] +[[package]] +name = "prompt-toolkit" +version = "3.0.50" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/e1/bd15cb8ffdcfeeb2bdc215de3c3cffca11408d829e4b8416dcfe71ba8854/prompt_toolkit-3.0.50.tar.gz", hash = "sha256:544748f3860a2623ca5cd6d2795e7a14f3d0e1c3c9728359013f79877fc89bab", size = 429087 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e4/ea/d836f008d33151c7a1f62caf3d8dd782e4d15f6a43897f64480c2b8de2ad/prompt_toolkit-3.0.50-py3-none-any.whl", hash = "sha256:9b6427eb19e479d98acff65196a307c555eb567989e6d88ebbb1b509d9779198", size = 387816 }, +] + [[package]] name = "propcache" version = "0.3.1" @@ -2672,6 +2705,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c8/19/4ec628951a74043532ca2cf5d97b7b14863931476d117c471e8e2b1eb39f/urllib3-2.3.0-py3-none-any.whl", hash = "sha256:1cee9ad369867bfdbbb48b7dd50374c0967a0bb7710050facf0dd6911440e3df", size = 128369 }, ] +[[package]] +name = "wasabi" +version = "1.1.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/f9/054e6e2f1071e963b5e746b48d1e3727470b2a490834d18ad92364929db3/wasabi-1.1.3.tar.gz", hash = "sha256:4bb3008f003809db0c3e28b4daf20906ea871a2bb43f9914197d540f4f2e0878", size = 30391 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/06/7c/34330a89da55610daa5f245ddce5aab81244321101614751e7537f125133/wasabi-1.1.3-py3-none-any.whl", hash = "sha256:f76e16e8f7e79f8c4c8be49b4024ac725713ab10cd7f19350ad18a8e3f71728c", size = 27880 }, +] + [[package]] name = "watchdog" version = "6.0.0" @@ -2709,6 +2754,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/33/e8/e40370e6d74ddba47f002a32919d91310d6074130fe4e17dabcafc15cbf1/watchdog-6.0.0-py3-none-win_ia64.whl", hash = "sha256:a1914259fa9e1454315171103c6a30961236f508b9b623eae470268bbcc6a22f", size = 79067 }, ] +[[package]] +name = "wcwidth" +version = "0.2.13" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6c/63/53559446a878410fc5a5974feb13d31d78d752eb18aeba59c7fef1af7598/wcwidth-0.2.13.tar.gz", hash = "sha256:72ea0c06399eb286d978fdedb6923a9eb47e1c486ce63e9b4e64fc18303972b5", size = 101301 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/84/fd2ba7aafacbad3c4201d395674fc6348826569da3c0937e75505ead3528/wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859", size = 34166 }, +] + [[package]] name = "websockets" version = "15.0.1"