1use serde::{Deserialize, Serialize};
4use tokio::sync::mpsc;
5
6use crate::ipc::client::ResponseDelta;
7
8pub enum BatchChatResult {
10 Complete(Vec<ClientResponse>),
12 Stream(mpsc::Receiver<ClientDelta>),
14}
15
16#[derive(Debug, Clone, Default, Serialize, Deserialize)]
18pub struct UsageStats {
19 pub prompt_tokens: u32,
20 pub completion_tokens: u32,
21 pub total_tokens: u32,
22}
23
24#[derive(Debug, Clone, Default, Serialize, Deserialize)]
26pub struct ClientDelta {
27 pub request_id: u64,
28 pub sequence_id: Option<u64>,
29 pub prompt_index: Option<u32>,
30 pub candidate_index: Option<u32>,
31 pub content: Option<String>,
32 pub content_len: Option<u32>,
33 pub inline_content_bytes: Option<u32>,
34 pub is_final: bool,
35 pub finish_reason: Option<String>,
36 pub error: Option<String>,
37 pub prompt_token_count: Option<u32>,
38 pub num_tokens_in_delta: Option<u32>,
39 pub generation_len: Option<u32>,
40 pub tokens: Vec<i32>,
41 pub top_logprobs: Vec<crate::ipc::client::TokenLogProb>,
42 pub cumulative_logprob: Option<f64>,
43 pub modal_decoder_id: Option<String>,
44 pub modal_bytes_b64: Option<String>,
45}
46
47impl From<ResponseDelta> for ClientDelta {
48 fn from(delta: ResponseDelta) -> Self {
49 Self {
50 request_id: delta.request_id,
51 sequence_id: delta.sequence_id,
52 prompt_index: delta.prompt_index,
53 candidate_index: delta.candidate_index,
54 content: delta.content,
55 content_len: delta.content_len,
56 inline_content_bytes: delta.inline_content_bytes,
57 is_final: delta.is_final_delta,
58 finish_reason: delta.finish_reason,
59 error: delta.error,
60 prompt_token_count: delta.prompt_token_count,
61 num_tokens_in_delta: delta.num_tokens_in_delta,
62 generation_len: delta.generation_len,
63 tokens: delta.tokens,
64 top_logprobs: delta.top_logprobs,
65 cumulative_logprob: delta.cumulative_logprob,
66 modal_decoder_id: delta.modal_decoder_id,
67 modal_bytes_b64: delta.modal_bytes_b64,
68 }
69 }
70}
71
72#[derive(Debug, Clone, Default, Serialize, Deserialize)]
74pub struct ClientResponse {
75 pub text: String,
76 pub finish_reason: Option<String>,
77 pub usage: UsageStats,
78 #[serde(skip_serializing_if = "Vec::is_empty")]
79 pub deltas: Vec<ClientDelta>,
80}
81
82#[cfg(test)]
83mod tests {
84 use super::*;
85
86 #[test]
87 fn test_usage_stats_default() {
88 let usage = UsageStats::default();
89 assert_eq!(usage.prompt_tokens, 0);
90 assert_eq!(usage.completion_tokens, 0);
91 assert_eq!(usage.total_tokens, 0);
92 }
93
94 #[test]
95 fn test_client_delta_from_response_delta() {
96 let response = ResponseDelta {
97 request_id: 123,
98 sequence_id: Some(1),
99 prompt_index: Some(0),
100 candidate_index: Some(0),
101 content: Some("Hello".to_string()),
102 content_len: Some(5),
103 inline_content_bytes: Some(5),
104 is_final_delta: false,
105 finish_reason: None,
106 error: None,
107 prompt_token_count: Some(10),
108 num_tokens_in_delta: Some(3),
109 generation_len: Some(5),
110 tokens: vec![1, 2, 3],
111 top_logprobs: vec![],
112 cumulative_logprob: Some(-1.5),
113 modal_decoder_id: Some("moondream3.coord".to_string()),
114 modal_bytes_b64: Some("AAAA".to_string()),
115 embedding_bytes: None,
116 state_events: vec![],
117 cached_token_count: Some(0),
118 reasoning_tokens: Some(0),
119 };
120
121 let delta = ClientDelta::from(response);
122 assert_eq!(delta.request_id, 123);
123 assert_eq!(delta.sequence_id, Some(1));
124 assert_eq!(delta.prompt_index, Some(0));
125 assert_eq!(delta.candidate_index, Some(0));
126 assert_eq!(delta.content, Some("Hello".to_string()));
127 assert_eq!(delta.content_len, Some(5));
128 assert_eq!(delta.num_tokens_in_delta, Some(3));
129 assert!(!delta.is_final);
130 assert_eq!(delta.tokens, vec![1, 2, 3]);
131 assert_eq!(delta.cumulative_logprob, Some(-1.5));
132 assert_eq!(delta.modal_decoder_id, Some("moondream3.coord".to_string()));
133 }
134}