1use async_trait::async_trait;
2use futures::stream::Stream;
3use serde::{Deserialize, Serialize};
4use std::pin::Pin;
5
6use crate::messaging::AgentMessage;
7use crate::tools::ToolSchema;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct LlmRequest {
12 pub system_prompt: String,
13 pub messages: Vec<AgentMessage>,
14 #[serde(default)]
16 pub tools: Vec<ToolSchema>,
17}
18
19impl LlmRequest {
20 pub fn new(system_prompt: impl Into<String>, messages: Vec<AgentMessage>) -> Self {
22 Self {
23 system_prompt: system_prompt.into(),
24 messages,
25 tools: Vec::new(),
26 }
27 }
28
29 pub fn with_tools(mut self, tools: Vec<ToolSchema>) -> Self {
31 self.tools = tools;
32 self
33 }
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct LlmResponse {
38 pub message: AgentMessage,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
43pub enum StreamChunk {
44 TextDelta(String),
46 Done {
48 message: AgentMessage,
50 },
51 Error(String),
53}
54
55pub type ChunkStream = Pin<Box<dyn Stream<Item = anyhow::Result<StreamChunk>> + Send>>;
57
58#[async_trait]
59pub trait LanguageModel: Send + Sync {
60 async fn generate(&self, request: LlmRequest) -> anyhow::Result<LlmResponse>;
62
63 async fn generate_stream(&self, request: LlmRequest) -> anyhow::Result<ChunkStream> {
66 let response = self.generate(request).await?;
68 Ok(Box::pin(futures::stream::once(async move {
69 Ok(StreamChunk::Done {
70 message: response.message,
71 })
72 })))
73 }
74}