use async_trait::async_trait;
use compact_str::CompactString;
use deepstrike_core::context::renderer::RenderedContext;
use deepstrike_core::runtime::session::ProviderReplay;
use deepstrike_core::types::message::{Content, Message, Role, ToolCall, ToolSchema};
use futures::{Stream, StreamExt};
pub mod anthropic;
pub mod openai;
pub type ProviderRunState = serde_json::Value;
#[derive(Debug, Clone, Default)]
pub struct RuntimePolicy {
pub max_turns: Option<u32>,
pub timeout_ms: Option<u64>,
}
#[derive(Debug, Clone)]
pub enum StreamEvent {
TextDelta {
delta: String,
},
ThinkingDelta {
delta: String,
},
ToolCall {
id: String,
name: String,
arguments: serde_json::Value,
},
Usage {
total_tokens: u32,
input_tokens: u32,
output_tokens: u32,
cache_read_input_tokens: u32,
cache_creation_input_tokens: u32,
cache_read_input_tokens_by_slot: Option<CacheReadBySlot>,
stop_reason: Option<String>,
},
Done,
}
#[derive(Debug, Clone, Default)]
pub struct CacheReadBySlot {
pub system: Option<u32>,
pub tools: Option<u32>,
pub messages: Option<u32>,
}
#[async_trait]
pub trait LLMProvider: Send + Sync {
fn create_run_state(&self) -> Option<ProviderRunState> {
None
}
fn runtime_policy(&self) -> RuntimePolicy {
RuntimePolicy::default()
}
fn peek_provider_replay(
&self,
_content: &str,
_tool_calls: &[ToolCall],
) -> Option<ProviderReplay> {
None
}
fn seed_provider_replay(
&self,
_content: &str,
_tool_calls: &[ToolCall],
_replay: &ProviderReplay,
) {
}
fn commit_stream_replay(&self, _content: &str, _tool_calls: &[ToolCall]) {}
async fn complete(
&self,
context: &RenderedContext,
tools: &[ToolSchema],
extensions: Option<&serde_json::Value>,
) -> crate::Result<Message> {
let mut stream = self.stream(context, tools, extensions, None).await?;
collect_message_from_stream(&mut stream).await
}
async fn stream(
&self,
context: &RenderedContext,
tools: &[ToolSchema],
extensions: Option<&serde_json::Value>,
state: Option<&ProviderRunState>,
) -> crate::Result<Box<dyn Stream<Item = crate::Result<StreamEvent>> + Send + Unpin>>;
}
pub async fn collect_message_from_stream(
stream: &mut (dyn Stream<Item = crate::Result<StreamEvent>> + Send + Unpin),
) -> crate::Result<Message> {
let mut content = String::new();
let mut tool_calls = Vec::new();
while let Some(evt) = stream.next().await {
match evt? {
StreamEvent::TextDelta { delta } => content.push_str(&delta),
StreamEvent::ThinkingDelta { .. } => {}
StreamEvent::ToolCall {
id,
name,
arguments,
} => {
tool_calls.push(ToolCall {
id: CompactString::new(&id),
name: CompactString::new(&name),
arguments,
});
}
StreamEvent::Usage { .. } | StreamEvent::Done => {}
}
}
Ok(Message {
role: Role::Assistant,
content: Content::Text(content),
tool_calls,
token_count: None,
})
}
#[derive(Debug, Clone, Default)]
pub struct TokenUsage {
pub input_tokens: u32,
pub output_tokens: u32,
pub cache_read_input_tokens: u32,
pub cache_creation_input_tokens: u32,
}
impl TokenUsage {
pub fn total_tokens(&self) -> u32 {
self.input_tokens + self.output_tokens
}
}
#[derive(Debug, Clone)]
pub struct ProviderToolSpec {
pub name: String,
pub description: String,
pub parameters: serde_json::Value,
}