feat: add ProxyClient with upstream forwarding and reply caching

openai-proxy/Cargo.lock

@@ -1314,6 +1314,7 @@ dependencies = [
  "sqlx",
  "tokio",
  "tokio-test",
+ "tower",
  "tower-http",
  "tracing",
  "tracing-subscriber",

openai-proxy/Cargo.toml

@@ -22,6 +22,7 @@ tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 anyhow = "1"
 tower-http = { version = "0.6", features = ["cors", "trace"] }
+tower = { version = "0.5", features = ["util"] }
 futures = "0.3"
 bytes = "1"
 dotenv = "0.15"

openai-proxy/src/proxy.rs

@@ -0,0 +1,180 @@
+use anyhow::{Context, Result};
+use bytes::Bytes;
+use reqwest::Client;
+use serde_json::Value;
+use tracing::{debug, warn};
+
+use crate::cache::CacheStore;
+use crate::config::Config;
+use crate::types::ChatCompletionRequest;
+
+#[derive(Clone)]
+pub struct ProxyClient {
+    client: Client,
+    config: Config,
+    cache: CacheStore,
+}
+
+pub enum ProxyResponse {
+    Buffered(Bytes),
+    Streamed(reqwest::Response),
+}
+
+impl ProxyClient {
+    pub fn new(config: Config, cache: CacheStore) -> Result<Self> {
+        let client = Client::builder()
+            .build()
+            .context("failed to build reqwest client")?;
+        Ok(Self { client, config, cache })
+    }
+
+    pub async fn forward_request(&self, req: ChatCompletionRequest) -> Result<ProxyResponse> {
+        self.cache.cache_messages(&req).await?;
+
+        let is_streaming = req.stream.unwrap_or(false);
+        let url = format!("{}/v1/chat/completions", self.config.upstream_url.trim_end_matches('/'));
+
+        debug!("Forwarding to upstream: {}", url);
+
+        let http_resp = self
+            .client
+            .post(&url)
+            .bearer_auth(&self.config.upstream_api_key)
+            .json(&req)
+            .send()
+            .await
+            .context("upstream request failed")?;
+
+        if !http_resp.status().is_success() {
+            let status = http_resp.status();
+            let body = http_resp.bytes().await.unwrap_or_default();
+            return Err(anyhow::anyhow!(
+                "upstream returned {}: {}",
+                status,
+                String::from_utf8_lossy(&body)
+            ));
+        }
+
+        if is_streaming {
+            return Ok(ProxyResponse::Streamed(http_resp));
+        }
+
+        let body_bytes = http_resp.bytes().await.context("failed to read upstream body")?;
+
+        if let Ok(json) = serde_json::from_slice::<Value>(&body_bytes) {
+            if let Some(content) = json
+                .get("choices")
+                .and_then(|c| c.get(0))
+                .and_then(|c| c.get("message"))
+                .and_then(|m| m.get("content"))
+                .and_then(|c| c.as_str())
+            {
+                if let Err(e) = self.cache.cache_assistant_reply(content).await {
+                    warn!("Failed to cache assistant reply: {}", e);
+                }
+            }
+        }
+
+        Ok(ProxyResponse::Buffered(body_bytes))
+    }
+
+    pub async fn forward_get(&self, path: &str) -> Result<Bytes> {
+        let url = format!(
+            "{}/{}",
+            self.config.upstream_url.trim_end_matches('/'),
+            path.trim_start_matches('/')
+        );
+        let resp = self
+            .client
+            .get(&url)
+            .bearer_auth(&self.config.upstream_api_key)
+            .send()
+            .await
+            .context("upstream GET failed")?;
+        let bytes = resp.bytes().await.context("failed to read upstream GET body")?;
+        Ok(bytes)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_url_construction() {
+        let base = "https://api.openai.com/";
+        let url = format!("{}/v1/chat/completions", base.trim_end_matches('/'));
+        assert_eq!(url, "https://api.openai.com/v1/chat/completions");
+    }
+
+    #[test]
+    fn test_url_construction_no_trailing_slash() {
+        let base = "https://api.openai.com";
+        let url = format!("{}/v1/chat/completions", base.trim_end_matches('/'));
+        assert_eq!(url, "https://api.openai.com/v1/chat/completions");
+    }
+
+    #[test]
+    fn test_get_path_construction() {
+        let base = "https://api.openai.com";
+        let path = "/v1/models";
+        let url = format!(
+            "{}/{}",
+            base.trim_end_matches('/'),
+            path.trim_start_matches('/')
+        );
+        assert_eq!(url, "https://api.openai.com/v1/models");
+    }
+
+    #[test]
+    fn test_assistant_content_extraction_from_response() {
+        let json = serde_json::json!({
+            "id": "chatcmpl-abc",
+            "object": "chat.completion",
+            "created": 1699000000,
+            "model": "gpt-4o",
+            "choices": [{
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "The answer is 42."
+                },
+                "finish_reason": "stop"
+            }],
+            "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        });
+
+        let content = json
+            .get("choices")
+            .and_then(|c| c.get(0))
+            .and_then(|c| c.get("message"))
+            .and_then(|m| m.get("content"))
+            .and_then(|c| c.as_str());
+
+        assert_eq!(content, Some("The answer is 42."));
+    }
+
+    #[test]
+    fn test_assistant_content_extraction_null_content() {
+        let json = serde_json::json!({
+            "choices": [{
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": null,
+                    "tool_calls": []
+                },
+                "finish_reason": "tool_calls"
+            }]
+        });
+
+        let content = json
+            .get("choices")
+            .and_then(|c| c.get(0))
+            .and_then(|c| c.get("message"))
+            .and_then(|m| m.get("content"))
+            .and_then(|c| c.as_str());
+
+        assert!(content.is_none(), "null content should not be cached");
+    }
+}