provisioning-platform/crates/rag/src/agent.rs

//! RAG Agent using Rig framework for Claude integration

use serde::{Deserialize, Serialize};

use crate::context::WorkspaceContext;
use crate::error::Result;
use crate::llm::LlmClient;
use crate::retrieval::{RetrieverEngine, SearchResult};

/// RAG Agent response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentResponse {
    /// The answer to the user's question
    pub answer: String,

    /// Source documents used for the answer
    pub sources: Vec<SearchResult>,

    /// Confidence score (0.0-1.0)
    pub confidence: f32,

    /// Retrieved context used
    pub context: String,
}

/// RAG Agent that answers questions using retrieved documents
pub struct RagAgent {
    retriever: RetrieverEngine,
    workspace_context: WorkspaceContext,
    llm_client: LlmClient,
}

impl RagAgent {
    /// Create a new RAG agent with Claude LLM
    pub fn new(
        retriever: RetrieverEngine,
        workspace_context: WorkspaceContext,
        llm_model: String,
    ) -> Result<Self> {
        let llm_client = LlmClient::new(llm_model)?;

        Ok(Self {
            retriever,
            workspace_context,
            llm_client,
        })
    }

    /// Ask a question and get an answer based on retrieved documents
    pub async fn ask(&self, question: &str) -> Result<AgentResponse> {
        // 1. Enrich query with workspace context
        let enriched_query = self.workspace_context.enrich_query(question);

        tracing::info!("Processing question: {}", question);
        tracing::debug!("Enriched query: {}", enriched_query);

        // 2. Retrieve relevant documents
        let sources = self.retriever.search(&enriched_query, None).await?;

        tracing::info!("Retrieved {} documents", sources.len());

        if sources.is_empty() {
            return Ok(AgentResponse {
                answer: "I could not find relevant information to answer your question. Please \
                         provide more context or check if the documentation has been indexed."
                    .to_string(),
                sources: vec![],
                confidence: 0.0,
                context: "No documents found".to_string(),
            });
        }

        // 3. Build context from retrieved documents
        let context = build_context(&sources);

        // 4. Generate answer using Claude API
        let answer = self
            .llm_client
            .generate_answer(&enriched_query, &context)
            .await?;

        // 5. Calculate confidence based on number and quality of sources
        let confidence = (sources.len() as f32 / 5.0).min(1.0);

        Ok(AgentResponse {
            answer,
            sources,
            confidence,
            context,
        })
    }

    /// Get agent metadata
    pub fn metadata(&self) -> AgentMetadata {
        AgentMetadata {
            model: self.llm_client.model.clone(),
            workspace: self.workspace_context.workspace_name.clone(),
            version: "0.1.0".to_string(),
        }
    }
}

/// Agent metadata
#[derive(Debug, Clone, Serialize)]
pub struct AgentMetadata {
    pub model: String,
    pub workspace: String,
    pub version: String,
}

/// Build context string from retrieved documents
fn build_context(sources: &[SearchResult]) -> String {
    let mut context = String::from("# Retrieved Context\n\n");

    for (idx, source) in sources.iter().enumerate() {
        context.push_str(&format!("## Document {}: {}\n", idx + 1, source.doc_id));
        context.push_str(&format!("**Type**: {}\n", source.doc_type));
        context.push_str(&format!("**Source**: {}\n", source.source_path));

        if !source.metadata.is_empty() {
            context.push_str("**Metadata**:\n");
            for (key, value) in &source.metadata {
                context.push_str(&format!("  - {}: {}\n", key, value));
            }
        }

        context.push_str(&format!("\n{}\n\n", source.content));
    }

    context
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_context_building() {
        let sources = vec![
            SearchResult {
                doc_id: "doc-1".to_string(),
                source_path: "docs/arch.md".to_string(),
                doc_type: "markdown".to_string(),
                content: "System architecture details".to_string(),
                similarity: 0.9,
                metadata: std::collections::HashMap::new(),
            },
            SearchResult {
                doc_id: "doc-2".to_string(),
                source_path: "docs/deploy.md".to_string(),
                doc_type: "markdown".to_string(),
                content: "Deployment procedures".to_string(),
                similarity: 0.85,
                metadata: std::collections::HashMap::new(),
            },
        ];

        let context = build_context(&sources);

        assert!(context.contains("Retrieved Context"));
        assert!(context.contains("doc-1"));
        assert!(context.contains("doc-2"));
        assert!(context.contains("markdown"));
    }
}