use std::collections::HashMap;
use std::fmt;
use std::pin::Pin;
use async_trait::async_trait;
use futures::stream::{Stream, StreamExt};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{error::LLMError, ToolCall};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
#[serde(alias = "input_tokens")]
pub prompt_tokens: u32,
#[serde(alias = "output_tokens")]
pub completion_tokens: u32,
pub total_tokens: u32,
#[serde(
skip_serializing_if = "Option::is_none",
alias = "output_tokens_details"
)]
pub completion_tokens_details: Option<CompletionTokensDetails>,
#[serde(
skip_serializing_if = "Option::is_none",
alias = "input_tokens_details"
)]
pub prompt_tokens_details: Option<PromptTokensDetails>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamResponse {
pub choices: Vec<StreamChoice>,
#[serde(skip_serializing_if = "Option::is_none")]
pub usage: Option<Usage>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamChoice {
pub delta: StreamDelta,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamDelta {
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
}
#[derive(Debug, Clone)]
pub enum StreamChunk {
Text(String),
ToolUseStart {
index: usize,
id: String,
name: String,
},
ToolUseInputDelta {
index: usize,
partial_json: String,
},
ToolUseComplete {
index: usize,
tool_call: ToolCall,
},
Done {
stop_reason: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CompletionTokensDetails {
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub audio_tokens: Option<u32>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptTokensDetails {
#[serde(skip_serializing_if = "Option::is_none")]
pub cached_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub audio_tokens: Option<u32>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChatRole {
User,
Assistant,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum ImageMime {
JPEG,
PNG,
GIF,
WEBP,
}
impl ImageMime {
pub fn mime_type(&self) -> &'static str {
match self {
ImageMime::JPEG => "image/jpeg",
ImageMime::PNG => "image/png",
ImageMime::GIF => "image/gif",
ImageMime::WEBP => "image/webp",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub enum MessageType {
#[default]
Text,
Image((ImageMime, Vec<u8>)),
Pdf(Vec<u8>),
ImageURL(String),
ToolUse(Vec<ToolCall>),
ToolResult(Vec<ToolCall>),
}
pub enum ReasoningEffort {
Low,
Medium,
High,
}
#[derive(Debug, Clone)]
pub struct ChatMessage {
pub role: ChatRole,
pub message_type: MessageType,
pub content: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct ParameterProperty {
#[serde(rename = "type")]
pub property_type: String,
pub description: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub items: Option<Box<ParameterProperty>>,
#[serde(skip_serializing_if = "Option::is_none", rename = "enum")]
pub enum_list: Option<Vec<String>>,
}
#[derive(Debug, Clone, Serialize)]
pub struct ParametersSchema {
#[serde(rename = "type")]
pub schema_type: String,
pub properties: HashMap<String, ParameterProperty>,
pub required: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct FunctionTool {
pub name: String,
pub description: String,
pub parameters: Value,
}
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct StructuredOutputFormat {
pub name: String,
pub description: Option<String>,
pub schema: Option<Value>,
pub strict: Option<bool>,
}
#[derive(Debug, Clone, Serialize)]
pub struct Tool {
#[serde(rename = "type")]
pub tool_type: String,
pub function: FunctionTool,
}
#[derive(Debug, Clone, Default)]
pub enum ToolChoice {
Any,
#[default]
Auto,
Tool(String),
None,
}
impl Serialize for ToolChoice {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
match self {
ToolChoice::Any => serializer.serialize_str("required"),
ToolChoice::Auto => serializer.serialize_str("auto"),
ToolChoice::None => serializer.serialize_str("none"),
ToolChoice::Tool(name) => {
use serde::ser::SerializeMap;
let mut map = serializer.serialize_map(Some(2))?;
map.serialize_entry("type", "function")?;
let mut function_obj = std::collections::HashMap::new();
function_obj.insert("name", name.as_str());
map.serialize_entry("function", &function_obj)?;
map.end()
}
}
}
}
pub trait ChatResponse: std::fmt::Debug + std::fmt::Display + Send + Sync {
fn text(&self) -> Option<String>;
fn tool_calls(&self) -> Option<Vec<ToolCall>>;
fn thinking(&self) -> Option<String> {
None
}
fn usage(&self) -> Option<Usage> {
None
}
}
#[async_trait]
pub trait ChatProvider: Sync + Send {
async fn chat(&self, messages: &[ChatMessage]) -> Result<Box<dyn ChatResponse>, LLMError> {
self.chat_with_tools(messages, None).await
}
async fn chat_with_tools(
&self,
messages: &[ChatMessage],
tools: Option<&[Tool]>,
) -> Result<Box<dyn ChatResponse>, LLMError>;
async fn chat_with_web_search(
&self,
_input: String,
) -> Result<Box<dyn ChatResponse>, LLMError> {
Err(LLMError::Generic(
"Web search not supported for this provider".to_string(),
))
}
async fn chat_stream(
&self,
_messages: &[ChatMessage],
) -> Result<std::pin::Pin<Box<dyn Stream<Item = Result<String, LLMError>> + Send>>, LLMError>
{
Err(LLMError::Generic(
"Streaming not supported for this provider".to_string(),
))
}
async fn chat_stream_struct(
&self,
_messages: &[ChatMessage],
) -> Result<
std::pin::Pin<Box<dyn Stream<Item = Result<StreamResponse, LLMError>> + Send>>,
LLMError,
> {
Err(LLMError::Generic(
"Structured streaming not supported for this provider".to_string(),
))
}
async fn chat_stream_with_tools(
&self,
_messages: &[ChatMessage],
_tools: Option<&[Tool]>,
) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, LLMError>> + Send>>, LLMError> {
Err(LLMError::Generic(
"Streaming with tools not supported for this provider".to_string(),
))
}
async fn memory_contents(&self) -> Option<Vec<ChatMessage>> {
None
}
async fn summarize_history(&self, msgs: &[ChatMessage]) -> Result<String, LLMError> {
let prompt = format!(
"Summarize in 2-3 sentences:\n{}",
msgs.iter()
.map(|m| format!("{:?}: {}", m.role, m.content))
.collect::<Vec<_>>()
.join("\n"),
);
let req = [ChatMessage::user().content(prompt).build()];
self.chat(&req)
.await?
.text()
.ok_or(LLMError::Generic("no text in summary response".into()))
}
}
impl fmt::Display for ReasoningEffort {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ReasoningEffort::Low => write!(f, "low"),
ReasoningEffort::Medium => write!(f, "medium"),
ReasoningEffort::High => write!(f, "high"),
}
}
}
impl ChatMessage {
pub fn user() -> ChatMessageBuilder {
ChatMessageBuilder::new(ChatRole::User)
}
pub fn assistant() -> ChatMessageBuilder {
ChatMessageBuilder::new(ChatRole::Assistant)
}
}
#[derive(Debug)]
pub struct ChatMessageBuilder {
role: ChatRole,
message_type: MessageType,
content: String,
}
impl ChatMessageBuilder {
pub fn new(role: ChatRole) -> Self {
Self {
role,
message_type: MessageType::default(),
content: String::new(),
}
}
pub fn content<S: Into<String>>(mut self, content: S) -> Self {
self.content = content.into();
self
}
pub fn image(mut self, image_mime: ImageMime, raw_bytes: Vec<u8>) -> Self {
self.message_type = MessageType::Image((image_mime, raw_bytes));
self
}
pub fn pdf(mut self, raw_bytes: Vec<u8>) -> Self {
self.message_type = MessageType::Pdf(raw_bytes);
self
}
pub fn image_url(mut self, url: impl Into<String>) -> Self {
self.message_type = MessageType::ImageURL(url.into());
self
}
pub fn tool_use(mut self, tools: Vec<ToolCall>) -> Self {
self.message_type = MessageType::ToolUse(tools);
self
}
pub fn tool_result(mut self, tools: Vec<ToolCall>) -> Self {
self.message_type = MessageType::ToolResult(tools);
self
}
pub fn build(self) -> ChatMessage {
ChatMessage {
role: self.role,
message_type: self.message_type,
content: self.content,
}
}
}
pub(crate) fn create_sse_stream<F>(
response: reqwest::Response,
parser: F,
) -> std::pin::Pin<Box<dyn Stream<Item = Result<String, LLMError>> + Send>>
where
F: Fn(&str) -> Result<Option<String>, LLMError> + Send + 'static,
{
let stream = response
.bytes_stream()
.scan(
(String::new(), Vec::new()),
move |(buffer, utf8_buffer), chunk| {
let result = match chunk {
Ok(bytes) => {
utf8_buffer.extend_from_slice(&bytes);
match String::from_utf8(utf8_buffer.clone()) {
Ok(text) => {
buffer.push_str(&text);
utf8_buffer.clear();
}
Err(e) => {
let valid_up_to = e.utf8_error().valid_up_to();
if valid_up_to > 0 {
let valid =
String::from_utf8_lossy(&utf8_buffer[..valid_up_to]);
buffer.push_str(&valid);
utf8_buffer.drain(..valid_up_to);
}
}
}
let mut results = Vec::new();
while let Some(pos) = buffer.find("\n\n") {
let event = buffer[..pos + 2].to_string();
buffer.drain(..pos + 2);
match parser(&event) {
Ok(Some(content)) => results.push(Ok(content)),
Ok(None) => {}
Err(e) => results.push(Err(e)),
}
}
Some(results)
}
Err(e) => Some(vec![Err(LLMError::HttpError(e.to_string()))]),
};
async move { result }
},
)
.flat_map(futures::stream::iter);
Box::pin(stream)
}
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
use futures::stream::StreamExt;
#[tokio::test]
async fn test_create_sse_stream_handles_split_utf8() {
let test_data = "data: Positive reactions\n\n".as_bytes();
let chunks: Vec<Result<Bytes, reqwest::Error>> = vec![
Ok(Bytes::from(&test_data[..10])),
Ok(Bytes::from(&test_data[10..])),
];
let mock_response = create_mock_response(chunks);
let parser = |event: &str| -> Result<Option<String>, LLMError> {
if let Some(content) = event.strip_prefix("data: ") {
let content = content.trim();
if content.is_empty() {
return Ok(None);
}
Ok(Some(content.to_string()))
} else {
Ok(None)
}
};
let mut stream = create_sse_stream(mock_response, parser);
let mut results = Vec::new();
while let Some(result) = stream.next().await {
results.push(result);
}
assert_eq!(results.len(), 1);
assert_eq!(results[0].as_ref().unwrap(), "Positive reactions");
}
#[tokio::test]
async fn test_create_sse_stream_handles_split_sse_events() {
let event1 = "data: First event\n\n";
let event2 = "data: Second event\n\n";
let combined = format!("{}{}", event1, event2);
let test_data = combined.as_bytes().to_vec();
let split_point = event1.len() + 5;
let chunks: Vec<Result<Bytes, reqwest::Error>> = vec![
Ok(Bytes::from(test_data[..split_point].to_vec())),
Ok(Bytes::from(test_data[split_point..].to_vec())),
];
let mock_response = create_mock_response(chunks);
let parser = |event: &str| -> Result<Option<String>, LLMError> {
if let Some(content) = event.strip_prefix("data: ") {
let content = content.trim();
if content.is_empty() {
return Ok(None);
}
Ok(Some(content.to_string()))
} else {
Ok(None)
}
};
let mut stream = create_sse_stream(mock_response, parser);
let mut results = Vec::new();
while let Some(result) = stream.next().await {
results.push(result);
}
assert_eq!(results.len(), 2);
assert_eq!(results[0].as_ref().unwrap(), "First event");
assert_eq!(results[1].as_ref().unwrap(), "Second event");
}
#[tokio::test]
async fn test_create_sse_stream_handles_multibyte_utf8_split() {
let multibyte_char = "✨";
let event = format!("data: Star {}\n\n", multibyte_char);
let test_data = event.as_bytes().to_vec();
let emoji_start = event.find(multibyte_char).unwrap();
let split_in_emoji = emoji_start + 1;
let chunks: Vec<Result<Bytes, reqwest::Error>> = vec![
Ok(Bytes::from(test_data[..split_in_emoji].to_vec())),
Ok(Bytes::from(test_data[split_in_emoji..].to_vec())),
];
let mock_response = create_mock_response(chunks);
let parser = |event: &str| -> Result<Option<String>, LLMError> {
if let Some(content) = event.strip_prefix("data: ") {
let content = content.trim();
if content.is_empty() {
return Ok(None);
}
Ok(Some(content.to_string()))
} else {
Ok(None)
}
};
let mut stream = create_sse_stream(mock_response, parser);
let mut results = Vec::new();
while let Some(result) = stream.next().await {
results.push(result);
}
assert_eq!(results.len(), 1);
assert_eq!(
results[0].as_ref().unwrap(),
&format!("Star {}", multibyte_char)
);
}
fn create_mock_response(chunks: Vec<Result<Bytes, reqwest::Error>>) -> reqwest::Response {
use http_body_util::StreamBody;
use reqwest::Body;
let frame_stream = futures::stream::iter(
chunks
.into_iter()
.map(|chunk| chunk.map(|bytes| hyper::body::Frame::data(bytes))),
);
let body = StreamBody::new(frame_stream);
let body = Body::wrap(body);
let http_response = http::Response::builder().status(200).body(body).unwrap();
http_response.into()
}
}