feat: rag-embedding-ai-chat (#1)

Co-authored-by: Sharang Parnerkar <parnerkarsharang@gmail.com>
Reviewed-on: #1

compliance-agent/src/llm/client.rs

@@ -8,6 +8,7 @@ pub struct LlmClient {
     base_url: String,
     api_key: SecretString,
     model: String,
+    embed_model: String,
     http: reqwest::Client,
 }
 
@@ -42,16 +43,46 @@ struct ChatResponseMessage {
     content: String,
 }
 
+/// Request body for the embeddings API
+#[derive(Serialize)]
+struct EmbeddingRequest {
+    model: String,
+    input: Vec<String>,
+}
+
+/// Response from the embeddings API
+#[derive(Deserialize)]
+struct EmbeddingResponse {
+    data: Vec<EmbeddingData>,
+}
+
+/// A single embedding result
+#[derive(Deserialize)]
+struct EmbeddingData {
+    embedding: Vec<f64>,
+    index: usize,
+}
+
 impl LlmClient {
-    pub fn new(base_url: String, api_key: SecretString, model: String) -> Self {
+    pub fn new(
+        base_url: String,
+        api_key: SecretString,
+        model: String,
+        embed_model: String,
+    ) -> Self {
         Self {
             base_url,
             api_key,
             model,
+            embed_model,
             http: reqwest::Client::new(),
         }
     }
 
+    pub fn embed_model(&self) -> &str {
+        &self.embed_model
+    }
+
     pub async fn chat(
         &self,
         system_prompt: &str,
@@ -169,4 +200,49 @@ impl LlmClient {
             .map(|c| c.message.content.clone())
             .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))
     }
+
+    /// Generate embeddings for a batch of texts
+    pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
+        let url = format!("{}/v1/embeddings", self.base_url.trim_end_matches('/'));
+
+        let request_body = EmbeddingRequest {
+            model: self.embed_model.clone(),
+            input: texts,
+        };
+
+        let mut req = self
+            .http
+            .post(&url)
+            .header("content-type", "application/json")
+            .json(&request_body);
+
+        let key = self.api_key.expose_secret();
+        if !key.is_empty() {
+            req = req.header("Authorization", format!("Bearer {key}"));
+        }
+
+        let resp = req
+            .send()
+            .await
+            .map_err(|e| AgentError::Other(format!("Embedding request failed: {e}")))?;
+
+        if !resp.status().is_success() {
+            let status = resp.status();
+            let body = resp.text().await.unwrap_or_default();
+            return Err(AgentError::Other(format!(
+                "Embedding API returned {status}: {body}"
+            )));
+        }
+
+        let body: EmbeddingResponse = resp
+            .json()
+            .await
+            .map_err(|e| AgentError::Other(format!("Failed to parse embedding response: {e}")))?;
+
+        // Sort by index to maintain input order
+        let mut data = body.data;
+        data.sort_by_key(|d| d.index);
+
+        Ok(data.into_iter().map(|d| d.embedding).collect())
+    }
 }