use crate::chat::tool_types::{StreamOutcome, ToolSpec};
use anyhow::Result;
use std::io::{BufRead, BufReader, Read};
use std::time::Duration;
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, serde::Serialize)]
pub struct ProviderId(pub String);
impl std::fmt::Display for ProviderId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<&str> for ProviderId {
fn from(value: &str) -> Self {
Self(value.to_string())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, serde::Serialize)]
pub struct ModelId(pub String);
impl std::fmt::Display for ModelId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<&str> for ModelId {
fn from(value: &str) -> Self {
Self(value.to_string())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum RuntimeKind {
Local,
Remote,
}
impl RuntimeKind {
pub fn label(self) -> &'static str {
match self {
Self::Local => "LOCAL",
Self::Remote => "REMOTE",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelInfo {
pub id: String,
pub context_length: Option<u64>,
pub size_bytes: Option<u64>,
pub quantization: Option<String>,
}
impl ModelInfo {
pub fn named(id: impl Into<String>) -> Self {
Self {
id: id.into(),
context_length: None,
size_bytes: None,
quantization: None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProviderHealth {
Checking,
Ready,
Unavailable(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
User,
Assistant,
}
#[derive(Debug, Clone)]
pub struct ChatMessage {
pub role: Role,
pub content: String,
pub tool_call: Option<crate::chat::tool_types::ToolCallRecord>,
pub tool_result: Option<crate::chat::tool_types::ToolResultRecord>,
}
impl ChatMessage {
pub fn text(role: Role, content: impl Into<String>) -> Self {
Self {
role,
content: content.into(),
tool_call: None,
tool_result: None,
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct ChatUsage {
pub input_tokens: u64,
pub output_tokens: u64,
}
impl ChatUsage {
pub fn merge(&mut self, other: ChatUsage) {
if other.input_tokens > 0 {
self.input_tokens = other.input_tokens;
}
if other.output_tokens > 0 {
self.output_tokens = other.output_tokens;
}
}
}
pub trait ChatProvider: Send + Sync {
fn name(&self) -> &str;
fn default_model(&self) -> &str;
fn requires_key(&self) -> bool;
fn env_var(&self) -> &str;
fn runtime_kind(&self) -> RuntimeKind {
RuntimeKind::Remote
}
fn supports_tool_calling(&self) -> bool {
true
}
fn list_models(&self, _api_key: Option<&str>) -> Result<Vec<ModelInfo>> {
Ok(vec![ModelInfo::named(self.default_model())])
}
fn health(&self, api_key: Option<&str>) -> ProviderHealth {
match self.list_models(api_key) {
Ok(_) => ProviderHealth::Ready,
Err(error) => ProviderHealth::Unavailable(error.to_string()),
}
}
fn stream_chat(
&self,
api_key: Option<&str>,
model: &str,
system: Option<&str>,
messages: &[ChatMessage],
tools: &[ToolSpec],
on_chunk: &mut dyn FnMut(&str) -> Result<()>,
) -> Result<(ChatUsage, StreamOutcome)>;
}
pub fn ask_once(
provider: &dyn ChatProvider,
api_key: Option<&str>,
model: &str,
system: &str,
user_message: &str,
) -> Result<String> {
let messages = [ChatMessage::text(Role::User, user_message)];
let mut full = String::new();
let (_, outcome) =
provider.stream_chat(api_key, model, Some(system), &messages, &[], &mut |chunk| {
full.push_str(chunk);
Ok(())
})?;
if matches!(outcome, StreamOutcome::ToolCalls(_)) {
anyhow::bail!("provider proposed a tool call in a non-tool-aware context (ask_once)");
}
Ok(full)
}
pub fn build_agent() -> ureq::Agent {
let config = ureq::Agent::config_builder()
.timeout_connect(Some(Duration::from_secs(10)))
.timeout_recv_response(Some(Duration::from_secs(30)))
.timeout_recv_body(Some(Duration::from_secs(300)))
.http_status_as_error(false)
.build();
ureq::Agent::new_with_config(config)
}
pub fn read_error_body(resp: &mut ureq::http::Response<ureq::Body>) -> String {
let mut buf = [0u8; 2048];
let n = resp.body_mut().as_reader().read(&mut buf).unwrap_or(0);
String::from_utf8_lossy(&buf[..n]).to_string()
}
pub fn read_sse_stream<R: Read>(
reader: R,
mut on_data: impl FnMut(&str) -> Result<()>,
) -> Result<()> {
let buf = BufReader::new(reader);
for line in buf.lines() {
let line = line?;
let Some(payload) = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))
else {
continue;
};
let payload = payload.trim();
if payload.is_empty() {
continue;
}
if payload == "[DONE]" {
break;
}
on_data(payload)?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
struct UnavailableProvider;
impl ChatProvider for UnavailableProvider {
fn name(&self) -> &str {
"offline"
}
fn default_model(&self) -> &str {
"none"
}
fn requires_key(&self) -> bool {
false
}
fn env_var(&self) -> &str {
""
}
fn list_models(&self, _api_key: Option<&str>) -> Result<Vec<ModelInfo>> {
anyhow::bail!("backend unavailable")
}
fn stream_chat(
&self,
_api_key: Option<&str>,
_model: &str,
_system: Option<&str>,
_messages: &[ChatMessage],
_tools: &[ToolSpec],
_on_chunk: &mut dyn FnMut(&str) -> Result<()>,
) -> Result<(ChatUsage, StreamOutcome)> {
anyhow::bail!("backend unavailable")
}
}
#[test]
fn unavailable_provider_health_is_actionable() {
assert_eq!(
UnavailableProvider.health(None),
ProviderHealth::Unavailable("backend unavailable".to_string())
);
}
#[test]
fn provider_id_and_model_id_display_their_inner_string() {
assert_eq!(ProviderId::from("anthropic").to_string(), "anthropic");
assert_eq!(ModelId::from("claude").to_string(), "claude");
}
}