Skip to main content

praxis_llm/
traits.rs

1use crate::openai::{ReasoningConfig, ResponsesResponse};
2use crate::streaming::StreamEvent;
3use crate::types::{Message, Tool, ToolChoice};
4use anyhow::Result;
5use async_trait::async_trait;
6use futures::Stream;
7use serde::{Deserialize, Serialize};
8use std::pin::Pin;
9
10/// Trait for chat-based LLM interactions (GPT-4, etc)
11/// 
12/// Provides both streaming and non-streaming completions for conversational use cases.
13#[async_trait]
14pub trait ChatClient: Send + Sync {
15    /// Non-streaming chat completion
16    async fn chat(&self, request: ChatRequest) -> Result<ChatResponse>;
17    
18    /// Streaming chat completion
19    async fn chat_stream(
20        &self,
21        request: ChatRequest,
22    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent>> + Send>>>;
23}
24
25/// Trait for reasoning-based LLM interactions (o1 models)
26/// 
27/// Provides access to models with extended reasoning capabilities.
28#[async_trait]
29pub trait ReasoningClient: Send + Sync {
30    /// Non-streaming reasoning completion
31    async fn reason(&self, request: ResponseRequest) -> Result<ResponseOutput>;
32    
33    /// Streaming reasoning completion
34    async fn reason_stream(
35        &self,
36        request: ResponseRequest,
37    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent>> + Send>>>;
38}
39
40/// Convenience trait for clients that support both chat and reasoning
41pub trait LLMClient: ChatClient + ReasoningClient {}
42
43#[derive(Debug, Clone)]
44pub struct ChatRequest {
45    pub model: String,
46    pub messages: Vec<Message>,
47    pub options: ChatOptions,
48}
49
50impl ChatRequest {
51    pub fn new(model: impl Into<String>, messages: Vec<Message>) -> Self {
52        Self {
53            model: model.into(),
54            messages,
55            options: ChatOptions::default(),
56        }
57    }
58    
59    pub fn with_options(mut self, options: ChatOptions) -> Self {
60        self.options = options;
61        self
62    }
63}
64
65#[derive(Debug, Clone, Default)]
66pub struct ChatOptions {
67    pub temperature: Option<f32>,
68    pub max_tokens: Option<u32>,
69    pub tools: Option<Vec<Tool>>,
70    pub tool_choice: Option<ToolChoice>,
71    pub reasoning_effort: Option<String>,
72}
73
74impl ChatOptions {
75    pub fn new() -> Self {
76        Self::default()
77    }
78    
79    pub fn temperature(mut self, temp: f32) -> Self {
80        self.temperature = Some(temp);
81        self
82    }
83    
84    pub fn max_tokens(mut self, tokens: u32) -> Self {
85        self.max_tokens = Some(tokens);
86        self
87    }
88    
89    pub fn tools(mut self, tools: Vec<Tool>) -> Self {
90        self.tools = Some(tools);
91        self
92    }
93    
94    pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
95        self.tool_choice = Some(choice);
96        self
97    }
98    
99    pub fn reasoning_effort(mut self, effort: impl Into<String>) -> Self {
100        self.reasoning_effort = Some(effort.into());
101        self
102    }
103}
104
105#[derive(Debug, Clone)]
106pub struct ChatResponse {
107    pub content: Option<String>,
108    pub tool_calls: Option<Vec<crate::types::ToolCall>>,
109    pub usage: Option<TokenUsage>,
110    pub finish_reason: Option<String>,
111    pub raw: serde_json::Value,
112}
113
114#[derive(Debug, Clone)]
115pub struct ResponseRequest {
116    pub model: String,
117    pub input: Vec<Message>,
118    pub reasoning: Option<ReasoningConfig>,
119    pub options: ResponseOptions,
120}
121
122impl ResponseRequest {
123    pub fn new(model: impl Into<String>, input: Vec<Message>) -> Self {
124        Self {
125            model: model.into(),
126            input,
127            reasoning: None,
128            options: ResponseOptions::default(),
129        }
130    }
131    
132    pub fn with_reasoning(mut self, reasoning: ReasoningConfig) -> Self {
133        self.reasoning = Some(reasoning);
134        self
135    }
136    
137    pub fn with_options(mut self, options: ResponseOptions) -> Self {
138        self.options = options;
139        self
140    }
141}
142
143#[derive(Debug, Clone, Default)]
144pub struct ResponseOptions {
145    pub temperature: Option<f32>,
146    pub max_output_tokens: Option<u32>,
147}
148
149impl ResponseOptions {
150    pub fn new() -> Self {
151        Self::default()
152    }
153    
154    pub fn temperature(mut self, temp: f32) -> Self {
155        self.temperature = Some(temp);
156        self
157    }
158    
159    pub fn max_output_tokens(mut self, tokens: u32) -> Self {
160        self.max_output_tokens = Some(tokens);
161        self
162    }
163}
164
165#[derive(Debug, Clone)]
166pub struct ResponseOutput {
167    pub reasoning: Option<String>,
168    pub message: Option<String>,
169    pub usage: Option<TokenUsage>,
170    pub status: Option<String>,
171    pub raw: ResponsesResponse,
172}
173
174#[derive(Debug, Clone, Serialize, Deserialize)]
175pub struct TokenUsage {
176    pub input_tokens: u32,
177    pub output_tokens: u32,
178    pub total_tokens: u32,
179    #[serde(skip_serializing_if = "Option::is_none")]
180    pub reasoning_tokens: Option<u32>,
181}
182