Initial commit: svibe agentic CLI tool
Features: - Multi-provider support (OpenAI, Ollama) with bearer/basic auth - Hierarchical agents with supervisor pattern - Model routing per task type (chat, code, planning, tool_use) - Built-in tools: file ops, shell commands, git operations - MCP protocol support for external tools - CLI with REPL and single-command modes - TOML configuration with environment variable overrides Architecture: - config/ - Configuration system - providers/ - AI model provider abstraction - tools/ - Built-in development tools - mcp/ - Model Context Protocol implementation - agents/ - Agent system with delegation - cli/ - Command-line interface Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
This commit is contained in:
1
.gitignore
vendored
Normal file
1
.gitignore
vendored
Normal file
@@ -0,0 +1 @@
|
||||
/target
|
||||
2161
Cargo.lock
generated
Normal file
2161
Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load Diff
44
Cargo.toml
Normal file
44
Cargo.toml
Normal file
@@ -0,0 +1,44 @@
|
||||
[package]
|
||||
name = "svibe"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
description = "Agentic CLI tool with hierarchical agents, multi-provider support, and MCP protocol"
|
||||
authors = ["svibe developers"]
|
||||
|
||||
[dependencies]
|
||||
# Async runtime
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
|
||||
# HTTP client with rustls for TLS
|
||||
reqwest = { version = "0.12", features = ["json", "rustls-tls"], default-features = false }
|
||||
|
||||
# Serialization
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
toml = "0.8"
|
||||
|
||||
# CLI
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
|
||||
# Logging
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
|
||||
# Error handling
|
||||
thiserror = "2"
|
||||
anyhow = "1"
|
||||
|
||||
# Async utilities
|
||||
futures = "0.3"
|
||||
|
||||
# UUID for session/tool IDs
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
||||
# Home directory
|
||||
dirs = "5"
|
||||
|
||||
# MCP protocol (JSON-RPC based)
|
||||
async-trait = "0.1"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
209
QWEN.md
Normal file
209
QWEN.md
Normal file
@@ -0,0 +1,209 @@
|
||||
# svibe - Agentic CLI Tool
|
||||
|
||||
A Rust-based agentic CLI tool with hierarchical agents, multi-provider support, tool usage, and MCP (Model Context Protocol) support.
|
||||
|
||||
## Project Overview
|
||||
|
||||
**svibe** is an intelligent command-line assistant that uses AI agents to help with development tasks. It supports:
|
||||
|
||||
- **Multiple Model Providers**: OpenAI (cloud) and Ollama (local)
|
||||
- **Hierarchical Agents**: Supervisor agent coordinates specialized sub-agents
|
||||
- **Model Routing**: Different models for different task types (chat, code, planning)
|
||||
- **Tool Support**: File operations, shell commands, git operations
|
||||
- **MCP Protocol**: Connect to external MCP servers for additional tools/resources
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
src/
|
||||
├── main.rs # CLI entry point
|
||||
├── config/ # Configuration system
|
||||
│ ├── types.rs # Config data structures
|
||||
│ ├── loader.rs # Config loading/saving
|
||||
│ └── error.rs # Config errors
|
||||
├── providers/ # AI model providers
|
||||
│ ├── traits.rs # Provider trait
|
||||
│ ├── types.rs # Common types (Message, Tool, etc.)
|
||||
│ ├── openai.rs # OpenAI client
|
||||
│ ├── ollama.rs # Ollama client
|
||||
│ └── registry.rs # Provider registry
|
||||
├── tools/ # Built-in tools
|
||||
│ ├── file_tools.rs # File operations
|
||||
│ ├── shell_tools.rs# Shell command execution
|
||||
│ ├── git_tools.rs # Git operations
|
||||
│ └── registry.rs # Tool registry
|
||||
├── mcp/ # MCP protocol support
|
||||
│ ├── client.rs # MCP client
|
||||
│ ├── types.rs # MCP types
|
||||
│ └── server.rs # MCP manager
|
||||
├── agents/ # Agent system
|
||||
│ ├── agent.rs # Base agent implementation
|
||||
│ ├── supervisor.rs # Supervisor agent
|
||||
│ └── types.rs # Agent types
|
||||
└── cli/ # CLI interface
|
||||
├── commands.rs # CLI commands
|
||||
├── repl.rs # Interactive REPL
|
||||
└── session.rs # Session management
|
||||
```
|
||||
|
||||
## Building and Running
|
||||
|
||||
```bash
|
||||
# Build the project
|
||||
cargo build
|
||||
|
||||
# Run in release mode
|
||||
cargo run --release
|
||||
|
||||
# Run specific command
|
||||
cargo run -- run "Help me fix this bug"
|
||||
|
||||
# Start REPL mode
|
||||
cargo run -- repl
|
||||
|
||||
# Initialize configuration
|
||||
cargo run -- init
|
||||
|
||||
# Show help
|
||||
cargo run -- --help
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
Run `svibe init` to create a configuration file at `~/.config/svibe/config.toml`:
|
||||
|
||||
```toml
|
||||
[providers.openai]
|
||||
api_key = "sk-..." # Or set OPENAI_API_KEY env var
|
||||
base_url = "https://api.openai.com/v1"
|
||||
default_model = "gpt-4o"
|
||||
|
||||
[providers.ollama]
|
||||
base_url = "http://localhost:11434"
|
||||
default_model = "llama3"
|
||||
# Bearer token auth (for cloud-hosted Ollama-compatible APIs)
|
||||
bearer_token = "your-token"
|
||||
# OR Basic auth
|
||||
# username = "user"
|
||||
# password = "pass"
|
||||
|
||||
[agents.default]
|
||||
description = "Default general-purpose agent"
|
||||
system_prompt = "You are a helpful AI assistant."
|
||||
|
||||
[agents.default.models]
|
||||
chat = "ollama/llama3"
|
||||
code = "openai/gpt-4o"
|
||||
planning = "openai/gpt-4o"
|
||||
tool_use = "openai/gpt-4o"
|
||||
default = "ollama/llama3"
|
||||
|
||||
[agents.default.tools]
|
||||
tools = ["file", "shell", "git"]
|
||||
|
||||
[agents.code]
|
||||
description = "Code development agent"
|
||||
# ... similar configuration
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Single Command Mode
|
||||
|
||||
```bash
|
||||
# Run a task
|
||||
svibe run "Fix the bug in main.rs"
|
||||
|
||||
# Use specific agent
|
||||
svibe -a code run "Write a function to parse JSON"
|
||||
|
||||
# Specify task type for model routing
|
||||
svibe run --task-type code "Implement binary search"
|
||||
```
|
||||
|
||||
### REPL Mode
|
||||
|
||||
```bash
|
||||
svibe repl
|
||||
```
|
||||
|
||||
REPL commands:
|
||||
- `/help` - Show help
|
||||
- `/clear` - Clear screen
|
||||
- `/status` - Show agent status
|
||||
- `/exit` or `quit` - Exit REPL
|
||||
|
||||
### Other Commands
|
||||
|
||||
```bash
|
||||
# Show agent status
|
||||
svibe status
|
||||
|
||||
# List available tools
|
||||
svibe tools
|
||||
|
||||
# List available agents
|
||||
svibe agents
|
||||
|
||||
# Health check providers
|
||||
svibe health
|
||||
|
||||
# MCP server management
|
||||
svibe mcp list
|
||||
svibe mcp connect myserver npx -y @modelcontextprotocol/server-example
|
||||
svibe mcp disconnect myserver
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
- `OPENAI_API_KEY` - OpenAI API key
|
||||
- `OLLAMA_HOST` - Ollama host (default: http://localhost:11434)
|
||||
- `RUST_LOG` - Logging level (e.g., `svibe=debug`)
|
||||
|
||||
## Development Conventions
|
||||
|
||||
- **Code Style**: Follow standard Rust idioms
|
||||
- **Formatting**: Run `cargo fmt` before committing
|
||||
- **Linting**: Run `cargo clippy` for lint checks
|
||||
- **Testing**: Add tests in module files with `#[cfg(test)]`
|
||||
|
||||
## Features
|
||||
|
||||
### Model Routing
|
||||
|
||||
Each agent can specify different models for different task types:
|
||||
- `chat` - General conversation
|
||||
- `code` - Code generation and analysis
|
||||
- `planning` - Complex reasoning and planning
|
||||
- `tool_use` - Structured tool invocation
|
||||
- `default` - Fallback model
|
||||
|
||||
### Hierarchical Agents
|
||||
|
||||
The supervisor agent coordinates sub-agents:
|
||||
- Routes tasks to appropriate sub-agents based on keywords
|
||||
- Sub-agents can specialize in specific domains (code, research, etc.)
|
||||
- Agents can delegate to other agents
|
||||
|
||||
### Built-in Tools
|
||||
|
||||
1. **File Tools**: read, write, edit, search, list directory
|
||||
2. **Shell Tools**: Execute shell commands with safety checks
|
||||
3. **Git Tools**: status, diff, commit, branch, log, add, checkout
|
||||
|
||||
### MCP Support
|
||||
|
||||
Connect to MCP servers for additional tools and resources:
|
||||
```bash
|
||||
svibe mcp connect filesystem npx -y @modelcontextprotocol/server-filesystem /path/to/allow
|
||||
```
|
||||
|
||||
## Future Enhancements
|
||||
|
||||
- Streaming responses
|
||||
- Proper MCP response correlation
|
||||
- More MCP transports (SSE)
|
||||
- Additional built-in agents (research, documentation)
|
||||
- Session persistence
|
||||
- Conversation history management
|
||||
- Tool output caching
|
||||
337
src/agents/agent.rs
Normal file
337
src/agents/agent.rs
Normal file
@@ -0,0 +1,337 @@
|
||||
use super::{
|
||||
AgentDelegation, AgentError, AgentMetrics, AgentRequest, AgentResponse, AgentRuntimeConfig,
|
||||
AgentState, AgentToolCall, Result,
|
||||
};
|
||||
use crate::config::AgentConfig;
|
||||
use crate::mcp::McpManager;
|
||||
use crate::providers::{
|
||||
ChatRequest, ChatResponse, Message, MessageRole, ProviderRegistry, TaskType, Tool,
|
||||
};
|
||||
use crate::tools::{ToolOutput, ToolRegistry};
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// Base agent trait
|
||||
#[async_trait]
|
||||
pub trait AgentTrait: Send + Sync {
|
||||
/// Get the agent name
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// Get the agent description
|
||||
fn description(&self) -> &str;
|
||||
|
||||
/// Process a request and return a response
|
||||
async fn process(&self, request: AgentRequest) -> Result<AgentResponse>;
|
||||
|
||||
/// Get the current state
|
||||
async fn get_state(&self) -> AgentState;
|
||||
|
||||
/// Get metrics
|
||||
fn get_metrics(&self) -> AgentMetrics;
|
||||
}
|
||||
|
||||
/// A concrete agent implementation
|
||||
pub struct Agent {
|
||||
config: AgentRuntimeConfig,
|
||||
provider_registry: Arc<ProviderRegistry>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
mcp_manager: Arc<McpManager>,
|
||||
state: Arc<Mutex<AgentState>>,
|
||||
metrics: Arc<Mutex<AgentMetrics>>,
|
||||
}
|
||||
|
||||
impl Agent {
|
||||
pub fn new(
|
||||
name: &str,
|
||||
config: AgentConfig,
|
||||
provider_registry: Arc<ProviderRegistry>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
mcp_manager: Arc<McpManager>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: AgentRuntimeConfig {
|
||||
name: name.to_string(),
|
||||
config,
|
||||
available_tools: vec![],
|
||||
available_sub_agents: vec![],
|
||||
},
|
||||
provider_registry,
|
||||
tool_registry,
|
||||
mcp_manager,
|
||||
state: Arc::new(Mutex::new(AgentState::default())),
|
||||
metrics: Arc::new(Mutex::new(AgentMetrics::default())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_tools(mut self, tools: Vec<String>) -> Self {
|
||||
self.config.available_tools = tools;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_sub_agents(mut self, agents: Vec<String>) -> Self {
|
||||
self.config.available_sub_agents = agents;
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the model for a specific task type
|
||||
fn get_model(&self, task_type: TaskType) -> Result<String> {
|
||||
self.config
|
||||
.get_model_for_task(task_type)
|
||||
.ok_or_else(|| AgentError::ConfigError(format!("No model configured for {:?}", task_type)))
|
||||
}
|
||||
|
||||
/// Get available tools as provider Tool definitions
|
||||
async fn get_provider_tools(&self) -> Vec<Tool> {
|
||||
let mut tools = Vec::new();
|
||||
|
||||
// Add built-in tools
|
||||
for tool_name in &self.config.available_tools {
|
||||
if let Some(tool) = self.tool_registry.get(tool_name) {
|
||||
tools.push(Tool {
|
||||
name: tool.name().to_string(),
|
||||
description: tool.description().to_string(),
|
||||
input_schema: tool.parameters(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Add MCP tools
|
||||
let mcp_tools = self.mcp_manager.list_all_tools().await;
|
||||
for (full_name, mcp_tool) in mcp_tools {
|
||||
tools.push(Tool {
|
||||
name: full_name,
|
||||
description: mcp_tool.description.unwrap_or_default(),
|
||||
input_schema: mcp_tool.input_schema.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
tools
|
||||
}
|
||||
|
||||
/// Execute a tool call
|
||||
async fn execute_tool(&self, tool_call: &AgentToolCall) -> Result<ToolOutput> {
|
||||
// Try built-in tools first
|
||||
if let Some(tool) = self.tool_registry.get(&tool_call.name) {
|
||||
return Ok(tool.execute(tool_call.arguments.clone()).await?);
|
||||
}
|
||||
|
||||
// Try MCP tools
|
||||
if tool_call.name.contains('/') {
|
||||
let result = self
|
||||
.mcp_manager
|
||||
.call_tool(&tool_call.name, Some(tool_call.arguments.clone()))
|
||||
.await?;
|
||||
|
||||
let content = result
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
crate::mcp::McpContent::Text { text } => Some(text.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
return Ok(ToolOutput::success(content));
|
||||
}
|
||||
|
||||
Err(AgentError::ToolError(
|
||||
crate::tools::ToolError::NotFound(tool_call.name.clone()),
|
||||
))
|
||||
}
|
||||
|
||||
/// Build the system prompt
|
||||
fn build_system_prompt(&self) -> String {
|
||||
let mut prompt = self.config.config.system_prompt.clone();
|
||||
|
||||
if !self.config.available_tools.is_empty() {
|
||||
prompt.push_str("\n\nYou have access to the following tools:\n");
|
||||
for tool_name in &self.config.available_tools {
|
||||
if let Some(tool) = self.tool_registry.get(tool_name) {
|
||||
prompt.push_str(&format!("- {}: {}\n", tool.name(), tool.description()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !self.config.available_sub_agents.is_empty() {
|
||||
prompt.push_str("\n\nYou can delegate tasks to these sub-agents:\n");
|
||||
for agent_name in &self.config.available_sub_agents {
|
||||
prompt.push_str(&format!("- {}\n", agent_name));
|
||||
}
|
||||
}
|
||||
|
||||
prompt
|
||||
}
|
||||
|
||||
/// Process tool calls in the response
|
||||
async fn process_tool_calls(
|
||||
&self,
|
||||
response: &ChatResponse,
|
||||
) -> Result<Vec<(AgentToolCall, ToolOutput)>> {
|
||||
let mut results = Vec::new();
|
||||
|
||||
for choice in &response.choices {
|
||||
for tool_call in &choice.message.tool_calls {
|
||||
let agent_tool_call = AgentToolCall {
|
||||
id: tool_call.id.clone(),
|
||||
name: tool_call.name.clone(),
|
||||
arguments: tool_call.arguments.clone(),
|
||||
};
|
||||
|
||||
let output = self.execute_tool(&agent_tool_call).await?;
|
||||
|
||||
results.push((agent_tool_call, output));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTrait for Agent {
|
||||
fn name(&self) -> &str {
|
||||
&self.config.name
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
&self.config.config.description
|
||||
}
|
||||
|
||||
async fn process(&self, request: AgentRequest) -> Result<AgentResponse> {
|
||||
let start_time = std::time::Instant::now();
|
||||
let mut metrics = self.metrics.lock().await;
|
||||
metrics.llm_calls += 1;
|
||||
|
||||
// Get the model for this task type
|
||||
let model = self.get_model(request.task_type)?;
|
||||
|
||||
// Build messages
|
||||
let mut messages = Vec::new();
|
||||
|
||||
// Add system prompt
|
||||
messages.push(Message {
|
||||
role: MessageRole::System,
|
||||
content: self.build_system_prompt(),
|
||||
tool_calls: vec![],
|
||||
tool_call_id: None,
|
||||
});
|
||||
|
||||
// Add context
|
||||
for context in request.context {
|
||||
messages.push(Message {
|
||||
role: MessageRole::User,
|
||||
content: format!("Context: {}", context),
|
||||
tool_calls: vec![],
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
// Add history
|
||||
messages.extend(request.history);
|
||||
|
||||
// Add current task
|
||||
messages.push(Message {
|
||||
role: MessageRole::User,
|
||||
content: request.task,
|
||||
tool_calls: vec![],
|
||||
tool_call_id: None,
|
||||
});
|
||||
|
||||
// Get available tools
|
||||
let tools = self.get_provider_tools().await;
|
||||
|
||||
// Resolve provider
|
||||
let (provider, _model_name) = self
|
||||
.provider_registry
|
||||
.resolve_provider(&model)
|
||||
.map_err(|e| AgentError::ProviderError(e))?;
|
||||
|
||||
// Make the request
|
||||
let chat_request = ChatRequest {
|
||||
model: model.clone(),
|
||||
messages: messages.clone(),
|
||||
tools: if tools.is_empty() { None } else { Some(tools) },
|
||||
max_tokens: Some(4096),
|
||||
temperature: Some(0.7),
|
||||
stream: Some(false),
|
||||
};
|
||||
|
||||
let response = provider
|
||||
.chat(chat_request)
|
||||
.await
|
||||
.map_err(|e| AgentError::ProviderError(e))?;
|
||||
|
||||
// Update metrics
|
||||
metrics.tokens_used += response.usage.total_tokens;
|
||||
|
||||
// Extract response
|
||||
let content = response
|
||||
.choices
|
||||
.first()
|
||||
.map(|c| c.message.content.clone())
|
||||
.unwrap_or_default();
|
||||
|
||||
// Process tool calls
|
||||
let mut tool_calls = Vec::new();
|
||||
let mut delegations = Vec::new();
|
||||
|
||||
for choice in &response.choices {
|
||||
for tc in &choice.message.tool_calls {
|
||||
// Check if this is a delegation (special tool call)
|
||||
if tc.name == "delegate_to_agent" {
|
||||
if let Some(target) = tc.arguments.get("target_agent").and_then(|v| v.as_str()) {
|
||||
delegations.push(AgentDelegation {
|
||||
target_agent: target.to_string(),
|
||||
task: tc
|
||||
.arguments
|
||||
.get("task")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
context: vec![],
|
||||
});
|
||||
}
|
||||
} else {
|
||||
tool_calls.push(AgentToolCall {
|
||||
id: tc.id.clone(),
|
||||
name: tc.name.clone(),
|
||||
arguments: tc.arguments.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
metrics.tool_calls += tool_calls.len() as u32;
|
||||
metrics.delegations += delegations.len() as u32;
|
||||
|
||||
// Update state
|
||||
let mut state = self.state.lock().await;
|
||||
state.messages = messages;
|
||||
state.task_complete = tool_calls.is_empty() && delegations.is_empty();
|
||||
state.task_result = if state.task_complete {
|
||||
Some(content.clone())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
metrics.execution_time_ms = start_time.elapsed().as_millis() as u64;
|
||||
|
||||
Ok(AgentResponse {
|
||||
content,
|
||||
complete: state.task_complete,
|
||||
error: None,
|
||||
tool_calls,
|
||||
delegations,
|
||||
})
|
||||
}
|
||||
|
||||
async fn get_state(&self) -> AgentState {
|
||||
self.state.lock().await.clone()
|
||||
}
|
||||
|
||||
fn get_metrics(&self) -> AgentMetrics {
|
||||
self.metrics.try_lock().map(|m| m.clone()).unwrap_or_default()
|
||||
}
|
||||
}
|
||||
33
src/agents/error.rs
Normal file
33
src/agents/error.rs
Normal file
@@ -0,0 +1,33 @@
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum AgentError {
|
||||
#[error("Agent not found: {0}")]
|
||||
NotFound(String),
|
||||
|
||||
#[error("Provider error: {0}")]
|
||||
ProviderError(#[from] crate::providers::ProviderError),
|
||||
|
||||
#[error("Tool error: {0}")]
|
||||
ToolError(#[from] crate::tools::ToolError),
|
||||
|
||||
#[error("MCP error: {0}")]
|
||||
McpError(#[from] crate::mcp::McpError),
|
||||
|
||||
#[error("Configuration error: {0}")]
|
||||
ConfigError(String),
|
||||
|
||||
#[error("Execution failed: {0}")]
|
||||
ExecutionFailed(String),
|
||||
|
||||
#[error("Delegation failed: {0}")]
|
||||
DelegationFailed(String),
|
||||
|
||||
#[error("Max iterations reached")]
|
||||
MaxIterationsReached,
|
||||
|
||||
#[error("Invalid state: {0}")]
|
||||
InvalidState(String),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, AgentError>;
|
||||
81
src/agents/executor.rs
Normal file
81
src/agents/executor.rs
Normal file
@@ -0,0 +1,81 @@
|
||||
use super::{AgentRequest, AgentResponse, AgentTrait, AgentMetrics, Result};
|
||||
use crate::config::AgentConfig;
|
||||
use crate::providers::{ProviderRegistry, TaskType};
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::mcp::McpManager;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Executor agent for direct task execution
|
||||
pub struct ExecutorAgent {
|
||||
name: String,
|
||||
config: AgentConfig,
|
||||
provider_registry: Arc<ProviderRegistry>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
mcp_manager: Arc<McpManager>,
|
||||
}
|
||||
|
||||
impl ExecutorAgent {
|
||||
pub fn new(
|
||||
name: &str,
|
||||
config: AgentConfig,
|
||||
provider_registry: Arc<ProviderRegistry>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
mcp_manager: Arc<McpManager>,
|
||||
) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
config,
|
||||
provider_registry,
|
||||
tool_registry,
|
||||
mcp_manager,
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a simple task directly
|
||||
pub async fn execute(&self, task: &str) -> Result<String> {
|
||||
let request = AgentRequest {
|
||||
task: task.to_string(),
|
||||
task_type: TaskType::Default,
|
||||
context: vec![],
|
||||
history: vec![],
|
||||
};
|
||||
|
||||
let response = self.process(request).await?;
|
||||
Ok(response.content)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AgentTrait for ExecutorAgent {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
&self.config.description
|
||||
}
|
||||
|
||||
async fn process(&self, request: AgentRequest) -> Result<AgentResponse> {
|
||||
// For now, delegate to a base Agent implementation
|
||||
// This could be enhanced with specific execution logic
|
||||
|
||||
let agent = super::Agent::new(
|
||||
&self.name,
|
||||
self.config.clone(),
|
||||
self.provider_registry.clone(),
|
||||
self.tool_registry.clone(),
|
||||
self.mcp_manager.clone(),
|
||||
)
|
||||
.with_tools(self.config.tools.clone());
|
||||
|
||||
agent.process(request).await
|
||||
}
|
||||
|
||||
async fn get_state(&self) -> crate::agents::AgentState {
|
||||
crate::agents::AgentState::default()
|
||||
}
|
||||
|
||||
fn get_metrics(&self) -> AgentMetrics {
|
||||
AgentMetrics::default()
|
||||
}
|
||||
}
|
||||
11
src/agents/mod.rs
Normal file
11
src/agents/mod.rs
Normal file
@@ -0,0 +1,11 @@
|
||||
mod error;
|
||||
mod types;
|
||||
mod agent;
|
||||
mod supervisor;
|
||||
mod executor;
|
||||
|
||||
pub use error::*;
|
||||
pub use types::*;
|
||||
pub use agent::{Agent, AgentTrait};
|
||||
pub use supervisor::SupervisorAgent;
|
||||
// pub use executor::ExecutorAgent;
|
||||
219
src/agents/supervisor.rs
Normal file
219
src/agents/supervisor.rs
Normal file
@@ -0,0 +1,219 @@
|
||||
use super::{
|
||||
Agent, AgentDelegation, AgentError, AgentMetrics, AgentRequest, AgentResponse, AgentTrait,
|
||||
Result,
|
||||
};
|
||||
use crate::config::{AgentConfig, Config};
|
||||
use crate::providers::ProviderRegistry;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::mcp::McpManager;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Supervisor agent that manages sub-agents
|
||||
pub struct SupervisorAgent {
|
||||
pub(crate) name: String,
|
||||
pub(crate) config: AgentConfig,
|
||||
pub(crate) sub_agents: Arc<RwLock<HashMap<String, Arc<dyn AgentTrait>>>>,
|
||||
pub(crate) provider_registry: Arc<ProviderRegistry>,
|
||||
pub(crate) tool_registry: Arc<ToolRegistry>,
|
||||
pub(crate) mcp_manager: Arc<McpManager>,
|
||||
pub(crate) metrics: Arc<RwLock<AgentMetrics>>,
|
||||
}
|
||||
|
||||
impl SupervisorAgent {
|
||||
pub fn new(
|
||||
name: &str,
|
||||
config: AgentConfig,
|
||||
provider_registry: Arc<ProviderRegistry>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
mcp_manager: Arc<McpManager>,
|
||||
) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
config,
|
||||
sub_agents: Arc::new(RwLock::new(HashMap::new())),
|
||||
provider_registry,
|
||||
tool_registry,
|
||||
mcp_manager,
|
||||
metrics: Arc::new(RwLock::new(AgentMetrics::default())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a sub-agent
|
||||
pub async fn add_sub_agent(&self, name: &str, agent: Arc<dyn AgentTrait>) {
|
||||
let mut sub_agents = self.sub_agents.write().await;
|
||||
sub_agents.insert(name.to_string(), agent);
|
||||
}
|
||||
|
||||
/// Remove a sub-agent
|
||||
pub async fn remove_sub_agent(&self, name: &str) {
|
||||
let mut sub_agents = self.sub_agents.write().await;
|
||||
sub_agents.remove(name);
|
||||
}
|
||||
|
||||
/// Get a sub-agent by name
|
||||
pub async fn get_sub_agent(&self, name: &str) -> Option<Arc<dyn AgentTrait>> {
|
||||
let sub_agents = self.sub_agents.read().await;
|
||||
sub_agents.get(name).cloned()
|
||||
}
|
||||
|
||||
/// List all sub-agents
|
||||
pub async fn list_sub_agents(&self) -> Vec<String> {
|
||||
let sub_agents = self.sub_agents.read().await;
|
||||
sub_agents.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// Delegate a task to a sub-agent
|
||||
pub async fn delegate(
|
||||
&self,
|
||||
delegation: &AgentDelegation,
|
||||
) -> Result<AgentResponse> {
|
||||
let sub_agent = self
|
||||
.get_sub_agent(&delegation.target_agent)
|
||||
.await
|
||||
.ok_or_else(|| {
|
||||
AgentError::DelegationFailed(format!(
|
||||
"Sub-agent '{}' not found",
|
||||
delegation.target_agent
|
||||
))
|
||||
})?;
|
||||
|
||||
let request = AgentRequest {
|
||||
task: delegation.task.clone(),
|
||||
task_type: crate::providers::TaskType::Default,
|
||||
context: delegation.context.clone(),
|
||||
history: vec![],
|
||||
};
|
||||
|
||||
let response = sub_agent.process(request).await?;
|
||||
|
||||
// Update metrics
|
||||
let mut metrics = self.metrics.write().await;
|
||||
metrics.delegations += 1;
|
||||
metrics.llm_calls += 1;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Create agents from configuration
|
||||
pub async fn create_from_config(
|
||||
config: &Config,
|
||||
provider_registry: Arc<ProviderRegistry>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
mcp_manager: Arc<McpManager>,
|
||||
) -> Result<Self> {
|
||||
let supervisor_config = config
|
||||
.agents
|
||||
.get("supervisor")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| AgentConfig {
|
||||
description: "Supervisor agent that coordinates sub-agents".to_string(),
|
||||
models: config
|
||||
.agents
|
||||
.get(&config.default_agent)
|
||||
.map(|a| a.models.clone())
|
||||
.unwrap_or_default(),
|
||||
system_prompt: "You are a supervisor agent. Your role is to coordinate tasks by delegating to specialized sub-agents.".to_string(),
|
||||
tools: vec![],
|
||||
sub_agents: config.agents.keys().filter(|k| *k != "supervisor").cloned().collect(),
|
||||
parent: None,
|
||||
});
|
||||
|
||||
let supervisor = Self::new(
|
||||
"supervisor",
|
||||
supervisor_config,
|
||||
provider_registry.clone(),
|
||||
tool_registry.clone(),
|
||||
mcp_manager.clone(),
|
||||
);
|
||||
|
||||
// Create sub-agents
|
||||
for (name, agent_config) in &config.agents {
|
||||
if name == "supervisor" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let agent = Agent::new(
|
||||
name,
|
||||
agent_config.clone(),
|
||||
provider_registry.clone(),
|
||||
tool_registry.clone(),
|
||||
mcp_manager.clone(),
|
||||
)
|
||||
.with_tools(agent_config.tools.clone())
|
||||
.with_sub_agents(agent_config.sub_agents.clone());
|
||||
|
||||
supervisor
|
||||
.add_sub_agent(name, Arc::new(agent))
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(supervisor)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AgentTrait for SupervisorAgent {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
&self.config.description
|
||||
}
|
||||
|
||||
async fn process(&self, request: AgentRequest) -> Result<AgentResponse> {
|
||||
// For now, simple routing based on task analysis
|
||||
let task_lower = request.task.to_lowercase();
|
||||
|
||||
// Try to find the best sub-agent for this task
|
||||
let sub_agents = self.sub_agents.read().await;
|
||||
|
||||
// Simple keyword-based routing - just get the agent name
|
||||
let target_agent_name = if task_lower.contains("code") || task_lower.contains("function") || task_lower.contains("bug") {
|
||||
sub_agents.get("code").map(|_| "code")
|
||||
} else if task_lower.contains("file") || task_lower.contains("read") || task_lower.contains("write") {
|
||||
sub_agents.get("code").map(|_| "code")
|
||||
} else if task_lower.contains("git") || task_lower.contains("commit") || task_lower.contains("branch") {
|
||||
sub_agents.get("code").map(|_| "code")
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Clone what we need before dropping the lock
|
||||
let should_use_code_agent = target_agent_name.is_some();
|
||||
drop(sub_agents);
|
||||
|
||||
if should_use_code_agent {
|
||||
if let Some(agent) = self.sub_agents.read().await.get("code").cloned() {
|
||||
let response = agent.process(request).await?;
|
||||
|
||||
let mut metrics = self.metrics.write().await;
|
||||
metrics.delegations += 1;
|
||||
|
||||
return Ok(response);
|
||||
}
|
||||
}
|
||||
|
||||
// No suitable sub-agent found, handle directly
|
||||
// Create a temporary agent
|
||||
let agent = Agent::new(
|
||||
&self.name,
|
||||
self.config.clone(),
|
||||
self.provider_registry.clone(),
|
||||
self.tool_registry.clone(),
|
||||
self.mcp_manager.clone(),
|
||||
);
|
||||
|
||||
agent.process(request).await
|
||||
}
|
||||
|
||||
async fn get_state(&self) -> crate::agents::AgentState {
|
||||
crate::agents::AgentState::default()
|
||||
}
|
||||
|
||||
fn get_metrics(&self) -> AgentMetrics {
|
||||
self.metrics.try_read().map(|m| m.clone()).unwrap_or_default()
|
||||
}
|
||||
}
|
||||
179
src/agents/types.rs
Normal file
179
src/agents/types.rs
Normal file
@@ -0,0 +1,179 @@
|
||||
use crate::config::AgentConfig;
|
||||
use crate::providers::{Message, TaskType};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Unique identifier for an agent
|
||||
#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AgentId(pub String);
|
||||
|
||||
impl AgentId {
|
||||
pub fn new() -> Self {
|
||||
Self(Uuid::new_v4().to_string())
|
||||
}
|
||||
|
||||
pub fn named(name: impl Into<String>) -> Self {
|
||||
Self(name.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AgentId {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Unique identifier for a session
|
||||
#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct SessionId(pub String);
|
||||
|
||||
impl SessionId {
|
||||
pub fn new() -> Self {
|
||||
Self(Uuid::new_v4().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SessionId {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent state
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct AgentState {
|
||||
/// Conversation history
|
||||
pub messages: Vec<Message>,
|
||||
|
||||
/// Pending tool calls
|
||||
pub pending_tool_calls: Vec<PendingToolCall>,
|
||||
|
||||
/// Current task description
|
||||
pub current_task: Option<String>,
|
||||
|
||||
/// Task completion status
|
||||
pub task_complete: bool,
|
||||
|
||||
/// Task result
|
||||
pub task_result: Option<String>,
|
||||
}
|
||||
|
||||
/// A pending tool call
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PendingToolCall {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub arguments: serde_json::Value,
|
||||
pub status: ToolCallStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub enum ToolCallStatus {
|
||||
Pending,
|
||||
Executing,
|
||||
Completed(String),
|
||||
Failed(String),
|
||||
}
|
||||
|
||||
/// Request to an agent
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentRequest {
|
||||
/// The task to perform
|
||||
pub task: String,
|
||||
|
||||
/// Task type for model routing
|
||||
pub task_type: TaskType,
|
||||
|
||||
/// Additional context
|
||||
#[serde(default)]
|
||||
pub context: Vec<String>,
|
||||
|
||||
/// Conversation history
|
||||
#[serde(default)]
|
||||
pub history: Vec<Message>,
|
||||
}
|
||||
|
||||
/// Response from an agent
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentResponse {
|
||||
/// The response content
|
||||
pub content: String,
|
||||
|
||||
/// Whether the task is complete
|
||||
pub complete: bool,
|
||||
|
||||
/// Any error that occurred
|
||||
#[serde(default)]
|
||||
pub error: Option<String>,
|
||||
|
||||
/// Tool calls that need to be executed
|
||||
#[serde(default)]
|
||||
pub tool_calls: Vec<AgentToolCall>,
|
||||
|
||||
/// Delegation requests to other agents
|
||||
#[serde(default)]
|
||||
pub delegations: Vec<AgentDelegation>,
|
||||
}
|
||||
|
||||
/// A tool call from an agent
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentToolCall {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub arguments: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A delegation request to another agent
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentDelegation {
|
||||
/// Target agent name
|
||||
pub target_agent: String,
|
||||
|
||||
/// Task to delegate
|
||||
pub task: String,
|
||||
|
||||
/// Context to pass
|
||||
pub context: Vec<String>,
|
||||
}
|
||||
|
||||
/// Agent configuration with runtime info
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AgentRuntimeConfig {
|
||||
pub name: String,
|
||||
pub config: AgentConfig,
|
||||
pub available_tools: Vec<String>,
|
||||
pub available_sub_agents: Vec<String>,
|
||||
}
|
||||
|
||||
impl AgentRuntimeConfig {
|
||||
pub fn get_model_for_task(&self, task_type: TaskType) -> Option<String> {
|
||||
let model_str = match task_type {
|
||||
TaskType::Chat => self.config.models.chat.as_ref(),
|
||||
TaskType::Code => self.config.models.code.as_ref(),
|
||||
TaskType::Planning => self.config.models.planning.as_ref(),
|
||||
TaskType::ToolUse => self.config.models.tool_use.as_ref(),
|
||||
TaskType::Default => self.config.models.default.as_ref(),
|
||||
};
|
||||
|
||||
model_str.cloned()
|
||||
}
|
||||
}
|
||||
|
||||
/// Metrics for agent execution
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct AgentMetrics {
|
||||
/// Number of LLM calls made
|
||||
pub llm_calls: u32,
|
||||
|
||||
/// Number of tool calls made
|
||||
pub tool_calls: u32,
|
||||
|
||||
/// Number of delegations made
|
||||
pub delegations: u32,
|
||||
|
||||
/// Total tokens used
|
||||
pub tokens_used: u32,
|
||||
|
||||
/// Execution time in milliseconds
|
||||
pub execution_time_ms: u64,
|
||||
}
|
||||
86
src/cli/commands.rs
Normal file
86
src/cli/commands.rs
Normal file
@@ -0,0 +1,86 @@
|
||||
use clap::{Parser, Subcommand};
|
||||
|
||||
/// svibe - Agentic CLI tool with hierarchical agents
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "svibe")]
|
||||
#[command(author, version, about, long_about = None)]
|
||||
pub struct Cli {
|
||||
/// Configuration file path
|
||||
#[arg(short, long, global = true)]
|
||||
pub config: Option<String>,
|
||||
|
||||
/// Enable verbose output
|
||||
#[arg(short, long, global = true)]
|
||||
pub verbose: bool,
|
||||
|
||||
/// Agent to use (default: from config)
|
||||
#[arg(short, long, global = true)]
|
||||
pub agent: Option<String>,
|
||||
|
||||
/// Subcommand to run
|
||||
#[command(subcommand)]
|
||||
pub command: Option<Commands>,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug)]
|
||||
pub enum Commands {
|
||||
/// Initialize configuration
|
||||
Init,
|
||||
|
||||
/// Run a single task
|
||||
Run {
|
||||
/// Task description
|
||||
#[arg(required = true)]
|
||||
task: String,
|
||||
|
||||
/// Task type for model routing
|
||||
#[arg(short, long)]
|
||||
task_type: Option<String>,
|
||||
},
|
||||
|
||||
/// Start interactive REPL mode
|
||||
Repl,
|
||||
|
||||
/// Show agent status
|
||||
Status,
|
||||
|
||||
/// List available tools
|
||||
Tools,
|
||||
|
||||
/// List available agents
|
||||
Agents,
|
||||
|
||||
/// Manage MCP servers
|
||||
Mcp {
|
||||
#[command(subcommand)]
|
||||
action: McpAction,
|
||||
},
|
||||
|
||||
/// Health check providers
|
||||
Health,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug)]
|
||||
pub enum McpAction {
|
||||
/// List connected MCP servers
|
||||
List,
|
||||
|
||||
/// Connect to an MCP server
|
||||
Connect {
|
||||
/// Server name
|
||||
name: String,
|
||||
|
||||
/// Command to run
|
||||
command: String,
|
||||
|
||||
/// Command arguments
|
||||
#[arg(last = true)]
|
||||
args: Vec<String>,
|
||||
},
|
||||
|
||||
/// Disconnect from an MCP server
|
||||
Disconnect {
|
||||
/// Server name
|
||||
name: String,
|
||||
},
|
||||
}
|
||||
7
src/cli/mod.rs
Normal file
7
src/cli/mod.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
mod commands;
|
||||
mod repl;
|
||||
mod session;
|
||||
|
||||
pub use commands::{Cli, Commands, McpAction};
|
||||
pub use repl::Repl;
|
||||
// pub use session::Session;
|
||||
138
src/cli/repl.rs
Normal file
138
src/cli/repl.rs
Normal file
@@ -0,0 +1,138 @@
|
||||
use crate::agents::{AgentRequest, AgentTrait};
|
||||
use crate::providers::TaskType;
|
||||
use std::io::{self, Write};
|
||||
|
||||
/// Interactive REPL for svibe
|
||||
pub struct Repl {
|
||||
agent: std::sync::Arc<dyn AgentTrait>,
|
||||
}
|
||||
|
||||
impl Repl {
|
||||
pub fn new(agent: std::sync::Arc<dyn AgentTrait>) -> Self {
|
||||
Self { agent }
|
||||
}
|
||||
|
||||
/// Run the REPL
|
||||
pub fn run(&self) -> Result<(), Box<dyn std::error::Error>> {
|
||||
println!("svibe REPL - Type 'exit' or 'quit' to leave");
|
||||
println!("Commands: /help, /clear, /status, /agents, /tools");
|
||||
println!();
|
||||
|
||||
let stdin = io::stdin();
|
||||
let mut stdout = io::stdout();
|
||||
let mut input = String::new();
|
||||
|
||||
loop {
|
||||
print!("> ");
|
||||
stdout.flush()?;
|
||||
|
||||
input.clear();
|
||||
let n = stdin.read_line(&mut input)?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
|
||||
let input = input.trim();
|
||||
|
||||
if input.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
if input == "exit" || input == "quit" || input == ":q" {
|
||||
println!("Goodbye!");
|
||||
break;
|
||||
}
|
||||
|
||||
if input.starts_with('/') {
|
||||
self.handle_command(input)?;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Process as a task - need to use tokio runtime
|
||||
let input = input.to_string();
|
||||
let agent = self.agent.clone();
|
||||
|
||||
let rt = tokio::runtime::Handle::current();
|
||||
let result = rt.block_on(async move {
|
||||
let request = AgentRequest {
|
||||
task: input,
|
||||
task_type: TaskType::Default,
|
||||
context: vec![],
|
||||
history: vec![],
|
||||
};
|
||||
agent.process(request).await
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(response) => {
|
||||
println!("\n{}", response.content);
|
||||
if !response.complete {
|
||||
println!("\n[Task incomplete - continuing...]");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("\nError: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
println!();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_command(&self, input: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let parts: Vec<&str> = input.split_whitespace().collect();
|
||||
|
||||
match parts.get(0).copied() {
|
||||
Some("help") | Some("/help") => {
|
||||
println!(
|
||||
r#"
|
||||
Available commands:
|
||||
/help Show this help message
|
||||
/clear Clear the screen
|
||||
/status Show agent status and metrics
|
||||
/agents List available agents
|
||||
/tools List available tools
|
||||
/exit, /quit, :q Exit the REPL
|
||||
|
||||
Just type your task to get started!
|
||||
"#
|
||||
);
|
||||
}
|
||||
Some("clear") | Some("/clear") => {
|
||||
print!("\x1B[2J\x1B[1;1H");
|
||||
io::stdout().flush()?;
|
||||
}
|
||||
Some("status") | Some("/status") => {
|
||||
let metrics = self.agent.get_metrics();
|
||||
println!(
|
||||
r#"
|
||||
Agent: {}
|
||||
Metrics:
|
||||
LLM calls: {}
|
||||
Tool calls: {}
|
||||
Delegations: {}
|
||||
Tokens used: {}
|
||||
"#,
|
||||
self.agent.name(),
|
||||
metrics.llm_calls,
|
||||
metrics.tool_calls,
|
||||
metrics.delegations,
|
||||
metrics.tokens_used,
|
||||
);
|
||||
}
|
||||
Some("agents") | Some("/agents") => {
|
||||
println!("Agent listing not available in this context");
|
||||
}
|
||||
Some("tools") | Some("/tools") => {
|
||||
println!("Tool listing not available in this context");
|
||||
}
|
||||
_ => {
|
||||
println!("Unknown command. Type /help for available commands.");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
102
src/cli/session.rs
Normal file
102
src/cli/session.rs
Normal file
@@ -0,0 +1,102 @@
|
||||
use crate::agents::{AgentRequest, AgentResponse, AgentTrait};
|
||||
use crate::providers::TaskType;
|
||||
use std::sync::Arc;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// A conversation session with an agent
|
||||
pub struct Session {
|
||||
id: String,
|
||||
agent: Arc<dyn AgentTrait>,
|
||||
history: Vec<AgentRequest>,
|
||||
responses: Vec<AgentResponse>,
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub fn new(agent: Arc<dyn AgentTrait>) -> Self {
|
||||
Self {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
agent,
|
||||
history: Vec::new(),
|
||||
responses: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn id(&self) -> &str {
|
||||
&self.id
|
||||
}
|
||||
|
||||
/// Process a task in this session
|
||||
pub async fn process(&mut self, task: &str) -> Result<AgentResponse, crate::agents::AgentError> {
|
||||
let request = AgentRequest {
|
||||
task: task.to_string(),
|
||||
task_type: TaskType::Default,
|
||||
context: vec![],
|
||||
history: self.get_recent_history(5),
|
||||
};
|
||||
|
||||
let response = self.agent.process(request.clone()).await?;
|
||||
|
||||
self.history.push(request);
|
||||
self.responses.push(response.clone());
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Get recent conversation history
|
||||
fn get_recent_history(&self, count: usize) -> Vec<crate::providers::Message> {
|
||||
use crate::providers::{Message, MessageRole};
|
||||
|
||||
let mut messages = Vec::new();
|
||||
|
||||
// Add recent exchanges
|
||||
let start = self.history.len().saturating_sub(count);
|
||||
for i in start..self.history.len() {
|
||||
if let Some(req) = self.history.get(i) {
|
||||
messages.push(Message {
|
||||
role: MessageRole::User,
|
||||
content: req.task.clone(),
|
||||
tool_calls: vec![],
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
if let Some(resp) = self.responses.get(i) {
|
||||
messages.push(Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: resp.content.clone(),
|
||||
tool_calls: vec![],
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
messages
|
||||
}
|
||||
|
||||
/// Clear session history
|
||||
pub fn clear(&mut self) {
|
||||
self.history.clear();
|
||||
self.responses.clear();
|
||||
}
|
||||
|
||||
/// Get session statistics
|
||||
pub fn stats(&self) -> SessionStats {
|
||||
let total_tokens: u32 = self
|
||||
.responses
|
||||
.iter()
|
||||
.map(|r| r.content.len() as u32 / 4)
|
||||
.sum();
|
||||
|
||||
SessionStats {
|
||||
message_count: self.history.len(),
|
||||
response_count: self.responses.len(),
|
||||
estimated_tokens: total_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Session statistics
|
||||
pub struct SessionStats {
|
||||
pub message_count: usize,
|
||||
pub response_count: usize,
|
||||
pub estimated_tokens: u32,
|
||||
}
|
||||
24
src/config/error.rs
Normal file
24
src/config/error.rs
Normal file
@@ -0,0 +1,24 @@
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ConfigError {
|
||||
#[error("Failed to read config file: {0}")]
|
||||
ReadError(#[from] std::io::Error),
|
||||
|
||||
#[error("Failed to parse config: {0}")]
|
||||
ParseError(#[from] toml::de::Error),
|
||||
|
||||
#[error("Failed to serialize config: {0}")]
|
||||
SerializeError(#[from] toml::ser::Error),
|
||||
|
||||
#[error("Config file not found at {0}. Run `svibe init` to create one.")]
|
||||
NotFound(String),
|
||||
|
||||
#[error("Missing required configuration: {0}")]
|
||||
MissingField(String),
|
||||
|
||||
#[error("Invalid configuration: {0}")]
|
||||
Invalid(String),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, ConfigError>;
|
||||
147
src/config/loader.rs
Normal file
147
src/config/loader.rs
Normal file
@@ -0,0 +1,147 @@
|
||||
use super::{Config, ConfigError, Result, OpenAiConfig, OllamaConfig, AgentConfig, AgentModels};
|
||||
use super::types::{default_openai_base_url, default_openai_model, default_ollama_base_url, default_ollama_model};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Configuration loader
|
||||
pub struct ConfigLoader {
|
||||
config_path: PathBuf,
|
||||
}
|
||||
|
||||
impl ConfigLoader {
|
||||
/// Create a new config loader with default path
|
||||
pub fn new() -> Self {
|
||||
let config_path = default_config_path();
|
||||
Self { config_path }
|
||||
}
|
||||
|
||||
/// Create a config loader with custom path
|
||||
pub fn with_path(path: impl AsRef<Path>) -> Self {
|
||||
Self {
|
||||
config_path: path.as_ref().to_path_buf(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Load configuration from file
|
||||
pub fn load(&self) -> Result<Config> {
|
||||
if !self.config_path.exists() {
|
||||
return Err(ConfigError::NotFound(
|
||||
self.config_path.display().to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let content = std::fs::read_to_string(&self.config_path)?;
|
||||
let mut config: Config = toml::from_str(&content)?;
|
||||
|
||||
// Merge with environment variables
|
||||
config.merge_env();
|
||||
|
||||
// Ensure default agent exists
|
||||
if config.agents.is_empty() {
|
||||
config.agents.insert("default".to_string(), default_agent_config());
|
||||
}
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Save configuration to file
|
||||
pub fn save(&self, config: &Config) -> Result<()> {
|
||||
if let Some(parent) = self.config_path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let content = toml::to_string_pretty(config)?;
|
||||
std::fs::write(&self.config_path, content)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the config path
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.config_path
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ConfigLoader {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the default configuration path
|
||||
fn default_config_path() -> PathBuf {
|
||||
dirs::config_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join("svibe")
|
||||
.join("config.toml")
|
||||
}
|
||||
|
||||
/// Merge environment variables into config
|
||||
impl Config {
|
||||
fn merge_env(&mut self) {
|
||||
// OpenAI API key from environment
|
||||
if let Some(ref mut openai) = self.providers.openai {
|
||||
if openai.api_key.is_none() {
|
||||
openai.api_key = std::env::var("OPENAI_API_KEY").ok();
|
||||
}
|
||||
}
|
||||
|
||||
// Ollama base URL from environment
|
||||
if let Some(ref mut ollama) = self.providers.ollama {
|
||||
if let Ok(url) = std::env::var("OLLAMA_HOST") {
|
||||
if !url.is_empty() {
|
||||
ollama.base_url = url;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a default agent configuration
|
||||
fn default_agent_config() -> AgentConfig {
|
||||
AgentConfig {
|
||||
description: "Default general-purpose agent".to_string(),
|
||||
models: AgentModels {
|
||||
chat: Some("ollama/llama3".to_string()),
|
||||
code: Some("openai/gpt-4o".to_string()),
|
||||
planning: Some("openai/gpt-4o".to_string()),
|
||||
tool_use: Some("openai/gpt-4o".to_string()),
|
||||
default: Some("ollama/llama3".to_string()),
|
||||
},
|
||||
system_prompt: "You are a helpful AI assistant. Help the user with their tasks.".to_string(),
|
||||
tools: vec!["file".to_string(), "shell".to_string(), "git".to_string()],
|
||||
sub_agents: vec![],
|
||||
parent: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a default configuration file
|
||||
pub fn create_default_config() -> Result<Config> {
|
||||
let mut config = Config::default();
|
||||
|
||||
config.providers.openai = Some(OpenAiConfig {
|
||||
api_key: None, // User must set this
|
||||
base_url: default_openai_base_url(),
|
||||
default_model: default_openai_model(),
|
||||
});
|
||||
|
||||
config.providers.ollama = Some(OllamaConfig {
|
||||
base_url: default_ollama_base_url(),
|
||||
default_model: default_ollama_model(),
|
||||
bearer_token: None,
|
||||
username: None,
|
||||
password: None,
|
||||
});
|
||||
|
||||
config.agents.insert("default".to_string(), default_agent_config());
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_config() {
|
||||
let config = create_default_config();
|
||||
assert!(config.is_ok());
|
||||
}
|
||||
}
|
||||
7
src/config/mod.rs
Normal file
7
src/config/mod.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
mod types;
|
||||
mod loader;
|
||||
mod error;
|
||||
|
||||
pub use types::*;
|
||||
pub use loader::*;
|
||||
pub use error::*;
|
||||
222
src/config/types.rs
Normal file
222
src/config/types.rs
Normal file
@@ -0,0 +1,222 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Main configuration structure
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Config {
|
||||
/// Provider configurations
|
||||
#[serde(default)]
|
||||
pub providers: ProvidersConfig,
|
||||
|
||||
/// Agent definitions
|
||||
#[serde(default)]
|
||||
pub agents: HashMap<String, AgentConfig>,
|
||||
|
||||
/// Default agent to use
|
||||
#[serde(default = "default_agent")]
|
||||
pub default_agent: String,
|
||||
|
||||
/// Global settings
|
||||
#[serde(default)]
|
||||
pub settings: Settings,
|
||||
}
|
||||
|
||||
fn default_agent() -> String {
|
||||
"default".to_string()
|
||||
}
|
||||
|
||||
/// Provider configurations
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct ProvidersConfig {
|
||||
/// OpenAI provider config
|
||||
#[serde(default)]
|
||||
pub openai: Option<OpenAiConfig>,
|
||||
|
||||
/// Ollama provider config
|
||||
#[serde(default)]
|
||||
pub ollama: Option<OllamaConfig>,
|
||||
}
|
||||
|
||||
/// OpenAI configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OpenAiConfig {
|
||||
/// API key (can also be set via OPENAI_API_KEY env var)
|
||||
pub api_key: Option<String>,
|
||||
|
||||
/// Base URL (for proxies or custom endpoints)
|
||||
#[serde(default = "default_openai_base_url")]
|
||||
pub base_url: String,
|
||||
|
||||
/// Default model for this provider
|
||||
#[serde(default = "default_openai_model")]
|
||||
pub default_model: String,
|
||||
}
|
||||
|
||||
pub fn default_openai_base_url() -> String {
|
||||
"https://api.openai.com/v1".to_string()
|
||||
}
|
||||
|
||||
pub fn default_openai_model() -> String {
|
||||
"gpt-4o".to_string()
|
||||
}
|
||||
|
||||
impl Default for OpenAiConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
api_key: None,
|
||||
base_url: default_openai_base_url(),
|
||||
default_model: default_openai_model(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Ollama configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OllamaConfig {
|
||||
/// Base URL
|
||||
#[serde(default = "default_ollama_base_url")]
|
||||
pub base_url: String,
|
||||
|
||||
/// Default model for this provider
|
||||
#[serde(default = "default_ollama_model")]
|
||||
pub default_model: String,
|
||||
|
||||
/// Bearer token for authentication (for cloud-hosted Ollama-compatible APIs)
|
||||
#[serde(default)]
|
||||
pub bearer_token: Option<String>,
|
||||
|
||||
/// Basic auth username
|
||||
#[serde(default)]
|
||||
pub username: Option<String>,
|
||||
|
||||
/// Basic auth password
|
||||
#[serde(default)]
|
||||
pub password: Option<String>,
|
||||
}
|
||||
|
||||
pub fn default_ollama_base_url() -> String {
|
||||
"http://localhost:11434".to_string()
|
||||
}
|
||||
|
||||
pub fn default_ollama_model() -> String {
|
||||
"llama3".to_string()
|
||||
}
|
||||
|
||||
impl Default for OllamaConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
base_url: default_ollama_base_url(),
|
||||
default_model: default_ollama_model(),
|
||||
bearer_token: None,
|
||||
username: None,
|
||||
password: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentConfig {
|
||||
/// Human-readable description
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
|
||||
/// Model routing configuration
|
||||
pub models: AgentModels,
|
||||
|
||||
/// System prompt for this agent
|
||||
#[serde(default)]
|
||||
pub system_prompt: String,
|
||||
|
||||
/// Tools this agent can use
|
||||
#[serde(default)]
|
||||
pub tools: Vec<String>,
|
||||
|
||||
/// Sub-agents this agent can delegate to
|
||||
#[serde(default)]
|
||||
pub sub_agents: Vec<String>,
|
||||
|
||||
/// Parent agent (for hierarchical structure)
|
||||
#[serde(default)]
|
||||
pub parent: Option<String>,
|
||||
}
|
||||
|
||||
/// Model routing for different task types
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct AgentModels {
|
||||
/// Model for general chat/conversation
|
||||
#[serde(default)]
|
||||
pub chat: Option<String>,
|
||||
|
||||
/// Model for code generation and analysis
|
||||
#[serde(default)]
|
||||
pub code: Option<String>,
|
||||
|
||||
/// Model for planning and reasoning
|
||||
#[serde(default)]
|
||||
pub planning: Option<String>,
|
||||
|
||||
/// Model for tool use (can be optimized for structured output)
|
||||
#[serde(default)]
|
||||
pub tool_use: Option<String>,
|
||||
|
||||
/// Fallback model when task type is unclear
|
||||
#[serde(default)]
|
||||
pub default: Option<String>,
|
||||
}
|
||||
|
||||
impl AgentModels {
|
||||
/// Get the model for a specific task type
|
||||
pub fn for_task(&self, task_type: &str) -> Option<&str> {
|
||||
match task_type {
|
||||
"chat" => self.chat.as_deref(),
|
||||
"code" => self.code.as_deref(),
|
||||
"planning" => self.planning.as_deref(),
|
||||
"tool_use" => self.tool_use.as_deref(),
|
||||
_ => self.default.as_deref(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Global settings
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Settings {
|
||||
/// Enable debug/verbose output
|
||||
#[serde(default)]
|
||||
pub verbose: bool,
|
||||
|
||||
/// Maximum tokens for responses
|
||||
#[serde(default = "default_max_tokens")]
|
||||
pub max_tokens: u32,
|
||||
|
||||
/// Temperature for generation
|
||||
#[serde(default = "default_temperature")]
|
||||
pub temperature: f32,
|
||||
|
||||
/// Working directory (defaults to current dir)
|
||||
#[serde(default)]
|
||||
pub working_dir: Option<String>,
|
||||
|
||||
/// Session storage directory
|
||||
#[serde(default)]
|
||||
pub session_dir: Option<String>,
|
||||
}
|
||||
|
||||
fn default_max_tokens() -> u32 {
|
||||
4096
|
||||
}
|
||||
|
||||
fn default_temperature() -> f32 {
|
||||
0.7
|
||||
}
|
||||
|
||||
impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
providers: ProvidersConfig::default(),
|
||||
agents: HashMap::new(),
|
||||
default_agent: default_agent(),
|
||||
settings: Settings::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
320
src/main.rs
Normal file
320
src/main.rs
Normal file
@@ -0,0 +1,320 @@
|
||||
// Allow dead code for future-use features
|
||||
#![allow(dead_code)]
|
||||
|
||||
mod agents;
|
||||
mod cli;
|
||||
mod config;
|
||||
mod mcp;
|
||||
mod providers;
|
||||
mod tools;
|
||||
|
||||
use agents::{AgentTrait, SupervisorAgent};
|
||||
use cli::{Cli, Commands, McpAction, Repl};
|
||||
use clap::Parser;
|
||||
use config::{Config, ConfigLoader, create_default_config};
|
||||
use mcp::McpManager;
|
||||
use providers::ProviderRegistry;
|
||||
use std::sync::Arc;
|
||||
use tools::{FileTools, GitTools, ShellTools, ToolRegistry};
|
||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// Initialize logging
|
||||
tracing_subscriber::registry()
|
||||
.with(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "svibe=info".into()),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.init();
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
// Load or create configuration
|
||||
let config = load_config(cli.config.as_deref())?;
|
||||
|
||||
// Update config from CLI args
|
||||
let mut config = config;
|
||||
if cli.verbose {
|
||||
config.settings.verbose = true;
|
||||
}
|
||||
|
||||
// Initialize components
|
||||
let provider_registry = Arc::new(ProviderRegistry::from_config(&config)?);
|
||||
let tool_registry = Arc::new(create_tool_registry());
|
||||
let mcp_manager = Arc::new(McpManager::new());
|
||||
|
||||
// Create supervisor agent with sub-agents
|
||||
let supervisor = SupervisorAgent::create_from_config(
|
||||
&config,
|
||||
provider_registry.clone(),
|
||||
tool_registry.clone(),
|
||||
mcp_manager.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Get the agent to use
|
||||
let agent: Arc<dyn AgentTrait> = if let Some(agent_name) = &cli.agent {
|
||||
supervisor
|
||||
.get_sub_agent(agent_name)
|
||||
.await
|
||||
.unwrap_or_else(|| Arc::new(supervisor.clone()))
|
||||
} else {
|
||||
Arc::new(supervisor.clone())
|
||||
};
|
||||
|
||||
// Handle commands
|
||||
match cli.command {
|
||||
None | Some(Commands::Repl) => {
|
||||
// Start REPL
|
||||
let repl = Repl::new(agent);
|
||||
if let Err(e) = repl.run() {
|
||||
eprintln!("REPL error: {}", e);
|
||||
}
|
||||
}
|
||||
Some(Commands::Init) => {
|
||||
cmd_init()?;
|
||||
}
|
||||
Some(Commands::Run { task, task_type }) => {
|
||||
cmd_run(&agent, &task, task_type.as_deref()).await?;
|
||||
}
|
||||
Some(Commands::Status) => {
|
||||
cmd_status(&agent, &provider_registry).await?;
|
||||
}
|
||||
Some(Commands::Tools) => {
|
||||
cmd_tools(&tool_registry).await?;
|
||||
}
|
||||
Some(Commands::Agents) => {
|
||||
cmd_agents(&supervisor).await?;
|
||||
}
|
||||
Some(Commands::Mcp { action }) => {
|
||||
cmd_mcp(&mcp_manager, action).await?;
|
||||
}
|
||||
Some(Commands::Health) => {
|
||||
cmd_health(&provider_registry).await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn load_config(config_path: Option<&str>) -> anyhow::Result<Config> {
|
||||
let loader = if let Some(path) = config_path {
|
||||
ConfigLoader::with_path(path)
|
||||
} else {
|
||||
ConfigLoader::new()
|
||||
};
|
||||
|
||||
match loader.load() {
|
||||
Ok(config) => Ok(config),
|
||||
Err(config::ConfigError::NotFound(_)) => {
|
||||
eprintln!("Config file not found. Run `svibe init` to create one.");
|
||||
eprintln!("Or set OPENAI_API_KEY environment variable for quick start.");
|
||||
|
||||
// Create a minimal config if env vars are set
|
||||
if std::env::var("OPENAI_API_KEY").is_ok() {
|
||||
eprintln!("Using OpenAI API key from environment.");
|
||||
return create_default_config().map_err(|e| e.into());
|
||||
}
|
||||
|
||||
std::process::exit(1);
|
||||
}
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn cmd_init() -> anyhow::Result<()> {
|
||||
let loader = ConfigLoader::new();
|
||||
let config = create_default_config()?;
|
||||
|
||||
println!("Creating configuration at: {}", loader.path().display());
|
||||
loader.save(&config)?;
|
||||
|
||||
println!(r#"
|
||||
Configuration created successfully!
|
||||
|
||||
Next steps:
|
||||
1. Set your OpenAI API key:
|
||||
export OPENAI_API_KEY=your-key-here
|
||||
|
||||
Or edit the config file and add your API key.
|
||||
|
||||
2. (Optional) Start Ollama for local models:
|
||||
ollama serve
|
||||
ollama pull llama3
|
||||
|
||||
3. Run svibe:
|
||||
svibe repl
|
||||
svibe run "Help me with this task"
|
||||
"#);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cmd_run(agent: &Arc<dyn AgentTrait>, task: &str, task_type: Option<&str>) -> anyhow::Result<()> {
|
||||
use agents::AgentRequest;
|
||||
use providers::TaskType;
|
||||
|
||||
let task_type = match task_type {
|
||||
Some("chat") => TaskType::Chat,
|
||||
Some("code") => TaskType::Code,
|
||||
Some("planning") => TaskType::Planning,
|
||||
Some("tool_use") => TaskType::ToolUse,
|
||||
_ => TaskType::Default,
|
||||
};
|
||||
|
||||
let request = AgentRequest {
|
||||
task: task.to_string(),
|
||||
task_type,
|
||||
context: vec![],
|
||||
history: vec![],
|
||||
};
|
||||
|
||||
println!("Processing task: {}\n", task);
|
||||
|
||||
let response = agent.process(request).await?;
|
||||
|
||||
println!("{}", response.content);
|
||||
|
||||
if !response.tool_calls.is_empty() {
|
||||
println!("\nTool calls to execute: {}", response.tool_calls.len());
|
||||
for tc in &response.tool_calls {
|
||||
println!(" - {}: {:?}", tc.name, tc.arguments);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cmd_status(agent: &Arc<dyn AgentTrait>, providers: &Arc<ProviderRegistry>) -> anyhow::Result<()> {
|
||||
println!("=== svibe status ===\n");
|
||||
|
||||
println!("Agent: {}", agent.name());
|
||||
println!("Description: {}", agent.description());
|
||||
|
||||
let metrics = agent.get_metrics();
|
||||
println!(
|
||||
"\nMetrics:
|
||||
LLM calls: {}
|
||||
Tool calls: {}
|
||||
Delegations: {}
|
||||
Tokens used: {}",
|
||||
metrics.llm_calls,
|
||||
metrics.tool_calls,
|
||||
metrics.delegations,
|
||||
metrics.tokens_used,
|
||||
);
|
||||
|
||||
println!("\nProviders:");
|
||||
let health = providers.health_check_all().await;
|
||||
for provider in providers.list() {
|
||||
let status = health.get(provider).copied().unwrap_or(false);
|
||||
println!(" - {}: {}", provider, if status { "✓" } else { "✗" });
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cmd_tools(tool_registry: &Arc<ToolRegistry>) -> anyhow::Result<()> {
|
||||
println!("=== Available tools ===\n");
|
||||
|
||||
for tool in tool_registry.list() {
|
||||
println!("{}: {}", tool.name, tool.description);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cmd_agents(supervisor: &SupervisorAgent) -> anyhow::Result<()> {
|
||||
println!("=== Available agents ===\n");
|
||||
|
||||
let sub_agents = supervisor.list_sub_agents().await;
|
||||
|
||||
if sub_agents.is_empty() {
|
||||
println!("No sub-agents configured.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
for name in sub_agents {
|
||||
if let Some(agent) = supervisor.get_sub_agent(&name).await {
|
||||
println!("{}: {}", name, agent.description());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cmd_mcp(mcp_manager: &Arc<McpManager>, action: McpAction) -> anyhow::Result<()> {
|
||||
match action {
|
||||
McpAction::List => {
|
||||
let servers = mcp_manager.list_servers().await;
|
||||
if servers.is_empty() {
|
||||
println!("No MCP servers connected.");
|
||||
} else {
|
||||
println!("Connected MCP servers:");
|
||||
for server in servers {
|
||||
println!(" - {}", server);
|
||||
}
|
||||
}
|
||||
}
|
||||
McpAction::Connect { name, command, args } => {
|
||||
let args: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
|
||||
mcp_manager.connect(&name, &command, &args).await?;
|
||||
println!("Connected to MCP server: {}", name);
|
||||
}
|
||||
McpAction::Disconnect { name } => {
|
||||
mcp_manager.disconnect(&name).await?;
|
||||
println!("Disconnected from MCP server: {}", name);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cmd_health(providers: &Arc<ProviderRegistry>) -> anyhow::Result<()> {
|
||||
println!("=== Provider health check ===\n");
|
||||
|
||||
let health = providers.health_check_all().await;
|
||||
|
||||
let mut all_healthy = true;
|
||||
for (provider, healthy) in health {
|
||||
if healthy {
|
||||
println!("✓ {} - healthy", provider);
|
||||
} else {
|
||||
println!("✗ {} - unhealthy", provider);
|
||||
all_healthy = false;
|
||||
}
|
||||
}
|
||||
|
||||
if !all_healthy {
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn create_tool_registry() -> ToolRegistry {
|
||||
let mut registry = ToolRegistry::new();
|
||||
|
||||
registry.register(Arc::new(FileTools::new()));
|
||||
registry.register(Arc::new(ShellTools::new()));
|
||||
registry.register(Arc::new(GitTools::new()));
|
||||
|
||||
registry
|
||||
}
|
||||
|
||||
// Need to implement Clone for SupervisorAgent
|
||||
impl Clone for SupervisorAgent {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
name: self.name.clone(),
|
||||
config: self.config.clone(),
|
||||
sub_agents: self.sub_agents.clone(),
|
||||
provider_registry: self.provider_registry.clone(),
|
||||
tool_registry: self.tool_registry.clone(),
|
||||
mcp_manager: self.mcp_manager.clone(),
|
||||
metrics: self.metrics.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
230
src/mcp/client.rs
Normal file
230
src/mcp/client.rs
Normal file
@@ -0,0 +1,230 @@
|
||||
use super::{
|
||||
JsonRpcRequest, JsonRpcResponse, McpCallToolParams, McpCallToolResult, McpClientCapabilities,
|
||||
McpError, McpInitializeParams, McpInitializeResult, McpTool, Result,
|
||||
RootsCapability, MCP_VERSION, ClientInfo,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::process::{Child, Command};
|
||||
use tokio::sync::{mpsc, Mutex};
|
||||
|
||||
/// MCP Client for communicating with MCP servers
|
||||
pub struct McpClient;
|
||||
|
||||
impl McpClient {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
|
||||
/// Connect to an MCP server via stdio (spawn a subprocess)
|
||||
pub async fn connect_stdio(
|
||||
&self,
|
||||
command: &str,
|
||||
args: &[&str],
|
||||
) -> Result<McpServerHandle> {
|
||||
let mut child = Command::new(command)
|
||||
.args(args)
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.map_err(|e| McpError::ConnectionFailed(format!("Failed to spawn: {}", e)))?;
|
||||
|
||||
let stdin = child
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| McpError::ConnectionFailed("No stdin".to_string()))?;
|
||||
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| McpError::ConnectionFailed("No stdout".to_string()))?;
|
||||
|
||||
// Create channels for communication
|
||||
let (tx, mut rx) = mpsc::channel::<Value>(100);
|
||||
|
||||
// Reader task
|
||||
let reader_handle = tokio::spawn(async move {
|
||||
let mut reader = BufReader::new(stdout);
|
||||
let mut line = String::new();
|
||||
|
||||
loop {
|
||||
line.clear();
|
||||
match reader.read_line(&mut line).await {
|
||||
Ok(0) => break, // EOF
|
||||
Ok(_) => {
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
match serde_json::from_str::<JsonRpcResponse>(trimmed) {
|
||||
Ok(_response) => {
|
||||
// For now, we don't correlate responses
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("MCP parse error: {} - {}", e, trimmed);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("MCP read error: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Writer task
|
||||
let stdin_arc = Arc::new(Mutex::new(stdin));
|
||||
let stdin_arc_clone = stdin_arc.clone();
|
||||
let writer_handle = tokio::spawn(async move {
|
||||
while let Some(value) = rx.recv().await {
|
||||
let mut stdin_guard = stdin_arc_clone.lock().await;
|
||||
let line = serde_json::to_string(&value).unwrap();
|
||||
let _ = stdin_guard.write_all(line.as_bytes()).await;
|
||||
let _ = stdin_guard.write_all(b"\n").await;
|
||||
let _ = stdin_guard.flush().await;
|
||||
}
|
||||
});
|
||||
|
||||
let stdin_mutex = Arc::new(Mutex::new(stdin_arc));
|
||||
|
||||
// Create handle
|
||||
let handle = McpServerHandle {
|
||||
name: command.to_string(),
|
||||
stdin: stdin_mutex,
|
||||
sender: tx,
|
||||
child: Arc::new(Mutex::new(Some(child))),
|
||||
reader_handle: Arc::new(Mutex::new(Some(reader_handle))),
|
||||
writer_handle: Arc::new(Mutex::new(Some(writer_handle))),
|
||||
};
|
||||
|
||||
// Initialize
|
||||
handle.initialize().await?;
|
||||
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
/// Connect to an MCP server via SSE
|
||||
pub async fn connect_sse(&self, _url: &str) -> Result<McpServerHandle> {
|
||||
Err(McpError::ConnectionFailed(
|
||||
"SSE transport not yet implemented".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for McpClient {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle to a connected MCP server
|
||||
#[derive(Clone)]
|
||||
pub struct McpServerHandle {
|
||||
name: String,
|
||||
stdin: Arc<Mutex<Arc<Mutex<tokio::process::ChildStdin>>>>,
|
||||
sender: mpsc::Sender<Value>,
|
||||
child: Arc<Mutex<Option<Child>>>,
|
||||
reader_handle: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>,
|
||||
writer_handle: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>,
|
||||
}
|
||||
|
||||
impl McpServerHandle {
|
||||
async fn send_request(&self, method: &str, params: Option<Value>) -> Result<Value> {
|
||||
let request = JsonRpcRequest::new(method, params, 1);
|
||||
let request_value = serde_json::to_value(&request)?;
|
||||
|
||||
self.sender
|
||||
.send(request_value)
|
||||
.await
|
||||
.map_err(|e| McpError::ProtocolError(format!("Send failed: {}", e)))?;
|
||||
|
||||
// For now, return success without waiting for response
|
||||
Ok(Value::Null)
|
||||
}
|
||||
|
||||
/// Initialize the MCP connection
|
||||
pub async fn initialize(&self) -> Result<McpInitializeResult> {
|
||||
let params = McpInitializeParams {
|
||||
protocol_version: MCP_VERSION.to_string(),
|
||||
capabilities: McpClientCapabilities {
|
||||
roots: Some(RootsCapability { list_changed: true }),
|
||||
sampling: None,
|
||||
},
|
||||
client_info: ClientInfo {
|
||||
name: "svibe".to_string(),
|
||||
version: env!("CARGO_PKG_VERSION").to_string(),
|
||||
},
|
||||
};
|
||||
|
||||
let params_value = serde_json::to_value(params)?;
|
||||
let _ = self.send_request("initialize", Some(params_value)).await?;
|
||||
|
||||
// Send initialized notification
|
||||
let _ = self.send_request("notifications/initialized", None).await;
|
||||
|
||||
// Return a default result for now
|
||||
Ok(McpInitializeResult {
|
||||
protocol_version: MCP_VERSION.to_string(),
|
||||
capabilities: super::McpServerCapabilities {
|
||||
tools: None,
|
||||
resources: None,
|
||||
prompts: None,
|
||||
},
|
||||
server_info: super::ServerInfo {
|
||||
name: self.name.clone(),
|
||||
version: "unknown".to_string(),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/// List available tools
|
||||
pub async fn list_tools(&self) -> Result<Vec<McpTool>> {
|
||||
let _result = self.send_request("tools/list", None).await?;
|
||||
|
||||
// For now return empty - proper implementation needs response correlation
|
||||
Ok(vec![])
|
||||
}
|
||||
|
||||
/// Call a tool
|
||||
pub async fn call_tool(
|
||||
&self,
|
||||
name: &str,
|
||||
arguments: Option<Value>,
|
||||
) -> Result<McpCallToolResult> {
|
||||
let params = McpCallToolParams {
|
||||
name: name.to_string(),
|
||||
arguments,
|
||||
};
|
||||
|
||||
let params_value = serde_json::to_value(params)?;
|
||||
let _result = self.send_request("tools/call", Some(params_value)).await?;
|
||||
|
||||
// For now return empty result
|
||||
Ok(McpCallToolResult {
|
||||
content: vec![],
|
||||
is_error: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the server name
|
||||
pub fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
/// Close the connection
|
||||
pub async fn close(&self) -> Result<()> {
|
||||
let _ = self.send_request("close", None).await;
|
||||
|
||||
let mut child_guard = self.child.lock().await;
|
||||
if let Some(mut child) = child_guard.take() {
|
||||
let _ = child.kill().await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
30
src/mcp/error.rs
Normal file
30
src/mcp/error.rs
Normal file
@@ -0,0 +1,30 @@
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum McpError {
|
||||
#[error("Connection failed: {0}")]
|
||||
ConnectionFailed(String),
|
||||
|
||||
#[error("Transport error: {0}")]
|
||||
TransportError(#[from] std::io::Error),
|
||||
|
||||
#[error("JSON parsing failed: {0}")]
|
||||
JsonError(#[from] serde_json::Error),
|
||||
|
||||
#[error("Server not found: {0}")]
|
||||
ServerNotFound(String),
|
||||
|
||||
#[error("Tool not found: {0}")]
|
||||
ToolNotFound(String),
|
||||
|
||||
#[error("Timeout: {0}")]
|
||||
Timeout(String),
|
||||
|
||||
#[error("Protocol error: {0}")]
|
||||
ProtocolError(String),
|
||||
|
||||
#[error("JSON-RPC error: {code}: {message}")]
|
||||
JsonRpcError { code: i32, message: String },
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, McpError>;
|
||||
9
src/mcp/mod.rs
Normal file
9
src/mcp/mod.rs
Normal file
@@ -0,0 +1,9 @@
|
||||
mod error;
|
||||
mod types;
|
||||
mod client;
|
||||
mod server;
|
||||
|
||||
pub use error::*;
|
||||
pub use types::*;
|
||||
pub use client::{McpClient, McpServerHandle};
|
||||
pub use server::McpManager;
|
||||
97
src/mcp/server.rs
Normal file
97
src/mcp/server.rs
Normal file
@@ -0,0 +1,97 @@
|
||||
use super::{McpClient, McpError, McpServerHandle, Result};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Manager for MCP server connections
|
||||
pub struct McpManager {
|
||||
client: McpClient,
|
||||
servers: Arc<RwLock<HashMap<String, McpServerHandle>>>,
|
||||
}
|
||||
|
||||
impl McpManager {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
client: McpClient::new(),
|
||||
servers: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Connect to an MCP server
|
||||
pub async fn connect(
|
||||
&self,
|
||||
name: &str,
|
||||
command: &str,
|
||||
args: &[&str],
|
||||
) -> Result<()> {
|
||||
let handle = self.client.connect_stdio(command, args).await?;
|
||||
|
||||
let mut servers = self.servers.write().await;
|
||||
servers.insert(name.to_string(), handle);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Disconnect from a server
|
||||
pub async fn disconnect(&self, name: &str) -> Result<()> {
|
||||
let mut servers = self.servers.write().await;
|
||||
|
||||
if let Some(handle) = servers.remove(name) {
|
||||
handle.close().await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get a server handle
|
||||
pub async fn get_server(&self, name: &str) -> Option<McpServerHandle> {
|
||||
let servers = self.servers.read().await;
|
||||
servers.get(name).cloned()
|
||||
}
|
||||
|
||||
/// List all connected servers
|
||||
pub async fn list_servers(&self) -> Vec<String> {
|
||||
let servers = self.servers.read().await;
|
||||
servers.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// List all tools from all servers
|
||||
pub async fn list_all_tools(&self) -> Vec<(String, super::McpTool)> {
|
||||
let servers = self.servers.read().await;
|
||||
let mut all_tools = Vec::new();
|
||||
|
||||
for (server_name, handle) in servers.iter() {
|
||||
if let Ok(tools) = handle.list_tools().await {
|
||||
for tool in tools {
|
||||
all_tools.push((format!("{}/{}", server_name, tool.name), tool));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
all_tools
|
||||
}
|
||||
|
||||
/// Call a tool on a specific server
|
||||
pub async fn call_tool(
|
||||
&self,
|
||||
server_tool: &str,
|
||||
arguments: Option<serde_json::Value>,
|
||||
) -> Result<super::McpCallToolResult> {
|
||||
let (server_name, tool_name) = server_tool
|
||||
.split_once('/')
|
||||
.ok_or_else(|| McpError::ToolNotFound(server_tool.to_string()))?;
|
||||
|
||||
let servers = self.servers.read().await;
|
||||
let handle = servers
|
||||
.get(server_name)
|
||||
.ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))?;
|
||||
|
||||
handle.call_tool(tool_name, arguments).await
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for McpManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
166
src/mcp/types.rs
Normal file
166
src/mcp/types.rs
Normal file
@@ -0,0 +1,166 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
/// JSON-RPC 2.0 request
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct JsonRpcRequest {
|
||||
pub jsonrpc: String,
|
||||
pub method: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub params: Option<Value>,
|
||||
pub id: Value,
|
||||
}
|
||||
|
||||
impl JsonRpcRequest {
|
||||
pub fn new(method: impl Into<String>, params: Option<Value>, id: u64) -> Self {
|
||||
Self {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
method: method.into(),
|
||||
params,
|
||||
id: Value::Number(id.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON-RPC 2.0 response
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct JsonRpcResponse {
|
||||
pub jsonrpc: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<JsonRpcErrorObject>,
|
||||
pub id: Value,
|
||||
}
|
||||
|
||||
/// JSON-RPC error object
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct JsonRpcErrorObject {
|
||||
pub code: i32,
|
||||
pub message: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub data: Option<Value>,
|
||||
}
|
||||
|
||||
/// MCP Protocol Version
|
||||
pub const MCP_VERSION: &str = "2024-11-05";
|
||||
|
||||
/// MCP Initialize request params
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpInitializeParams {
|
||||
#[serde(rename = "protocolVersion")]
|
||||
pub protocol_version: String,
|
||||
pub capabilities: McpClientCapabilities,
|
||||
#[serde(rename = "clientInfo")]
|
||||
pub client_info: ClientInfo,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ClientInfo {
|
||||
pub name: String,
|
||||
pub version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpClientCapabilities {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub roots: Option<RootsCapability>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub sampling: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RootsCapability {
|
||||
#[serde(rename = "listChanged")]
|
||||
pub list_changed: bool,
|
||||
}
|
||||
|
||||
/// MCP Initialize result
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpInitializeResult {
|
||||
#[serde(rename = "protocolVersion")]
|
||||
pub protocol_version: String,
|
||||
pub capabilities: McpServerCapabilities,
|
||||
#[serde(rename = "serverInfo")]
|
||||
pub server_info: ServerInfo,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ServerInfo {
|
||||
pub name: String,
|
||||
pub version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpServerCapabilities {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<ToolsCapability>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub resources: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub prompts: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolsCapability {
|
||||
#[serde(rename = "listChanged")]
|
||||
pub list_changed: Option<bool>,
|
||||
}
|
||||
|
||||
/// MCP Tool definition
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpTool {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
#[serde(rename = "inputSchema")]
|
||||
pub input_schema: Value,
|
||||
}
|
||||
|
||||
/// MCP List tools result
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpListToolsResult {
|
||||
pub tools: Vec<McpTool>,
|
||||
}
|
||||
|
||||
/// MCP Call tool params
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpCallToolParams {
|
||||
pub name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub arguments: Option<Value>,
|
||||
}
|
||||
|
||||
/// MCP Call tool result
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpCallToolResult {
|
||||
pub content: Vec<McpContent>,
|
||||
#[serde(rename = "isError", skip_serializing_if = "Option::is_none")]
|
||||
pub is_error: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum McpContent {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
#[serde(rename = "image")]
|
||||
Image { data: String, mime_type: String },
|
||||
#[serde(rename = "resource")]
|
||||
Resource { resource: Value },
|
||||
}
|
||||
|
||||
/// MCP List resources result
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpListResourcesResult {
|
||||
pub resources: Vec<McpResource>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpResource {
|
||||
pub uri: String,
|
||||
pub name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mime_type: Option<String>,
|
||||
}
|
||||
30
src/providers/error.rs
Normal file
30
src/providers/error.rs
Normal file
@@ -0,0 +1,30 @@
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ProviderError {
|
||||
#[error("HTTP request failed: {0}")]
|
||||
HttpError(#[from] reqwest::Error),
|
||||
|
||||
#[error("JSON parsing failed: {0}")]
|
||||
JsonError(#[from] serde_json::Error),
|
||||
|
||||
#[error("Provider API error: {0}")]
|
||||
ApiError(String),
|
||||
|
||||
#[error("Rate limit exceeded. Try again later.")]
|
||||
RateLimit,
|
||||
|
||||
#[error("Authentication failed: {0}")]
|
||||
AuthError(String),
|
||||
|
||||
#[error("Model not found: {0}")]
|
||||
ModelNotFound(String),
|
||||
|
||||
#[error("Invalid request: {0}")]
|
||||
InvalidRequest(String),
|
||||
|
||||
#[error("Unknown provider: {0}")]
|
||||
UnknownProvider(String),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, ProviderError>;
|
||||
14
src/providers/mod.rs
Normal file
14
src/providers/mod.rs
Normal file
@@ -0,0 +1,14 @@
|
||||
mod error;
|
||||
mod traits;
|
||||
mod types;
|
||||
mod openai;
|
||||
mod ollama;
|
||||
mod registry;
|
||||
|
||||
pub use error::*;
|
||||
pub use traits::*;
|
||||
pub use types::*;
|
||||
pub use openai::OpenAiProvider;
|
||||
pub use ollama::OllamaProvider;
|
||||
// pub use ollama::OllamaAuth;
|
||||
pub use registry::ProviderRegistry;
|
||||
230
src/providers/ollama.rs
Normal file
230
src/providers/ollama.rs
Normal file
@@ -0,0 +1,230 @@
|
||||
use super::{
|
||||
ChatRequest, ChatResponse, Message, MessageRole, Provider, ProviderError, Result, Tool,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Authentication method for Ollama provider
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub enum OllamaAuth {
|
||||
#[default]
|
||||
None,
|
||||
Bearer(String),
|
||||
Basic(String, String),
|
||||
}
|
||||
|
||||
/// Ollama API provider
|
||||
pub struct OllamaProvider {
|
||||
client: Client,
|
||||
base_url: String,
|
||||
default_model: String,
|
||||
auth: OllamaAuth,
|
||||
}
|
||||
|
||||
impl OllamaProvider {
|
||||
pub fn new(base_url: impl Into<String>, default_model: impl Into<String>) -> Self {
|
||||
Self {
|
||||
client: Client::new(),
|
||||
base_url: base_url.into(),
|
||||
default_model: default_model.into(),
|
||||
auth: OllamaAuth::None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
|
||||
self.auth = OllamaAuth::Bearer(token.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_basic_auth(mut self, username: impl Into<String>, password: impl Into<String>) -> Self {
|
||||
self.auth = OllamaAuth::Basic(username.into(), password.into());
|
||||
self
|
||||
}
|
||||
|
||||
fn parse_model(&self, model_spec: &str) -> String {
|
||||
// Handle "ollama/model-name" format
|
||||
if let Some(stripped) = model_spec.strip_prefix("ollama/") {
|
||||
stripped.to_string()
|
||||
} else {
|
||||
model_spec.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply authentication headers to a request builder
|
||||
fn apply_auth(&self, request: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
|
||||
match &self.auth {
|
||||
OllamaAuth::None => request,
|
||||
OllamaAuth::Bearer(token) => request.bearer_auth(token),
|
||||
OllamaAuth::Basic(username, password) => request.basic_auth(username, Some(password)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct OllamaChatRequest {
|
||||
model: String,
|
||||
messages: Vec<OllamaMessage>,
|
||||
stream: bool,
|
||||
options: Option<OllamaOptions>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct OllamaMessage {
|
||||
role: String,
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct OllamaOptions {
|
||||
temperature: Option<f32>,
|
||||
num_predict: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct OllamaChatResponse {
|
||||
message: OllamaMessage,
|
||||
done: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for OllamaProvider {
|
||||
fn name(&self) -> &str {
|
||||
"ollama"
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> Result<bool> {
|
||||
let url = format!("{}/api/tags", self.base_url);
|
||||
let response = self.apply_auth(self.client.get(&url)).send().await?;
|
||||
Ok(response.status().is_success())
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>> {
|
||||
let url = format!("{}/api/tags", self.base_url);
|
||||
let response = self.apply_auth(self.client.get(&url)).send().await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(ProviderError::ApiError(format!(
|
||||
"Failed to list models: {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ModelList {
|
||||
models: Vec<ModelInfo>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ModelInfo {
|
||||
name: String,
|
||||
}
|
||||
|
||||
let body: ModelList = response.json().await?;
|
||||
Ok(body.models.iter().map(|m| m.name.clone()).collect())
|
||||
}
|
||||
|
||||
async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
|
||||
let url = format!("{}/api/chat", self.base_url);
|
||||
let model = self.parse_model(&request.model);
|
||||
|
||||
let ollama_request = OllamaChatRequest {
|
||||
model,
|
||||
messages: request
|
||||
.messages
|
||||
.iter()
|
||||
.map(|m| OllamaMessage {
|
||||
role: match m.role {
|
||||
MessageRole::User => "user".to_string(),
|
||||
MessageRole::Assistant => "assistant".to_string(),
|
||||
MessageRole::System => "system".to_string(),
|
||||
MessageRole::Tool => "tool".to_string(),
|
||||
},
|
||||
content: m.content.clone(),
|
||||
})
|
||||
.collect(),
|
||||
stream: false,
|
||||
options: Some(OllamaOptions {
|
||||
temperature: request.temperature,
|
||||
num_predict: request.max_tokens,
|
||||
}),
|
||||
};
|
||||
|
||||
let response = self
|
||||
.apply_auth(
|
||||
self.client
|
||||
.post(&url)
|
||||
.header("Content-Type", "application/json")
|
||||
)
|
||||
.json(&ollama_request)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let error_text = response.text().await?;
|
||||
return Err(ProviderError::ApiError(error_text));
|
||||
}
|
||||
|
||||
let ollama_response: OllamaChatResponse = response.json().await?;
|
||||
|
||||
let chat_response = ChatResponse {
|
||||
choices: vec![super::Choice {
|
||||
message: Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: ollama_response.message.content,
|
||||
tool_calls: vec![],
|
||||
tool_call_id: None,
|
||||
},
|
||||
finish_reason: if ollama_response.done {
|
||||
Some("stop".to_string())
|
||||
} else {
|
||||
Some("length".to_string())
|
||||
},
|
||||
}],
|
||||
usage: super::Usage::default(),
|
||||
};
|
||||
|
||||
Ok(chat_response)
|
||||
}
|
||||
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
request: ChatRequest,
|
||||
) -> Result<tokio::sync::mpsc::Receiver<Result<String>>> {
|
||||
// For now, implement as non-streaming wrapped in a channel
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(100);
|
||||
|
||||
let this = self.clone();
|
||||
let request_clone = request;
|
||||
|
||||
tokio::spawn(async move {
|
||||
match this.chat(request_clone).await {
|
||||
Ok(response) => {
|
||||
for choice in response.choices {
|
||||
let _ = tx.send(Ok(choice.message.content)).await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(e)).await;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
|
||||
fn get_tools(&self) -> Vec<Tool> {
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for OllamaProvider {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
client: self.client.clone(),
|
||||
base_url: self.base_url.clone(),
|
||||
default_model: self.default_model.clone(),
|
||||
auth: self.auth.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
216
src/providers/openai.rs
Normal file
216
src/providers/openai.rs
Normal file
@@ -0,0 +1,216 @@
|
||||
use super::{
|
||||
ChatRequest, ChatResponse, MessageRole, Provider, ProviderError, Result, Tool,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use reqwest::Client;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// OpenAI API provider
|
||||
pub struct OpenAiProvider {
|
||||
client: Client,
|
||||
base_url: String,
|
||||
api_key: String,
|
||||
default_model: String,
|
||||
}
|
||||
|
||||
impl OpenAiProvider {
|
||||
pub fn new(base_url: impl Into<String>, api_key: impl Into<String>, default_model: impl Into<String>) -> Self {
|
||||
Self {
|
||||
client: Client::new(),
|
||||
base_url: base_url.into(),
|
||||
api_key: api_key.into(),
|
||||
default_model: default_model.into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_model(&self, model_spec: &str) -> String {
|
||||
// Handle "openai/model-name" format
|
||||
if let Some(stripped) = model_spec.strip_prefix("openai/") {
|
||||
stripped.to_string()
|
||||
} else {
|
||||
model_spec.to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for OpenAiProvider {
|
||||
fn name(&self) -> &str {
|
||||
"openai"
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> Result<bool> {
|
||||
let url = format!("{}/models", self.base_url);
|
||||
let response = self
|
||||
.client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
Ok(response.status().is_success())
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>> {
|
||||
let url = format!("{}/models", self.base_url);
|
||||
let response = self
|
||||
.client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(ProviderError::ApiError(format!(
|
||||
"Failed to list models: {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
|
||||
let body: serde_json::Value = response.json().await?;
|
||||
let models = body
|
||||
.get("data")
|
||||
.and_then(|d| d.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|m| m.get("id").and_then(|id| id.as_str()))
|
||||
.map(|s| s.to_string())
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(models)
|
||||
}
|
||||
|
||||
async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
let model = self.parse_model(&request.model);
|
||||
|
||||
let mut req_body = serde_json::json!({
|
||||
"model": model,
|
||||
"messages": request.messages.iter().map(|m| {
|
||||
let mut msg = serde_json::json!({
|
||||
"role": match m.role {
|
||||
MessageRole::User => "user",
|
||||
MessageRole::Assistant => "assistant",
|
||||
MessageRole::System => "system",
|
||||
MessageRole::Tool => "tool",
|
||||
},
|
||||
"content": m.content,
|
||||
});
|
||||
|
||||
if m.role == MessageRole::Tool {
|
||||
if let Some(ref id) = m.tool_call_id {
|
||||
msg["tool_call_id"] = serde_json::json!(id);
|
||||
}
|
||||
}
|
||||
|
||||
if !m.tool_calls.is_empty() {
|
||||
msg["tool_calls"] = serde_json::json!(m.tool_calls.iter().map(|tc| {
|
||||
serde_json::json!({
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.name,
|
||||
"arguments": tc.arguments,
|
||||
}
|
||||
})
|
||||
}).collect::<Vec<_>>());
|
||||
}
|
||||
|
||||
msg
|
||||
}).collect::<Vec<_>>(),
|
||||
});
|
||||
|
||||
if let Some(ref tools) = request.tools {
|
||||
req_body["tools"] = serde_json::json!(tools.iter().map(|t| {
|
||||
serde_json::json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": t.name,
|
||||
"description": t.description,
|
||||
"parameters": t.input_schema,
|
||||
}
|
||||
})
|
||||
}).collect::<Vec<_>>());
|
||||
}
|
||||
|
||||
if let Some(max_tokens) = request.max_tokens {
|
||||
req_body["max_tokens"] = serde_json::json!(max_tokens);
|
||||
}
|
||||
|
||||
if let Some(temperature) = request.temperature {
|
||||
req_body["temperature"] = serde_json::json!(temperature);
|
||||
}
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&req_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if response.status() == 429 {
|
||||
return Err(ProviderError::RateLimit);
|
||||
}
|
||||
|
||||
if !response.status().is_success() {
|
||||
let error_text = response.text().await?;
|
||||
return Err(ProviderError::ApiError(error_text));
|
||||
}
|
||||
|
||||
let response: ChatResponse = response.json().await?;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
request: ChatRequest,
|
||||
) -> Result<tokio::sync::mpsc::Receiver<Result<String>>> {
|
||||
// For now, implement as non-streaming wrapped in a channel
|
||||
// TODO: Implement proper SSE streaming
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(100);
|
||||
|
||||
let this = Arc::new(self.clone_box());
|
||||
let request_clone = request;
|
||||
|
||||
tokio::spawn(async move {
|
||||
match this.chat(request_clone).await {
|
||||
Ok(response) => {
|
||||
for choice in response.choices {
|
||||
let _ = tx.send(Ok(choice.message.content)).await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(e)).await;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
|
||||
fn get_tools(&self) -> Vec<Tool> {
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for OpenAiProvider {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
client: self.client.clone(),
|
||||
base_url: self.base_url.clone(),
|
||||
api_key: self.api_key.clone(),
|
||||
default_model: self.default_model.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Helper for cloning trait objects
|
||||
impl OpenAiProvider {
|
||||
fn clone_box(&self) -> Box<dyn Provider> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
}
|
||||
134
src/providers/registry.rs
Normal file
134
src/providers/registry.rs
Normal file
@@ -0,0 +1,134 @@
|
||||
use super::{Provider, ProviderError, Result, OpenAiProvider, OllamaProvider};
|
||||
use crate::config::Config;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Registry for managing multiple providers
|
||||
pub struct ProviderRegistry {
|
||||
providers: HashMap<String, Arc<dyn Provider>>,
|
||||
}
|
||||
|
||||
impl ProviderRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
providers: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a registry from configuration
|
||||
pub fn from_config(config: &Config) -> Result<Self> {
|
||||
let mut registry = Self::new();
|
||||
|
||||
// Register OpenAI provider if configured
|
||||
if let Some(ref openai_config) = config.providers.openai {
|
||||
let api_key = openai_config
|
||||
.api_key
|
||||
.clone()
|
||||
.ok_or_else(|| ProviderError::AuthError("OpenAI API key not configured".to_string()))?;
|
||||
|
||||
let provider = OpenAiProvider::new(
|
||||
&openai_config.base_url,
|
||||
&api_key,
|
||||
&openai_config.default_model,
|
||||
);
|
||||
|
||||
registry.register("openai", Arc::new(provider));
|
||||
}
|
||||
|
||||
// Register Ollama provider if configured
|
||||
if let Some(ref ollama_config) = config.providers.ollama {
|
||||
let mut provider = OllamaProvider::new(
|
||||
&ollama_config.base_url,
|
||||
&ollama_config.default_model,
|
||||
);
|
||||
|
||||
// Apply authentication
|
||||
if let Some(ref token) = ollama_config.bearer_token {
|
||||
provider = provider.with_bearer_token(token);
|
||||
} else if let (Some(ref username), Some(ref password)) =
|
||||
(&ollama_config.username, &ollama_config.password) {
|
||||
provider = provider.with_basic_auth(username, password);
|
||||
}
|
||||
|
||||
registry.register("ollama", Arc::new(provider));
|
||||
}
|
||||
|
||||
Ok(registry)
|
||||
}
|
||||
|
||||
/// Register a provider
|
||||
pub fn register(&mut self, name: impl Into<String>, provider: Arc<dyn Provider>) {
|
||||
self.providers.insert(name.into(), provider);
|
||||
}
|
||||
|
||||
/// Get a provider by name
|
||||
pub fn get(&self, name: &str) -> Option<Arc<dyn Provider>> {
|
||||
self.providers.get(name).cloned()
|
||||
}
|
||||
|
||||
/// Get a provider or return an error
|
||||
pub fn get_or_error(&self, name: &str) -> Result<Arc<dyn Provider>> {
|
||||
self.get(name)
|
||||
.ok_or_else(|| ProviderError::UnknownProvider(name.to_string()))
|
||||
}
|
||||
|
||||
/// Resolve a model specification to a provider
|
||||
/// Format: "provider/model-name" or just "model-name" (uses default)
|
||||
pub fn resolve_provider(&self, model_spec: &str) -> Result<(Arc<dyn Provider>, String)> {
|
||||
if let Some((provider_name, model_name)) = model_spec.split_once('/') {
|
||||
let provider = self.get_or_error(provider_name)?;
|
||||
Ok((provider, model_name.to_string()))
|
||||
} else {
|
||||
// Try to find a provider that has this model as default
|
||||
// For now, default to openai if available
|
||||
if let Some(provider) = self.get("openai") {
|
||||
return Ok((provider, model_spec.to_string()));
|
||||
}
|
||||
if let Some(provider) = self.get("ollama") {
|
||||
return Ok((provider, model_spec.to_string()));
|
||||
}
|
||||
Err(ProviderError::UnknownProvider(
|
||||
"No providers available".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// List all registered providers
|
||||
pub fn list(&self) -> Vec<&str> {
|
||||
self.providers.keys().map(|s| s.as_str()).collect()
|
||||
}
|
||||
|
||||
/// Check health of all providers
|
||||
pub async fn health_check_all(&self) -> HashMap<String, bool> {
|
||||
let mut results = HashMap::new();
|
||||
|
||||
for (name, provider) in &self.providers {
|
||||
let healthy = provider.health_check().await.unwrap_or(false);
|
||||
results.insert(name.clone(), healthy);
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ProviderRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_resolve_provider_format() {
|
||||
let registry = ProviderRegistry::new();
|
||||
|
||||
// Test parsing format
|
||||
assert_eq!(
|
||||
"openai/gpt-4".split_once('/'),
|
||||
Some(("openai", "gpt-4"))
|
||||
);
|
||||
}
|
||||
}
|
||||
27
src/providers/traits.rs
Normal file
27
src/providers/traits.rs
Normal file
@@ -0,0 +1,27 @@
|
||||
use super::{ChatRequest, ChatResponse, Result, Tool};
|
||||
use async_trait::async_trait;
|
||||
|
||||
/// Trait for AI model providers
|
||||
#[async_trait]
|
||||
pub trait Provider: Send + Sync {
|
||||
/// Get the provider name
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// Check if the provider is available
|
||||
async fn health_check(&self) -> Result<bool>;
|
||||
|
||||
/// List available models
|
||||
async fn list_models(&self) -> Result<Vec<String>>;
|
||||
|
||||
/// Send a chat request and get a response
|
||||
async fn chat(&self, request: ChatRequest) -> Result<ChatResponse>;
|
||||
|
||||
/// Send a chat request with streaming response
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
request: ChatRequest,
|
||||
) -> Result<tokio::sync::mpsc::Receiver<Result<String>>>;
|
||||
|
||||
/// Get available tools
|
||||
fn get_tools(&self) -> Vec<Tool>;
|
||||
}
|
||||
139
src/providers/types.rs
Normal file
139
src/providers/types.rs
Normal file
@@ -0,0 +1,139 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// A message in a conversation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Message {
|
||||
/// Role of the message sender
|
||||
pub role: MessageRole,
|
||||
|
||||
/// Content of the message
|
||||
pub content: String,
|
||||
|
||||
/// Optional tool calls in the message
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
|
||||
/// Optional tool call result (for tool responses)
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tool_call_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Role of a message sender
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum MessageRole {
|
||||
User,
|
||||
Assistant,
|
||||
System,
|
||||
Tool,
|
||||
}
|
||||
|
||||
/// A tool call from the model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCall {
|
||||
/// Unique ID for this tool call
|
||||
pub id: String,
|
||||
|
||||
/// Name of the tool to call
|
||||
pub name: String,
|
||||
|
||||
/// Arguments to pass to the tool
|
||||
pub arguments: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A tool definition
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Tool {
|
||||
/// Name of the tool
|
||||
pub name: String,
|
||||
|
||||
/// Description of what the tool does
|
||||
pub description: String,
|
||||
|
||||
/// JSON schema for the tool's parameters
|
||||
#[serde(rename = "parameters")]
|
||||
pub input_schema: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Request to a chat completion API
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChatRequest {
|
||||
/// Model to use
|
||||
pub model: String,
|
||||
|
||||
/// Messages in the conversation
|
||||
pub messages: Vec<Message>,
|
||||
|
||||
/// Tools available for use
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<Vec<Tool>>,
|
||||
|
||||
/// Maximum tokens to generate
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_tokens: Option<u32>,
|
||||
|
||||
/// Temperature for sampling
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f32>,
|
||||
|
||||
/// Whether to stream the response
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stream: Option<bool>,
|
||||
}
|
||||
|
||||
/// Response from a chat completion API
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChatResponse {
|
||||
/// The model's response
|
||||
pub choices: Vec<Choice>,
|
||||
|
||||
/// Usage statistics
|
||||
#[serde(default)]
|
||||
pub usage: Usage,
|
||||
}
|
||||
|
||||
/// A single choice in a response
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Choice {
|
||||
/// The message content
|
||||
pub message: Message,
|
||||
|
||||
/// Finish reason
|
||||
#[serde(default)]
|
||||
pub finish_reason: Option<String>,
|
||||
}
|
||||
|
||||
/// Token usage statistics
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct Usage {
|
||||
#[serde(default)]
|
||||
pub prompt_tokens: u32,
|
||||
|
||||
#[serde(default)]
|
||||
pub completion_tokens: u32,
|
||||
|
||||
#[serde(default)]
|
||||
pub total_tokens: u32,
|
||||
}
|
||||
|
||||
/// Task type for model routing
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum TaskType {
|
||||
Chat,
|
||||
Code,
|
||||
Planning,
|
||||
ToolUse,
|
||||
Default,
|
||||
}
|
||||
|
||||
impl TaskType {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
TaskType::Chat => "chat",
|
||||
TaskType::Code => "code",
|
||||
TaskType::Planning => "planning",
|
||||
TaskType::ToolUse => "tool_use",
|
||||
TaskType::Default => "default",
|
||||
}
|
||||
}
|
||||
}
|
||||
30
src/tools/error.rs
Normal file
30
src/tools/error.rs
Normal file
@@ -0,0 +1,30 @@
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ToolError {
|
||||
#[error("Tool not found: {0}")]
|
||||
NotFound(String),
|
||||
|
||||
#[error("Invalid arguments: {0}")]
|
||||
InvalidArguments(String),
|
||||
|
||||
#[error("Execution failed: {0}")]
|
||||
ExecutionFailed(String),
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
IoError(#[from] std::io::Error),
|
||||
|
||||
#[error("Permission denied: {0}")]
|
||||
PermissionDenied(String),
|
||||
|
||||
#[error("File not found: {0}")]
|
||||
FileNotFound(String),
|
||||
|
||||
#[error("Directory not found: {0}")]
|
||||
DirectoryNotFound(String),
|
||||
|
||||
#[error("Git error: {0}")]
|
||||
GitError(String),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, ToolError>;
|
||||
245
src/tools/file_tools.rs
Normal file
245
src/tools/file_tools.rs
Normal file
@@ -0,0 +1,245 @@
|
||||
use super::{Tool, ToolError, ToolOutput, ToolResult};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::path::Path;
|
||||
|
||||
/// File operation tools
|
||||
pub struct FileTools;
|
||||
|
||||
impl FileTools {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FileTools {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for FileTools {
|
||||
fn name(&self) -> &str {
|
||||
"file"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"File operations: read, write, edit, search, and list directory contents"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["read", "write", "edit", "search", "list"],
|
||||
"description": "The file operation to perform"
|
||||
},
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "File or directory path"
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "Content to write (for write/edit operations)"
|
||||
},
|
||||
"pattern": {
|
||||
"type": "string",
|
||||
"description": "Search pattern (for search operation)"
|
||||
},
|
||||
"line_number": {
|
||||
"type": "integer",
|
||||
"description": "Line number for targeted edit"
|
||||
}
|
||||
},
|
||||
"required": ["operation", "path"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value) -> ToolResult<ToolOutput> {
|
||||
let operation = args
|
||||
.get("operation")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("Missing 'operation' argument".to_string()))?;
|
||||
|
||||
let path = args
|
||||
.get("path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("Missing 'path' argument".to_string()))?;
|
||||
|
||||
match operation {
|
||||
"read" => self.read_file(path).await,
|
||||
"write" => {
|
||||
let content = args
|
||||
.get("content")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("Missing 'content' argument".to_string()))?;
|
||||
self.write_file(path, content).await
|
||||
}
|
||||
"edit" => {
|
||||
let content = args
|
||||
.get("content")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("Missing 'content' argument".to_string()))?;
|
||||
let line_number = args.get("line_number").and_then(|v| v.as_u64());
|
||||
self.edit_file(path, content, line_number).await
|
||||
}
|
||||
"search" => {
|
||||
let pattern = args
|
||||
.get("pattern")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("Missing 'pattern' argument".to_string()))?;
|
||||
self.search_file(path, pattern).await
|
||||
}
|
||||
"list" => self.list_directory(path).await,
|
||||
_ => Err(ToolError::InvalidArguments(format!(
|
||||
"Unknown operation: {}",
|
||||
operation
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FileTools {
|
||||
async fn read_file(&self, path: &str) -> ToolResult<ToolOutput> {
|
||||
let full_path = Path::new(path);
|
||||
|
||||
if !full_path.exists() {
|
||||
return Err(ToolError::FileNotFound(path.to_string()));
|
||||
}
|
||||
|
||||
let content = tokio::fs::read_to_string(full_path).await?;
|
||||
Ok(ToolOutput::success(format!(
|
||||
"File contents ({} bytes):\n\n{}",
|
||||
content.len(),
|
||||
content
|
||||
)))
|
||||
}
|
||||
|
||||
async fn write_file(&self, path: &str, content: &str) -> ToolResult<ToolOutput> {
|
||||
let full_path = Path::new(path);
|
||||
|
||||
if let Some(parent) = full_path.parent() {
|
||||
if !parent.exists() {
|
||||
tokio::fs::create_dir_all(parent).await?;
|
||||
}
|
||||
}
|
||||
|
||||
tokio::fs::write(full_path, content).await?;
|
||||
Ok(ToolOutput::success(format!(
|
||||
"Successfully wrote {} bytes to {}",
|
||||
content.len(),
|
||||
path
|
||||
)))
|
||||
}
|
||||
|
||||
async fn edit_file(&self, path: &str, content: &str, line_number: Option<u64>) -> ToolResult<ToolOutput> {
|
||||
let full_path = Path::new(path);
|
||||
|
||||
if !full_path.exists() {
|
||||
return Err(ToolError::FileNotFound(path.to_string()));
|
||||
}
|
||||
|
||||
let existing = tokio::fs::read_to_string(full_path).await?;
|
||||
|
||||
let new_content = if let Some(line_num) = line_number {
|
||||
let mut lines: Vec<&str> = existing.lines().collect();
|
||||
let idx = (line_num - 1) as usize;
|
||||
if idx >= lines.len() {
|
||||
lines.push(content);
|
||||
} else {
|
||||
lines[idx] = content;
|
||||
}
|
||||
lines.join("\n")
|
||||
} else {
|
||||
// Append by default
|
||||
format!("{}{}", existing, content)
|
||||
};
|
||||
|
||||
tokio::fs::write(full_path, &new_content).await?;
|
||||
Ok(ToolOutput::success(format!(
|
||||
"Successfully edited {} ({} bytes)",
|
||||
path,
|
||||
new_content.len()
|
||||
)))
|
||||
}
|
||||
|
||||
async fn search_file(&self, path: &str, pattern: &str) -> ToolResult<ToolOutput> {
|
||||
let full_path = Path::new(path);
|
||||
|
||||
if !full_path.exists() {
|
||||
return Err(ToolError::FileNotFound(path.to_string()));
|
||||
}
|
||||
|
||||
let content = tokio::fs::read_to_string(full_path).await?;
|
||||
let mut matches = Vec::new();
|
||||
|
||||
for (line_num, line) in content.lines().enumerate() {
|
||||
if line.contains(pattern) {
|
||||
matches.push(format!("{}: {}", line_num + 1, line));
|
||||
}
|
||||
}
|
||||
|
||||
if matches.is_empty() {
|
||||
Ok(ToolOutput::success(format!(
|
||||
"No matches found for '{}' in {}",
|
||||
pattern, path
|
||||
)))
|
||||
} else {
|
||||
Ok(ToolOutput::success(format!(
|
||||
"Found {} matches for '{}' in {}:\n\n{}",
|
||||
matches.len(),
|
||||
pattern,
|
||||
path,
|
||||
matches.join("\n")
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_directory(&self, path: &str) -> ToolResult<ToolOutput> {
|
||||
let full_path = Path::new(path);
|
||||
|
||||
if !full_path.exists() {
|
||||
return Err(ToolError::DirectoryNotFound(path.to_string()));
|
||||
}
|
||||
|
||||
if !full_path.is_dir() {
|
||||
return Err(ToolError::InvalidArguments(format!(
|
||||
"{} is not a directory",
|
||||
path
|
||||
)));
|
||||
}
|
||||
|
||||
let mut entries: Vec<String> = Vec::new();
|
||||
let mut read_dir = tokio::fs::read_dir(full_path).await?;
|
||||
|
||||
while let Some(entry) = read_dir.next_entry().await? {
|
||||
let entry_path = entry.path();
|
||||
let file_name = entry_path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("<unknown>");
|
||||
|
||||
let file_type = entry.file_type().await?;
|
||||
let type_str = if file_type.is_dir() {
|
||||
"[DIR]"
|
||||
} else if file_type.is_symlink() {
|
||||
"[LINK]"
|
||||
} else {
|
||||
"[FILE]"
|
||||
};
|
||||
|
||||
entries.push(format!("{} {}", type_str, file_name));
|
||||
}
|
||||
|
||||
entries.sort();
|
||||
|
||||
Ok(ToolOutput::success(format!(
|
||||
"Directory contents of {}:\n\n{}",
|
||||
path,
|
||||
entries.join("\n")
|
||||
)))
|
||||
}
|
||||
}
|
||||
179
src/tools/git_tools.rs
Normal file
179
src/tools/git_tools.rs
Normal file
@@ -0,0 +1,179 @@
|
||||
use super::{Tool, ToolError, ToolOutput, ToolResult};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
/// Git operation tools
|
||||
pub struct GitTools;
|
||||
|
||||
impl GitTools {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for GitTools {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GitTools {
|
||||
fn name(&self) -> &str {
|
||||
"git"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Git version control operations: status, diff, commit, branch, log, and more"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["status", "diff", "commit", "branch", "log", "add", "checkout", "merge", "rebase"],
|
||||
"description": "The git operation to perform"
|
||||
},
|
||||
"args": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Additional arguments for the git command"
|
||||
},
|
||||
"message": {
|
||||
"type": "string",
|
||||
"description": "Commit message (for commit operation)"
|
||||
},
|
||||
"branch_name": {
|
||||
"type": "string",
|
||||
"description": "Branch name (for branch/checkout operations)"
|
||||
}
|
||||
},
|
||||
"required": ["operation"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value) -> ToolResult<ToolOutput> {
|
||||
let operation = args
|
||||
.get("operation")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("Missing 'operation' argument".to_string()))?;
|
||||
|
||||
match operation {
|
||||
"status" => self.run_git(&["status"]).await,
|
||||
"diff" => {
|
||||
let unstaged = args.get("unstaged").and_then(|v| v.as_bool()).unwrap_or(true);
|
||||
if unstaged {
|
||||
self.run_git(&["diff"]).await
|
||||
} else {
|
||||
self.run_git(&["diff", "--staged"]).await
|
||||
}
|
||||
}
|
||||
"commit" => {
|
||||
let message = args
|
||||
.get("message")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("Missing 'message' argument for commit".to_string()))?;
|
||||
self.run_git(&["commit", "-m", message]).await
|
||||
}
|
||||
"branch" => {
|
||||
if let Some(branch_name) = args.get("branch_name").and_then(|v| v.as_str()) {
|
||||
self.run_git(&["branch", branch_name]).await
|
||||
} else {
|
||||
self.run_git(&["branch"]).await
|
||||
}
|
||||
}
|
||||
"log" => {
|
||||
let count = args.get("count").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
|
||||
self.run_git(&["log", "-n", &count.to_string()]).await
|
||||
}
|
||||
"add" => {
|
||||
let paths: Vec<&str> = args
|
||||
.get("paths")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
|
||||
.unwrap_or_else(|| vec!["."]);
|
||||
|
||||
let mut cmd_args = vec!["add"];
|
||||
cmd_args.extend(paths);
|
||||
self.run_git(&cmd_args).await
|
||||
}
|
||||
"checkout" => {
|
||||
if let Some(branch_name) = args.get("branch_name").and_then(|v| v.as_str()) {
|
||||
self.run_git(&["checkout", branch_name]).await
|
||||
} else {
|
||||
Err(ToolError::InvalidArguments("Missing 'branch_name' for checkout".to_string()))
|
||||
}
|
||||
}
|
||||
"merge" => {
|
||||
if let Some(branch_name) = args.get("branch_name").and_then(|v| v.as_str()) {
|
||||
self.run_git(&["merge", branch_name]).await
|
||||
} else {
|
||||
Err(ToolError::InvalidArguments("Missing 'branch_name' for merge".to_string()))
|
||||
}
|
||||
}
|
||||
"rebase" => {
|
||||
if let Some(branch_name) = args.get("branch_name").and_then(|v| v.as_str()) {
|
||||
self.run_git(&["rebase", branch_name]).await
|
||||
} else {
|
||||
Err(ToolError::InvalidArguments("Missing 'branch_name' for rebase".to_string()))
|
||||
}
|
||||
}
|
||||
_ => Err(ToolError::InvalidArguments(format!(
|
||||
"Unknown git operation: {}",
|
||||
operation
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GitTools {
|
||||
async fn run_git(&self, args: &[&str]) -> ToolResult<ToolOutput> {
|
||||
// Check if we're in a git repository
|
||||
let git_check = tokio::process::Command::new("git")
|
||||
.arg("rev-parse")
|
||||
.arg("--git-dir")
|
||||
.output()
|
||||
.await;
|
||||
|
||||
if let Ok(output) = git_check {
|
||||
if !output.status.success() {
|
||||
return Err(ToolError::GitError(
|
||||
"Not in a git repository".to_string(),
|
||||
));
|
||||
}
|
||||
} else {
|
||||
return Err(ToolError::GitError(
|
||||
"Git is not installed or not in PATH".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut cmd = tokio::process::Command::new("git");
|
||||
cmd.args(args);
|
||||
|
||||
let output = cmd.output().await?;
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
let status = output.status;
|
||||
|
||||
let mut result = String::new();
|
||||
result.push_str(&format!("git {}\nExit code: {}\n", args.join(" "), status));
|
||||
|
||||
if !stdout.is_empty() {
|
||||
result.push_str("\n");
|
||||
result.push_str(&stdout);
|
||||
}
|
||||
|
||||
if !stderr.is_empty() {
|
||||
result.push_str("\n--- stderr ---\n");
|
||||
result.push_str(&stderr);
|
||||
}
|
||||
|
||||
if status.success() {
|
||||
Ok(ToolOutput::success(result.trim().to_string()))
|
||||
} else {
|
||||
Ok(ToolOutput::success(format!("Git command failed:\n{}", result)))
|
||||
}
|
||||
}
|
||||
}
|
||||
15
src/tools/mod.rs
Normal file
15
src/tools/mod.rs
Normal file
@@ -0,0 +1,15 @@
|
||||
mod error;
|
||||
mod traits;
|
||||
mod types;
|
||||
mod registry;
|
||||
mod file_tools;
|
||||
mod shell_tools;
|
||||
mod git_tools;
|
||||
|
||||
pub use error::*;
|
||||
pub use traits::*;
|
||||
pub use types::*;
|
||||
pub use registry::ToolRegistry;
|
||||
pub use file_tools::FileTools;
|
||||
pub use shell_tools::ShellTools;
|
||||
pub use git_tools::GitTools;
|
||||
58
src/tools/registry.rs
Normal file
58
src/tools/registry.rs
Normal file
@@ -0,0 +1,58 @@
|
||||
use super::{Tool, ToolDefinition, ToolError, ToolOutput, ToolResult};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Registry for managing tools
|
||||
pub struct ToolRegistry {
|
||||
tools: HashMap<String, Arc<dyn Tool>>,
|
||||
}
|
||||
|
||||
impl ToolRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tools: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a tool
|
||||
pub fn register(&mut self, tool: Arc<dyn Tool>) {
|
||||
self.tools.insert(tool.name().to_string(), tool);
|
||||
}
|
||||
|
||||
/// Get a tool by name
|
||||
pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
|
||||
self.tools.get(name).cloned()
|
||||
}
|
||||
|
||||
/// Execute a tool by name
|
||||
pub async fn execute(&self, name: &str, args: Value) -> ToolResult<ToolOutput> {
|
||||
let tool = self
|
||||
.get(name)
|
||||
.ok_or_else(|| ToolError::NotFound(name.to_string()))?;
|
||||
tool.execute(args).await
|
||||
}
|
||||
|
||||
/// List all registered tools
|
||||
pub fn list(&self) -> Vec<ToolDefinition> {
|
||||
self.tools
|
||||
.values()
|
||||
.map(|tool| ToolDefinition {
|
||||
name: tool.name().to_string(),
|
||||
description: tool.description().to_string(),
|
||||
parameters: tool.parameters(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get tool names
|
||||
pub fn tool_names(&self) -> Vec<&str> {
|
||||
self.tools.keys().map(|s| s.as_str()).collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ToolRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
166
src/tools/shell_tools.rs
Normal file
166
src/tools/shell_tools.rs
Normal file
@@ -0,0 +1,166 @@
|
||||
use super::{Tool, ToolError, ToolOutput, ToolResult};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
/// Shell command execution tools
|
||||
pub struct ShellTools {
|
||||
/// Whether to require confirmation for dangerous commands
|
||||
require_confirmation: bool,
|
||||
}
|
||||
|
||||
impl ShellTools {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
require_confirmation: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_confirmation(require: bool) -> Self {
|
||||
Self {
|
||||
require_confirmation: require,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ShellTools {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
// List of dangerous command patterns that should require confirmation
|
||||
const DANGEROUS_PATTERNS: &[&str] = &[
|
||||
"rm -rf",
|
||||
"rm -r /",
|
||||
"rm -rf /",
|
||||
"dd if=",
|
||||
"mkfs",
|
||||
"chmod -R 777",
|
||||
"chmod -R 000",
|
||||
"sudo rm",
|
||||
"sudo chmod",
|
||||
":(){:|:&};:", // Fork bomb
|
||||
"> /dev/sda",
|
||||
"echo > /dev/sda",
|
||||
];
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ShellTools {
|
||||
fn name(&self) -> &str {
|
||||
"shell"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Execute shell commands with safety checks. Use for running build commands, git operations, and other CLI tasks."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "The shell command to execute"
|
||||
},
|
||||
"working_dir": {
|
||||
"type": "string",
|
||||
"description": "Working directory for command execution"
|
||||
},
|
||||
"timeout": {
|
||||
"type": "integer",
|
||||
"description": "Timeout in milliseconds (default: 60000)"
|
||||
}
|
||||
},
|
||||
"required": ["command"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value) -> ToolResult<ToolOutput> {
|
||||
let command = args
|
||||
.get("command")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidArguments("Missing 'command' argument".to_string()))?;
|
||||
|
||||
let working_dir = args
|
||||
.get("working_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let timeout = args
|
||||
.get("timeout")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(60000);
|
||||
|
||||
// Check for dangerous patterns
|
||||
if self.is_dangerous(command) {
|
||||
return Err(ToolError::PermissionDenied(
|
||||
"Command contains dangerous patterns and requires manual execution".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
self.run_command(command, working_dir.as_deref(), timeout)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl ShellTools {
|
||||
fn is_dangerous(&self, command: &str) -> bool {
|
||||
let cmd_lower = command.to_lowercase();
|
||||
DANGEROUS_PATTERNS.iter().any(|pattern| cmd_lower.contains(pattern))
|
||||
}
|
||||
|
||||
async fn run_command(
|
||||
&self,
|
||||
command: &str,
|
||||
working_dir: Option<&str>,
|
||||
timeout: u64,
|
||||
) -> ToolResult<ToolOutput> {
|
||||
let mut cmd = tokio::process::Command::new("sh");
|
||||
cmd.arg("-c").arg(command);
|
||||
|
||||
if let Some(dir) = working_dir {
|
||||
cmd.current_dir(dir);
|
||||
}
|
||||
|
||||
// Set timeout
|
||||
let timeout_duration = std::time::Duration::from_millis(timeout);
|
||||
|
||||
// Run with timeout
|
||||
let result = tokio::time::timeout(timeout_duration, cmd.output()).await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(output)) => {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
let status = output.status;
|
||||
|
||||
let mut result = String::new();
|
||||
result.push_str(&format!("Exit code: {}\n", status));
|
||||
|
||||
if !stdout.is_empty() {
|
||||
result.push_str("\n--- stdout ---\n");
|
||||
result.push_str(&stdout);
|
||||
}
|
||||
|
||||
if !stderr.is_empty() {
|
||||
result.push_str("\n--- stderr ---\n");
|
||||
result.push_str(&stderr);
|
||||
}
|
||||
|
||||
if status.success() {
|
||||
Ok(ToolOutput::success(result))
|
||||
} else {
|
||||
Ok(ToolOutput::success(format!("Command failed:\n{}", result)))
|
||||
}
|
||||
}
|
||||
Ok(Err(e)) => Err(ToolError::ExecutionFailed(format!(
|
||||
"Failed to execute command: {}",
|
||||
e
|
||||
))),
|
||||
Err(_) => Err(ToolError::ExecutionFailed(format!(
|
||||
"Command timed out after {}ms",
|
||||
timeout
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
48
src/tools/traits.rs
Normal file
48
src/tools/traits.rs
Normal file
@@ -0,0 +1,48 @@
|
||||
use super::ToolError;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
/// Tool output
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolOutput {
|
||||
/// Text content of the output
|
||||
pub content: String,
|
||||
|
||||
/// Whether the output is an error
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
impl ToolOutput {
|
||||
pub fn success(content: impl Into<String>) -> Self {
|
||||
Self {
|
||||
content: content.into(),
|
||||
is_error: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn error(content: impl Into<String>) -> Self {
|
||||
Self {
|
||||
content: content.into(),
|
||||
is_error: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result type for tool operations
|
||||
pub type ToolResult<T> = std::result::Result<T, ToolError>;
|
||||
|
||||
/// Trait for all tools
|
||||
#[async_trait]
|
||||
pub trait Tool: Send + Sync {
|
||||
/// Get the tool name
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// Get the tool description
|
||||
fn description(&self) -> &str;
|
||||
|
||||
/// Get the JSON schema for tool parameters
|
||||
fn parameters(&self) -> Value;
|
||||
|
||||
/// Execute the tool with given arguments
|
||||
async fn execute(&self, args: Value) -> ToolResult<ToolOutput>;
|
||||
}
|
||||
24
src/tools/types.rs
Normal file
24
src/tools/types.rs
Normal file
@@ -0,0 +1,24 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Tool definition for the provider API
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolDefinition {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub parameters: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Tool call request
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCallRequest {
|
||||
pub tool_name: String,
|
||||
pub arguments: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Tool execution result
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCallResult {
|
||||
pub tool_name: String,
|
||||
pub output: String,
|
||||
pub is_error: bool,
|
||||
}
|
||||
Reference in New Issue
Block a user