Skip to main content

neuron_types/
types.rs

1//! Core message and request/response types.
2
3use std::collections::HashMap;
4use std::path::PathBuf;
5use std::sync::Arc;
6
7use serde::{Deserialize, Serialize};
8use tokio_util::sync::CancellationToken;
9
10use crate::wasm::{WasmCompatSend, WasmCompatSync};
11
12/// The role of a message participant.
13#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
14pub enum Role {
15    /// A human user.
16    User,
17    /// An AI assistant.
18    Assistant,
19    /// A system message.
20    System,
21}
22
23/// A content block within a message.
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub enum ContentBlock {
26    /// Plain text content.
27    Text(String),
28    /// Extended thinking from reasoning models.
29    Thinking {
30        /// The thinking text.
31        thinking: String,
32        /// Cryptographic signature for verification.
33        signature: String,
34    },
35    /// Redacted thinking (not visible to user).
36    RedactedThinking {
37        /// Opaque data blob.
38        data: String,
39    },
40    /// A tool invocation request from the assistant.
41    ToolUse {
42        /// Unique identifier for this tool call.
43        id: String,
44        /// Name of the tool to invoke.
45        name: String,
46        /// JSON input arguments.
47        input: serde_json::Value,
48    },
49    /// Result of a tool invocation.
50    ToolResult {
51        /// References the `id` from the corresponding `ToolUse`.
52        tool_use_id: String,
53        /// Content items in the result.
54        content: Vec<ContentItem>,
55        /// Whether this result represents an error.
56        is_error: bool,
57    },
58    /// An image content block.
59    Image {
60        /// The image source.
61        source: ImageSource,
62    },
63    /// A document content block.
64    Document {
65        /// The document source.
66        source: DocumentSource,
67    },
68}
69
70/// A content item within a tool result.
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub enum ContentItem {
73    /// Plain text content.
74    Text(String),
75    /// An image.
76    Image {
77        /// The image source.
78        source: ImageSource,
79    },
80}
81
82/// Source of an image.
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub enum ImageSource {
85    /// Base64-encoded image data.
86    Base64 {
87        /// MIME type (e.g. "image/png").
88        media_type: String,
89        /// Base64-encoded data.
90        data: String,
91    },
92    /// URL to an image.
93    Url {
94        /// The image URL.
95        url: String,
96    },
97}
98
99/// Source of a document.
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub enum DocumentSource {
102    /// Base64-encoded PDF.
103    Base64Pdf {
104        /// Base64-encoded PDF data.
105        data: String,
106    },
107    /// Plain text document.
108    PlainText {
109        /// The text content.
110        data: String,
111    },
112    /// URL to a document.
113    Url {
114        /// The document URL.
115        url: String,
116    },
117}
118
119/// A message in a conversation.
120#[derive(Debug, Clone, Serialize, Deserialize)]
121pub struct Message {
122    /// The role of the message author.
123    pub role: Role,
124    /// The content blocks of this message.
125    pub content: Vec<ContentBlock>,
126}
127
128// --- Completion request/response types ---
129
130/// System prompt configuration.
131#[derive(Debug, Clone, Serialize, Deserialize)]
132pub enum SystemPrompt {
133    /// A simple text system prompt.
134    Text(String),
135    /// Structured system prompt blocks with optional cache control.
136    Blocks(Vec<SystemBlock>),
137}
138
139/// A block within a structured system prompt.
140#[derive(Debug, Clone, Serialize, Deserialize)]
141pub struct SystemBlock {
142    /// The text content of this block.
143    pub text: String,
144    /// Optional cache control for this block.
145    pub cache_control: Option<CacheControl>,
146}
147
148/// Cache control configuration for prompt caching.
149#[derive(Debug, Clone, Serialize, Deserialize)]
150pub struct CacheControl {
151    /// Time-to-live for the cached content.
152    pub ttl: Option<CacheTtl>,
153}
154
155/// Cache time-to-live options.
156#[derive(Debug, Clone, Serialize, Deserialize)]
157pub enum CacheTtl {
158    /// Cache for 5 minutes.
159    FiveMinutes,
160    /// Cache for 1 hour.
161    OneHour,
162}
163
164/// Tool selection strategy for the model.
165#[derive(Debug, Clone, Serialize, Deserialize)]
166pub enum ToolChoice {
167    /// Model decides whether to use tools.
168    Auto,
169    /// Model must not use tools.
170    None,
171    /// Model must use at least one tool.
172    Required,
173    /// Model must use the specified tool.
174    Specific {
175        /// Name of the required tool.
176        name: String,
177    },
178}
179
180/// Response format configuration.
181#[derive(Debug, Clone, Serialize, Deserialize)]
182pub enum ResponseFormat {
183    /// Plain text response.
184    Text,
185    /// JSON object response.
186    JsonObject,
187    /// Structured JSON response with schema validation.
188    JsonSchema {
189        /// Name of the schema.
190        name: String,
191        /// The JSON Schema.
192        schema: serde_json::Value,
193        /// Whether to enforce strict schema adherence.
194        strict: bool,
195    },
196}
197
198/// Extended thinking configuration for reasoning models.
199#[derive(Debug, Clone, Serialize, Deserialize)]
200pub enum ThinkingConfig {
201    /// Enable thinking with a token budget.
202    Enabled {
203        /// Maximum tokens for thinking.
204        budget_tokens: usize,
205    },
206    /// Disable thinking.
207    Disabled,
208    /// Let the model decide.
209    Adaptive,
210}
211
212/// Reasoning effort level for reasoning models.
213#[derive(Debug, Clone, Serialize, Deserialize)]
214pub enum ReasoningEffort {
215    /// No reasoning.
216    None,
217    /// Minimal reasoning.
218    Low,
219    /// Moderate reasoning.
220    Medium,
221    /// Maximum reasoning.
222    High,
223}
224
225/// A completion request to an LLM provider.
226#[derive(Debug, Clone, Default, Serialize, Deserialize)]
227pub struct CompletionRequest {
228    /// The model identifier.
229    pub model: String,
230    /// The conversation messages.
231    pub messages: Vec<Message>,
232    /// Optional system prompt.
233    pub system: Option<SystemPrompt>,
234    /// Tool definitions available to the model.
235    pub tools: Vec<ToolDefinition>,
236    /// Maximum tokens to generate.
237    pub max_tokens: Option<usize>,
238    /// Sampling temperature (0.0 to 1.0).
239    pub temperature: Option<f32>,
240    /// Nucleus sampling parameter.
241    pub top_p: Option<f32>,
242    /// Sequences that cause generation to stop.
243    pub stop_sequences: Vec<String>,
244    /// Tool selection strategy.
245    pub tool_choice: Option<ToolChoice>,
246    /// Response format constraint.
247    pub response_format: Option<ResponseFormat>,
248    /// Extended thinking configuration.
249    pub thinking: Option<ThinkingConfig>,
250    /// Reasoning effort level.
251    pub reasoning_effort: Option<ReasoningEffort>,
252    /// Provider-specific extra fields forwarded verbatim.
253    pub extra: Option<serde_json::Value>,
254}
255
256/// A completion response from an LLM provider.
257#[derive(Debug, Clone, Serialize, Deserialize)]
258pub struct CompletionResponse {
259    /// Provider-assigned message ID.
260    pub id: String,
261    /// The model that generated this response.
262    pub model: String,
263    /// The response message.
264    pub message: Message,
265    /// Token usage statistics.
266    pub usage: TokenUsage,
267    /// Why the model stopped generating.
268    pub stop_reason: StopReason,
269}
270
271/// Reason the model stopped generating.
272#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
273pub enum StopReason {
274    /// Model reached a natural end.
275    EndTurn,
276    /// Model wants to use a tool.
277    ToolUse,
278    /// Hit the max token limit.
279    MaxTokens,
280    /// Hit a stop sequence.
281    StopSequence,
282    /// Content was filtered.
283    ContentFilter,
284}
285
286/// Token usage statistics for a completion.
287#[derive(Debug, Clone, Default, Serialize, Deserialize)]
288pub struct TokenUsage {
289    /// Tokens in the input/prompt.
290    pub input_tokens: usize,
291    /// Tokens in the output/completion.
292    pub output_tokens: usize,
293    /// Tokens read from cache.
294    pub cache_read_tokens: Option<usize>,
295    /// Tokens written to cache.
296    pub cache_creation_tokens: Option<usize>,
297    /// Tokens used for reasoning/thinking.
298    pub reasoning_tokens: Option<usize>,
299}
300
301// --- Tool definition types (needed by CompletionRequest) ---
302
303/// Definition of a tool available to the model.
304#[derive(Debug, Clone, Serialize, Deserialize)]
305pub struct ToolDefinition {
306    /// The tool name (unique identifier).
307    pub name: String,
308    /// Optional human-readable title.
309    pub title: Option<String>,
310    /// Description of what the tool does.
311    pub description: String,
312    /// JSON Schema for the tool's input parameters.
313    pub input_schema: serde_json::Value,
314    /// Optional JSON Schema for the tool's output.
315    pub output_schema: Option<serde_json::Value>,
316    /// Optional behavioral annotations (MCP spec).
317    pub annotations: Option<ToolAnnotations>,
318    /// Optional cache control for this tool definition.
319    pub cache_control: Option<CacheControl>,
320}
321
322/// Behavioral annotations for a tool (from MCP spec).
323#[derive(Debug, Clone, Serialize, Deserialize)]
324pub struct ToolAnnotations {
325    /// Whether the tool only reads data.
326    pub read_only_hint: Option<bool>,
327    /// Whether the tool performs destructive operations.
328    pub destructive_hint: Option<bool>,
329    /// Whether repeated calls with same args produce same result.
330    pub idempotent_hint: Option<bool>,
331    /// Whether the tool interacts with external systems.
332    pub open_world_hint: Option<bool>,
333}
334
335/// Output from a tool execution.
336#[derive(Debug, Clone, Serialize, Deserialize)]
337pub struct ToolOutput {
338    /// Human-readable content items.
339    pub content: Vec<ContentItem>,
340    /// Optional structured JSON output for programmatic consumption.
341    pub structured_content: Option<serde_json::Value>,
342    /// Whether this output represents an error.
343    pub is_error: bool,
344}
345
346/// Runtime context provided to tools during execution.
347pub struct ToolContext {
348    /// Current working directory.
349    pub cwd: PathBuf,
350    /// Session identifier.
351    pub session_id: String,
352    /// Environment variables available to the tool.
353    pub environment: HashMap<String, String>,
354    /// Token for cooperative cancellation.
355    pub cancellation_token: CancellationToken,
356    /// Optional progress reporter for long-running tools.
357    pub progress_reporter: Option<Arc<dyn ProgressReporter>>,
358}
359
360/// Reports progress for long-running tool operations.
361pub trait ProgressReporter: WasmCompatSend + WasmCompatSync {
362    /// Report progress.
363    ///
364    /// # Arguments
365    /// * `progress` - Current progress value.
366    /// * `total` - Optional total value (for percentage calculation).
367    /// * `message` - Optional status message.
368    fn report(&self, progress: f64, total: Option<f64>, message: Option<&str>);
369}
370
371impl From<String> for SystemPrompt {
372    fn from(s: String) -> Self {
373        SystemPrompt::Text(s)
374    }
375}
376
377impl From<&str> for SystemPrompt {
378    fn from(s: &str) -> Self {
379        SystemPrompt::Text(s.to_string())
380    }
381}