Skip to main content

pe_core/
llm.rs

1//! LLM provider abstraction — the interface every execution path depends on.
2//!
3//! `LlmProvider` is object-safe (`Box<dyn LlmProvider>` works) via explicit
4//! `Pin<Box<dyn Future>>` return types. This is intentional — the trait must
5//! be storable in runtime structs without generics.
6//!
7//! Based on Group 15.1 of the pre-plan.
8
9use std::collections::HashMap;
10use std::future::Future;
11use std::pin::Pin;
12
13use futures::Stream;
14use serde::{Deserialize, Serialize};
15
16use crate::error::PeError;
17use crate::message::{AiMessage, Message};
18
19/// Future returned by [`LlmProvider::stream`].
20pub type StreamFuture<'a> = Pin<
21    Box<
22        dyn Future<Output = Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send>>, PeError>>
23            + Send
24            + 'a,
25    >,
26>;
27
28/// Core LLM abstraction. Every execution path in the library goes through this.
29///
30/// Implement this trait to plug in any model provider (OpenAI, Anthropic,
31/// Ollama, local models, etc.).
32///
33/// The trait is object-safe: `Box<dyn LlmProvider>` compiles.
34/// Uses `Pin<Box<dyn Future>>` instead of `async fn` for object safety.
35///
36/// # Example
37///
38/// ```ignore
39/// let provider: Box<dyn LlmProvider> = Box::new(MyProvider::new());
40/// let response = provider.complete(messages, tools).await?;
41/// ```
42pub trait LlmProvider: Send + Sync + 'static {
43    /// Non-streaming completion. Returns the full response.
44    ///
45    /// `tools` is empty when no tool calling is needed.
46    fn complete(
47        &self,
48        messages: &[Message],
49        tools: &[ToolSchema],
50    ) -> Pin<Box<dyn Future<Output = Result<LlmResponse, PeError>> + Send + '_>>;
51
52    /// Streaming completion. Yields tokens (and tool call deltas) as they arrive.
53    ///
54    /// The final item in the stream is `StreamChunk::Done` carrying the full response.
55    /// `tools` enables streaming tool calls (supported by Anthropic and OpenAI).
56    fn stream(&self, messages: &[Message], tools: &[ToolSchema]) -> StreamFuture<'_>;
57
58    /// Embed text into a vector. Used for semantic routing and memory search.
59    fn embed(
60        &self,
61        text: &str,
62    ) -> Pin<Box<dyn Future<Output = Result<Vec<f32>, PeError>> + Send + '_>>;
63
64    /// Human-readable provider name for logging ("openai", "anthropic", "mock").
65    fn provider_name(&self) -> &'static str;
66}
67
68/// Response from an LLM completion call.
69#[derive(Debug, Clone)]
70pub struct LlmResponse {
71    /// The AI message produced — contains content, tool_calls, usage.
72    pub message: AiMessage,
73    /// Raw provider metadata (model used, stop reason, etc.).
74    pub provider_metadata: HashMap<String, serde_json::Value>,
75}
76
77/// A chunk from a streaming LLM response.
78///
79#[derive(Debug, Clone)]
80#[non_exhaustive]
81pub enum StreamChunk {
82    /// A single token from the LLM output.
83    Token(String),
84    /// Streaming complete — carries the assembled full response.
85    Done(LlmResponse),
86    /// Error mid-stream.
87    Error(PeError),
88}
89
90/// JSON schema description for a tool, sent to the LLM in completion calls.
91///
92/// The LLM uses this to understand what tools are available and how to call them.
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct ToolSchema {
95    /// Tool name — must match the tool's registered name.
96    pub name: String,
97    /// Human-readable description of what the tool does.
98    pub description: String,
99    /// Parameters schema in JSON Schema format.
100    pub parameters: serde_json::Value,
101    /// Whether the LLM must strictly follow the schema (provider-dependent).
102    #[serde(default)]
103    pub strict: bool,
104}