Refactoring the structure a bit
This commit is contained in:
Dave
378
server/src/llm/chat.rs
Normal file
378
server/src/llm/chat.rs
Normal file
@@ -0,0 +1,378 @@
|
||||
use crate::llm::prompts::SYSTEM_PROMPT;
|
||||
use crate::llm::types::{Message, Role, ToolCall, ToolDefinition, ToolFunctionDefinition};
|
||||
use crate::state::SessionState;
|
||||
use crate::store::StoreOps;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
|
||||
const MAX_TURNS: usize = 30;
|
||||
const KEY_ANTHROPIC_API_KEY: &str = "anthropic_api_key";
|
||||
|
||||
#[derive(Deserialize, Clone)]
|
||||
pub struct ProviderConfig {
|
||||
pub provider: String,
|
||||
pub model: String,
|
||||
pub base_url: Option<String>,
|
||||
pub enable_tools: Option<bool>,
|
||||
}
|
||||
|
||||
fn get_anthropic_api_key_exists_impl(store: &dyn StoreOps) -> bool {
|
||||
match store.get(KEY_ANTHROPIC_API_KEY) {
|
||||
Some(value) => value.as_str().map(|k| !k.is_empty()).unwrap_or(false),
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn set_anthropic_api_key_impl(store: &dyn StoreOps, api_key: &str) -> Result<(), String> {
|
||||
store.set(KEY_ANTHROPIC_API_KEY, json!(api_key));
|
||||
store.save()?;
|
||||
|
||||
match store.get(KEY_ANTHROPIC_API_KEY) {
|
||||
Some(value) => {
|
||||
if let Some(retrieved) = value.as_str() {
|
||||
if retrieved != api_key {
|
||||
return Err("Retrieved key does not match saved key".to_string());
|
||||
}
|
||||
} else {
|
||||
return Err("Stored value is not a string".to_string());
|
||||
}
|
||||
}
|
||||
None => {
|
||||
return Err("API key was saved but cannot be retrieved".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_anthropic_api_key_impl(store: &dyn StoreOps) -> Result<String, String> {
|
||||
match store.get(KEY_ANTHROPIC_API_KEY) {
|
||||
Some(value) => {
|
||||
if let Some(key) = value.as_str() {
|
||||
if key.is_empty() {
|
||||
Err("Anthropic API key is empty. Please set your API key.".to_string())
|
||||
} else {
|
||||
Ok(key.to_string())
|
||||
}
|
||||
} else {
|
||||
Err("Stored API key is not a string".to_string())
|
||||
}
|
||||
}
|
||||
None => Err("Anthropic API key not found. Please set your API key.".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_tool_arguments(args_str: &str) -> Result<serde_json::Value, String> {
|
||||
serde_json::from_str(args_str).map_err(|e| format!("Error parsing arguments: {e}"))
|
||||
}
|
||||
|
||||
pub fn get_tool_definitions() -> Vec<ToolDefinition> {
|
||||
vec![
|
||||
ToolDefinition {
|
||||
kind: "function".to_string(),
|
||||
function: ToolFunctionDefinition {
|
||||
name: "read_file".to_string(),
|
||||
description: "Reads the complete content of a file from the project. Use this to understand existing code before making changes.".to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": { "type": "string", "description": "Relative path to the file from project root" }
|
||||
},
|
||||
"required": ["path"]
|
||||
}),
|
||||
},
|
||||
},
|
||||
ToolDefinition {
|
||||
kind: "function".to_string(),
|
||||
function: ToolFunctionDefinition {
|
||||
name: "write_file".to_string(),
|
||||
description: "Creates or completely overwrites a file with new content. YOU MUST USE THIS to implement code changes - do not suggest code to the user. The content parameter must contain the COMPLETE file including all imports, functions, and unchanged code.".to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": { "type": "string", "description": "Relative path to the file from project root" },
|
||||
"content": { "type": "string", "description": "The complete file content to write (not a diff or partial code)" }
|
||||
},
|
||||
"required": ["path", "content"]
|
||||
}),
|
||||
},
|
||||
},
|
||||
ToolDefinition {
|
||||
kind: "function".to_string(),
|
||||
function: ToolFunctionDefinition {
|
||||
name: "list_directory".to_string(),
|
||||
description: "Lists all files and directories at a given path. Use this to explore the project structure.".to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": { "type": "string", "description": "Relative path to list (use '.' for project root)" }
|
||||
},
|
||||
"required": ["path"]
|
||||
}),
|
||||
},
|
||||
},
|
||||
ToolDefinition {
|
||||
kind: "function".to_string(),
|
||||
function: ToolFunctionDefinition {
|
||||
name: "search_files".to_string(),
|
||||
description: "Searches for text patterns across all files in the project. Use this to find functions, variables, or code patterns when you don't know which file they're in."
|
||||
.to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": { "type": "string", "description": "The text pattern to search for across all files" }
|
||||
},
|
||||
"required": ["query"]
|
||||
}),
|
||||
},
|
||||
},
|
||||
ToolDefinition {
|
||||
kind: "function".to_string(),
|
||||
function: ToolFunctionDefinition {
|
||||
name: "exec_shell".to_string(),
|
||||
description: "Executes a shell command in the project root directory. Use this to run tests, build commands, git operations, or any command-line tool. Examples: cargo check, npm test, git status.".to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "The command binary to execute (e.g., 'git', 'cargo', 'npm', 'ls')"
|
||||
},
|
||||
"args": {
|
||||
"type": "array",
|
||||
"items": { "type": "string" },
|
||||
"description": "Array of arguments to pass to the command (e.g., ['status'] for git status)"
|
||||
}
|
||||
},
|
||||
"required": ["command", "args"]
|
||||
}),
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
pub async fn get_ollama_models(base_url: Option<String>) -> Result<Vec<String>, String> {
|
||||
use crate::llm::providers::ollama::OllamaProvider;
|
||||
let url = base_url.unwrap_or_else(|| "http://localhost:11434".to_string());
|
||||
OllamaProvider::get_models(&url).await
|
||||
}
|
||||
|
||||
pub fn get_anthropic_api_key_exists(store: &dyn StoreOps) -> Result<bool, String> {
|
||||
Ok(get_anthropic_api_key_exists_impl(store))
|
||||
}
|
||||
|
||||
pub fn set_anthropic_api_key(store: &dyn StoreOps, api_key: String) -> Result<(), String> {
|
||||
set_anthropic_api_key_impl(store, &api_key)
|
||||
}
|
||||
|
||||
pub async fn chat<F, U>(
|
||||
messages: Vec<Message>,
|
||||
config: ProviderConfig,
|
||||
state: &SessionState,
|
||||
store: &dyn StoreOps,
|
||||
mut on_update: F,
|
||||
mut on_token: U,
|
||||
) -> Result<Vec<Message>, String>
|
||||
where
|
||||
F: FnMut(&[Message]) + Send,
|
||||
U: FnMut(&str) + Send,
|
||||
{
|
||||
use crate::llm::providers::anthropic::AnthropicProvider;
|
||||
use crate::llm::providers::ollama::OllamaProvider;
|
||||
|
||||
let _ = state.cancel_tx.send(false);
|
||||
let mut cancel_rx = state.cancel_rx.clone();
|
||||
cancel_rx.borrow_and_update();
|
||||
|
||||
let base_url = config
|
||||
.base_url
|
||||
.clone()
|
||||
.unwrap_or_else(|| "http://localhost:11434".to_string());
|
||||
|
||||
let is_claude = config.model.starts_with("claude-");
|
||||
|
||||
if !is_claude && config.provider.as_str() != "ollama" {
|
||||
return Err(format!("Unsupported provider: {}", config.provider));
|
||||
}
|
||||
|
||||
let tool_defs = get_tool_definitions();
|
||||
let tools = if config.enable_tools.unwrap_or(true) {
|
||||
tool_defs.as_slice()
|
||||
} else {
|
||||
&[]
|
||||
};
|
||||
|
||||
let mut current_history = messages.clone();
|
||||
|
||||
current_history.insert(
|
||||
0,
|
||||
Message {
|
||||
role: Role::System,
|
||||
content: SYSTEM_PROMPT.to_string(),
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
},
|
||||
);
|
||||
|
||||
current_history.insert(
|
||||
1,
|
||||
Message {
|
||||
role: Role::System,
|
||||
content: "REMINDER: Distinguish between showing examples (use code blocks in chat) vs implementing changes (use write_file tool). Keywords like 'show me', 'example', 'how does' = chat response. Keywords like 'create', 'add', 'implement', 'fix' = use tools."
|
||||
.to_string(),
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
},
|
||||
);
|
||||
|
||||
let mut new_messages: Vec<Message> = Vec::new();
|
||||
let mut turn_count = 0;
|
||||
|
||||
loop {
|
||||
if *cancel_rx.borrow() {
|
||||
return Err("Chat cancelled by user".to_string());
|
||||
}
|
||||
|
||||
if turn_count >= MAX_TURNS {
|
||||
return Err("Max conversation turns reached.".to_string());
|
||||
}
|
||||
turn_count += 1;
|
||||
|
||||
let response = if is_claude {
|
||||
let api_key = get_anthropic_api_key_impl(store)?;
|
||||
let anthropic_provider = AnthropicProvider::new(api_key);
|
||||
anthropic_provider
|
||||
.chat_stream(
|
||||
&config.model,
|
||||
¤t_history,
|
||||
tools,
|
||||
&mut cancel_rx,
|
||||
|token| on_token(token),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| format!("Anthropic Error: {e}"))?
|
||||
} else {
|
||||
let ollama_provider = OllamaProvider::new(base_url.clone());
|
||||
ollama_provider
|
||||
.chat_stream(
|
||||
&config.model,
|
||||
¤t_history,
|
||||
tools,
|
||||
&mut cancel_rx,
|
||||
|token| on_token(token),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| format!("Ollama Error: {e}"))?
|
||||
};
|
||||
|
||||
if let Some(tool_calls) = response.tool_calls {
|
||||
let assistant_msg = Message {
|
||||
role: Role::Assistant,
|
||||
content: response.content.unwrap_or_default(),
|
||||
tool_calls: Some(tool_calls.clone()),
|
||||
tool_call_id: None,
|
||||
};
|
||||
|
||||
current_history.push(assistant_msg.clone());
|
||||
new_messages.push(assistant_msg);
|
||||
on_update(¤t_history[2..]);
|
||||
|
||||
for call in tool_calls {
|
||||
if *cancel_rx.borrow() {
|
||||
return Err("Chat cancelled before tool execution".to_string());
|
||||
}
|
||||
|
||||
let output = execute_tool(&call, state).await;
|
||||
|
||||
let tool_msg = Message {
|
||||
role: Role::Tool,
|
||||
content: output,
|
||||
tool_calls: None,
|
||||
tool_call_id: call.id,
|
||||
};
|
||||
|
||||
current_history.push(tool_msg.clone());
|
||||
new_messages.push(tool_msg);
|
||||
on_update(¤t_history[2..]);
|
||||
}
|
||||
} else {
|
||||
let assistant_msg = Message {
|
||||
role: Role::Assistant,
|
||||
content: response.content.unwrap_or_default(),
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
};
|
||||
|
||||
new_messages.push(assistant_msg.clone());
|
||||
current_history.push(assistant_msg);
|
||||
on_update(¤t_history[2..]);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(new_messages)
|
||||
}
|
||||
|
||||
async fn execute_tool(call: &ToolCall, state: &SessionState) -> String {
|
||||
use crate::io::{fs, search, shell};
|
||||
|
||||
let name = call.function.name.as_str();
|
||||
let args: serde_json::Value = match parse_tool_arguments(&call.function.arguments) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return e,
|
||||
};
|
||||
|
||||
match name {
|
||||
"read_file" => {
|
||||
let path = args["path"].as_str().unwrap_or("").to_string();
|
||||
match fs::read_file(path, state).await {
|
||||
Ok(content) => content,
|
||||
Err(e) => format!("Error: {e}"),
|
||||
}
|
||||
}
|
||||
"write_file" => {
|
||||
let path = args["path"].as_str().unwrap_or("").to_string();
|
||||
let content = args["content"].as_str().unwrap_or("").to_string();
|
||||
match fs::write_file(path, content, state).await {
|
||||
Ok(()) => "File written successfully.".to_string(),
|
||||
Err(e) => format!("Error: {e}"),
|
||||
}
|
||||
}
|
||||
"list_directory" => {
|
||||
let path = args["path"].as_str().unwrap_or("").to_string();
|
||||
match fs::list_directory(path, state).await {
|
||||
Ok(entries) => serde_json::to_string(&entries).unwrap_or_default(),
|
||||
Err(e) => format!("Error: {e}"),
|
||||
}
|
||||
}
|
||||
"search_files" => {
|
||||
let query = args["query"].as_str().unwrap_or("").to_string();
|
||||
match search::search_files(query, state).await {
|
||||
Ok(results) => serde_json::to_string(&results).unwrap_or_default(),
|
||||
Err(e) => format!("Error: {e}"),
|
||||
}
|
||||
}
|
||||
"exec_shell" => {
|
||||
let command = args["command"].as_str().unwrap_or("").to_string();
|
||||
let args_vec: Vec<String> = args["args"]
|
||||
.as_array()
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.map(|v| v.as_str().unwrap_or("").to_string())
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
match shell::exec_shell(command, args_vec, state).await {
|
||||
Ok(output) => serde_json::to_string(&output).unwrap_or_default(),
|
||||
Err(e) => format!("Error: {e}"),
|
||||
}
|
||||
}
|
||||
_ => format!("Unknown tool: {name}"),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cancel_chat(state: &SessionState) -> Result<(), String> {
|
||||
state.cancel_tx.send(true).map_err(|e| e.to_string())?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod chat;
|
||||
pub mod prompts;
|
||||
pub mod providers;
|
||||
pub mod types;
|
||||
|
||||
Reference in New Issue
Block a user