starweaver_model/adapter/
context.rs1use serde::{Deserialize, Serialize};
2use serde_json::{Map, Value};
3use starweaver_core::{CancellationToken, ConversationId, RunId, TraceContext};
4
5#[derive(Clone, Deserialize, Serialize)]
7pub struct ModelRequestContext {
8 pub run_id: RunId,
10 pub conversation_id: ConversationId,
12 #[serde(default, skip_serializing_if = "TraceContext::is_empty")]
14 pub trace_context: TraceContext,
15 #[serde(default, skip_serializing_if = "Map::is_empty")]
17 pub llm_trace_metadata: Map<String, Value>,
18 #[serde(skip, default)]
20 pub cancellation_token: CancellationToken,
21}
22
23impl std::fmt::Debug for ModelRequestContext {
24 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25 f.debug_struct("ModelRequestContext")
26 .field("run_id", &self.run_id)
27 .field("conversation_id", &self.conversation_id)
28 .field("trace_context", &self.trace_context)
29 .field("cancelled", &self.cancellation_token.is_cancelled())
30 .finish_non_exhaustive()
31 }
32}
33
34impl PartialEq for ModelRequestContext {
35 fn eq(&self, other: &Self) -> bool {
36 self.run_id == other.run_id
37 && self.conversation_id == other.conversation_id
38 && self.trace_context == other.trace_context
39 && self.llm_trace_metadata == other.llm_trace_metadata
40 && self.cancellation_token == other.cancellation_token
41 }
42}
43
44impl Eq for ModelRequestContext {}
45
46impl ModelRequestContext {
47 #[must_use]
49 pub fn new(run_id: RunId, conversation_id: ConversationId) -> Self {
50 Self {
51 run_id,
52 conversation_id,
53 trace_context: TraceContext::default(),
54 llm_trace_metadata: Map::new(),
55 cancellation_token: CancellationToken::default(),
56 }
57 }
58
59 #[must_use]
61 pub fn with_trace_context(mut self, trace_context: TraceContext) -> Self {
62 self.trace_context = trace_context;
63 self
64 }
65
66 #[must_use]
68 pub fn with_llm_trace_metadata(mut self, metadata: Map<String, Value>) -> Self {
69 self.llm_trace_metadata = metadata;
70 self
71 }
72
73 #[must_use]
75 pub fn with_cancellation_token(mut self, token: CancellationToken) -> Self {
76 self.cancellation_token = token;
77 self
78 }
79
80 #[must_use]
82 pub fn cancellation_token(&self) -> CancellationToken {
83 self.cancellation_token.clone()
84 }
85}