feat(provider): provider support client

Author: bao-cn

Cargo.lock

@@ -22,6 +22,7 @@ futures = "0.3.32"
 futures-util = "0.3.32"
 serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0.150"
+serde_urlencoded = "0.7.1"
 rmcp = "1.7.0"
 rusqlite = { version = "0.40.0", features = ["bundled"] }
 ratatui = "0.30.0"

Cargo.toml

@@ -0,0 +1,109 @@
+use anyhow::{anyhow, Result};
+use bytes::Bytes;
+use futures_util::Stream;
+use reqwest::Method;
+use serde::{de::DeserializeOwned, Serialize};
+use std::collections::HashMap;
+use std::pin::Pin;
+
+use crate::core::client::base_client::BaseClient;
+use crate::core::client::config::ClientConfig;
+use crate::core::transport::http_transoprt::HttpTransport;
+use crate::core::transport::retry::RetryPolicy;
+use crate::core::transport::TransportRequest;
+
+#[derive(Clone)]
+pub struct AiClient {
+    transport: HttpTransport,
+}
+
+impl AiClient {
+    pub fn new(config: ClientConfig, retry: RetryPolicy) -> Result<Self> {
+        let base_client = BaseClient::new(config)?;
+        let transport = HttpTransport::new(base_client, retry);
+        
+        Ok(Self { transport })
+    }
+
+    pub async fn post_json<T>(
+        &self,
+        url: &str,
+        headers: HashMap<String, String>,
+        body: &impl Serialize,
+    ) -> Result<T>
+    where
+        T: DeserializeOwned,
+    {
+        let body_bytes = serde_json::to_vec(body)
+            .map_err(|e| anyhow!("Failed to serialize request body: {}", e))?;
+
+        let request = TransportRequest {
+            method: Method::POST,
+            url: url.to_string(),
+            headers,
+            body: Some(Bytes::from(body_bytes)),
+        };
+
+        let response = self.transport.execute(request).await?;
+        
+        serde_json::from_slice(&response.body)
+            .map_err(|e| anyhow!("Failed to deserialize response: {}", e))
+    }
+
+    pub async fn get_json<T>(
+        &self,
+        url: &str,
+        headers: HashMap<String, String>,
+        query: &impl Serialize,
+    ) -> Result<T>
+    where
+        T: DeserializeOwned,
+    {
+        let query_string = serde_urlencoded::to_string(query)
+            .map_err(|e| anyhow!("Failed to serialize query: {}", e))?;
+        
+        let url_with_query = if query_string.is_empty() {
+            url.to_string()
+        } else {
+            format!("{}?{}", url, query_string)
+        };
+
+        let request = TransportRequest {
+            method: Method::GET,
+            url: url_with_query,
+            headers,
+            body: None,
+        };
+
+        let response = self.transport.execute(request).await?;
+        
+        serde_json::from_slice(&response.body)
+            .map_err(|e| anyhow!("Failed to deserialize response: {}", e))
+    }
+
+    pub async fn post_stream<T>(
+        &self,
+        url: &str,
+        headers: HashMap<String, String>,
+        body: &impl Serialize,
+    ) -> Result<Pin<Box<dyn Stream<Item = Result<T>> + Send + 'static>>>
+    where
+        T: DeserializeOwned + Send + 'static,
+    {
+        let body_bytes = serde_json::to_vec(body)
+            .map_err(|e| anyhow!("Failed to serialize request body: {}", e))?;
+
+        let request = TransportRequest {
+            method: Method::POST,
+            url: url.to_string(),
+            headers,
+            body: Some(Bytes::from(body_bytes)),
+        };
+
+        self.transport.execute_stream(request).await.map_err(|e| anyhow!(e))
+    }
+
+    pub fn transport(&self) -> &HttpTransport {
+        &self.transport
+    }
+}

src/core/client/ai_client.rs

@@ -0,0 +1,2 @@
+pub mod provider;
+pub use provider::GeminiProvider;

src/core/providers/gemini/mod.rs

