agents_core/
llm.rs

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/// Minimal request structure passed to a language model.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct LlmRequest {
12    pub system_prompt: String,
13    pub messages: Vec<AgentMessage>,
14    /// Available tools that the LLM can invoke
15    #[serde(default)]
16    pub tools: Vec<ToolSchema>,
17}
18
19impl LlmRequest {
20    /// Create a new LLM request
21    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    /// Add tools to the request
30    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/// A chunk of streaming response from the LLM
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub enum StreamChunk {
44    /// A text delta to append to the response
45    TextDelta(String),
46    /// The stream has finished
47    Done {
48        /// The complete final message
49        message: AgentMessage,
50    },
51    /// An error occurred during streaming
52    Error(String),
53}
54
55/// Type alias for a pinned boxed stream of chunks
56pub type ChunkStream = Pin<Box<dyn Stream<Item = anyhow::Result<StreamChunk>> + Send>>;
57
58#[async_trait]
59pub trait LanguageModel: Send + Sync {
60    /// Generate a complete response (non-streaming)
61    async fn generate(&self, request: LlmRequest) -> anyhow::Result<LlmResponse>;
62
63    /// Generate a streaming response
64    /// Default implementation falls back to non-streaming generate()
65    async fn generate_stream(&self, request: LlmRequest) -> anyhow::Result<ChunkStream> {
66        // Default implementation: call generate() and return complete response as a single chunk
67        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}