Skip to main content

entelix_runnable/
structured.rs

1//! `StructuredOutputAdapter` — bridges
2//! [`entelix_core::chat::ChatModel`]'s `complete_typed::<O>` into the
3//! `Runnable<Vec<Message>, O>` composition contract so typed
4//! structured outputs drop into `.pipe()` chains alongside
5//! [`ToolToRunnableAdapter`](crate::ToolToRunnableAdapter) and the
6//! built-in [`ChatModel`] `Runnable<Vec<Message>, Message>` impl.
7//!
8//! ```ignore
9//! use entelix::prelude::*;
10//! use entelix_runnable::ChatModelExt;
11//!
12//! let chain = prompt.pipe(model.with_structured_output::<Order>());
13//! let order: Order = chain.invoke(vars, &ctx).await?;
14//! ```
15//!
16//! The adapter is `Clone`-cheap (internal `Arc`) and forwards every
17//! `invoke` through `ChatModel::complete_typed`, inheriting the
18//! validation-retry budget configured on the underlying `ChatModel`
19//! (CLAUDE.md invariant 20).
20
21use std::marker::PhantomData;
22use std::sync::Arc;
23
24use entelix_core::chat::ChatModel;
25use entelix_core::codecs::Codec;
26use entelix_core::ir::Message;
27use entelix_core::transports::Transport;
28use entelix_core::{ExecutionContext, Result};
29
30use crate::runnable::Runnable;
31
32/// Adapts a `ChatModel<C, T>` to `Runnable<Vec<Message>, O>` by
33/// routing every `invoke` through
34/// [`ChatModel::complete_typed::<O>`](entelix_core::chat::ChatModel::complete_typed).
35///
36/// `O` is constrained to the same trait set as `complete_typed`:
37/// `schemars::JsonSchema + DeserializeOwned + Send + 'static`. The
38/// validation-retry budget configured on the underlying `ChatModel`
39/// (`ChatModelConfig::validation_retries`) flows through unchanged
40/// — schema-mismatch retries reflect the parser diagnostic to the
41/// model and re-invoke (CLAUDE.md invariant 20).
42///
43/// Construct via [`Self::new`] (consumes the model), [`Self::from_arc`]
44/// (shares an existing `Arc<ChatModel>`), or — most ergonomic — via
45/// [`crate::ChatModelExt::with_structured_output`].
46pub struct StructuredOutputAdapter<O, C: Codec, T: Transport> {
47    inner: Arc<ChatModel<C, T>>,
48    _phantom: PhantomData<fn() -> O>,
49}
50
51impl<O, C: Codec, T: Transport> StructuredOutputAdapter<O, C, T> {
52    /// Wrap a concrete `ChatModel`.
53    pub fn new(model: ChatModel<C, T>) -> Self {
54        Self {
55            inner: Arc::new(model),
56            _phantom: PhantomData,
57        }
58    }
59
60    /// Wrap an already-shared `Arc<ChatModel>` — avoids a second
61    /// `Arc::new` for operators that hold the model behind an `Arc`
62    /// already.
63    pub const fn from_arc(model: Arc<ChatModel<C, T>>) -> Self {
64        Self {
65            inner: model,
66            _phantom: PhantomData,
67        }
68    }
69
70    /// Borrow the wrapped `ChatModel` — useful for inspecting config
71    /// or threading the same model through multiple adapters.
72    pub const fn inner(&self) -> &Arc<ChatModel<C, T>> {
73        &self.inner
74    }
75}
76
77impl<O, C: Codec, T: Transport> Clone for StructuredOutputAdapter<O, C, T> {
78    fn clone(&self) -> Self {
79        Self {
80            inner: Arc::clone(&self.inner),
81            _phantom: PhantomData,
82        }
83    }
84}
85
86impl<O, C: Codec, T: Transport> std::fmt::Debug for StructuredOutputAdapter<O, C, T> {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.debug_struct("StructuredOutputAdapter")
89            .field("output", &std::any::type_name::<O>())
90            .finish()
91    }
92}
93
94#[async_trait::async_trait]
95impl<O, C, T> Runnable<Vec<Message>, O> for StructuredOutputAdapter<O, C, T>
96where
97    O: schemars::JsonSchema + serde::de::DeserializeOwned + Send + 'static,
98    C: Codec,
99    T: Transport,
100{
101    async fn invoke(&self, input: Vec<Message>, ctx: &ExecutionContext) -> Result<O> {
102        self.inner.complete_typed::<O>(input, ctx).await
103    }
104}
105
106/// Extension methods on [`ChatModel`] that produce typed `Runnable`
107/// adapters.
108///
109/// `model.with_structured_output::<Order>()` returns a
110/// `Runnable<Vec<Message>, Order>` ready for `.pipe()` composition —
111/// the [`with_structured_output`](Self::with_structured_output)
112/// ergonomic, typed.
113pub trait ChatModelExt<C: Codec, T: Transport>: Sized {
114    /// Adapt this `ChatModel` into a
115    /// `Runnable<Vec<Message>, O>` that routes every invocation
116    /// through `complete_typed::<O>`. The original model is consumed;
117    /// operators sharing the model behind an `Arc` reach for
118    /// [`StructuredOutputAdapter::from_arc`] directly.
119    fn with_structured_output<O>(self) -> StructuredOutputAdapter<O, C, T>
120    where
121        O: schemars::JsonSchema + serde::de::DeserializeOwned + Send + 'static;
122}
123
124impl<C, T> ChatModelExt<C, T> for ChatModel<C, T>
125where
126    C: Codec,
127    T: Transport,
128{
129    fn with_structured_output<O>(self) -> StructuredOutputAdapter<O, C, T>
130    where
131        O: schemars::JsonSchema + serde::de::DeserializeOwned + Send + 'static,
132    {
133        StructuredOutputAdapter::new(self)
134    }
135}