@@ -0,0 +1,72 @@
+use anyhow::Result;
+use async_trait::async_trait;
+use std::collections::HashMap;
+use crate::core::client::{AiClient, ClientConfig};
+use crate::core::transport::retry::RetryPolicy;
+use crate::core::providers::provider::Provider;
+use crate::core::providers::types::{CompletionRequest, CompletionResponse};
+
+pub struct GeminiProvider {
+    client: AiClient,
+    api_key: String,
+    model: String,
+}
+
+impl GeminiProvider {
+    pub fn new(api_key: String, model: String) -> Result<Self> {
+        let config = ClientConfig::default();
+        let retry = RetryPolicy::default();
+        let client = AiClient::new(config, retry)?;
+        
+        Ok(Self {
+            client,
+            api_key,
+            model,
+        })
+    }
+}
+
+#[async_trait]
+impl Provider for GeminiProvider {
+    fn client(&self) -> &AiClient {
+        &self.client
+    }
+    
+    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse> {
+        let mut headers = HashMap::new();
+        headers.insert("Content-Type".to_string(), "application/json".to_string());
+        
+        let url = format!("https://generativelanguage.googleapis.com/v1/models/{model}:generateContent?key={api_key}", 
+            model = self.model,
+            api_key = self.api_key
+        );
+        
+        let response: serde_json::Value = self.client
+            .post_json(&url, headers, &serde_json::json!({
+                "contents": request.messages.iter().map(|m| {
+                    serde_json::json!({
+                        "role": match m.role {
+                            crate::core::providers::types::Role::System => "system",
+                            crate::core::providers::types::Role::User => "user",
+                            crate::core::providers::types::Role::Assistant => "model",
+                        },
+                        "parts": [serde_json::json!({"text": m.content})]
+                    })
+                }).collect::<Vec<_>>(),
+                "generationConfig": {
+                    "temperature": request.temperature.unwrap_or(0.7),
+                    "maxOutputTokens": request.max_tokens,
+                }
+            }))
+            .await
+            .map_err(|e| anyhow::anyhow!("Gemini API error: {}", e))?;
+
+        let content = response["candidates"][0]["content"]["parts"][0]["text"].as_str().unwrap_or("").to_string();
+        
+        Ok(CompletionResponse {
+            content,
+            finish_reason: None,
+            usage: None,
+        })
+    }
+}

src/core/providers/gemini/provider.rs

@@ -0,0 +1,2 @@
+pub mod provider;
+pub use provider::OllamaProvider;

src/core/providers/ollama/mod.rs

@@ -0,0 +1,62 @@
+use anyhow::Result;
+use async_trait::async_trait;
+use std::collections::HashMap;
+use crate::core::client::{AiClient, ClientConfig};
+use crate::core::transport::retry::RetryPolicy;
+use crate::core::providers::provider::Provider;
+use crate::core::providers::types::{CompletionRequest, CompletionResponse};
+
+pub struct OllamaProvider {
+    client: AiClient,
+    base_url: String,
+}
+
+impl OllamaProvider {
+    pub fn new(base_url: String) -> Result<Self> {
+        let config = ClientConfig::default();
+        let retry = RetryPolicy::default();
+        let client = AiClient::new(config, retry)?;
+        
+        Ok(Self {
+            client,
+            base_url,
+        })
+    }
+}
+
+#[async_trait]
+impl Provider for OllamaProvider {
+    fn client(&self) -> &AiClient {
+        &self.client
+    }
+    
+    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse> {
+        let mut headers = HashMap::new();
+        headers.insert("Content-Type".to_string(), "application/json".to_string());
+        
+        let url = format!("{}/api/chat", self.base_url);
+        
+        let response: serde_json::Value = self.client
+            .post_json(&url, headers, &serde_json::json!({
+                "model": request.model,
+                "messages": request.messages,
+                "stream": request.stream,
+                "temperature": request.temperature,
+                "max_tokens": request.max_tokens,
+            }))
+            .await
+            .map_err(|e| anyhow::anyhow!("Ollama API error: {}", e))?;
+
+        let content = response["message"]["content"].as_str().unwrap_or("").to_string();
+        let finish_reason = match response["done"].as_bool() {
+            Some(true) => Some(crate::core::providers::types::FinishReason::Stop),
+            _ => None,
+        };
+
+        Ok(CompletionResponse {
+            content,
+            finish_reason,
+            usage: None,
+        })
+    }
+}

src/core/providers/ollama/provider.rs

@@ -1,11 +1,16 @@
 use std::time::Duration;
+use std::pin::Pin;
 
+use futures_util::Stream;
 use tokio::time::sleep;
 use crate::core::transport::errors::TransportError;
 use crate::core::transport::request::TransportRequest;
 use crate::core::transport::response::TransportResponse;
 use crate::core::transport::retry::RetryPolicy;
+use crate::core::transport::sse::{SseStream, sse_stream_to_json};
 use crate::core::client::base_client::BaseClient;
+use serde::de::DeserializeOwned;
+
 #[derive(Clone)]
 pub struct HttpTransport {
     client: BaseClient,
@@ -100,4 +105,26 @@ impl HttpTransport {
     fn backoff(&self, retry: usize) -> Duration {
         self.retry.base_delay * retry as u32
     }
+
+    pub async fn execute_stream<T>(
+        &self,
+        req: TransportRequest,
+    ) -> Result<Pin<Box<dyn Stream<Item = anyhow::Result<T>> + Send + 'static>>, TransportError>
+    where
+        T: DeserializeOwned + Send + 'static,
+    {
+        let request = self.build_request(req)?;
+        
+        let response = self.client.execute(request).await?;
+        
+        let status = response.status();
+        if status.is_client_error() || status.is_server_error() {
+            return Err(TransportError::Server(status.to_string()));
+        }
+
+        let sse_stream = SseStream::new(response);
+        let json_stream = sse_stream_to_json(sse_stream).await;
+        
+        Ok(json_stream)
+    }
 }