nxtgauge-backend-rust/crates/cache/src/ollama.rs

//! Ollama client for AI-powered text generation
//!
//! Used for generating job descriptions, resume analysis, and other AI features

use reqwest::{Client, Error as ReqwestError};
use serde::{Deserialize, Serialize};
use std::time::Duration;

const OLLAMA_URL: &str = "http://nxtgauge-ai-assistant:11434";
const DEFAULT_MODEL: &str = "gemma3:270m";
const REQUEST_TIMEOUT: Duration = Duration::from_secs(120);

#[derive(Debug, Clone)]
pub struct OllamaClient {
    http_client: Client,
    base_url: String,
    model: String,
}

#[derive(Debug, Serialize)]
struct GenerateRequest {
    model: String,
    prompt: String,
    stream: bool,
    options: Option<GenerationOptions>,
}

#[derive(Debug, Serialize, Default)]
struct GenerationOptions {
    temperature: Option<f32>,
    top_p: Option<f32>,
    top_k: Option<i32>,
    num_predict: Option<i32>,
}

#[derive(Debug, Deserialize)]
pub struct GenerateResponse {
    pub model: String,
    pub created_at: String,
    pub response: String,
    pub done: bool,
    pub context: Option<Vec<i32>>,
    pub total_duration: Option<u64>,
    pub load_duration: Option<u64>,
    pub prompt_eval_count: Option<i32>,
    pub prompt_eval_duration: Option<u64>,
    pub eval_count: Option<i32>,
    pub eval_duration: Option<u64>,
}

#[derive(Debug, Deserialize)]
#[allow(dead_code)]
struct OllamaErrorResponse {
    error: String,
}

#[derive(Debug, thiserror::Error)]
pub enum OllamaError {
    #[error("HTTP request failed: {0}")]
    RequestFailed(#[from] ReqwestError),
    
    #[error("Ollama API error: {0}")]
    ApiError(String),
    
    #[error("Failed to parse response: {0}")]
    ParseError(String),
    
    #[error("Connection timeout")]
    Timeout,
    
    #[error("Model not found: {0}")]
    ModelNotFound(String),
}

impl OllamaClient {
    pub fn new() -> Self {
        let http_client = Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .build()
            .expect("Failed to create HTTP client");
        
        Self {
            http_client,
            base_url: OLLAMA_URL.to_string(),
            model: DEFAULT_MODEL.to_string(),
        }
    }
    
    pub fn with_url(base_url: impl Into<String>) -> Self {
        let http_client = Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .build()
            .expect("Failed to create HTTP client");
        
        Self {
            http_client,
            base_url: base_url.into(),
            model: DEFAULT_MODEL.to_string(),
        }
    }
    
    pub fn with_model(mut self, model: impl Into<String>) -> Self {
        self.model = model.into();
        self
    }
    
    pub fn get_model(&self) -> &str {
        &self.model
    }
    
    /// Generate text using the configured model and prompt
    pub async fn generate(&self, prompt: impl Into<String>) -> Result<GenerateResponse, OllamaError> {
        let request = GenerateRequest {
            model: self.model.clone(),
            prompt: prompt.into(),
            stream: false,
            options: None,
        };
        
        let url = format!("{}/api/generate", self.base_url);
        
        let response = self.http_client
            .post(&url)
            .json(&request)
            .send()
            .await
            .map_err(|e| {
                if e.is_timeout() {
                    OllamaError::Timeout
                } else {
                    OllamaError::RequestFailed(e)
                }
            })?;
        
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            
            if status.as_u16() == 404 {
                return Err(OllamaError::ModelNotFound(self.model.clone()));
            }
            
            return Err(OllamaError::ApiError(format!("{}: {}", status, error_text)));
        }
        
        let result = response.json::<GenerateResponse>()
            .await
            .map_err(|e| OllamaError::ParseError(e.to_string()))?;
        
        Ok(result)
    }
    
    /// Generate a job description based on a prompt
    pub async fn generate_job_description(&self, prompt: &str) -> Result<String, OllamaError> {
        let enhanced_prompt = format!(
            "Generate a professional job description based on the following prompt:\n\n{}\n\n\
             Provide a well-structured description with clear responsibilities and requirements.",
            prompt
        );
        
        let response = self.generate(enhanced_prompt).await?;
        Ok(response.response)
    }
    
    /// Analyze a resume and provide feedback
    pub async fn analyze_resume(&self, resume_content: &str, job_description: &str) -> Result<String, OllamaError> {
        let prompt = format!(
            "Analyze the following resume against this job description:\n\n\
             Job Description:\n{}\n\n\
             Resume:\n{}\n\n\
             Provide specific feedback on:\n\
             1. How well the resume matches the job requirements\n\
             2. Missing skills or experience\n\
             3. Suggestions for improvement\n\
             4. Overall match percentage",
            job_description, resume_content
        );
        
        let response = self.generate(prompt).await?;
        Ok(response.response)
    }
    
    /// Generate a cover letter
    pub async fn generate_cover_letter(
        &self,
        candidate_info: &str,
        job_description: &str,
        tone: &str,
    ) -> Result<String, OllamaError> {
        let prompt = format!(
            "Write a {} cover letter for a candidate with the following background:\n\n\
             Candidate: {}\n\n\
             Job Description: {}\n\n\
             The cover letter should be professional and highlight relevant experience.",
            tone, candidate_info, job_description
        );
        
        let response = self.generate(prompt).await?;
        Ok(response.response)
    }
}

impl Default for OllamaClient {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_client_creation() {
        let client = OllamaClient::new();
        assert_eq!(client.get_model(), DEFAULT_MODEL);
    }
    
    #[test]
    fn test_client_with_custom_model() {
        let client = OllamaClient::new()
            .with_model("gemma:4b");
        assert_eq!(client.get_model(), "gemma:4b");
    }
    
    #[test]
    fn test_client_with_custom_url() {
        let client = OllamaClient::with_url("http://custom:11434");
        assert_eq!(client.get_model(), DEFAULT_MODEL);
    }
}