refactor: modularize codebase and add 404 unit tests (#13)

compliance-agent/src/llm/client.rs

@@ -1,147 +1,17 @@
 use secrecy::{ExposeSecret, SecretString};
-use serde::{Deserialize, Serialize};
 
+use super::types::*;
 use crate::error::AgentError;
 
 #[derive(Clone)]
 pub struct LlmClient {
-    base_url: String,
-    api_key: SecretString,
-    model: String,
-    embed_model: String,
-    http: reqwest::Client,
+    pub(crate) base_url: String,
+    pub(crate) api_key: SecretString,
+    pub(crate) model: String,
+    pub(crate) embed_model: String,
+    pub(crate) http: reqwest::Client,
 }
 
-// ── Request types ──────────────────────────────────────────────
-
-#[derive(Serialize, Clone, Debug)]
-pub struct ChatMessage {
-    pub role: String,
-    pub content: Option<String>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub tool_calls: Option<Vec<ToolCallRequest>>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub tool_call_id: Option<String>,
-}
-
-#[derive(Serialize)]
-struct ChatCompletionRequest {
-    model: String,
-    messages: Vec<ChatMessage>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    temperature: Option<f64>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    max_tokens: Option<u32>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    tools: Option<Vec<ToolDefinitionPayload>>,
-}
-
-#[derive(Serialize)]
-struct ToolDefinitionPayload {
-    r#type: String,
-    function: ToolFunctionPayload,
-}
-
-#[derive(Serialize)]
-struct ToolFunctionPayload {
-    name: String,
-    description: String,
-    parameters: serde_json::Value,
-}
-
-// ── Response types ─────────────────────────────────────────────
-
-#[derive(Deserialize)]
-struct ChatCompletionResponse {
-    choices: Vec<ChatChoice>,
-}
-
-#[derive(Deserialize)]
-struct ChatChoice {
-    message: ChatResponseMessage,
-}
-
-#[derive(Deserialize)]
-struct ChatResponseMessage {
-    #[serde(default)]
-    content: Option<String>,
-    #[serde(default)]
-    tool_calls: Option<Vec<ToolCallResponse>>,
-}
-
-#[derive(Deserialize)]
-struct ToolCallResponse {
-    id: String,
-    function: ToolCallFunction,
-}
-
-#[derive(Deserialize)]
-struct ToolCallFunction {
-    name: String,
-    arguments: String,
-}
-
-// ── Public types for tool calling ──────────────────────────────
-
-/// Definition of a tool that the LLM can invoke
-#[derive(Debug, Clone, Serialize)]
-pub struct ToolDefinition {
-    pub name: String,
-    pub description: String,
-    pub parameters: serde_json::Value,
-}
-
-/// A tool call request from the LLM
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct LlmToolCall {
-    pub id: String,
-    pub name: String,
-    pub arguments: serde_json::Value,
-}
-
-/// A tool call in the request message format (for sending back tool_calls in assistant messages)
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ToolCallRequest {
-    pub id: String,
-    pub r#type: String,
-    pub function: ToolCallRequestFunction,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ToolCallRequestFunction {
-    pub name: String,
-    pub arguments: String,
-}
-
-/// Response from the LLM — either content or tool calls
-#[derive(Debug, Clone)]
-pub enum LlmResponse {
-    Content(String),
-    /// Tool calls with optional reasoning text from the LLM
-    ToolCalls { calls: Vec<LlmToolCall>, reasoning: String },
-}
-
-// ── Embedding types ────────────────────────────────────────────
-
-#[derive(Serialize)]
-struct EmbeddingRequest {
-    model: String,
-    input: Vec<String>,
-}
-
-#[derive(Deserialize)]
-struct EmbeddingResponse {
-    data: Vec<EmbeddingData>,
-}
-
-#[derive(Deserialize)]
-struct EmbeddingData {
-    embedding: Vec<f64>,
-    index: usize,
-}
-
-// ── Implementation ─────────────────────────────────────────────
-
 impl LlmClient {
     pub fn new(
         base_url: String,
@@ -158,18 +28,14 @@ impl LlmClient {
         }
     }
 
-    pub fn embed_model(&self) -> &str {
-        &self.embed_model
-    }
-
-    fn chat_url(&self) -> String {
+    pub(crate) fn chat_url(&self) -> String {
         format!(
             "{}/v1/chat/completions",
             self.base_url.trim_end_matches('/')
         )
     }
 
-    fn auth_header(&self) -> Option<String> {
+    pub(crate) fn auth_header(&self) -> Option<String> {
         let key = self.api_key.expose_secret();
         if key.is_empty() {
             None
@@ -241,12 +107,12 @@ impl LlmClient {
             tools: None,
         };
 
-        self.send_chat_request(&request_body).await.map(|resp| {
-            match resp {
+        self.send_chat_request(&request_body)
+            .await
+            .map(|resp| match resp {
                 LlmResponse::Content(c) => c,
                 LlmResponse::ToolCalls { .. } => String::new(),
-            }
-        })
+            })
     }
 
     /// Chat with tool definitions — returns either content or tool calls.
@@ -292,7 +158,7 @@ impl LlmClient {
     ) -> Result<LlmResponse, AgentError> {
         let mut req = self
             .http
-            .post(&self.chat_url())
+            .post(self.chat_url())
             .header("content-type", "application/json")
             .json(request_body);
 
@@ -345,54 +211,7 @@ impl LlmClient {
         }
 
         // Otherwise return content
-        let content = choice
-            .message
-            .content
-            .clone()
-            .unwrap_or_default();
+        let content = choice.message.content.clone().unwrap_or_default();
         Ok(LlmResponse::Content(content))
     }
-
-    /// 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);
-
-        if let Some(auth) = self.auth_header() {
-            req = req.header("Authorization", auth);
-        }
-
-        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}")))?;
-
-        let mut data = body.data;
-        data.sort_by_key(|d| d.index);
-
-        Ok(data.into_iter().map(|d| d.embedding).collect())
-    }
 }