Files
certifai/src/infrastructure/provider_client.rs
Sharang Parnerkar 50237f5377
All checks were successful
CI / Format (push) Successful in 2s
CI / Clippy (push) Successful in 2m13s
CI / Security Audit (push) Successful in 1m37s
CI / Tests (push) Successful in 2m52s
CI / Deploy (push) Successful in 2s
feat(chat): added chat interface and connection to ollama (#10)
Co-authored-by: Sharang Parnerkar <parnerkarsharang@gmail.com>
Reviewed-on: #10
2026-02-20 19:40:25 +00:00

149 lines
4.7 KiB
Rust

//! Unified LLM provider dispatch.
//!
//! Routes chat completion requests to Ollama, OpenAI, Anthropic, or
//! HuggingFace based on the session's provider setting. All providers
//! except Anthropic use the OpenAI-compatible chat completions format.
use reqwest::Client;
use serde::{Deserialize, Serialize};
use super::server_state::ServerState;
/// OpenAI-compatible chat message used for request bodies.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderMessage {
pub role: String,
pub content: String,
}
/// Send a chat completion request to the configured provider.
///
/// # Arguments
///
/// * `state` - Server state (for default Ollama URL/model)
/// * `provider` - Provider name (`"ollama"`, `"openai"`, `"anthropic"`, `"huggingface"`)
/// * `model` - Model ID
/// * `messages` - Conversation history
/// * `api_key` - API key (required for non-Ollama providers)
/// * `stream` - Whether to request streaming
///
/// # Returns
///
/// The raw `reqwest::Response` for the caller to consume (streaming or not).
///
/// # Errors
///
/// Returns an error if the HTTP request fails.
pub async fn send_chat_request(
state: &ServerState,
provider: &str,
model: &str,
messages: &[ProviderMessage],
api_key: Option<&str>,
stream: bool,
) -> Result<reqwest::Response, reqwest::Error> {
let client = Client::new();
match provider {
"openai" => {
let body = serde_json::json!({
"model": model,
"messages": messages,
"stream": stream,
});
client
.post("https://api.openai.com/v1/chat/completions")
.header("content-type", "application/json")
.header(
"Authorization",
format!("Bearer {}", api_key.unwrap_or_default()),
)
.json(&body)
.send()
.await
}
"anthropic" => {
// Anthropic uses a different API format -- translate.
// Extract system message separately, convert roles.
let system_msg: String = messages
.iter()
.filter(|m| m.role == "system")
.map(|m| m.content.clone())
.collect::<Vec<_>>()
.join("\n");
let anthropic_msgs: Vec<serde_json::Value> = messages
.iter()
.filter(|m| m.role != "system")
.map(|m| {
serde_json::json!({
"role": m.role,
"content": m.content,
})
})
.collect();
let mut body = serde_json::json!({
"model": model,
"messages": anthropic_msgs,
"max_tokens": 4096,
"stream": stream,
});
if !system_msg.is_empty() {
body["system"] = serde_json::Value::String(system_msg);
}
client
.post("https://api.anthropic.com/v1/messages")
.header("content-type", "application/json")
.header("x-api-key", api_key.unwrap_or_default())
.header("anthropic-version", "2023-06-01")
.json(&body)
.send()
.await
}
"huggingface" => {
let url = format!(
"https://api-inference.huggingface.co/models/{}/v1/chat/completions",
model
);
let body = serde_json::json!({
"model": model,
"messages": messages,
"stream": stream,
});
client
.post(&url)
.header("content-type", "application/json")
.header(
"Authorization",
format!("Bearer {}", api_key.unwrap_or_default()),
)
.json(&body)
.send()
.await
}
// Default: Ollama (OpenAI-compatible endpoint)
_ => {
let base_url = &state.services.ollama_url;
let resolved_model = if model.is_empty() {
&state.services.ollama_model
} else {
model
};
let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/'));
let body = serde_json::json!({
"model": resolved_model,
"messages": messages,
"stream": stream,
});
client
.post(&url)
.header("content-type", "application/json")
.json(&body)
.send()
.await
}
}
}