feat: AI-driven automated penetration testing (#12)

compliance-agent/src/llm/client.rs

@@ -12,10 +12,16 @@ pub struct LlmClient {
     http: reqwest::Client,
 }
 
-#[derive(Serialize)]
-struct ChatMessage {
-    role: String,
-    content: String,
+// ── 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)]
@@ -26,8 +32,25 @@ struct ChatCompletionRequest {
     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>,
@@ -40,29 +63,85 @@ struct ChatChoice {
 
 #[derive(Deserialize)]
 struct ChatResponseMessage {
-    content: String,
+    #[serde(default)]
+    content: Option<String>,
+    #[serde(default)]
+    tool_calls: Option<Vec<ToolCallResponse>>,
 }
 
-/// Request body for the embeddings API
+#[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>,
 }
 
-/// 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,
 }
 
+// ── Implementation ─────────────────────────────────────────────
+
 impl LlmClient {
     pub fn new(
         base_url: String,
@@ -83,98 +162,142 @@ impl LlmClient {
         &self.embed_model
     }
 
+    fn chat_url(&self) -> String {
+        format!(
+            "{}/v1/chat/completions",
+            self.base_url.trim_end_matches('/')
+        )
+    }
+
+    fn auth_header(&self) -> Option<String> {
+        let key = self.api_key.expose_secret();
+        if key.is_empty() {
+            None
+        } else {
+            Some(format!("Bearer {key}"))
+        }
+    }
+
+    /// Simple chat: system + user prompt → text response
     pub async fn chat(
         &self,
         system_prompt: &str,
         user_prompt: &str,
         temperature: Option<f64>,
     ) -> Result<String, AgentError> {
-        let url = format!(
-            "{}/v1/chat/completions",
-            self.base_url.trim_end_matches('/')
-        );
+        let messages = vec![
+            ChatMessage {
+                role: "system".to_string(),
+                content: Some(system_prompt.to_string()),
+                tool_calls: None,
+                tool_call_id: None,
+            },
+            ChatMessage {
+                role: "user".to_string(),
+                content: Some(user_prompt.to_string()),
+                tool_calls: None,
+                tool_call_id: None,
+            },
+        ];
 
         let request_body = ChatCompletionRequest {
             model: self.model.clone(),
-            messages: vec![
-                ChatMessage {
-                    role: "system".to_string(),
-                    content: system_prompt.to_string(),
-                },
-                ChatMessage {
-                    role: "user".to_string(),
-                    content: user_prompt.to_string(),
-                },
-            ],
+            messages,
             temperature,
             max_tokens: Some(4096),
+            tools: None,
         };
 
-        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!("LiteLLM 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!(
-                "LiteLLM returned {status}: {body}"
-            )));
-        }
-
-        let body: ChatCompletionResponse = resp
-            .json()
-            .await
-            .map_err(|e| AgentError::Other(format!("Failed to parse LiteLLM response: {e}")))?;
-
-        body.choices
-            .first()
-            .map(|c| c.message.content.clone())
-            .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))
+        self.send_chat_request(&request_body).await.map(|resp| {
+            match resp {
+                LlmResponse::Content(c) => c,
+                LlmResponse::ToolCalls { .. } => String::new(), // shouldn't happen without tools
+            }
+        })
     }
 
+    /// Chat with a list of (role, content) messages → text response
     #[allow(dead_code)]
     pub async fn chat_with_messages(
         &self,
         messages: Vec<(String, String)>,
         temperature: Option<f64>,
     ) -> Result<String, AgentError> {
-        let url = format!(
-            "{}/v1/chat/completions",
-            self.base_url.trim_end_matches('/')
-        );
+        let messages = messages
+            .into_iter()
+            .map(|(role, content)| ChatMessage {
+                role,
+                content: Some(content),
+                tool_calls: None,
+                tool_call_id: None,
+            })
+            .collect();
 
         let request_body = ChatCompletionRequest {
             model: self.model.clone(),
-            messages: messages
-                .into_iter()
-                .map(|(role, content)| ChatMessage { role, content })
-                .collect(),
+            messages,
             temperature,
             max_tokens: Some(4096),
+            tools: None,
         };
 
+        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.
+    /// Use this for the AI pentest orchestrator loop.
+    pub async fn chat_with_tools(
+        &self,
+        messages: Vec<ChatMessage>,
+        tools: &[ToolDefinition],
+        temperature: Option<f64>,
+        max_tokens: Option<u32>,
+    ) -> Result<LlmResponse, AgentError> {
+        let tool_payloads: Vec<ToolDefinitionPayload> = tools
+            .iter()
+            .map(|t| ToolDefinitionPayload {
+                r#type: "function".to_string(),
+                function: ToolFunctionPayload {
+                    name: t.name.clone(),
+                    description: t.description.clone(),
+                    parameters: t.parameters.clone(),
+                },
+            })
+            .collect();
+
+        let request_body = ChatCompletionRequest {
+            model: self.model.clone(),
+            messages,
+            temperature,
+            max_tokens: Some(max_tokens.unwrap_or(8192)),
+            tools: if tool_payloads.is_empty() {
+                None
+            } else {
+                Some(tool_payloads)
+            },
+        };
+
+        self.send_chat_request(&request_body).await
+    }
+
+    /// Internal method to send a chat completion request and parse the response
+    async fn send_chat_request(
+        &self,
+        request_body: &ChatCompletionRequest,
+    ) -> Result<LlmResponse, AgentError> {
         let mut req = self
             .http
-            .post(&url)
+            .post(&self.chat_url())
             .header("content-type", "application/json")
-            .json(&request_body);
+            .json(request_body);
 
-        let key = self.api_key.expose_secret();
-        if !key.is_empty() {
-            req = req.header("Authorization", format!("Bearer {key}"));
+        if let Some(auth) = self.auth_header() {
+            req = req.header("Authorization", auth);
         }
 
         let resp = req
@@ -195,10 +318,39 @@ impl LlmClient {
             .await
             .map_err(|e| AgentError::Other(format!("Failed to parse LiteLLM response: {e}")))?;
 
-        body.choices
+        let choice = body
+            .choices
             .first()
-            .map(|c| c.message.content.clone())
-            .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))
+            .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))?;
+
+        // Check for tool calls first
+        if let Some(tool_calls) = &choice.message.tool_calls {
+            if !tool_calls.is_empty() {
+                let calls: Vec<LlmToolCall> = tool_calls
+                    .iter()
+                    .map(|tc| {
+                        let arguments = serde_json::from_str(&tc.function.arguments)
+                            .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
+                        LlmToolCall {
+                            id: tc.id.clone(),
+                            name: tc.function.name.clone(),
+                            arguments,
+                        }
+                    })
+                    .collect();
+                // Capture any reasoning text the LLM included alongside tool calls
+                let reasoning = choice.message.content.clone().unwrap_or_default();
+                return Ok(LlmResponse::ToolCalls { calls, reasoning });
+            }
+        }
+
+        // Otherwise return content
+        let content = choice
+            .message
+            .content
+            .clone()
+            .unwrap_or_default();
+        Ok(LlmResponse::Content(content))
     }
 
     /// Generate embeddings for a batch of texts
@@ -216,9 +368,8 @@ impl LlmClient {
             .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}"));
+        if let Some(auth) = self.auth_header() {
+            req = req.header("Authorization", auth);
         }
 
         let resp = req
@@ -239,7 +390,6 @@ impl LlmClient {
             .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);