synaptic_core/
lib.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use futures::Stream;
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use thiserror::Error;
11
12#[cfg(feature = "schemars")]
13pub use schemars;
14
15pub mod context_budget;
16pub mod token_counter;
17
18pub use context_budget::{ContextBudget, ContextSlot, Priority};
19pub use token_counter::{HeuristicTokenCounter, TokenCounter};
20
21// ---------------------------------------------------------------------------
22// ContentBlock — multimodal message content
23// ---------------------------------------------------------------------------
24
25/// A block of content within a message, supporting multimodal inputs.
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(tag = "type", rename_all = "snake_case")]
28pub enum ContentBlock {
29    Text {
30        text: String,
31    },
32    Image {
33        url: String,
34        #[serde(default, skip_serializing_if = "Option::is_none")]
35        detail: Option<String>,
36    },
37    Audio {
38        url: String,
39    },
40    Video {
41        url: String,
42    },
43    File {
44        url: String,
45        #[serde(default, skip_serializing_if = "Option::is_none")]
46        mime_type: Option<String>,
47    },
48    Data {
49        data: Value,
50    },
51    Reasoning {
52        content: String,
53    },
54}
55
56// ---------------------------------------------------------------------------
57// Message
58// ---------------------------------------------------------------------------
59
60/// Represents a chat message. Tagged enum with System, Human, AI, and Tool variants.
61#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
62#[serde(tag = "role")]
63pub enum Message {
64    #[serde(rename = "system")]
65    System {
66        content: String,
67        #[serde(default, skip_serializing_if = "Option::is_none")]
68        id: Option<String>,
69        #[serde(default, skip_serializing_if = "Option::is_none")]
70        name: Option<String>,
71        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
72        additional_kwargs: HashMap<String, Value>,
73        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
74        response_metadata: HashMap<String, Value>,
75        #[serde(default, skip_serializing_if = "Vec::is_empty")]
76        content_blocks: Vec<ContentBlock>,
77    },
78    #[serde(rename = "human")]
79    Human {
80        content: String,
81        #[serde(default, skip_serializing_if = "Option::is_none")]
82        id: Option<String>,
83        #[serde(default, skip_serializing_if = "Option::is_none")]
84        name: Option<String>,
85        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
86        additional_kwargs: HashMap<String, Value>,
87        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
88        response_metadata: HashMap<String, Value>,
89        #[serde(default, skip_serializing_if = "Vec::is_empty")]
90        content_blocks: Vec<ContentBlock>,
91    },
92    #[serde(rename = "assistant")]
93    AI {
94        content: String,
95        #[serde(default, skip_serializing_if = "Vec::is_empty")]
96        tool_calls: Vec<ToolCall>,
97        #[serde(default, skip_serializing_if = "Option::is_none")]
98        id: Option<String>,
99        #[serde(default, skip_serializing_if = "Option::is_none")]
100        name: Option<String>,
101        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
102        additional_kwargs: HashMap<String, Value>,
103        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
104        response_metadata: HashMap<String, Value>,
105        #[serde(default, skip_serializing_if = "Vec::is_empty")]
106        content_blocks: Vec<ContentBlock>,
107        #[serde(default, skip_serializing_if = "Option::is_none")]
108        usage_metadata: Option<TokenUsage>,
109        #[serde(default, skip_serializing_if = "Vec::is_empty")]
110        invalid_tool_calls: Vec<InvalidToolCall>,
111    },
112    #[serde(rename = "tool")]
113    Tool {
114        content: String,
115        tool_call_id: String,
116        #[serde(default, skip_serializing_if = "Option::is_none")]
117        id: Option<String>,
118        #[serde(default, skip_serializing_if = "Option::is_none")]
119        name: Option<String>,
120        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
121        additional_kwargs: HashMap<String, Value>,
122        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
123        response_metadata: HashMap<String, Value>,
124        #[serde(default, skip_serializing_if = "Vec::is_empty")]
125        content_blocks: Vec<ContentBlock>,
126    },
127    #[serde(rename = "chat")]
128    Chat {
129        custom_role: String,
130        content: String,
131        #[serde(default, skip_serializing_if = "Option::is_none")]
132        id: Option<String>,
133        #[serde(default, skip_serializing_if = "Option::is_none")]
134        name: Option<String>,
135        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
136        additional_kwargs: HashMap<String, Value>,
137        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
138        response_metadata: HashMap<String, Value>,
139        #[serde(default, skip_serializing_if = "Vec::is_empty")]
140        content_blocks: Vec<ContentBlock>,
141    },
142    /// A special message that signals removal of a message by its ID.
143    /// Used in message history management.
144    #[serde(rename = "remove")]
145    Remove {
146        /// ID of the message to remove.
147        id: String,
148    },
149}
150
151/// Helper macro to set a shared field across all Message variants.
152/// Note: Remove variant has no common fields, so it is a no-op.
153macro_rules! set_message_field {
154    ($self:expr, $field:ident, $value:expr) => {
155        match $self {
156            Message::System { $field, .. } => *$field = $value,
157            Message::Human { $field, .. } => *$field = $value,
158            Message::AI { $field, .. } => *$field = $value,
159            Message::Tool { $field, .. } => *$field = $value,
160            Message::Chat { $field, .. } => *$field = $value,
161            Message::Remove { .. } => { /* Remove has no common fields */ }
162        }
163    };
164}
165
166/// Helper macro to get a shared field from all Message variants.
167/// Note: Remove variant panics — callers handle Remove before using this macro.
168macro_rules! get_message_field {
169    ($self:expr, $field:ident) => {
170        match $self {
171            Message::System { $field, .. } => $field,
172            Message::Human { $field, .. } => $field,
173            Message::AI { $field, .. } => $field,
174            Message::Tool { $field, .. } => $field,
175            Message::Chat { $field, .. } => $field,
176            Message::Remove { .. } => unreachable!("get_message_field called on Remove variant"),
177        }
178    };
179}
180
181impl Message {
182    // -- Factory methods -----------------------------------------------------
183
184    pub fn system(content: impl Into<String>) -> Self {
185        Message::System {
186            content: content.into(),
187            id: None,
188            name: None,
189            additional_kwargs: HashMap::new(),
190            response_metadata: HashMap::new(),
191            content_blocks: Vec::new(),
192        }
193    }
194
195    pub fn human(content: impl Into<String>) -> Self {
196        Message::Human {
197            content: content.into(),
198            id: None,
199            name: None,
200            additional_kwargs: HashMap::new(),
201            response_metadata: HashMap::new(),
202            content_blocks: Vec::new(),
203        }
204    }
205
206    pub fn ai(content: impl Into<String>) -> Self {
207        Message::AI {
208            content: content.into(),
209            tool_calls: vec![],
210            id: None,
211            name: None,
212            additional_kwargs: HashMap::new(),
213            response_metadata: HashMap::new(),
214            content_blocks: Vec::new(),
215            usage_metadata: None,
216            invalid_tool_calls: Vec::new(),
217        }
218    }
219
220    pub fn ai_with_tool_calls(content: impl Into<String>, tool_calls: Vec<ToolCall>) -> Self {
221        Message::AI {
222            content: content.into(),
223            tool_calls,
224            id: None,
225            name: None,
226            additional_kwargs: HashMap::new(),
227            response_metadata: HashMap::new(),
228            content_blocks: Vec::new(),
229            usage_metadata: None,
230            invalid_tool_calls: Vec::new(),
231        }
232    }
233
234    pub fn tool(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
235        Message::Tool {
236            content: content.into(),
237            tool_call_id: tool_call_id.into(),
238            id: None,
239            name: None,
240            additional_kwargs: HashMap::new(),
241            response_metadata: HashMap::new(),
242            content_blocks: Vec::new(),
243        }
244    }
245
246    pub fn chat(role: impl Into<String>, content: impl Into<String>) -> Self {
247        Message::Chat {
248            custom_role: role.into(),
249            content: content.into(),
250            id: None,
251            name: None,
252            additional_kwargs: HashMap::new(),
253            response_metadata: HashMap::new(),
254            content_blocks: Vec::new(),
255        }
256    }
257
258    /// Create a Remove message that signals removal of a message by its ID.
259    pub fn remove(id: impl Into<String>) -> Self {
260        Message::Remove { id: id.into() }
261    }
262
263    // -- Builder methods -----------------------------------------------------
264
265    pub fn with_id(mut self, value: impl Into<String>) -> Self {
266        set_message_field!(&mut self, id, Some(value.into()));
267        self
268    }
269
270    pub fn with_name(mut self, value: impl Into<String>) -> Self {
271        set_message_field!(&mut self, name, Some(value.into()));
272        self
273    }
274
275    pub fn with_additional_kwarg(mut self, key: impl Into<String>, value: Value) -> Self {
276        match &mut self {
277            Message::System {
278                additional_kwargs, ..
279            }
280            | Message::Human {
281                additional_kwargs, ..
282            }
283            | Message::AI {
284                additional_kwargs, ..
285            }
286            | Message::Tool {
287                additional_kwargs, ..
288            }
289            | Message::Chat {
290                additional_kwargs, ..
291            } => {
292                additional_kwargs.insert(key.into(), value);
293            }
294            Message::Remove { .. } => { /* Remove has no additional_kwargs */ }
295        }
296        self
297    }
298
299    pub fn with_response_metadata_entry(mut self, key: impl Into<String>, value: Value) -> Self {
300        match &mut self {
301            Message::System {
302                response_metadata, ..
303            }
304            | Message::Human {
305                response_metadata, ..
306            }
307            | Message::AI {
308                response_metadata, ..
309            }
310            | Message::Tool {
311                response_metadata, ..
312            }
313            | Message::Chat {
314                response_metadata, ..
315            } => {
316                response_metadata.insert(key.into(), value);
317            }
318            Message::Remove { .. } => { /* Remove has no response_metadata */ }
319        }
320        self
321    }
322
323    pub fn with_content_blocks(mut self, blocks: Vec<ContentBlock>) -> Self {
324        set_message_field!(&mut self, content_blocks, blocks);
325        self
326    }
327
328    pub fn with_usage_metadata(mut self, usage: TokenUsage) -> Self {
329        if let Message::AI { usage_metadata, .. } = &mut self {
330            *usage_metadata = Some(usage);
331        }
332        self
333    }
334
335    // -- Accessor methods ----------------------------------------------------
336
337    pub fn content(&self) -> &str {
338        match self {
339            Message::Remove { .. } => "",
340            other => get_message_field!(other, content),
341        }
342    }
343
344    pub fn role(&self) -> &str {
345        match self {
346            Message::System { .. } => "system",
347            Message::Human { .. } => "human",
348            Message::AI { .. } => "assistant",
349            Message::Tool { .. } => "tool",
350            Message::Chat { custom_role, .. } => custom_role,
351            Message::Remove { .. } => "remove",
352        }
353    }
354
355    pub fn is_system(&self) -> bool {
356        matches!(self, Message::System { .. })
357    }
358
359    pub fn is_human(&self) -> bool {
360        matches!(self, Message::Human { .. })
361    }
362
363    pub fn is_ai(&self) -> bool {
364        matches!(self, Message::AI { .. })
365    }
366
367    pub fn is_tool(&self) -> bool {
368        matches!(self, Message::Tool { .. })
369    }
370
371    pub fn is_chat(&self) -> bool {
372        matches!(self, Message::Chat { .. })
373    }
374
375    pub fn is_remove(&self) -> bool {
376        matches!(self, Message::Remove { .. })
377    }
378
379    pub fn tool_calls(&self) -> &[ToolCall] {
380        match self {
381            Message::AI { tool_calls, .. } => tool_calls,
382            _ => &[],
383        }
384    }
385
386    pub fn tool_call_id(&self) -> Option<&str> {
387        match self {
388            Message::Tool { tool_call_id, .. } => Some(tool_call_id),
389            _ => None,
390        }
391    }
392
393    pub fn id(&self) -> Option<&str> {
394        match self {
395            Message::Remove { id } => Some(id),
396            other => get_message_field!(other, id).as_deref(),
397        }
398    }
399
400    pub fn name(&self) -> Option<&str> {
401        match self {
402            Message::Remove { .. } => None,
403            other => get_message_field!(other, name).as_deref(),
404        }
405    }
406
407    pub fn additional_kwargs(&self) -> &HashMap<String, Value> {
408        match self {
409            Message::System {
410                additional_kwargs, ..
411            }
412            | Message::Human {
413                additional_kwargs, ..
414            }
415            | Message::AI {
416                additional_kwargs, ..
417            }
418            | Message::Tool {
419                additional_kwargs, ..
420            }
421            | Message::Chat {
422                additional_kwargs, ..
423            } => additional_kwargs,
424            Message::Remove { .. } => {
425                static EMPTY: std::sync::OnceLock<HashMap<String, Value>> =
426                    std::sync::OnceLock::new();
427                EMPTY.get_or_init(HashMap::new)
428            }
429        }
430    }
431
432    pub fn response_metadata(&self) -> &HashMap<String, Value> {
433        match self {
434            Message::System {
435                response_metadata, ..
436            }
437            | Message::Human {
438                response_metadata, ..
439            }
440            | Message::AI {
441                response_metadata, ..
442            }
443            | Message::Tool {
444                response_metadata, ..
445            }
446            | Message::Chat {
447                response_metadata, ..
448            } => response_metadata,
449            Message::Remove { .. } => {
450                static EMPTY: std::sync::OnceLock<HashMap<String, Value>> =
451                    std::sync::OnceLock::new();
452                EMPTY.get_or_init(HashMap::new)
453            }
454        }
455    }
456
457    pub fn content_blocks(&self) -> &[ContentBlock] {
458        match self {
459            Message::Remove { .. } => &[],
460            other => get_message_field!(other, content_blocks),
461        }
462    }
463
464    /// Return the remove ID if this is a Remove message.
465    pub fn remove_id(&self) -> Option<&str> {
466        match self {
467            Message::Remove { id } => Some(id),
468            _ => None,
469        }
470    }
471
472    pub fn usage_metadata(&self) -> Option<&TokenUsage> {
473        match self {
474            Message::AI { usage_metadata, .. } => usage_metadata.as_ref(),
475            _ => None,
476        }
477    }
478
479    pub fn invalid_tool_calls(&self) -> &[InvalidToolCall] {
480        match self {
481            Message::AI {
482                invalid_tool_calls, ..
483            } => invalid_tool_calls,
484            _ => &[],
485        }
486    }
487
488    /// Set the content of this message. No-op for Remove variant.
489    pub fn set_content(&mut self, new_content: impl Into<String>) {
490        let new_content = new_content.into();
491        set_message_field!(self, content, new_content);
492    }
493}
494
495// ---------------------------------------------------------------------------
496// Message utility functions
497// ---------------------------------------------------------------------------
498
499/// Filter messages by type, name, or id.
500pub fn filter_messages(
501    messages: &[Message],
502    include_types: Option<&[&str]>,
503    exclude_types: Option<&[&str]>,
504    include_names: Option<&[&str]>,
505    exclude_names: Option<&[&str]>,
506    include_ids: Option<&[&str]>,
507    exclude_ids: Option<&[&str]>,
508) -> Vec<Message> {
509    messages
510        .iter()
511        .filter(|msg| {
512            if let Some(include) = include_types {
513                if !include.contains(&msg.role()) {
514                    return false;
515                }
516            }
517            if let Some(exclude) = exclude_types {
518                if exclude.contains(&msg.role()) {
519                    return false;
520                }
521            }
522            if let Some(include) = include_names {
523                match msg.name() {
524                    Some(name) => {
525                        if !include.contains(&name) {
526                            return false;
527                        }
528                    }
529                    None => return false,
530                }
531            }
532            if let Some(exclude) = exclude_names {
533                if let Some(name) = msg.name() {
534                    if exclude.contains(&name) {
535                        return false;
536                    }
537                }
538            }
539            if let Some(include) = include_ids {
540                match msg.id() {
541                    Some(id) => {
542                        if !include.contains(&id) {
543                            return false;
544                        }
545                    }
546                    None => return false,
547                }
548            }
549            if let Some(exclude) = exclude_ids {
550                if let Some(id) = msg.id() {
551                    if exclude.contains(&id) {
552                        return false;
553                    }
554                }
555            }
556            true
557        })
558        .cloned()
559        .collect()
560}
561
562/// Strategy for trimming messages.
563#[derive(Debug, Clone, Copy, PartialEq, Eq)]
564pub enum TrimStrategy {
565    /// Keep the first messages that fit within the token budget.
566    First,
567    /// Keep the last messages that fit within the token budget.
568    Last,
569}
570
571/// Trim messages to fit within a token budget.
572///
573/// `token_counter` receives a single message and returns its token count.
574/// When `include_system` is true and `strategy` is `Last`, the leading system
575/// message is always preserved.
576pub fn trim_messages(
577    messages: Vec<Message>,
578    max_tokens: usize,
579    token_counter: impl Fn(&Message) -> usize,
580    strategy: TrimStrategy,
581    include_system: bool,
582) -> Vec<Message> {
583    if messages.is_empty() {
584        return messages;
585    }
586
587    match strategy {
588        TrimStrategy::First => {
589            let mut result = Vec::new();
590            let mut total = 0;
591            for msg in messages {
592                let count = token_counter(&msg);
593                if total + count > max_tokens {
594                    break;
595                }
596                total += count;
597                result.push(msg);
598            }
599            result
600        }
601        TrimStrategy::Last => {
602            let (system_msg, rest) = if include_system && messages[0].is_system() {
603                (Some(messages[0].clone()), &messages[1..])
604            } else {
605                (None, messages.as_slice())
606            };
607
608            let system_tokens = system_msg.as_ref().map(&token_counter).unwrap_or(0);
609            let budget = max_tokens.saturating_sub(system_tokens);
610
611            let mut selected = Vec::new();
612            let mut total = 0;
613            for msg in rest.iter().rev() {
614                let count = token_counter(msg);
615                if total + count > budget {
616                    break;
617                }
618                total += count;
619                selected.push(msg.clone());
620            }
621            selected.reverse();
622
623            let mut result = Vec::new();
624            if let Some(sys) = system_msg {
625                result.push(sys);
626            }
627            result.extend(selected);
628            result
629        }
630    }
631}
632
633/// Merge consecutive messages of the same role into a single message.
634pub fn merge_message_runs(messages: Vec<Message>) -> Vec<Message> {
635    if messages.is_empty() {
636        return messages;
637    }
638
639    let mut result: Vec<Message> = Vec::new();
640
641    for msg in messages {
642        let should_merge = result
643            .last()
644            .map(|last| last.role() == msg.role())
645            .unwrap_or(false);
646
647        if should_merge {
648            let last = result.last_mut().unwrap();
649            // Merge content
650            let merged_content = format!("{}\n{}", last.content(), msg.content());
651            match last {
652                Message::System { content, .. } => *content = merged_content,
653                Message::Human { content, .. } => *content = merged_content,
654                Message::AI {
655                    content,
656                    tool_calls,
657                    invalid_tool_calls,
658                    ..
659                } => {
660                    *content = merged_content;
661                    tool_calls.extend(msg.tool_calls().to_vec());
662                    invalid_tool_calls.extend(msg.invalid_tool_calls().to_vec());
663                }
664                Message::Tool { content, .. } => *content = merged_content,
665                Message::Chat { content, .. } => *content = merged_content,
666                Message::Remove { .. } => { /* Remove messages are not merged */ }
667            }
668        } else {
669            result.push(msg);
670        }
671    }
672
673    result
674}
675
676/// Convert messages to a human-readable buffer string.
677pub fn get_buffer_string(messages: &[Message], human_prefix: &str, ai_prefix: &str) -> String {
678    messages
679        .iter()
680        .map(|msg| {
681            let prefix = match msg {
682                Message::System { .. } => "System",
683                Message::Human { .. } => human_prefix,
684                Message::AI { .. } => ai_prefix,
685                Message::Tool { .. } => "Tool",
686                Message::Chat { custom_role, .. } => custom_role.as_str(),
687                Message::Remove { .. } => "Remove",
688            };
689            format!("{prefix}: {}", msg.content())
690        })
691        .collect::<Vec<_>>()
692        .join("\n")
693}
694
695// ---------------------------------------------------------------------------
696// AIMessageChunk
697// ---------------------------------------------------------------------------
698
699/// A streaming chunk from an AI model response. Supports merge via `+`/`+=` operators and conversion to `Message` via `into_message()`.
700#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
701pub struct AIMessageChunk {
702    pub content: String,
703    #[serde(default, skip_serializing_if = "Vec::is_empty")]
704    pub tool_calls: Vec<ToolCall>,
705    #[serde(default, skip_serializing_if = "Option::is_none")]
706    pub usage: Option<TokenUsage>,
707    #[serde(default, skip_serializing_if = "Option::is_none")]
708    pub id: Option<String>,
709    #[serde(default, skip_serializing_if = "Vec::is_empty")]
710    pub tool_call_chunks: Vec<ToolCallChunk>,
711    #[serde(default, skip_serializing_if = "Vec::is_empty")]
712    pub invalid_tool_calls: Vec<InvalidToolCall>,
713}
714
715impl AIMessageChunk {
716    pub fn into_message(self) -> Message {
717        Message::ai_with_tool_calls(self.content, self.tool_calls)
718    }
719}
720
721impl std::ops::Add for AIMessageChunk {
722    type Output = Self;
723
724    fn add(mut self, rhs: Self) -> Self {
725        self += rhs;
726        self
727    }
728}
729
730impl std::ops::AddAssign for AIMessageChunk {
731    fn add_assign(&mut self, rhs: Self) {
732        self.content.push_str(&rhs.content);
733        self.tool_calls.extend(rhs.tool_calls);
734        self.tool_call_chunks.extend(rhs.tool_call_chunks);
735        self.invalid_tool_calls.extend(rhs.invalid_tool_calls);
736        if self.id.is_none() {
737            self.id = rhs.id;
738        }
739        match (&mut self.usage, rhs.usage) {
740            (Some(u), Some(rhs_u)) => {
741                u.input_tokens += rhs_u.input_tokens;
742                u.output_tokens += rhs_u.output_tokens;
743                u.total_tokens += rhs_u.total_tokens;
744            }
745            (None, Some(rhs_u)) => {
746                self.usage = Some(rhs_u);
747            }
748            _ => {}
749        }
750    }
751}
752
753// ---------------------------------------------------------------------------
754// Tool-related types
755// ---------------------------------------------------------------------------
756
757/// Represents a tool invocation requested by an AI model, with an ID, function name, and JSON arguments.
758#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
759pub struct ToolCall {
760    pub id: String,
761    pub name: String,
762    pub arguments: Value,
763}
764
765/// A tool call that failed to parse correctly.
766#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
767pub struct InvalidToolCall {
768    #[serde(default, skip_serializing_if = "Option::is_none")]
769    pub id: Option<String>,
770    #[serde(default, skip_serializing_if = "Option::is_none")]
771    pub name: Option<String>,
772    #[serde(default, skip_serializing_if = "Option::is_none")]
773    pub arguments: Option<String>,
774    pub error: String,
775}
776
777/// A partial tool call chunk received during streaming.
778#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
779pub struct ToolCallChunk {
780    #[serde(default, skip_serializing_if = "Option::is_none")]
781    pub id: Option<String>,
782    #[serde(default, skip_serializing_if = "Option::is_none")]
783    pub name: Option<String>,
784    #[serde(default, skip_serializing_if = "Option::is_none")]
785    pub arguments: Option<String>,
786    #[serde(default, skip_serializing_if = "Option::is_none")]
787    pub index: Option<usize>,
788}
789
790/// Schema definition for a tool, including its name, description, and JSON Schema for parameters.
791#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
792pub struct ToolDefinition {
793    pub name: String,
794    pub description: String,
795    pub parameters: Value,
796    /// Provider-specific parameters (e.g., Anthropic's `cache_control`).
797    #[serde(default, skip_serializing_if = "Option::is_none")]
798    pub extras: Option<HashMap<String, Value>>,
799}
800
801/// Controls how the model selects tools: Auto, Required, None, or a Specific named tool.
802#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
803#[serde(rename_all = "lowercase")]
804pub enum ToolChoice {
805    Auto,
806    Required,
807    None,
808    Specific(String),
809}
810
811// ---------------------------------------------------------------------------
812// Chat request / response
813// ---------------------------------------------------------------------------
814
815/// A request to a chat model containing messages, optional tool definitions, and tool choice configuration.
816#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
817pub struct ChatRequest {
818    pub messages: Vec<Message>,
819    #[serde(default, skip_serializing_if = "Vec::is_empty")]
820    pub tools: Vec<ToolDefinition>,
821    #[serde(default, skip_serializing_if = "Option::is_none")]
822    pub tool_choice: Option<ToolChoice>,
823}
824
825impl ChatRequest {
826    pub fn new(messages: Vec<Message>) -> Self {
827        Self {
828            messages,
829            tools: vec![],
830            tool_choice: None,
831        }
832    }
833
834    pub fn with_tools(mut self, tools: Vec<ToolDefinition>) -> Self {
835        self.tools = tools;
836        self
837    }
838
839    pub fn with_tool_choice(mut self, choice: ToolChoice) -> Self {
840        self.tool_choice = Some(choice);
841        self
842    }
843}
844
845/// A response from a chat model containing the AI message and optional token usage statistics.
846#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
847pub struct ChatResponse {
848    pub message: Message,
849    pub usage: Option<TokenUsage>,
850}
851
852// ---------------------------------------------------------------------------
853// Token usage
854// ---------------------------------------------------------------------------
855
856#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
857pub struct TokenUsage {
858    pub input_tokens: u32,
859    pub output_tokens: u32,
860    pub total_tokens: u32,
861    #[serde(default, skip_serializing_if = "Option::is_none")]
862    pub input_details: Option<InputTokenDetails>,
863    #[serde(default, skip_serializing_if = "Option::is_none")]
864    pub output_details: Option<OutputTokenDetails>,
865}
866
867/// Detailed breakdown of input token usage.
868#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
869pub struct InputTokenDetails {
870    #[serde(default)]
871    pub cached: u32,
872    #[serde(default)]
873    pub audio: u32,
874}
875
876/// Detailed breakdown of output token usage.
877#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
878pub struct OutputTokenDetails {
879    #[serde(default)]
880    pub reasoning: u32,
881    #[serde(default)]
882    pub audio: u32,
883}
884
885// ---------------------------------------------------------------------------
886// Events
887// ---------------------------------------------------------------------------
888
889/// Lifecycle events emitted during agent execution, used by `CallbackHandler` implementations.
890#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
891pub enum RunEvent {
892    RunStarted {
893        run_id: String,
894        session_id: String,
895    },
896    RunStep {
897        run_id: String,
898        step: usize,
899    },
900    LlmCalled {
901        run_id: String,
902        message_count: usize,
903    },
904    ToolCalled {
905        run_id: String,
906        tool_name: String,
907    },
908    RunFinished {
909        run_id: String,
910        output: String,
911    },
912    RunFailed {
913        run_id: String,
914        error: String,
915    },
916}
917
918// ---------------------------------------------------------------------------
919// Errors
920// ---------------------------------------------------------------------------
921
922/// Unified error type for the Synaptic framework with variants covering all subsystems.
923#[derive(Debug, Error)]
924pub enum SynapticError {
925    #[error("prompt error: {0}")]
926    Prompt(String),
927    #[error("model error: {0}")]
928    Model(String),
929    #[error("tool error: {0}")]
930    Tool(String),
931    #[error("tool not found: {0}")]
932    ToolNotFound(String),
933    #[error("memory error: {0}")]
934    Memory(String),
935    #[error("rate limit: {0}")]
936    RateLimit(String),
937    #[error("timeout: {0}")]
938    Timeout(String),
939    #[error("validation error: {0}")]
940    Validation(String),
941    #[error("parsing error: {0}")]
942    Parsing(String),
943    #[error("callback error: {0}")]
944    Callback(String),
945    #[error("max steps exceeded: {max_steps}")]
946    MaxStepsExceeded { max_steps: usize },
947    #[error("embedding error: {0}")]
948    Embedding(String),
949    #[error("vector store error: {0}")]
950    VectorStore(String),
951    #[error("retriever error: {0}")]
952    Retriever(String),
953    #[error("loader error: {0}")]
954    Loader(String),
955    #[error("splitter error: {0}")]
956    Splitter(String),
957    #[error("graph error: {0}")]
958    Graph(String),
959    #[error("cache error: {0}")]
960    Cache(String),
961    #[error("store error: {0}")]
962    Store(String),
963    #[error("config error: {0}")]
964    Config(String),
965    #[error("mcp error: {0}")]
966    Mcp(String),
967}
968
969// ---------------------------------------------------------------------------
970// Core traits
971// ---------------------------------------------------------------------------
972
973/// Type alias for a pinned, boxed async stream of `AIMessageChunk` results.
974pub type ChatStream<'a> =
975    Pin<Box<dyn Stream<Item = Result<AIMessageChunk, SynapticError>> + Send + 'a>>;
976
977/// Describes a model's capabilities and limits.
978#[derive(Debug, Clone, Serialize, Deserialize)]
979pub struct ModelProfile {
980    pub name: String,
981    pub provider: String,
982    pub supports_tool_calling: bool,
983    pub supports_structured_output: bool,
984    pub supports_streaming: bool,
985    pub max_input_tokens: Option<usize>,
986    pub max_output_tokens: Option<usize>,
987}
988
989/// The core trait for language model providers. Implementations provide `chat()` for single responses and optionally `stream_chat()` for streaming.
990#[async_trait]
991pub trait ChatModel: Send + Sync {
992    async fn chat(&self, request: ChatRequest) -> Result<ChatResponse, SynapticError>;
993
994    /// Return the model's capability profile, if known.
995    fn profile(&self) -> Option<ModelProfile> {
996        None
997    }
998
999    fn stream_chat(&self, request: ChatRequest) -> ChatStream<'_> {
1000        Box::pin(async_stream::stream! {
1001            match self.chat(request).await {
1002                Ok(response) => {
1003                    yield Ok(AIMessageChunk {
1004                        content: response.message.content().to_string(),
1005                        tool_calls: response.message.tool_calls().to_vec(),
1006                        usage: response.usage,
1007                        ..Default::default()
1008                    });
1009                }
1010                Err(e) => yield Err(e),
1011            }
1012        })
1013    }
1014}
1015
1016/// Defines an executable tool that can be called by an AI model. Each tool has a name, description, JSON schema for parameters, and an async `call()` method.
1017#[async_trait]
1018pub trait Tool: Send + Sync {
1019    fn name(&self) -> &'static str;
1020    fn description(&self) -> &'static str;
1021
1022    fn parameters(&self) -> Option<Value> {
1023        None
1024    }
1025
1026    async fn call(&self, args: Value) -> Result<Value, SynapticError>;
1027
1028    fn as_tool_definition(&self) -> ToolDefinition {
1029        ToolDefinition {
1030            name: self.name().to_string(),
1031            description: self.description().to_string(),
1032            parameters: self
1033                .parameters()
1034                .unwrap_or(serde_json::json!({"type": "object", "properties": {}})),
1035            extras: None,
1036        }
1037    }
1038}
1039
1040// ---------------------------------------------------------------------------
1041// ToolContext — context-aware tool execution
1042// ---------------------------------------------------------------------------
1043
1044/// Context passed to tools during graph execution.
1045///
1046/// Provides access to the current graph state (serialized as JSON),
1047/// the tool call ID, and an optional key-value store reference.
1048#[derive(Debug, Clone, Default)]
1049pub struct ToolContext {
1050    /// The current graph state, serialized as JSON.
1051    pub state: Option<Value>,
1052    /// The ID of the tool call being executed.
1053    pub tool_call_id: String,
1054}
1055
1056/// A tool that receives execution context from the graph.
1057///
1058/// This extends the basic `Tool` trait with graph-level context
1059/// (current state, store, tool call ID). Implement this for tools
1060/// that need to read or modify graph state.
1061#[async_trait]
1062pub trait ContextAwareTool: Send + Sync {
1063    fn name(&self) -> &'static str;
1064    fn description(&self) -> &'static str;
1065    async fn call_with_context(
1066        &self,
1067        args: Value,
1068        ctx: ToolContext,
1069    ) -> Result<Value, SynapticError>;
1070}
1071
1072/// Wrapper that adapts a `ContextAwareTool` into a standard `Tool`.
1073///
1074/// When used outside a graph context, the tool receives a default
1075/// (empty) `ToolContext`.
1076pub struct ContextAwareToolAdapter {
1077    inner: Arc<dyn ContextAwareTool>,
1078}
1079
1080impl ContextAwareToolAdapter {
1081    pub fn new(inner: Arc<dyn ContextAwareTool>) -> Self {
1082        Self { inner }
1083    }
1084}
1085
1086#[async_trait]
1087impl Tool for ContextAwareToolAdapter {
1088    fn name(&self) -> &'static str {
1089        self.inner.name()
1090    }
1091
1092    fn description(&self) -> &'static str {
1093        self.inner.description()
1094    }
1095
1096    async fn call(&self, args: Value) -> Result<Value, SynapticError> {
1097        self.inner
1098            .call_with_context(args, ToolContext::default())
1099            .await
1100    }
1101}
1102
1103// ---------------------------------------------------------------------------
1104// MemoryStore
1105// ---------------------------------------------------------------------------
1106
1107/// Persistent storage for conversation message history, keyed by session ID.
1108#[async_trait]
1109pub trait MemoryStore: Send + Sync {
1110    async fn append(&self, session_id: &str, message: Message) -> Result<(), SynapticError>;
1111    async fn load(&self, session_id: &str) -> Result<Vec<Message>, SynapticError>;
1112    async fn clear(&self, session_id: &str) -> Result<(), SynapticError>;
1113}
1114
1115/// Handler for lifecycle events during agent execution. Receives `RunEvent` notifications at each stage.
1116#[async_trait]
1117pub trait CallbackHandler: Send + Sync {
1118    async fn on_event(&self, event: RunEvent) -> Result<(), SynapticError>;
1119}
1120
1121// ---------------------------------------------------------------------------
1122// RunnableConfig
1123// ---------------------------------------------------------------------------
1124
1125/// Runtime configuration passed through runnable chains, including tags, metadata, concurrency limits, and run identification.
1126#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1127pub struct RunnableConfig {
1128    #[serde(default)]
1129    pub tags: Vec<String>,
1130    #[serde(default)]
1131    pub metadata: HashMap<String, Value>,
1132    #[serde(default)]
1133    pub max_concurrency: Option<usize>,
1134    #[serde(default)]
1135    pub recursion_limit: Option<usize>,
1136    #[serde(default)]
1137    pub run_id: Option<String>,
1138    #[serde(default)]
1139    pub run_name: Option<String>,
1140}
1141
1142impl RunnableConfig {
1143    pub fn with_tags(mut self, tags: Vec<String>) -> Self {
1144        self.tags = tags;
1145        self
1146    }
1147
1148    pub fn with_run_name(mut self, name: impl Into<String>) -> Self {
1149        self.run_name = Some(name.into());
1150        self
1151    }
1152
1153    pub fn with_run_id(mut self, id: impl Into<String>) -> Self {
1154        self.run_id = Some(id.into());
1155        self
1156    }
1157
1158    pub fn with_max_concurrency(mut self, max: usize) -> Self {
1159        self.max_concurrency = Some(max);
1160        self
1161    }
1162
1163    pub fn with_recursion_limit(mut self, limit: usize) -> Self {
1164        self.recursion_limit = Some(limit);
1165        self
1166    }
1167
1168    pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
1169        self.metadata.insert(key.into(), value);
1170        self
1171    }
1172}
1173
1174// ---------------------------------------------------------------------------
1175// Store trait (forward-declared in core, implemented in synaptic-store)
1176// ---------------------------------------------------------------------------
1177
1178/// A stored item in the key-value store.
1179#[derive(Debug, Clone, Serialize, Deserialize)]
1180pub struct Item {
1181    pub namespace: Vec<String>,
1182    pub key: String,
1183    pub value: Value,
1184    pub created_at: String,
1185    pub updated_at: String,
1186    /// Relevance score from a search operation (e.g., similarity score).
1187    #[serde(default, skip_serializing_if = "Option::is_none")]
1188    pub score: Option<f64>,
1189}
1190
1191/// Persistent key-value store trait for cross-invocation state.
1192///
1193/// Namespaces are hierarchical (represented as slices of strings) and
1194/// keys are strings within a namespace. Values are arbitrary JSON.
1195#[async_trait]
1196pub trait Store: Send + Sync {
1197    /// Get an item by namespace and key.
1198    async fn get(&self, namespace: &[&str], key: &str) -> Result<Option<Item>, SynapticError>;
1199
1200    /// Search items within a namespace.
1201    async fn search(
1202        &self,
1203        namespace: &[&str],
1204        query: Option<&str>,
1205        limit: usize,
1206    ) -> Result<Vec<Item>, SynapticError>;
1207
1208    /// Put (upsert) an item.
1209    async fn put(&self, namespace: &[&str], key: &str, value: Value) -> Result<(), SynapticError>;
1210
1211    /// Delete an item.
1212    async fn delete(&self, namespace: &[&str], key: &str) -> Result<(), SynapticError>;
1213
1214    /// List all namespaces, optionally filtered by prefix.
1215    async fn list_namespaces(&self, prefix: &[&str]) -> Result<Vec<Vec<String>>, SynapticError>;
1216}
1217
1218// ---------------------------------------------------------------------------
1219// Embeddings trait (forward-declared here, implemented in synaptic-embeddings)
1220// ---------------------------------------------------------------------------
1221
1222/// Trait for embedding text into vectors.
1223#[async_trait]
1224pub trait Embeddings: Send + Sync {
1225    /// Embed multiple texts (for batch document embedding).
1226    async fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, SynapticError>;
1227
1228    /// Embed a single query text.
1229    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, SynapticError>;
1230}
1231
1232// ---------------------------------------------------------------------------
1233// StreamWriter
1234// ---------------------------------------------------------------------------
1235
1236/// Custom stream writer that nodes can use to emit custom events.
1237pub type StreamWriter = Arc<dyn Fn(Value) + Send + Sync>;
1238
1239// ---------------------------------------------------------------------------
1240// Runtime types
1241// ---------------------------------------------------------------------------
1242
1243/// Graph execution runtime context passed to nodes and middleware.
1244#[derive(Clone)]
1245pub struct Runtime {
1246    pub store: Option<Arc<dyn Store>>,
1247    pub stream_writer: Option<StreamWriter>,
1248}
1249
1250/// Tool execution runtime context.
1251#[derive(Clone)]
1252pub struct ToolRuntime {
1253    pub store: Option<Arc<dyn Store>>,
1254    pub stream_writer: Option<StreamWriter>,
1255    pub state: Option<Value>,
1256    pub tool_call_id: String,
1257    pub config: Option<RunnableConfig>,
1258}
1259
1260// ---------------------------------------------------------------------------
1261// RuntimeAwareTool
1262// ---------------------------------------------------------------------------
1263
1264/// Context-aware tool that receives runtime information.
1265///
1266/// This extends the basic `Tool` trait with runtime context
1267/// (current state, store, stream writer, tool call ID). Implement this
1268/// for tools that need to read or modify graph state.
1269#[async_trait]
1270pub trait RuntimeAwareTool: Send + Sync {
1271    fn name(&self) -> &'static str;
1272    fn description(&self) -> &'static str;
1273
1274    fn parameters(&self) -> Option<Value> {
1275        None
1276    }
1277
1278    async fn call_with_runtime(
1279        &self,
1280        args: Value,
1281        runtime: ToolRuntime,
1282    ) -> Result<Value, SynapticError>;
1283
1284    fn as_tool_definition(&self) -> ToolDefinition {
1285        ToolDefinition {
1286            name: self.name().to_string(),
1287            description: self.description().to_string(),
1288            parameters: self
1289                .parameters()
1290                .unwrap_or(serde_json::json!({"type": "object", "properties": {}})),
1291            extras: None,
1292        }
1293    }
1294}
1295
1296/// Adapter that wraps a `RuntimeAwareTool` into a standard `Tool`.
1297///
1298/// When used outside a graph context, the tool receives a default
1299/// (empty) `ToolRuntime`.
1300pub struct RuntimeAwareToolAdapter {
1301    inner: Arc<dyn RuntimeAwareTool>,
1302    runtime: Arc<tokio::sync::RwLock<Option<ToolRuntime>>>,
1303}
1304
1305impl RuntimeAwareToolAdapter {
1306    pub fn new(tool: Arc<dyn RuntimeAwareTool>) -> Self {
1307        Self {
1308            inner: tool,
1309            runtime: Arc::new(tokio::sync::RwLock::new(None)),
1310        }
1311    }
1312
1313    pub async fn set_runtime(&self, runtime: ToolRuntime) {
1314        *self.runtime.write().await = Some(runtime);
1315    }
1316}
1317
1318#[async_trait]
1319impl Tool for RuntimeAwareToolAdapter {
1320    fn name(&self) -> &'static str {
1321        self.inner.name()
1322    }
1323
1324    fn description(&self) -> &'static str {
1325        self.inner.description()
1326    }
1327
1328    fn parameters(&self) -> Option<Value> {
1329        self.inner.parameters()
1330    }
1331
1332    async fn call(&self, args: Value) -> Result<Value, SynapticError> {
1333        let runtime = self.runtime.read().await.clone().unwrap_or(ToolRuntime {
1334            store: None,
1335            stream_writer: None,
1336            state: None,
1337            tool_call_id: String::new(),
1338            config: None,
1339        });
1340        self.inner.call_with_runtime(args, runtime).await
1341    }
1342}
1343
1344// ---------------------------------------------------------------------------
1345// Document
1346// ---------------------------------------------------------------------------
1347
1348/// A document with content and metadata, used throughout the retrieval pipeline.
1349#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1350pub struct Document {
1351    pub id: String,
1352    pub content: String,
1353    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
1354    pub metadata: HashMap<String, Value>,
1355}
1356
1357impl Document {
1358    pub fn new(id: impl Into<String>, content: impl Into<String>) -> Self {
1359        Self {
1360            id: id.into(),
1361            content: content.into(),
1362            metadata: HashMap::new(),
1363        }
1364    }
1365
1366    pub fn with_metadata(
1367        id: impl Into<String>,
1368        content: impl Into<String>,
1369        metadata: HashMap<String, Value>,
1370    ) -> Self {
1371        Self {
1372            id: id.into(),
1373            content: content.into(),
1374            metadata,
1375        }
1376    }
1377}
1378
1379// ---------------------------------------------------------------------------
1380// Retriever trait (forward-declared here, implementations in synaptic-retrieval)
1381// ---------------------------------------------------------------------------
1382
1383/// Trait for retrieving relevant documents given a query string.
1384#[async_trait]
1385pub trait Retriever: Send + Sync {
1386    async fn retrieve(&self, query: &str, top_k: usize) -> Result<Vec<Document>, SynapticError>;
1387}
1388
1389// ---------------------------------------------------------------------------
1390// VectorStore trait (forward-declared here, implementations in synaptic-vectorstores)
1391// ---------------------------------------------------------------------------
1392
1393/// Trait for vector storage backends.
1394#[async_trait]
1395pub trait VectorStore: Send + Sync {
1396    /// Add documents to the store, computing their embeddings.
1397    async fn add_documents(
1398        &self,
1399        docs: Vec<Document>,
1400        embeddings: &dyn Embeddings,
1401    ) -> Result<Vec<String>, SynapticError>;
1402
1403    /// Search for similar documents by query string.
1404    async fn similarity_search(
1405        &self,
1406        query: &str,
1407        k: usize,
1408        embeddings: &dyn Embeddings,
1409    ) -> Result<Vec<Document>, SynapticError>;
1410
1411    /// Search with similarity scores (higher = more similar).
1412    async fn similarity_search_with_score(
1413        &self,
1414        query: &str,
1415        k: usize,
1416        embeddings: &dyn Embeddings,
1417    ) -> Result<Vec<(Document, f32)>, SynapticError>;
1418
1419    /// Search by pre-computed embedding vector instead of text query.
1420    async fn similarity_search_by_vector(
1421        &self,
1422        embedding: &[f32],
1423        k: usize,
1424    ) -> Result<Vec<Document>, SynapticError>;
1425
1426    /// Delete documents by ID.
1427    async fn delete(&self, ids: &[&str]) -> Result<(), SynapticError>;
1428}
1429
1430// ---------------------------------------------------------------------------
1431// Loader trait (forward-declared here, implementations in synaptic-loaders)
1432// ---------------------------------------------------------------------------
1433
1434/// Trait for loading documents from various sources.
1435#[async_trait]
1436pub trait Loader: Send + Sync {
1437    /// Load all documents from this source.
1438    async fn load(&self) -> Result<Vec<Document>, SynapticError>;
1439
1440    /// Stream documents lazily. Default implementation wraps load().
1441    fn lazy_load(
1442        &self,
1443    ) -> Pin<Box<dyn Stream<Item = Result<Document, SynapticError>> + Send + '_>> {
1444        Box::pin(async_stream::stream! {
1445            match self.load().await {
1446                Ok(docs) => {
1447                    for doc in docs {
1448                        yield Ok(doc);
1449                    }
1450                }
1451                Err(e) => yield Err(e),
1452            }
1453        })
1454    }
1455}
1456
1457// ---------------------------------------------------------------------------
1458// LlmCache trait (forward-declared here, implementations in synaptic-cache)
1459// ---------------------------------------------------------------------------
1460
1461/// Trait for caching LLM responses.
1462#[async_trait]
1463pub trait LlmCache: Send + Sync {
1464    /// Look up a cached response by cache key.
1465    async fn get(&self, key: &str) -> Result<Option<ChatResponse>, SynapticError>;
1466    /// Store a response in the cache.
1467    async fn put(&self, key: &str, response: &ChatResponse) -> Result<(), SynapticError>;
1468    /// Clear all entries from the cache.
1469    async fn clear(&self) -> Result<(), SynapticError>;
1470}
1471
1472// ---------------------------------------------------------------------------
1473// Entrypoint / Task metadata (used by proc macros)
1474// ---------------------------------------------------------------------------
1475
1476/// Configuration for an `#[entrypoint]`-decorated function.
1477#[derive(Debug, Clone)]
1478pub struct EntrypointConfig {
1479    pub name: &'static str,
1480    pub checkpointer: Option<&'static str>,
1481}
1482
1483/// An entrypoint wrapping an async function as a runnable workflow.
1484///
1485/// The `invoke_fn` field is a type-erased async function (`Value -> Result<Value, SynapticError>`).
1486/// Type alias for the async entrypoint function signature.
1487pub type EntrypointFn = dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, SynapticError>> + Send>>
1488    + Send
1489    + Sync;
1490
1491pub struct Entrypoint {
1492    pub config: EntrypointConfig,
1493    pub invoke_fn: Box<EntrypointFn>,
1494}
1495
1496impl Entrypoint {
1497    pub async fn invoke(&self, input: Value) -> Result<Value, SynapticError> {
1498        (self.invoke_fn)(input).await
1499    }
1500}
synaptic_core/lib.rs

synaptic_core/
lib.rs
