diff --git a/CLAUDE.md b/CLAUDE.md index 32a89fec..ba2887cb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -211,6 +211,19 @@ NEARAI_BASE_URL=https://private.near.ai # Agent settings AGENT_NAME=near-agent MAX_PARALLEL_JOBS=5 + +# Embeddings (for semantic memory search) +OPENAI_API_KEY=sk-... # For OpenAI embeddings +# Or use NEAR AI embeddings: +# EMBEDDING_PROVIDER=nearai +# EMBEDDING_ENABLED=true +EMBEDDING_MODEL=text-embedding-3-small # or text-embedding-3-large + +# Heartbeat (proactive periodic execution) +HEARTBEAT_ENABLED=true +HEARTBEAT_INTERVAL_SECS=1800 # 30 minutes +HEARTBEAT_NOTIFY_CHANNEL=tui +HEARTBEAT_NOTIFY_USER=default ``` ### NEAR AI Provider @@ -276,8 +289,9 @@ Key test patterns: 2. **Domain-specific tools** - `marketplace.rs`, `restaurant.rs`, `taskrabbit.rs`, `ecommerce.rs` return placeholder responses; need real API integrations 3. **Integration tests** - Need testcontainers setup for PostgreSQL 4. **MCP stdio transport** - Only HTTP transport implemented -5. **Embedding backfill** - Background job to generate embeddings for chunks missing them -6. **Auto-context compaction** - Context monitor exists but doesn't auto-trigger (requires manual `/compact`) +5. **WIT bindgen integration** - Auto-extract tool description/schema from WASM modules (stubbed) +6. **Capability granting after tool build** - Built tools get empty capabilities; need UX for granting HTTP/secrets access +7. **Tool versioning workflow** - No version tracking or rollback for dynamically built tools ### Completed @@ -285,11 +299,15 @@ Key test patterns: - ✅ **WASM sandboxing** - Full implementation in `tools/wasm/` with fuel metering, memory limits, capabilities - ✅ **Dynamic tool building** - `tools/builder/` has LlmSoftwareBuilder with iterative build loop - ✅ **HTTP webhook security** - Secret validation implemented, proper error handling (no panics) - -### Known Clippy Warnings (non-blocking) - -- `too_many_arguments` on Agent::new, Worker::new, Store::record_llm_call (refactor to config structs if desired) -- `implicit_saturating_sub` in CLI render (use `.saturating_sub()` instead of manual check) +- ✅ **Embeddings integration** - OpenAI and NEAR AI providers wired to workspace for semantic search +- ✅ **Workspace system prompt** - Identity files (AGENTS.md, SOUL.md, USER.md, IDENTITY.md) injected into LLM context +- ✅ **Heartbeat notifications** - Route through channel manager (broadcast API) instead of logging-only +- ✅ **Auto-context compaction** - Triggers automatically when context exceeds threshold +- ✅ **Embedding backfill** - Runs on startup when embeddings provider is enabled +- ✅ **Clippy clean** - All warnings addressed via config struct refactoring +- ✅ **Tool approval enforcement** - Tools with `requires_approval()` (shell, http, file write/patch, build_software) now gate execution, track auto-approved tools per session +- ✅ **Tool definition refresh** - Tool definitions refreshed each iteration so newly built tools become visible in same session +- ✅ **Worker tool call handling** - Uses `respond_with_tools()` to properly execute tool calls when `select_tools()` returns empty ## Adding a New Tool diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 9c61b671..41d13235 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -10,7 +10,7 @@ use crate::agent::compaction::ContextCompactor; use crate::agent::context_monitor::ContextMonitor; use crate::agent::heartbeat::spawn_heartbeat; use crate::agent::self_repair::DefaultSelfRepair; -use crate::agent::session::{Session, ThreadState}; +use crate::agent::session::{PendingApproval, Session, ThreadState}; use crate::agent::session_manager::SessionManager; use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult}; use crate::agent::{ @@ -27,20 +27,38 @@ use crate::safety::SafetyLayer; use crate::tools::ToolRegistry; use crate::workspace::Workspace; +/// Result of the agentic loop execution. +enum AgenticLoopResult { + /// Completed with a response. + Response(String), + /// A tool requires approval before continuing. + NeedApproval { + /// The pending approval request to store. + pending: PendingApproval, + }, +} + +/// Core dependencies for the agent. +/// +/// Bundles the shared components to reduce argument count. +pub struct AgentDeps { + pub store: Option>, + pub llm: Arc, + pub safety: Arc, + pub tools: Arc, + pub workspace: Option>, +} + /// The main agent that coordinates all components. pub struct Agent { config: AgentConfig, - store: Option>, - llm: Arc, - safety: Arc, - tools: Arc, - channels: ChannelManager, + deps: AgentDeps, + channels: Arc, context_manager: Arc, scheduler: Arc, router: Router, session_manager: Arc, context_monitor: ContextMonitor, - workspace: Option>, heartbeat_config: Option, } @@ -48,12 +66,8 @@ impl Agent { /// Create a new agent. pub fn new( config: AgentConfig, - store: Option>, - llm: Arc, - safety: Arc, - tools: Arc, + deps: AgentDeps, channels: ChannelManager, - workspace: Option>, heartbeat_config: Option, ) -> Self { let context_manager = Arc::new(ContextManager::new(config.max_parallel_jobs)); @@ -61,29 +75,46 @@ impl Agent { let scheduler = Arc::new(Scheduler::new( config.clone(), context_manager.clone(), - llm.clone(), - safety.clone(), - tools.clone(), - store.clone(), + deps.llm.clone(), + deps.safety.clone(), + deps.tools.clone(), + deps.store.clone(), )); Self { config, - store, - llm, - safety, - tools, - channels, + deps, + channels: Arc::new(channels), context_manager, scheduler, router: Router::new(), session_manager: Arc::new(SessionManager::new()), context_monitor: ContextMonitor::new(), - workspace, heartbeat_config, } } + // Convenience accessors + fn store(&self) -> Option<&Arc> { + self.deps.store.as_ref() + } + + fn llm(&self) -> &Arc { + &self.deps.llm + } + + fn safety(&self) -> &Arc { + &self.deps.safety + } + + fn tools(&self) -> &Arc { + &self.deps.tools + } + + fn workspace(&self) -> Option<&Arc> { + self.deps.workspace.as_ref() + } + /// Run the agent main loop. pub async fn run(self) -> Result<(), Error> { // Start channels @@ -104,30 +135,61 @@ impl Agent { // Spawn heartbeat if enabled let heartbeat_handle = if let Some(ref hb_config) = self.heartbeat_config { if hb_config.enabled { - if let Some(ref workspace) = self.workspace { + if let Some(workspace) = self.workspace() { let config = AgentHeartbeatConfig::default() .with_interval(std::time::Duration::from_secs(hb_config.interval_secs)); - // Set up notification channel if configured + // Set up notification channel let (notify_tx, mut notify_rx) = tokio::sync::mpsc::channel::(16); - // Spawn notification forwarder - // We can't clone ChannelManager directly, so we just log the notifications - // The heartbeat system will handle notifications via the response_tx + // Spawn notification forwarder that routes through channel manager let notify_channel = hb_config.notify_channel.clone(); let notify_user = hb_config.notify_user.clone(); + let channels = self.channels.clone(); tokio::spawn(async move { while let Some(response) = notify_rx.recv().await { - if let (Some(ch), Some(user)) = (¬ify_channel, ¬ify_user) { - // Log the heartbeat notification - // In a full implementation, we'd route this through a shared channel reference - tracing::info!( - "Heartbeat notification for {}/{}: {}", - ch, - user, - &response.content - ); + // Route notification to configured channel/user, or broadcast to all + match (¬ify_channel, ¬ify_user) { + (Some(channel), Some(user)) => { + // Send to specific channel and user + if let Err(e) = + channels.broadcast(channel, user, response.clone()).await + { + tracing::warn!( + "Failed to send heartbeat to {}/{}: {}", + channel, + user, + e + ); + } else { + tracing::debug!( + "Heartbeat notification sent to {}/{}", + channel, + user + ); + } + } + (None, Some(user)) => { + // Broadcast to all channels for this user + let results = channels.broadcast_all(user, response).await; + for (ch, result) in results { + if let Err(e) = result { + tracing::warn!( + "Failed to broadcast heartbeat to {}: {}", + ch, + e + ); + } + } + } + _ => { + // No target configured, just log + tracing::info!( + "Heartbeat notification (no target configured): {}", + &response.content + ); + } } } }); @@ -139,7 +201,7 @@ impl Agent { Some(spawn_heartbeat( config, workspace.clone(), - self.llm.clone(), + self.llm().clone(), Some(notify_tx), )) } else { @@ -230,11 +292,24 @@ impl Agent { Submission::Resume { checkpoint_id } => { self.process_resume(session, thread_id, checkpoint_id).await } - Submission::ExecApproval { .. } => { - // Not supported in simple chat flow - Ok(SubmissionResult::error( - "Approval flow not supported in this context", - )) + Submission::ExecApproval { + request_id, + approved, + always, + } => { + self.process_approval( + message, + session, + thread_id, + Some(request_id), + approved, + always, + ) + .await + } + Submission::ApprovalResponse { approved, always } => { + self.process_approval(message, session, thread_id, None, approved, always) + .await } }; @@ -244,8 +319,32 @@ impl Agent { SubmissionResult::Ok { message } => Ok(message), SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))), SubmissionResult::Interrupted => Ok(Some("Interrupted.".into())), - SubmissionResult::NeedApproval { .. } => { - Ok(Some("Approval required but not supported.".into())) + SubmissionResult::NeedApproval { + request_id, + tool_name, + description, + parameters, + } => { + // Format approval request for user + let params_preview = serde_json::to_string_pretty(¶meters) + .unwrap_or_else(|_| parameters.to_string()); + let params_truncated = if params_preview.len() > 200 { + format!("{}...", ¶ms_preview[..200]) + } else { + params_preview + }; + Ok(Some(format!( + "🔒 Tool requires approval:\n\n\ + **Tool:** {}\n\ + **Description:** {}\n\ + **Parameters:** ```\n{}\n```\n\n\ + Reply with:\n\ + - `yes` or `approve` to allow this tool\n\ + - `always` to always allow this tool in this session\n\ + - `no` or `deny` to reject\n\n\ + Request ID: {}", + tool_name, description, params_truncated, request_id + ))) } } } @@ -322,9 +421,9 @@ impl Agent { "Context at {:.1}% capacity, auto-compacting", self.context_monitor.usage_percent(&messages) ); - let compactor = ContextCompactor::new(self.llm.clone()); + let compactor = ContextCompactor::new(self.llm().clone()); if let Err(e) = compactor - .compact(thread, strategy, self.workspace.as_deref()) + .compact(thread, strategy, self.workspace().map(|w| w.as_ref())) .await { tracing::warn!("Auto-compaction failed: {}", e); @@ -389,9 +488,9 @@ impl Agent { return Ok(SubmissionResult::Interrupted); } - // Complete or fail the turn + // Complete, fail, or request approval match result { - Ok(response) => { + Ok(AgenticLoopResult::Response(response)) => { thread.complete_turn(&response); let _ = self .channels @@ -399,6 +498,27 @@ impl Agent { .await; Ok(SubmissionResult::response(response)) } + Ok(AgenticLoopResult::NeedApproval { pending }) => { + // Store pending approval in thread and update state + let request_id = pending.request_id; + let tool_name = pending.tool_name.clone(); + let description = pending.description.clone(); + let parameters = pending.parameters.clone(); + thread.await_approval(pending); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Awaiting approval".into()), + ) + .await; + Ok(SubmissionResult::NeedApproval { + request_id, + tool_name, + description, + parameters, + }) + } Err(e) => { thread.fail_turn(e.to_string()); Ok(SubmissionResult::error(e.to_string())) @@ -407,15 +527,34 @@ impl Agent { } /// Run the agentic loop: call LLM, execute tools, repeat until text response. + /// + /// Returns `AgenticLoopResult::Response` on completion, or + /// `AgenticLoopResult::NeedApproval` if a tool requires user approval. async fn run_agentic_loop( &self, message: &IncomingMessage, session: Arc>, thread_id: Uuid, initial_messages: Vec, - ) -> Result { - let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone()); - let tool_defs = self.tools.tool_definitions().await; + ) -> Result { + // Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.) + let system_prompt = if let Some(ws) = self.workspace() { + match ws.system_prompt().await { + Ok(prompt) if !prompt.is_empty() => Some(prompt), + Ok(_) => None, + Err(e) => { + tracing::debug!("Could not load workspace system prompt: {}", e); + None + } + } + } else { + None + }; + + let mut reasoning = Reasoning::new(self.llm().clone(), self.safety().clone()); + if let Some(prompt) = system_prompt { + reasoning = reasoning.with_system_prompt(prompt); + } // Build context with messages that we'll mutate during the loop let mut context_messages = initial_messages; @@ -425,6 +564,7 @@ impl Agent { const MAX_TOOL_ITERATIONS: usize = 10; let mut iteration = 0; + let mut tools_executed = false; loop { iteration += 1; @@ -450,19 +590,38 @@ impl Agent { } } + // Refresh tool definitions each iteration so newly built tools become visible + let tool_defs = self.tools().tool_definitions().await; + // Call LLM with current context let context = ReasoningContext::new() .with_messages(context_messages.clone()) - .with_tools(tool_defs.clone()); + .with_tools(tool_defs); let result = reasoning.respond_with_tools(&context).await?; match result { RespondResult::Text(text) => { - // Final response, return it - return Ok(text); + // If no tools have been executed yet, prompt the LLM to use tools + // This handles the case where the model explains what it will do + // instead of actually calling tools + if !tools_executed && iteration < 3 { + tracing::debug!( + "No tools executed yet (iteration {}), prompting for tool use", + iteration + ); + context_messages.push(ChatMessage::assistant(&text)); + context_messages.push(ChatMessage::user( + "Please proceed and use the available tools to complete this task.", + )); + continue; + } + + // Tools have been executed or we've tried multiple times, return response + return Ok(AgenticLoopResult::Response(text)); } RespondResult::ToolCalls(tool_calls) => { + tools_executed = true; // Execute tools and add results to context let _ = self .channels @@ -487,8 +646,33 @@ impl Agent { } } - // Execute each tool + // Execute each tool (with approval checking) for tc in tool_calls { + // Check if tool requires approval + if let Some(tool) = self.tools().get(&tc.name).await { + if tool.requires_approval() { + // Check if auto-approved for this session + let is_auto_approved = { + let sess = session.lock().await; + sess.is_tool_auto_approved(&tc.name) + }; + + if !is_auto_approved { + // Need approval - store pending request and return + let pending = PendingApproval { + request_id: Uuid::new_v4(), + tool_name: tc.name.clone(), + parameters: tc.arguments.clone(), + description: tool.description().to_string(), + tool_call_id: tc.id.clone(), + context_messages: context_messages.clone(), + }; + + return Ok(AgenticLoopResult::NeedApproval { pending }); + } + } + } + let tool_result = self .execute_chat_tool(&tc.name, &tc.arguments, &job_ctx) .await; @@ -514,8 +698,9 @@ impl Agent { let result_content = match tool_result { Ok(output) => { // Sanitize output before showing to LLM - let sanitized = self.safety.sanitize_tool_output(&tc.name, &output); - self.safety.wrap_for_llm( + let sanitized = + self.safety().sanitize_tool_output(&tc.name, &output); + self.safety().wrap_for_llm( &tc.name, &sanitized.content, sanitized.was_modified, @@ -543,7 +728,7 @@ impl Agent { job_ctx: &JobContext, ) -> Result { let tool = - self.tools + self.tools() .get(tool_name) .await .ok_or_else(|| crate::error::ToolError::NotFound { @@ -718,9 +903,9 @@ impl Agent { crate::agent::context_monitor::CompactionStrategy::Summarize { keep_recent: 5 }, ); - let compactor = ContextCompactor::new(self.llm.clone()); + let compactor = ContextCompactor::new(self.llm().clone()); match compactor - .compact(thread, strategy, self.workspace.as_deref()) + .compact(thread, strategy, self.workspace().map(|w| w.as_ref())) .await { Ok(result) => { @@ -757,6 +942,190 @@ impl Agent { Ok(SubmissionResult::ok_with_message("Thread cleared.")) } + /// Process an approval or rejection of a pending tool execution. + async fn process_approval( + &self, + message: &IncomingMessage, + session: Arc>, + thread_id: Uuid, + request_id: Option, + approved: bool, + always: bool, + ) -> Result { + // Get thread state and pending approval + let (_thread_state, pending) = { + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + if thread.state != ThreadState::AwaitingApproval { + return Ok(SubmissionResult::error("No pending approval request.")); + } + + let pending = thread.take_pending_approval(); + (thread.state, pending) + }; + + let pending = match pending { + Some(p) => p, + None => return Ok(SubmissionResult::error("No pending approval request.")), + }; + + // Verify request ID if provided + if let Some(req_id) = request_id { + if req_id != pending.request_id { + // Put it back and return error + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.await_approval(pending); + } + return Ok(SubmissionResult::error( + "Request ID mismatch. Use the correct request ID.", + )); + } + } + + if approved { + // If always, add to auto-approved set + if always { + let mut sess = session.lock().await; + sess.auto_approve_tool(&pending.tool_name); + tracing::info!( + "Auto-approved tool '{}' for session {}", + pending.tool_name, + sess.id + ); + } + + // Reset thread state to processing + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.state = ThreadState::Processing; + } + } + + // Execute the approved tool and continue the loop + let job_ctx = + JobContext::with_user(&message.user_id, "chat", "Interactive chat session"); + + let tool_result = self + .execute_chat_tool(&pending.tool_name, &pending.parameters, &job_ctx) + .await; + + // Build context including the tool result + let mut context_messages = pending.context_messages; + + // Record result in thread + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + if let Some(turn) = thread.last_turn_mut() { + match &tool_result { + Ok(output) => { + turn.record_tool_result(serde_json::json!(output)); + } + Err(e) => { + turn.record_tool_error(e.to_string()); + } + } + } + } + } + + // Add tool result to context + let result_content = match tool_result { + Ok(output) => { + let sanitized = self + .safety() + .sanitize_tool_output(&pending.tool_name, &output); + self.safety().wrap_for_llm( + &pending.tool_name, + &sanitized.content, + sanitized.was_modified, + ) + } + Err(e) => format!("Error: {}", e), + }; + + context_messages.push(ChatMessage::tool_result( + &pending.tool_call_id, + &pending.tool_name, + result_content, + )); + + // Continue the agentic loop + let result = self + .run_agentic_loop(message, session.clone(), thread_id, context_messages) + .await; + + // Handle the result + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + match result { + Ok(AgenticLoopResult::Response(response)) => { + thread.complete_turn(&response); + let _ = self + .channels + .send_status(&message.channel, StatusUpdate::Status("Done".into())) + .await; + Ok(SubmissionResult::response(response)) + } + Ok(AgenticLoopResult::NeedApproval { + pending: new_pending, + }) => { + let request_id = new_pending.request_id; + let tool_name = new_pending.tool_name.clone(); + let description = new_pending.description.clone(); + let parameters = new_pending.parameters.clone(); + thread.await_approval(new_pending); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Awaiting approval".into()), + ) + .await; + Ok(SubmissionResult::NeedApproval { + request_id, + tool_name, + description, + parameters, + }) + } + Err(e) => { + thread.fail_turn(e.to_string()); + Ok(SubmissionResult::error(e.to_string())) + } + } + } else { + // Rejected - clear approval and return to idle + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.clear_pending_approval(); + } + } + + let _ = self + .channels + .send_status(&message.channel, StatusUpdate::Status("Rejected".into())) + .await; + + Ok(SubmissionResult::response(format!( + "Tool '{}' was rejected. The agent will not execute this tool.\n\n\ + You can continue the conversation or try a different approach.", + pending.tool_name + ))) + } + } + async fn process_new_thread( &self, message: &IncomingMessage, @@ -842,7 +1211,7 @@ impl Agent { } // Persist new job to database (fire-and-forget) - if let Some(ref store) = self.store { + if let Some(store) = self.store() { if let Ok(ctx) = self.context_manager.get_context(job_id).await { let store = store.clone(); tokio::spawn(async move { @@ -990,7 +1359,7 @@ impl Agent { ))), "tools" => { - let tools = self.tools.list().await; + let tools = self.tools().list().await; Ok(Some(format!("Available tools: {}", tools.join(", ")))) } diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 7b9b506d..2467eb24 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -21,18 +21,18 @@ mod session_manager; pub mod submission; pub mod task; pub mod undo; -mod worker; +pub mod worker; -pub use agent_loop::Agent; +pub use agent_loop::{Agent, AgentDeps}; pub use compaction::{CompactionResult, ContextCompactor}; pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor}; pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat}; pub use router::{MessageIntent, Router}; pub use scheduler::Scheduler; pub use self_repair::{BrokenTool, RepairResult, RepairTask, SelfRepair, StuckJob}; -pub use session::{Session, Thread, ThreadState, Turn, TurnState}; +pub use session::{PendingApproval, Session, Thread, ThreadState, Turn, TurnState}; pub use session_manager::SessionManager; pub use submission::{Submission, SubmissionParser, SubmissionResult}; pub use task::{Task, TaskContext, TaskHandler, TaskOutput, TaskStatus}; pub use undo::{Checkpoint, UndoManager}; -pub use worker::Worker; +pub use worker::{Worker, WorkerDeps}; diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 2e7cd2fd..6c3adf8b 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -8,8 +8,8 @@ use tokio::sync::{RwLock, mpsc, oneshot}; use tokio::task::JoinHandle; use uuid::Uuid; -use crate::agent::Worker; use crate::agent::task::{Task, TaskContext, TaskOutput}; +use crate::agent::worker::{Worker, WorkerDeps}; use crate::config::AgentConfig; use crate::context::{ContextManager, JobContext, JobState}; use crate::error::{Error, JobError}; @@ -109,17 +109,17 @@ impl Scheduler { // Create worker channel let (tx, rx) = mpsc::channel(16); - // Create worker - let worker = Worker::new( - job_id, - self.context_manager.clone(), - self.llm.clone(), - self.safety.clone(), - self.tools.clone(), - self.store.clone(), - self.config.job_timeout, - self.config.use_planning, - ); + // Create worker with shared dependencies + let deps = WorkerDeps { + context_manager: self.context_manager.clone(), + llm: self.llm.clone(), + safety: self.safety.clone(), + tools: self.tools.clone(), + store: self.store.clone(), + timeout: self.config.job_timeout, + use_planning: self.config.use_planning, + }; + let worker = Worker::new(job_id, deps); // Spawn worker task let handle = tokio::spawn(async move { diff --git a/src/agent/session.rs b/src/agent/session.rs index 70846be1..3cf9faf0 100644 --- a/src/agent/session.rs +++ b/src/agent/session.rs @@ -10,7 +10,7 @@ //! - Compaction: Summarize old turns to save context //! - Resume: Continue from a saved checkpoint -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; @@ -35,6 +35,9 @@ pub struct Session { pub last_active_at: DateTime, /// Session metadata. pub metadata: serde_json::Value, + /// Tools that have been auto-approved for this session ("always approve"). + #[serde(default)] + pub auto_approved_tools: HashSet, } impl Session { @@ -49,9 +52,20 @@ impl Session { created_at: now, last_active_at: now, metadata: serde_json::Value::Null, + auto_approved_tools: HashSet::new(), } } + /// Check if a tool has been auto-approved for this session. + pub fn is_tool_auto_approved(&self, tool_name: &str) -> bool { + self.auto_approved_tools.contains(tool_name) + } + + /// Add a tool to the auto-approved set. + pub fn auto_approve_tool(&mut self, tool_name: impl Into) { + self.auto_approved_tools.insert(tool_name.into()); + } + /// Create a new thread in this session. pub fn create_thread(&mut self) -> &mut Thread { let thread = Thread::new(self.id); @@ -107,6 +121,23 @@ pub enum ThreadState { Interrupted, } +/// Pending tool approval request stored on a thread. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PendingApproval { + /// Unique request ID. + pub request_id: Uuid, + /// Tool name requiring approval. + pub tool_name: String, + /// Tool parameters. + pub parameters: serde_json::Value, + /// Description of what the tool will do. + pub description: String, + /// Tool call ID from LLM (for proper context continuation). + pub tool_call_id: String, + /// Context messages at the time of the request (to resume from). + pub context_messages: Vec, +} + /// A conversation thread within a session. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Thread { @@ -124,6 +155,9 @@ pub struct Thread { pub updated_at: DateTime, /// Thread metadata (e.g., title, tags). pub metadata: serde_json::Value, + /// Pending approval request (when state is AwaitingApproval). + #[serde(default)] + pub pending_approval: Option, } impl Thread { @@ -138,6 +172,7 @@ impl Thread { created_at: now, updated_at: now, metadata: serde_json::Value::Null, + pending_approval: None, } } @@ -184,9 +219,22 @@ impl Thread { self.updated_at = Utc::now(); } - /// Mark the thread as awaiting approval. - pub fn await_approval(&mut self) { + /// Mark the thread as awaiting approval with pending request details. + pub fn await_approval(&mut self, pending: PendingApproval) { self.state = ThreadState::AwaitingApproval; + self.pending_approval = Some(pending); + self.updated_at = Utc::now(); + } + + /// Take the pending approval (clearing it from the thread). + pub fn take_pending_approval(&mut self) -> Option { + self.pending_approval.take() + } + + /// Clear pending approval and return to idle state. + pub fn clear_pending_approval(&mut self) { + self.pending_approval = None; + self.state = ThreadState::Idle; self.updated_at = Utc::now(); } diff --git a/src/agent/submission.rs b/src/agent/submission.rs index 15598eae..f79c0b03 100644 --- a/src/agent/submission.rs +++ b/src/agent/submission.rs @@ -52,6 +52,30 @@ impl SubmissionParser { } } + // Approval responses (simple yes/no/always for pending approvals) + // These are short enough to check explicitly + match lower.as_str() { + "yes" | "y" | "approve" | "ok" => { + return Submission::ApprovalResponse { + approved: true, + always: false, + }; + } + "always" | "yes always" | "approve always" => { + return Submission::ApprovalResponse { + approved: true, + always: true, + }; + } + "no" | "n" | "deny" | "reject" | "cancel" => { + return Submission::ApprovalResponse { + approved: false, + always: false, + }; + } + _ => {} + } + // Default: user input Submission::UserInput { content: content.to_string(), @@ -68,7 +92,7 @@ pub enum Submission { content: String, }, - /// Response to an execution approval request. + /// Response to an execution approval request (with explicit request ID). ExecApproval { /// ID of the approval request being responded to. request_id: Uuid, @@ -78,6 +102,14 @@ pub enum Submission { always: bool, }, + /// Simple approval response (yes/no/always) for the current pending approval. + ApprovalResponse { + /// Whether the execution was approved. + approved: bool, + /// If true, auto-approve this tool for the rest of the session. + always: bool, + }, + /// Interrupt the current turn. Interrupt, diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 94dbd7f6..d49f8ac8 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -13,22 +13,30 @@ use crate::context::{ContextManager, JobState}; use crate::error::Error; use crate::history::Store; use crate::llm::{ - ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, ToolSelection, + ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection, }; use crate::safety::SafetyLayer; use crate::tools::ToolRegistry; +/// Shared dependencies for worker execution. +/// +/// This bundles the dependencies that are shared across all workers, +/// reducing the number of arguments to `Worker::new`. +#[derive(Clone)] +pub struct WorkerDeps { + pub context_manager: Arc, + pub llm: Arc, + pub safety: Arc, + pub tools: Arc, + pub store: Option>, + pub timeout: Duration, + pub use_planning: bool, +} + /// Worker that executes a single job. pub struct Worker { job_id: Uuid, - context_manager: Arc, - llm: Arc, - safety: Arc, - tools: Arc, - store: Option>, - timeout: Duration, - /// Whether to use planning before tool execution. - use_planning: bool, + deps: WorkerDeps, } /// Result of a tool execution with metadata for context building. @@ -37,32 +45,43 @@ struct ToolExecResult { } impl Worker { - /// Create a new worker. - pub fn new( - job_id: Uuid, - context_manager: Arc, - llm: Arc, - safety: Arc, - tools: Arc, - store: Option>, - timeout: Duration, - use_planning: bool, - ) -> Self { - Self { - job_id, - context_manager, - llm, - safety, - tools, - store, - timeout, - use_planning, - } + /// Create a new worker for a specific job. + pub fn new(job_id: Uuid, deps: WorkerDeps) -> Self { + Self { job_id, deps } + } + + // Convenience accessors to avoid deps.field everywhere + fn context_manager(&self) -> &Arc { + &self.deps.context_manager + } + + fn llm(&self) -> &Arc { + &self.deps.llm + } + + fn safety(&self) -> &Arc { + &self.deps.safety + } + + fn tools(&self) -> &Arc { + &self.deps.tools + } + + fn store(&self) -> Option<&Arc> { + self.deps.store.as_ref() + } + + fn timeout(&self) -> Duration { + self.deps.timeout + } + + fn use_planning(&self) -> bool { + self.deps.use_planning } /// Fire-and-forget persistence of job status. fn persist_status(&self, status: JobState, reason: Option) { - if let Some(ref store) = self.store { + if let Some(store) = self.store() { let store = store.clone(); let job_id = self.job_id; tokio::spawn(async move { @@ -91,16 +110,13 @@ impl Worker { } // Get job context - let job_ctx = self.context_manager.get_context(self.job_id).await?; + let job_ctx = self.context_manager().get_context(self.job_id).await?; // Create reasoning engine - let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone()); + let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone()); - // Build initial reasoning context - let tool_defs = self.tools.tool_definitions().await; - let mut reason_ctx = ReasoningContext::new() - .with_job(&job_ctx.description) - .with_tools(tool_defs); + // Build initial reasoning context (tool definitions refreshed each iteration in execution_loop) + let mut reason_ctx = ReasoningContext::new().with_job(&job_ctx.description); // Add system message reason_ctx.messages.push(ChatMessage::system(format!( @@ -116,7 +132,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# ))); // Main execution loop with timeout - let result = tokio::time::timeout(self.timeout, async { + let result = tokio::time::timeout(self.timeout(), async { self.execution_loop(&mut rx, &reasoning, &mut reason_ctx) .await }) @@ -148,8 +164,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."# let max_iterations = 50; let mut iteration = 0; + // Initial tool definitions for planning (will be refreshed in loop) + reason_ctx.available_tools = self.tools().tool_definitions().await; + // Generate plan if planning is enabled - let plan = if self.use_planning { + let plan = if self.use_planning() { match reasoning.plan(reason_ctx).await { Ok(p) => { tracing::info!( @@ -213,29 +232,61 @@ Report when the job is complete or if you encounter issues you cannot resolve."# return Ok(()); } + // Refresh tool definitions so newly built tools become visible + reason_ctx.available_tools = self.tools().tool_definitions().await; + // Select next tool(s) to use let selections = reasoning.select_tools(reason_ctx).await?; if selections.is_empty() { - // No tools selected, ask LLM for next steps - let response = reasoning.respond(reason_ctx).await?; + // No tools from select_tools, ask LLM directly (may still return tool calls) + let respond_result = reasoning.respond_with_tools(reason_ctx).await?; - if response.to_lowercase().contains("complete") - || response.to_lowercase().contains("finished") - || response.to_lowercase().contains("done") - { - self.mark_completed().await?; - return Ok(()); - } + match respond_result { + RespondResult::Text(response) => { + // Check for completion keywords + let response_lower = response.to_lowercase(); + if response_lower.contains("complete") + || response_lower.contains("finished") + || response_lower.contains("done") + { + self.mark_completed().await?; + return Ok(()); + } - // Add assistant response to context - reason_ctx.messages.push(ChatMessage::assistant(&response)); + // Add assistant response to context + reason_ctx.messages.push(ChatMessage::assistant(&response)); - // Give it one more chance to select a tool - if iteration > 3 && iteration % 5 == 0 { - reason_ctx.messages.push(ChatMessage::user( - "Are you stuck? Do you need help completing this job?", - )); + // Give it one more chance to select a tool + if iteration > 3 && iteration % 5 == 0 { + reason_ctx.messages.push(ChatMessage::user( + "Are you stuck? Do you need help completing this job?", + )); + } + } + RespondResult::ToolCalls(tool_calls) => { + // Model returned tool calls - execute them + tracing::debug!( + "Job {} respond_with_tools returned {} tool calls", + self.job_id, + tool_calls.len() + ); + + for tc in tool_calls { + let result = self.execute_tool(&tc.name, &tc.arguments).await; + + // Create synthetic selection for process_tool_result + let selection = ToolSelection { + tool_name: tc.name.clone(), + parameters: tc.arguments.clone(), + reasoning: String::new(), + alternatives: vec![], + }; + + self.process_tool_result(reason_ctx, &selection, result) + .await?; + } + } } } else if selections.len() == 1 { // Single tool: execute directly @@ -282,10 +333,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."# .map(|selection| { let tool_name = selection.tool_name.clone(); let params = selection.parameters.clone(); - let tools = self.tools.clone(); - let context_manager = self.context_manager.clone(); + let tools = self.tools().clone(); + let context_manager = self.context_manager().clone(); let job_id = self.job_id; - let store = self.store.clone(); + let store = self.deps.store.clone(); async move { let result = Self::execute_tool_inner( @@ -321,6 +372,15 @@ Report when the job is complete or if you encounter issues you cannot resolve."# name: tool_name.to_string(), })?; + // Log warning if tool requires approval (autonomous jobs auto-approve for now) + if tool.requires_approval() { + tracing::warn!( + job_id = %job_id, + tool = %tool_name, + "Executing sensitive tool in autonomous job (auto-approved)" + ); + } + // Get job context for the tool let job_ctx = context_manager.get_context(job_id).await?; @@ -412,11 +472,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."# Ok(output) => { // Sanitize output let sanitized = self - .safety + .safety() .sanitize_tool_output(&selection.tool_name, &output); // Add to context - let wrapped = self.safety.wrap_for_llm( + let wrapped = self.safety().wrap_for_llm( &selection.tool_name, &sanitized.content, sanitized.was_modified, @@ -445,7 +505,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# ); // Record failure for self-repair tracking - if let Some(ref store) = self.store { + if let Some(store) = self.store() { let store = store.clone(); let tool_name = selection.tool_name.clone(); let error_msg = e.to_string(); @@ -563,9 +623,9 @@ Report when the job is complete or if you encounter issues you cannot resolve."# params: &serde_json::Value, ) -> Result { Self::execute_tool_inner( - self.tools.clone(), - self.context_manager.clone(), - self.store.clone(), + self.tools().clone(), + self.context_manager().clone(), + self.deps.store.clone(), self.job_id, tool_name, params, @@ -574,7 +634,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } async fn mark_completed(&self) -> Result<(), Error> { - self.context_manager + self.context_manager() .update_context(self.job_id, |ctx| { ctx.transition_to( JobState::Completed, @@ -595,7 +655,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } async fn mark_failed(&self, reason: &str) -> Result<(), Error> { - self.context_manager + self.context_manager() .update_context(self.job_id, |ctx| { ctx.transition_to(JobState::Failed, Some(reason.to_string())) }) @@ -610,7 +670,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } async fn mark_stuck(&self, reason: &str) -> Result<(), Error> { - self.context_manager + self.context_manager() .update_context(self.job_id, |ctx| ctx.mark_stuck(reason)) .await? .map_err(|s| crate::error::JobError::ContextError { diff --git a/src/channels/channel.rs b/src/channels/channel.rs index cd9e8527..68ec714d 100644 --- a/src/channels/channel.rs +++ b/src/channels/channel.rs @@ -146,6 +146,20 @@ pub trait Channel: Send + Sync { Ok(()) } + /// Send a proactive message without a prior incoming message. + /// + /// Used for alerts, heartbeat notifications, and other agent-initiated communication. + /// The user_id helps target a specific user within the channel. + /// + /// Default implementation does nothing (for channels that don't support broadcast). + async fn broadcast( + &self, + _user_id: &str, + _response: OutgoingResponse, + ) -> Result<(), ChannelError> { + Ok(()) + } + /// Check if the channel is healthy. async fn health_check(&self) -> Result<(), ChannelError>; diff --git a/src/channels/cli/mod.rs b/src/channels/cli/mod.rs index cd80607e..0515e9e9 100644 --- a/src/channels/cli/mod.rs +++ b/src/channels/cli/mod.rs @@ -128,6 +128,22 @@ impl Channel for TuiChannel { Ok(()) } + async fn broadcast( + &self, + _user_id: &str, + response: OutgoingResponse, + ) -> Result<(), ChannelError> { + // For TUI, broadcasts appear as regular agent responses with a notification indicator + self.event_tx + .send(AppEvent::Response(response.content)) + .await + .map_err(|e| ChannelError::SendFailed { + name: "tui".to_string(), + reason: e.to_string(), + })?; + Ok(()) + } + async fn health_check(&self) -> Result<(), ChannelError> { // Channel is healthy if we haven't been closed if self.event_tx.is_closed() { diff --git a/src/channels/cli/render.rs b/src/channels/cli/render.rs index 9da4f8c1..d07332a2 100644 --- a/src/channels/cli/render.rs +++ b/src/channels/cli/render.rs @@ -89,11 +89,7 @@ fn render_messages(frame: &mut Frame, app: &AppState, area: Rect) { // Calculate scroll - show most recent messages let visible_height = area.height.saturating_sub(2) as usize; // Account for borders let total_lines = lines.len(); - let scroll_offset = if total_lines > visible_height { - total_lines - visible_height - } else { - 0 - }; + let scroll_offset = total_lines.saturating_sub(visible_height); let text = Text::from(lines); let messages = Paragraph::new(text) diff --git a/src/channels/manager.rs b/src/channels/manager.rs index 35cf3d73..6d8f230c 100644 --- a/src/channels/manager.rs +++ b/src/channels/manager.rs @@ -97,6 +97,45 @@ impl ChannelManager { } } + /// Broadcast a message to a specific user on a specific channel. + /// + /// Used for proactive notifications like heartbeat alerts. + pub async fn broadcast( + &self, + channel_name: &str, + user_id: &str, + response: OutgoingResponse, + ) -> Result<(), ChannelError> { + let channels = self.channels.read().await; + if let Some(channel) = channels.get(channel_name) { + channel.broadcast(user_id, response).await + } else { + Err(ChannelError::SendFailed { + name: channel_name.to_string(), + reason: "Channel not found".to_string(), + }) + } + } + + /// Broadcast a message to all channels. + /// + /// Sends to the specified user on every registered channel. + pub async fn broadcast_all( + &self, + user_id: &str, + response: OutgoingResponse, + ) -> Vec<(String, Result<(), ChannelError>)> { + let channels = self.channels.read().await; + let mut results = Vec::new(); + + for (name, channel) in channels.iter() { + let result = channel.broadcast(user_id, response.clone()).await; + results.push((name.clone(), result)); + } + + results + } + /// Check health of all channels. pub async fn health_check_all(&self) -> HashMap> { let channels = self.channels.read().await; diff --git a/src/config.rs b/src/config.rs index 15f2b56b..074cbb70 100644 --- a/src/config.rs +++ b/src/config.rs @@ -12,6 +12,7 @@ use crate::error::ConfigError; pub struct Config { pub database: DatabaseConfig, pub llm: LlmConfig, + pub embeddings: EmbeddingsConfig, pub channels: ChannelsConfig, pub agent: AgentConfig, pub safety: SafetyConfig, @@ -30,6 +31,7 @@ impl Config { Ok(Self { database: DatabaseConfig::from_env()?, llm: LlmConfig::from_env()?, + embeddings: EmbeddingsConfig::from_env()?, channels: ChannelsConfig::from_env()?, agent: AgentConfig::from_env()?, safety: SafetyConfig::from_env()?, @@ -106,6 +108,62 @@ impl LlmConfig { } } +/// Embeddings provider configuration. +#[derive(Debug, Clone)] +pub struct EmbeddingsConfig { + /// Whether embeddings are enabled. + pub enabled: bool, + /// Provider to use: "openai" or "nearai" + pub provider: String, + /// OpenAI API key (for OpenAI provider). + pub openai_api_key: Option, + /// Model to use for embeddings. + /// For OpenAI: "text-embedding-3-small", "text-embedding-3-large", "text-embedding-ada-002" + /// For NEAR AI: Uses the configured session for auth. + pub model: String, +} + +impl Default for EmbeddingsConfig { + fn default() -> Self { + Self { + enabled: false, + provider: "openai".to_string(), + openai_api_key: None, + model: "text-embedding-3-small".to_string(), + } + } +} + +impl EmbeddingsConfig { + fn from_env() -> Result { + let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from); + let provider = optional_env("EMBEDDING_PROVIDER")?.unwrap_or_else(|| "openai".to_string()); + + // Auto-enable if we have an API key + let enabled = optional_env("EMBEDDING_ENABLED")? + .map(|s| s.parse()) + .transpose() + .map_err(|e| ConfigError::InvalidValue { + key: "EMBEDDING_ENABLED".to_string(), + message: format!("must be 'true' or 'false': {e}"), + })? + .unwrap_or(openai_api_key.is_some()); + + Ok(Self { + enabled, + provider, + openai_api_key, + model: optional_env("EMBEDDING_MODEL")? + .unwrap_or_else(|| "text-embedding-3-small".to_string()), + }) + } + + /// Get the OpenAI API key if configured. + pub fn openai_api_key(&self) -> Option<&str> { + self.openai_api_key.as_ref().map(|s| s.expose_secret()) + } +} + /// Get the default session file path (~/.near-agent/session.json). fn default_session_path() -> PathBuf { dirs::home_dir() diff --git a/src/history/mod.rs b/src/history/mod.rs index ea321821..e8254479 100644 --- a/src/history/mod.rs +++ b/src/history/mod.rs @@ -9,4 +9,4 @@ mod analytics; mod store; pub use analytics::{JobStats, ToolStats}; -pub use store::Store; +pub use store::{LlmCallRecord, Store}; diff --git a/src/history/store.rs b/src/history/store.rs index 1e347fee..0850701d 100644 --- a/src/history/store.rs +++ b/src/history/store.rs @@ -9,6 +9,19 @@ use crate::config::DatabaseConfig; use crate::context::{ActionRecord, JobContext, JobState}; use crate::error::DatabaseError; +/// Record for an LLM call to be persisted. +#[derive(Debug, Clone)] +pub struct LlmCallRecord<'a> { + pub job_id: Option, + pub conversation_id: Option, + pub provider: &'a str, + pub model: &'a str, + pub input_tokens: u32, + pub output_tokens: u32, + pub cost: Decimal, + pub purpose: Option<&'a str>, +} + /// Database store for the agent. pub struct Store { pool: Pool, @@ -335,17 +348,7 @@ impl Store { // ==================== LLM Calls ==================== /// Record an LLM call. - pub async fn record_llm_call( - &self, - job_id: Option, - conversation_id: Option, - provider: &str, - model: &str, - input_tokens: u32, - output_tokens: u32, - cost: Decimal, - purpose: Option<&str>, - ) -> Result { + pub async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result { let conn = self.conn().await?; let id = Uuid::new_v4(); @@ -356,14 +359,14 @@ impl Store { "#, &[ &id, - &job_id, - &conversation_id, - &provider, - &model, - &(input_tokens as i32), - &(output_tokens as i32), - &cost, - &purpose, + &record.job_id, + &record.conversation_id, + &record.provider, + &record.model, + &(record.input_tokens as i32), + &(record.output_tokens as i32), + &record.cost, + &record.purpose, ], ) .await?; diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index 2206b17e..efad83e3 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -121,12 +121,29 @@ pub struct Reasoning { llm: Arc, #[allow(dead_code)] // Will be used for sanitizing tool outputs safety: Arc, + /// Optional workspace for loading identity/system prompts. + workspace_system_prompt: Option, } impl Reasoning { /// Create a new reasoning engine. pub fn new(llm: Arc, safety: Arc) -> Self { - Self { llm, safety } + Self { + llm, + safety, + workspace_system_prompt: None, + } + } + + /// Set a custom system prompt from workspace identity files. + /// + /// This is typically loaded from workspace.system_prompt() which combines + /// AGENTS.md, SOUL.md, USER.md, and IDENTITY.md into a unified prompt. + pub fn with_system_prompt(mut self, prompt: String) -> Self { + if !prompt.is_empty() { + self.workspace_system_prompt = Some(prompt); + } + self } /// Generate a plan for completing a goal. @@ -361,11 +378,18 @@ Respond with a JSON plan in this format: .map(|t| format!(" - {}: {}", t.name, t.description)) .collect(); format!( - "\n\n## Available Tools\nYou have access to these tools:\n{}\n\nCall tools directly when needed.", + "\n\n## Available Tools\nYou have access to these tools:\n{}\n\nCall tools when they would help accomplish the task.", tool_list.join("\n") ) }; + // Include workspace identity prompt if available + let identity_section = if let Some(ref identity) = self.workspace_system_prompt { + format!("\n\n---\n\n{}", identity) + } else { + String::new() + }; + format!( r#"You are NEAR AI Agent, an autonomous assistant. @@ -388,8 +412,8 @@ Here's the solution: [actual response to user] - For code, use appropriate code blocks with language tags - Call tools when they would help accomplish the task{} -The user sees ONLY content outside tags."#, - tools_section +The user sees ONLY content outside tags.{}"#, + tools_section, identity_section ) } diff --git a/src/main.rs b/src/main.rs index fc830e2e..2516fe2c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,7 +6,7 @@ use clap::Parser; use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; use near_agent::{ - agent::Agent, + agent::{Agent, AgentDeps}, channels::{ChannelManager, HttpChannel, TuiChannel}, cli::{Cli, Command, run_tool_command}, config::Config, @@ -17,7 +17,7 @@ use near_agent::{ ToolRegistry, wasm::{WasmToolLoader, WasmToolRuntime}, }, - workspace::Workspace, + workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace}, }; #[tokio::main] @@ -93,8 +93,8 @@ async fn main() -> anyhow::Result<()> { Some(Arc::new(store)) }; - // Initialize LLM provider - let llm = create_llm_provider(&config.llm, session)?; + // Initialize LLM provider (clone session so we can reuse it for embeddings) + let llm = create_llm_provider(&config.llm, session.clone())?; tracing::info!("LLM provider initialized: {}", llm.model_name()); // Initialize safety layer @@ -106,9 +106,52 @@ async fn main() -> anyhow::Result<()> { tools.register_builtin_tools(); tracing::info!("Registered {} built-in tools", tools.count()); + // Create embeddings provider if configured + let embeddings: Option> = if config.embeddings.enabled { + match config.embeddings.provider.as_str() { + "nearai" => { + tracing::info!( + "Embeddings enabled via NEAR AI (model: {})", + config.embeddings.model + ); + Some(Arc::new( + NearAiEmbeddings::new(&config.llm.nearai.base_url, session.clone()) + .with_model(&config.embeddings.model, 1536), + )) + } + _ => { + // Default to OpenAI for unknown providers + if let Some(api_key) = config.embeddings.openai_api_key() { + tracing::info!( + "Embeddings enabled via OpenAI (model: {})", + config.embeddings.model + ); + Some(Arc::new(OpenAiEmbeddings::with_model( + api_key, + &config.embeddings.model, + match config.embeddings.model.as_str() { + "text-embedding-3-large" => 3072, + _ => 1536, // text-embedding-3-small and ada-002 + }, + ))) + } else { + tracing::warn!("Embeddings configured but OPENAI_API_KEY not set"); + None + } + } + } + } else { + tracing::info!("Embeddings disabled (set OPENAI_API_KEY or EMBEDDING_ENABLED=true)"); + None + }; + // Register memory tools if database is available if let Some(ref store) = store { - let workspace = Arc::new(Workspace::new("default", store.pool())); + let mut workspace = Workspace::new("default", store.pool()); + if let Some(ref emb) = embeddings { + workspace = workspace.with_embeddings(emb.clone()); + } + let workspace = Arc::new(workspace); tools.register_memory_tools(workspace); } @@ -181,19 +224,39 @@ async fn main() -> anyhow::Result<()> { } // Create workspace for agent (shared with memory tools) - let workspace = store - .as_ref() - .map(|s| Arc::new(Workspace::new("default", s.pool()))); + let workspace = store.as_ref().map(|s| { + let mut ws = Workspace::new("default", s.pool()); + if let Some(ref emb) = embeddings { + ws = ws.with_embeddings(emb.clone()); + } + Arc::new(ws) + }); + + // Backfill embeddings if we just enabled the provider + if let (Some(ws), Some(_)) = (&workspace, &embeddings) { + match ws.backfill_embeddings().await { + Ok(count) if count > 0 => { + tracing::info!("Backfilled embeddings for {} chunks", count); + } + Ok(_) => {} + Err(e) => { + tracing::warn!("Failed to backfill embeddings: {}", e); + } + } + } // Create and run the agent - let agent = Agent::new( - config.agent.clone(), + let deps = AgentDeps { store, llm, safety, tools, - channels, workspace, + }; + let agent = Agent::new( + config.agent.clone(), + deps, + channels, Some(config.heartbeat.clone()), ); diff --git a/src/workspace/embeddings.rs b/src/workspace/embeddings.rs index 6cfee5a2..320eca30 100644 --- a/src/workspace/embeddings.rs +++ b/src/workspace/embeddings.rs @@ -213,6 +213,147 @@ impl EmbeddingProvider for OpenAiEmbeddings { } } +/// NEAR AI embedding provider using the NEAR AI API. +/// +/// Uses the same session-based auth as the LLM provider. +pub struct NearAiEmbeddings { + client: reqwest::Client, + base_url: String, + session: std::sync::Arc, + model: String, + dimension: usize, +} + +impl NearAiEmbeddings { + /// Create a new NEAR AI embedding provider. + /// + /// Uses the same session manager as the LLM provider for auth. + pub fn new( + base_url: impl Into, + session: std::sync::Arc, + ) -> Self { + Self { + client: reqwest::Client::new(), + base_url: base_url.into(), + session, + model: "text-embedding-3-small".to_string(), + dimension: 1536, + } + } + + /// Use a specific model. + pub fn with_model(mut self, model: impl Into, dimension: usize) -> Self { + self.model = model.into(); + self.dimension = dimension; + self + } +} + +#[derive(Debug, Serialize)] +struct NearAiEmbeddingRequest<'a> { + model: &'a str, + input: &'a [String], +} + +#[derive(Debug, Deserialize)] +struct NearAiEmbeddingResponse { + data: Vec, +} + +#[derive(Debug, Deserialize)] +struct NearAiEmbeddingData { + embedding: Vec, +} + +#[async_trait] +impl EmbeddingProvider for NearAiEmbeddings { + fn dimension(&self) -> usize { + self.dimension + } + + fn model_name(&self) -> &str { + &self.model + } + + fn max_input_length(&self) -> usize { + 32_000 + } + + async fn embed(&self, text: &str) -> Result, EmbeddingError> { + if text.len() > self.max_input_length() { + return Err(EmbeddingError::TextTooLong { + length: text.len(), + max: self.max_input_length(), + }); + } + + let embeddings = self.embed_batch(&[text.to_string()]).await?; + embeddings + .into_iter() + .next() + .ok_or_else(|| EmbeddingError::InvalidResponse("No embedding returned".to_string())) + } + + async fn embed_batch(&self, texts: &[String]) -> Result>, EmbeddingError> { + use secrecy::ExposeSecret; + + if texts.is_empty() { + return Ok(Vec::new()); + } + + let request = NearAiEmbeddingRequest { + model: &self.model, + input: texts, + }; + + let token = self + .session + .get_token() + .await + .map_err(|_| EmbeddingError::AuthFailed)?; + + let url = format!("{}/v1/embeddings", self.base_url); + + let response = self + .client + .post(&url) + .header("Authorization", format!("Bearer {}", token.expose_secret())) + .json(&request) + .send() + .await?; + + let status = response.status(); + + if status == reqwest::StatusCode::UNAUTHORIZED { + return Err(EmbeddingError::AuthFailed); + } + + if status == reqwest::StatusCode::TOO_MANY_REQUESTS { + let retry_after = response + .headers() + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + .map(std::time::Duration::from_secs); + return Err(EmbeddingError::RateLimited { retry_after }); + } + + if !status.is_success() { + let error_text = response.text().await.unwrap_or_default(); + return Err(EmbeddingError::HttpError(format!( + "Status {}: {}", + status, error_text + ))); + } + + let result: NearAiEmbeddingResponse = response.json().await.map_err(|e| { + EmbeddingError::InvalidResponse(format!("Failed to parse response: {}", e)) + })?; + + Ok(result.data.into_iter().map(|d| d.embedding).collect()) + } +} + /// A mock embedding provider for testing. /// /// Generates deterministic embeddings based on text hash. diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs index 841b9431..32cc2121 100644 --- a/src/workspace/mod.rs +++ b/src/workspace/mod.rs @@ -48,7 +48,7 @@ mod search; pub use chunker::{ChunkConfig, chunk_document}; pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths}; -pub use embeddings::{EmbeddingProvider, MockEmbeddings, OpenAiEmbeddings}; +pub use embeddings::{EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OpenAiEmbeddings}; pub use repository::Repository; pub use search::{SearchConfig, SearchResult};