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::domain::errors::Result;
10use crate::ports::outbound::ai_chat_client::AiChatClient;
11use crate::ports::outbound::ai_chat_response_cache::AiChatResponseCache;
12use crate::ports::outbound::ai_chat_tool_interceptor::{AiChatToolInterceptor, AiToolCallEnvelope};
13use crate::ports::outbound::runtime::runtime_metrics::RuntimeMetrics;
14
15pub const CHAT_REQUESTS_TOTAL: &str = "runtime.chat.requests.total";
16pub const CHAT_ERRORS_TOTAL: &str = "runtime.chat.errors.total";
17pub const CHAT_DURATION_MS: &str = "runtime.chat.duration_ms";
18pub const CHAT_CACHE_HIT_TOTAL: &str = "runtime.chat.cache.hit.total";
19pub const CHAT_CACHE_MISS_TOTAL: &str = "runtime.chat.cache.miss.total";
20pub const CHAT_TOOL_CALLS_TOTAL: &str = "runtime.chat.tool_calls.total";
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}
74
75impl TelemetryChatMiddleware {
76    pub fn new(metrics: Arc<dyn RuntimeMetrics>) -> Self {
77        Self { metrics }
78    }
79}
80
81impl ChatClientMiddleware for TelemetryChatMiddleware {
82    fn wrap(&self, inner: Arc<dyn AiChatClient>) -> Arc<dyn AiChatClient> {
83        Arc::new(TelemetryChatClient {
84            inner,
85            metrics: self.metrics.clone(),
86        })
87    }
88}
89
90#[derive(Clone)]
91struct TelemetryChatClient {
92    inner: Arc<dyn AiChatClient>,
93    metrics: Arc<dyn RuntimeMetrics>,
94}
95
96#[async_trait]
97impl AiChatClient for TelemetryChatClient {
98    async fn complete(
99        &self,
100        request: ChatRequest,
101        options: Option<&ChatOptions>,
102    ) -> Result<ChatResponse> {
103        self.metrics.incr_counter(CHAT_REQUESTS_TOTAL, 1);
104        let started = Instant::now();
105        match self.inner.complete(request, options).await {
106            Ok(response) => {
107                self.metrics
108                    .observe_duration_ms(CHAT_DURATION_MS, started.elapsed().as_millis() as u64);
109                Ok(response)
110            }
111            Err(err) => {
112                self.metrics.incr_counter(CHAT_ERRORS_TOTAL, 1);
113                self.metrics
114                    .observe_duration_ms(CHAT_DURATION_MS, started.elapsed().as_millis() as u64);
115                Err(err)
116            }
117        }
118    }
119}
120
121#[derive(Clone)]
122pub struct CacheChatMiddleware {
123    cache: Arc<dyn AiChatResponseCache>,
124    metrics: Option<Arc<dyn RuntimeMetrics>>,
125}
126
127impl CacheChatMiddleware {
128    pub fn new(cache: Arc<dyn AiChatResponseCache>) -> Self {
129        Self {
130            cache,
131            metrics: None,
132        }
133    }
134
135    pub fn with_metrics(mut self, metrics: Arc<dyn RuntimeMetrics>) -> Self {
136        self.metrics = Some(metrics);
137        self
138    }
139}
140
141impl ChatClientMiddleware for CacheChatMiddleware {
142    fn wrap(&self, inner: Arc<dyn AiChatClient>) -> Arc<dyn AiChatClient> {
143        Arc::new(CacheChatClient {
144            inner,
145            cache: self.cache.clone(),
146            metrics: self.metrics.clone(),
147        })
148    }
149}
150
151#[derive(Clone)]
152struct CacheChatClient {
153    inner: Arc<dyn AiChatClient>,
154    cache: Arc<dyn AiChatResponseCache>,
155    metrics: Option<Arc<dyn RuntimeMetrics>>,
156}
157
158#[async_trait]
159impl AiChatClient for CacheChatClient {
160    async fn complete(
161        &self,
162        request: ChatRequest,
163        options: Option<&ChatOptions>,
164    ) -> Result<ChatResponse> {
165        let cache_key = deterministic_cache_key(&request, options);
166        if let Some(cached) = self.cache.get(&cache_key) {
167            if let Some(metrics) = &self.metrics {
168                metrics.incr_counter(CHAT_CACHE_HIT_TOTAL, 1);
169            }
170            return Ok(cached);
171        }
172        if let Some(metrics) = &self.metrics {
173            metrics.incr_counter(CHAT_CACHE_MISS_TOTAL, 1);
174        }
175
176        let response = self.inner.complete(request, options).await?;
177        self.cache.set(&cache_key, response.clone());
178        Ok(response)
179    }
180}
181
182pub fn deterministic_cache_key(request: &ChatRequest, options: Option<&ChatOptions>) -> String {
183    let basis = format!("request={request:?}|options={options:?}");
184    let mut hasher = Sha256::new();
185    hasher.update(basis.as_bytes());
186    format!("chat:{}", hex::encode(hasher.finalize()))
187}
188
189#[derive(Clone)]
190pub struct ToolCallInterceptionChatMiddleware {
191    interceptor: Arc<dyn AiChatToolInterceptor>,
192    metrics: Option<Arc<dyn RuntimeMetrics>>,
193}
194
195impl ToolCallInterceptionChatMiddleware {
196    pub fn new(interceptor: Arc<dyn AiChatToolInterceptor>) -> Self {
197        Self {
198            interceptor,
199            metrics: None,
200        }
201    }
202
203    pub fn with_metrics(mut self, metrics: Arc<dyn RuntimeMetrics>) -> Self {
204        self.metrics = Some(metrics);
205        self
206    }
207}
208
209impl ChatClientMiddleware for ToolCallInterceptionChatMiddleware {
210    fn wrap(&self, inner: Arc<dyn AiChatClient>) -> Arc<dyn AiChatClient> {
211        Arc::new(ToolCallInterceptionChatClient {
212            inner,
213            interceptor: self.interceptor.clone(),
214            metrics: self.metrics.clone(),
215        })
216    }
217}
218
219#[derive(Clone)]
220struct ToolCallInterceptionChatClient {
221    inner: Arc<dyn AiChatClient>,
222    interceptor: Arc<dyn AiChatToolInterceptor>,
223    metrics: Option<Arc<dyn RuntimeMetrics>>,
224}
225
226#[async_trait]
227impl AiChatClient for ToolCallInterceptionChatClient {
228    async fn complete(
229        &self,
230        request: ChatRequest,
231        options: Option<&ChatOptions>,
232    ) -> Result<ChatResponse> {
233        let request_fingerprint = deterministic_cache_key(&request, options);
234        let response = self.inner.complete(request, options).await?;
235
236        let tool_calls = response.clone().into_tool_calls();
237        if !tool_calls.is_empty() {
238            let tool_call_count = tool_calls.len();
239            let tool_names = tool_calls.into_iter().map(|call| call.fn_name).collect();
240            self.interceptor.on_tool_calls(AiToolCallEnvelope {
241                request_fingerprint,
242                tool_call_count,
243                tool_names,
244            });
245            if let Some(metrics) = &self.metrics {
246                metrics.incr_counter(CHAT_TOOL_CALLS_TOTAL, tool_call_count as u64);
247            }
248        }
249
250        Ok(response)
251    }
252}
stasis/application/runtime/default_chat_middlewares.rs

stasis/application/runtime/
default_chat_middlewares.rs
