//! Core AI service implementation with RAG integration use std::net::SocketAddr; use std::sync::Arc; use anyhow::Result; use serde::{Deserialize, Serialize}; use tokio::sync::RwLock; use tracing::{debug, info}; /// MCP tool invocation request #[derive(Debug, Clone, Serialize, Deserialize)] pub struct McpToolRequest { /// Tool name (e.g., "execute_provisioning_plan") pub tool_name: String, /// Tool arguments as JSON pub args: serde_json::Value, } /// MCP tool invocation response #[derive(Debug, Clone, Serialize, Deserialize)] pub struct McpToolResponse { /// Tool execution result pub result: serde_json::Value, /// Execution time in milliseconds pub duration_ms: u64, } /// RAG-powered question request with optional hybrid tool execution #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AskRequest { /// User question pub question: String, /// Optional context pub context: Option, /// Enable automatic tool execution (hybrid mode) /// When true, the RAG answer may trigger tool calls automatically /// Default: false (explicit tool calls only) #[serde(default)] pub enable_tool_execution: Option, /// Maximum number of tools to execute automatically /// Default: 3 #[serde(default)] pub max_tool_calls: Option, } /// RAG-powered question response with optional tool execution results #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AskResponse { /// Answer from AI pub answer: String, /// Source documents used pub sources: Vec, /// Confidence level (0-100) pub confidence: u8, /// Reasoning explanation pub reasoning: String, /// Tool executions performed (if hybrid mode enabled) #[serde(skip_serializing_if = "Option::is_none")] pub tool_executions: Option>, } /// Extension DAG node #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DagNode { /// Extension/component name pub name: String, /// Dependencies on other nodes pub dependencies: Vec, /// Component version pub version: String, } /// Extension DAG response #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DagResponse { /// DAG nodes (extensions) pub nodes: Vec, /// DAG edges (dependencies) pub edges: Vec<(String, String)>, } /// Best practice entry #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BestPractice { /// Practice title pub title: String, /// Practice description pub description: String, /// Category (e.g., "deployment", "security") pub category: String, /// Relevance score (0-100) pub relevance: u8, } /// Knowledge base document (from RAG ingestion) #[derive(Debug, Clone, Serialize, Deserialize)] pub struct KnowledgeDocument { /// Document ID pub id: String, /// Document type #[serde(rename = "type")] pub doc_type: String, /// Document title pub title: Option, /// Document name (for extensions) pub name: Option, /// Full content pub content: String, /// Document category pub category: String, /// Tags pub tags: Vec, /// Relevance/importance pub relevance: Option, /// Dependencies (for extensions) pub dependencies: Option>, } /// Knowledge base with documents and relationships #[derive(Debug, Clone)] pub struct KnowledgeBase { /// All documents indexed by ID pub documents: std::collections::HashMap, /// Document relationships pub relationships: Vec, } /// Document relationship #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Relationship { /// Source document ID pub source_id: String, /// Target document ID pub target_id: String, /// Relationship type pub relationship_type: String, /// Strength (0-1) pub strength: f32, } /// Core AI service pub struct AiService { /// Server address addr: SocketAddr, /// Knowledge base knowledge_base: Arc>, /// MCP tool registry tool_registry: crate::mcp::ToolRegistry, } impl AiService { /// Create new AI service pub fn new(addr: SocketAddr) -> Self { Self { addr, knowledge_base: Arc::new(RwLock::new(KnowledgeBase { documents: std::collections::HashMap::new(), relationships: Vec::new(), })), tool_registry: crate::mcp::ToolRegistry::new(), } } /// Load knowledge base from JSON files pub async fn load_knowledge_base(&self, knowledge_base_dir: &str) -> Result<()> { info!("Loading knowledge base from: {}", knowledge_base_dir); // Load best practice documents let bp_path = format!("{}/best-practices-docs.json", knowledge_base_dir); let bp_content = std::fs::read_to_string(&bp_path) .map_err(|e| anyhow::anyhow!("Failed to read best practices: {}", e))?; let bp_docs: Vec = serde_json::from_str(&bp_content)?; // Load extension documents let ext_path = format!("{}/extension-docs.json", knowledge_base_dir); let ext_content = std::fs::read_to_string(&ext_path) .map_err(|e| anyhow::anyhow!("Failed to read extensions: {}", e))?; let ext_docs: Vec = serde_json::from_str(&ext_content)?; // Load relationships let rel_path = format!("{}/relationships.json", knowledge_base_dir); let rel_content = std::fs::read_to_string(&rel_path) .map_err(|e| anyhow::anyhow!("Failed to read relationships: {}", e))?; let relationships: Vec = serde_json::from_str(&rel_content)?; // Build document index let mut documents = std::collections::HashMap::new(); for doc in bp_docs.into_iter().chain(ext_docs.into_iter()) { documents.insert(doc.id.clone(), doc); } // Update knowledge base let mut kb = self.knowledge_base.write().await; kb.documents = documents; kb.relationships = relationships; info!( "Knowledge base loaded: {} documents, {} relationships", kb.documents.len(), kb.relationships.len() ); Ok(()) } /// Get service address pub fn addr(&self) -> SocketAddr { self.addr } /// Search knowledge base by keyword and category async fn search_knowledge( &self, query: &str, category: Option<&str>, ) -> Vec { let kb = self.knowledge_base.read().await; let query_lower = query.to_lowercase(); kb.documents .values() .filter(|doc| { let matches_query = doc.content.to_lowercase().contains(&query_lower) || doc .tags .iter() .any(|t| t.to_lowercase().contains(&query_lower)); let matches_category = category.map(|c| doc.category == c).unwrap_or(true); matches_query && matches_category }) .cloned() .collect() } /// Call an MCP tool via registry pub async fn call_mcp_tool(&self, req: McpToolRequest) -> Result { let start = std::time::Instant::now(); debug!("Calling MCP tool: {}", req.tool_name); let result = self.execute_mcp_tool(&req.tool_name, &req.args).await?; let duration_ms = start.elapsed().as_millis() as u64; Ok(McpToolResponse { result, duration_ms, }) } /// Execute MCP tool with arguments via registry async fn execute_mcp_tool( &self, tool_name: &str, args: &serde_json::Value, ) -> Result { // Check if tool exists in registry if !self.tool_registry.has_tool(tool_name) { return Err(anyhow::anyhow!("Unknown tool: {}", tool_name)); } // Execute tool through registry match self.tool_registry.execute(tool_name, args).await { Ok(result) => Ok(result), Err(e) => Err(anyhow::anyhow!("Tool execution failed: {}", e)), } } /// Execute suggested tools and collect results async fn execute_tools( &self, suggestions: &[crate::tool_integration::ToolSuggestion], max_tools: usize, ) -> Option<( Vec, Vec<(String, serde_json::Value)>, )> { if suggestions.is_empty() { return None; } debug!( "Executing {} suggested tools in hybrid mode", suggestions.len().min(max_tools) ); let mut executions = Vec::new(); let mut results = Vec::new(); for suggestion in suggestions.iter().take(max_tools) { match self .tool_registry .execute(&suggestion.tool_name, &suggestion.args) .await { Ok(result) => { debug!("Tool {} executed successfully", suggestion.tool_name); results.push((suggestion.tool_name.clone(), result.clone())); executions.push(crate::mcp::ToolExecution { tool_name: suggestion.tool_name.clone(), result: result.clone(), duration_ms: 0, }); } Err(e) => { debug!("Tool {} execution failed: {}", suggestion.tool_name, e); } } } if executions.is_empty() { None } else { Some((executions, results)) } } /// Ask AI a question using RAG with optional hybrid tool execution pub async fn ask(&self, req: AskRequest) -> Result { debug!("Processing RAG question: {}", req.question); // Search knowledge base for relevant documents let search_results = self.search_knowledge(&req.question, None).await; if search_results.is_empty() { return Ok(AskResponse { answer: "I couldn't find any relevant information in the knowledge base for this \ question." .to_string(), sources: vec![], confidence: 20, reasoning: "No matching documents found".to_string(), tool_executions: None, }); } // Sort by relevance (best practices have explicit relevance scores) let mut results = search_results; results.sort_by(|a, b| { let a_rel = a.relevance.unwrap_or(50); let b_rel = b.relevance.unwrap_or(50); b_rel.cmp(&a_rel) }); // Get top 3 most relevant documents let top_results: Vec<_> = results.iter().take(3).collect(); // Build answer from top results let mut answer_parts = vec!["Based on the knowledge base, here's what I found:".to_string()]; for doc in &top_results { if let Some(title) = &doc.title { answer_parts.push(format!( "\n- **{}**: {}", title, &doc.content[..std::cmp::min(150, doc.content.len())] )); } else if let Some(name) = &doc.name { answer_parts.push(format!( "\n- **{}**: {}", name, &doc.content[..std::cmp::min(150, doc.content.len())] )); } } let mut answer = answer_parts.join("\n"); let sources: Vec = top_results .iter() .filter_map(|d| d.title.clone().or_else(|| d.name.clone())) .collect(); let confidence = (top_results.iter().filter_map(|d| d.relevance).sum::() as f32 / top_results.len() as f32) as u8; // Handle hybrid execution mode (auto-trigger tools if enabled) let mut tool_executions = None; if req.enable_tool_execution.unwrap_or(false) { let max_tools = req.max_tool_calls.unwrap_or(3) as usize; let tool_suggestions = crate::tool_integration::analyze_for_tools(&answer, &req.question); if let Some((executions, results)) = self.execute_tools(&tool_suggestions, max_tools).await { if !results.is_empty() { answer = crate::tool_integration::enrich_answer_with_results(answer, &results); tool_executions = Some(executions); } } } Ok(AskResponse { answer, sources, confidence, reasoning: format!( "Retrieved {} relevant documents using keyword search across {} total documents", top_results.len(), results.len() ), tool_executions, }) } /// Get extension dependency DAG pub async fn get_extension_dag(&self) -> Result { debug!("Building extension DAG"); let kb = self.knowledge_base.read().await; // Build nodes from extension documents let nodes: Vec = kb .documents .values() .filter(|doc| doc.doc_type == "extension_metadata") .map(|doc| DagNode { name: doc.name.clone().unwrap_or_else(|| doc.id.clone()), dependencies: doc.dependencies.clone().unwrap_or_default(), version: "1.0.0".to_string(), }) .collect(); // Build edges from dependency relationships let edges: Vec<(String, String)> = kb .relationships .iter() .filter(|rel| rel.relationship_type == "depends_on") .map(|rel| { let source = rel.source_id.strip_prefix("ext_").unwrap_or(&rel.source_id); let target = rel.target_id.strip_prefix("ext_").unwrap_or(&rel.target_id); (source.to_string(), target.to_string()) }) .collect(); Ok(DagResponse { nodes, edges }) } /// Get best practices for a category pub async fn get_best_practices(&self, category: &str) -> Result> { debug!("Retrieving best practices for category: {}", category); let kb = self.knowledge_base.read().await; // Filter documents by category and type let mut practices: Vec = kb .documents .values() .filter(|doc| doc.category == category && doc.doc_type == "best_practice") .map(|doc| BestPractice { title: doc.title.clone().unwrap_or_else(|| doc.id.clone()), description: doc.content.clone(), category: doc.category.clone(), relevance: doc.relevance.unwrap_or(70), }) .collect(); // Sort by relevance descending practices.sort_by(|a, b| b.relevance.cmp(&a.relevance)); Ok(practices) } /// Health check endpoint pub async fn health_check(&self) -> Result<()> { Ok(()) } /// Get all available tools pub fn list_all_tools(&self) -> Vec { self.tool_registry.list_tools() } /// Get tools by category pub fn tools_by_category( &self, category: crate::mcp::ToolCategory, ) -> Vec { self.tool_registry.tools_by_category(category) } /// Check if a tool exists pub fn has_tool(&self, name: &str) -> bool { self.tool_registry.has_tool(name) } } impl Default for AiService { fn default() -> Self { Self::new("127.0.0.1:8083".parse().unwrap()) } }