stasis/application/runtime/
default_chat_middlewares.rs

1use std::sync::Arc;
2use std::time::Instant;
3
4use async_trait::async_trait;
5use genai::chat::{ChatOptions, ChatRequest, ChatResponse};
6use sha2::{Digest, Sha256};
7
8use crate::application::runtime::chat_client_middleware::ChatClientMiddleware;
9use crate::application::telemetry::spans as span_names;
10use crate::domain::errors::Result;
11use crate::ports::outbound::ai_chat_client::AiChatClient;
12use crate::ports::outbound::ai_chat_response_cache::AiChatResponseCache;
13use crate::ports::outbound::ai_chat_tool_interceptor::{AiChatToolInterceptor, AiToolCallEnvelope};
14use crate::ports::outbound::runtime::runtime_metrics::RuntimeMetrics;
15use crate::ports::outbound::runtime::runtime_tracing::{OtelAttribute, RuntimeTracing};
16
17pub use crate::application::telemetry::keys::{
18    CHAT_CACHE_HIT_TOTAL, CHAT_CACHE_MISS_TOTAL, CHAT_DURATION_MS, CHAT_ERRORS_TOTAL,
19    CHAT_REQUESTS_TOTAL, CHAT_TOOL_CALLS_TOTAL,
20};
21
22#[derive(Clone, Default)]
23pub struct LoggingChatMiddleware;
24
25impl ChatClientMiddleware for LoggingChatMiddleware {
26    fn wrap(&self, inner: Arc<dyn AiChatClient>) -> Arc<dyn AiChatClient> {
27        Arc::new(LoggingChatClient { inner })
28    }
29}
30
31#[derive(Clone)]
32struct LoggingChatClient {
33    inner: Arc<dyn AiChatClient>,
34}
35
36#[async_trait]
37impl AiChatClient for LoggingChatClient {
38    async fn complete(
39        &self,
40        request: ChatRequest,
41        options: Option<&ChatOptions>,
42    ) -> Result<ChatResponse> {
43        let started = Instant::now();
44        eprintln!(
45            "stasis.chat request messages={} options_present={}",
46            request.messages.len(),
47            options.is_some()
48        );
49
50        match self.inner.complete(request, options).await {
51            Ok(response) => {
52                eprintln!(
53                    "stasis.chat response ok elapsed_ms={}",
54                    started.elapsed().as_millis()
55                );
56                Ok(response)
57            }
58            Err(err) => {
59                eprintln!(
60                    "stasis.chat response error elapsed_ms={} error={}",
61                    started.elapsed().as_millis(),
62                    err
63                );
64                Err(err)
65            }
66        }
67    }
68}
69
70#[derive(Clone)]
71pub struct TelemetryChatMiddleware {
72    metrics: Arc<dyn RuntimeMetrics>,
73    tracing: Option<Arc<dyn RuntimeTracing>>,
74}
75
76impl TelemetryChatMiddleware {
77    pub fn new(metrics: Arc<dyn RuntimeMetrics>) -> Self {
78        Self {
79            metrics,
80            tracing: None,
81        }
82    }
83
84    pub fn with_tracing(mut self, tracing: Arc<dyn RuntimeTracing>) -> Self {
85        self.tracing = Some(tracing);
86        self
87    }
88}
89
90impl ChatClientMiddleware for TelemetryChatMiddleware {
91    fn wrap(&self, inner: Arc<dyn AiChatClient>) -> Arc<dyn AiChatClient> {
92        Arc::new(TelemetryChatClient {
93            inner,
94            metrics: self.metrics.clone(),
95            tracing: self.tracing.clone(),
96        })
97    }
98}
99
100#[derive(Clone)]
101struct TelemetryChatClient {
102    inner: Arc<dyn AiChatClient>,
103    metrics: Arc<dyn RuntimeMetrics>,
104    tracing: Option<Arc<dyn RuntimeTracing>>,
105}
106
107#[async_trait]
108impl AiChatClient for TelemetryChatClient {
109    async fn complete(
110        &self,
111        request: ChatRequest,
112        options: Option<&ChatOptions>,
113    ) -> Result<ChatResponse> {
114        let _span = self.tracing.as_ref().map(|tracing| {
115            tracing.start_span(
116                span_names::CHAT_COMPLETE,
117                &[OtelAttribute::int(
118                    "stasis.chat.messages",
119                    request.messages.len() as i64,
120                )],
121            )
122        });
123
124        self.metrics.incr_counter(CHAT_REQUESTS_TOTAL, 1);
125        let started = Instant::now();
126        match self.inner.complete(request, options).await {
127            Ok(response) => {
128                self.metrics
129                    .observe_duration_ms(CHAT_DURATION_MS, started.elapsed().as_millis() as u64);
130                Ok(response)
131            }
132            Err(err) => {
133                self.metrics.incr_counter(CHAT_ERRORS_TOTAL, 1);
134                self.metrics
135                    .observe_duration_ms(CHAT_DURATION_MS, started.elapsed().as_millis() as u64);
136                Err(err)
137            }
138        }
139    }
140}
141
142#[derive(Clone)]
143pub struct CacheChatMiddleware {
144    cache: Arc<dyn AiChatResponseCache>,
145    metrics: Option<Arc<dyn RuntimeMetrics>>,
146}
147
148impl CacheChatMiddleware {
149    pub fn new(cache: Arc<dyn AiChatResponseCache>) -> Self {
150        Self {
151            cache,
152            metrics: None,
153        }
154    }
155
156    pub fn with_metrics(mut self, metrics: Arc<dyn RuntimeMetrics>) -> Self {
157        self.metrics = Some(metrics);
158        self
159    }
160}
161
162impl ChatClientMiddleware for CacheChatMiddleware {
163    fn wrap(&self, inner: Arc<dyn AiChatClient>) -> Arc<dyn AiChatClient> {
164        Arc::new(CacheChatClient {
165            inner,
166            cache: self.cache.clone(),
167            metrics: self.metrics.clone(),
168        })
169    }
170}
171
172#[derive(Clone)]
173struct CacheChatClient {
174    inner: Arc<dyn AiChatClient>,
175    cache: Arc<dyn AiChatResponseCache>,
176    metrics: Option<Arc<dyn RuntimeMetrics>>,
177}
178
179#[async_trait]
180impl AiChatClient for CacheChatClient {
181    async fn complete(
182        &self,
183        request: ChatRequest,
184        options: Option<&ChatOptions>,
185    ) -> Result<ChatResponse> {
186        let cache_key = deterministic_cache_key(&request, options);
187        if let Some(cached) = self.cache.get(&cache_key) {
188            if let Some(metrics) = &self.metrics {
189                metrics.incr_counter(CHAT_CACHE_HIT_TOTAL, 1);
190            }
191            return Ok(cached);
192        }
193        if let Some(metrics) = &self.metrics {
194            metrics.incr_counter(CHAT_CACHE_MISS_TOTAL, 1);
195        }
196
197        let response = self.inner.complete(request, options).await?;
198        self.cache.set(&cache_key, response.clone());
199        Ok(response)
200    }
201}
202
203pub fn deterministic_cache_key(request: &ChatRequest, options: Option<&ChatOptions>) -> String {
204    let basis = format!("request={request:?}|options={options:?}");
205    let mut hasher = Sha256::new();
206    hasher.update(basis.as_bytes());
207    format!("chat:{}", hex::encode(hasher.finalize()))
208}
209
210#[derive(Clone)]
211pub struct ToolCallInterceptionChatMiddleware {
212    interceptor: Arc<dyn AiChatToolInterceptor>,
213    metrics: Option<Arc<dyn RuntimeMetrics>>,
214}
215
216impl ToolCallInterceptionChatMiddleware {
217    pub fn new(interceptor: Arc<dyn AiChatToolInterceptor>) -> Self {
218        Self {
219            interceptor,
220            metrics: None,
221        }
222    }
223
224    pub fn with_metrics(mut self, metrics: Arc<dyn RuntimeMetrics>) -> Self {
225        self.metrics = Some(metrics);
226        self
227    }
228}
229
230impl ChatClientMiddleware for ToolCallInterceptionChatMiddleware {
231    fn wrap(&self, inner: Arc<dyn AiChatClient>) -> Arc<dyn AiChatClient> {
232        Arc::new(ToolCallInterceptionChatClient {
233            inner,
234            interceptor: self.interceptor.clone(),
235            metrics: self.metrics.clone(),
236        })
237    }
238}
239
240#[derive(Clone)]
241struct ToolCallInterceptionChatClient {
242    inner: Arc<dyn AiChatClient>,
243    interceptor: Arc<dyn AiChatToolInterceptor>,
244    metrics: Option<Arc<dyn RuntimeMetrics>>,
245}
246
247#[async_trait]
248impl AiChatClient for ToolCallInterceptionChatClient {
249    async fn complete(
250        &self,
251        request: ChatRequest,
252        options: Option<&ChatOptions>,
253    ) -> Result<ChatResponse> {
254        let request_fingerprint = deterministic_cache_key(&request, options);
255        let response = self.inner.complete(request, options).await?;
256
257        let tool_calls = response.clone().into_tool_calls();
258        if !tool_calls.is_empty() {
259            let tool_call_count = tool_calls.len();
260            let tool_names = tool_calls.into_iter().map(|call| call.fn_name).collect();
261            self.interceptor.on_tool_calls(AiToolCallEnvelope {
262                request_fingerprint,
263                tool_call_count,
264                tool_names,
265            });
266            if let Some(metrics) = &self.metrics {
267                metrics.incr_counter(CHAT_TOOL_CALLS_TOTAL, tool_call_count as u64);
268            }
269        }
270
271        Ok(response)
272    }
273}
stasis/application/runtime/default_chat_middlewares.rs

stasis/application/runtime/
default_chat_middlewares.rs
