llm_bridge/
response.rs

1use std::fmt;
2use serde::{Deserialize, Serialize};
3
4
5
6#[derive(Serialize, Deserialize, Debug)]
7pub struct OpenAIResponse {
8    pub id: String,
9    pub object: String,
10    pub created: i64,
11    pub model: String,
12    pub choices: Vec<OpenAIChoice>,
13    pub usage: OpenAIUsage,
14}
15
16#[derive(Serialize, Deserialize, Debug, Default)]
17pub(crate) struct OpenAIUsage {
18    pub prompt_tokens: usize,
19    pub completion_tokens: usize,
20    pub total_tokens: usize,
21}
22#[derive(Serialize, Deserialize, Debug)]
23pub struct AnthropicResponse {
24    pub id: String,
25    pub role: String,
26    pub content: Vec<AnthropicContentBlock>,
27    pub model: String,
28    pub stop_reason: String,
29    pub stop_sequence: Option<String>,
30    pub usage: AnthropicUsage,
31}
32
33/// Represents a block of content in the API response.
34#[derive(Serialize, Deserialize, Debug)]
35#[serde(untagged)]
36pub enum AnthropicContentBlock {
37    /// Represents a text content block in the Anthropic API response.
38    Text {
39        /// The actual text content of the response.
40        text: String,
41        /// The type of the content block, always "text" for this variant.
42        #[serde(rename = "type")]
43        block_type: String,
44    },
45    /// Represents a tool use content block in the Anthropic API response.
46    /// This is used when the model decides to use a tool.
47    ToolUse {
48        /// The type of the content block, always "tool_use" for this variant.
49        #[serde(rename = "type")]
50        block_type: String,
51        /// A unique identifier for this tool use instance.
52        id: String,
53        /// The name of the tool being used.
54        name: String,
55        /// The input provided to the tool, represented as a JSON value.
56        /// This allows for flexibility in the structure of tool inputs.
57        input: serde_json::Value,
58    },
59}
60
61/// Represents the response message received from an LLM API.
62///
63/// The `ResponseMessage` enum encapsulates the different response types from various LLM APIs,
64/// providing a unified interface for accessing common fields and methods.
65#[derive(Serialize, Deserialize, Debug)]
66#[serde(untagged)]
67pub enum ResponseMessage {
68    Anthropic(AnthropicResponse),
69    OpenAI(OpenAIResponse),
70}
71
72impl ResponseMessage {
73    /// Returns the text content of the first message in the response.
74    ///
75    /// # Examples
76    ///
77    /// ```
78    /// # use llm_bridge::response::{AnthropicResponse, ResponseMessage};
79    /// let response = ResponseMessage::Anthropic(AnthropicResponse {
80    ///     id: "".to_string(),
81    ///     role: "".to_string(),
82    ///     content: vec![],
83    ///     model: "".to_string(),
84    ///     stop_reason: "".to_string(),
85    ///     stop_sequence: None,
86    ///     usage: Default::default(),}
87    /// );
88    /// let first_message = response.first_message();
89    /// println!("First message: {}", first_message);
90    /// ```
91    pub fn first_message(&self) -> String {
92        match self {
93            ResponseMessage::Anthropic(response) => {
94                if let Some(content) = response.content.first() {
95                    match content {
96                        AnthropicContentBlock::Text { text, .. } => text.clone(),
97                        AnthropicContentBlock::ToolUse { .. } => String::new(), // or handle tool use as needed
98                    }
99                } else {
100                    String::new()
101                }
102            }
103            ResponseMessage::OpenAI(response) => {
104                if let Some(choice) = response.choices.first() {
105                    choice.message.content.clone().unwrap_or_else(|| {
106                        // If content is None, it might be a function call response
107                        if let Some(tool_calls) = &choice.message.tool_calls {
108                            if let Some(first_tool) = tool_calls.first() {
109                                format!("Function call: {}", first_tool.function.name)
110                            } else {
111                                String::new()
112                            }
113                        } else {
114                            String::new()
115                        }
116                    })
117                } else {
118                    String::new()
119                }
120            }
121        }
122    }
123
124    pub fn tools(&self) -> Option<Vec<ToolResponse>> {
125        match self {
126            ResponseMessage::Anthropic(response) => {
127                let tool_uses: Vec<ToolResponse> = response.content.iter()
128                    .filter_map(|block| {
129                        if let AnthropicContentBlock::ToolUse { id, name, input, .. } = block {
130                            Some(ToolResponse {
131                                id: id.clone(),
132                                name: name.clone(),
133                                input: input.clone(),
134                            })
135                        } else {
136                            None
137                        }
138                    })
139                    .collect();
140                if tool_uses.is_empty() { None } else { Some(tool_uses) }
141            },
142            ResponseMessage::OpenAI(response) => {
143                let tool_calls: Vec<ToolResponse> = response.choices.iter()
144                    .filter_map(|choice| choice.message.tool_calls.as_ref())
145                    .flatten()
146                    .map(|tool_call| ToolResponse {
147                        id: tool_call.id.clone(),
148                        name: tool_call.function.name.clone(),
149                        input: serde_json::from_str(&tool_call.function.arguments).unwrap_or(serde_json::Value::Null),
150                    })
151                    .collect();
152                if tool_calls.is_empty() { None } else { Some(tool_calls) }
153            },
154        }
155    }
156
157    /// Returns the role of the sender in the response.
158    ///
159    /// # Examples
160    ///
161    /// ```
162    /// # use llm_bridge::response::{AnthropicResponse, ResponseMessage};
163    /// let response = ResponseMessage::Anthropic(AnthropicResponse {
164    ///     id: "".to_string(),
165    ///     role: "".to_string(),
166    ///     content: vec![],
167    ///     model: "".to_string(),
168    ///     stop_reason: "".to_string(),
169    ///     stop_sequence: None,
170    ///     usage: Default::default(),}
171    /// );
172    /// let role = response.role();
173    /// println!("Role: {}", role);
174    /// ```
175    pub fn role(&self) -> &str {
176        match self {
177            ResponseMessage::Anthropic(response) => &response.role,
178            ResponseMessage::OpenAI(response) => {
179                if let Some(choice) = response.choices.first() {
180                    &choice.message.role
181                } else {
182                    ""
183                }
184            }
185        }
186    }
187
188    /// Returns the name of the model used for generating the response.
189    ///
190    /// # Examples
191    ///
192    /// ```
193    /// # use llm_bridge::response::{AnthropicResponse, ResponseMessage};
194    /// let response = ResponseMessage::Anthropic(AnthropicResponse {
195    ///     id: "".to_string(),
196    ///     role: "".to_string(),
197    ///     content: vec![],
198    ///     model: "".to_string(),
199    ///     stop_reason: "".to_string(),
200    ///     stop_sequence: None,
201    ///     usage: Default::default(),}
202    /// );
203    /// let model = response.model();
204    /// println!("Model: {}", model);
205    /// ```
206    pub fn model(&self) -> &str {
207        match self {
208            ResponseMessage::Anthropic(response) => &response.model,
209            ResponseMessage::OpenAI(response) => &response.model,
210        }
211    }
212
213    /// Returns the stop reason for the generated response.
214    ///
215    /// # Examples
216    ///
217    /// ```
218    /// # use llm_bridge::response::{AnthropicResponse, ResponseMessage};
219    /// let response = ResponseMessage::Anthropic(AnthropicResponse {
220    ///     id: "".to_string(),
221    ///     role: "".to_string(),
222    ///     content: vec![],
223    ///     model: "".to_string(),
224    ///     stop_reason: "".to_string(),
225    ///     stop_sequence: None,
226    ///     usage: Default::default(),}
227    /// );
228    /// let stop_reason = response.stop_reason();
229    /// println!("Stop reason: {}", stop_reason);
230    /// ```
231    pub fn stop_reason(&self) -> &str {
232        match self {
233            ResponseMessage::Anthropic(response) => &response.stop_reason,
234            ResponseMessage::OpenAI(response) => {
235                if let Some(choice) = response.choices.first() {
236                    &choice.finish_reason
237                } else {
238                    ""
239                }
240            }
241        }
242    }
243
244    /// Returns the usage information for the generated response.
245    ///
246    /// # Examples
247    ///
248    /// ```
249    /// # use llm_bridge::response::{AnthropicResponse, ResponseMessage};
250    /// let response = ResponseMessage::Anthropic(AnthropicResponse {
251    ///     id: "".to_string(),
252    ///     role: "".to_string(),
253    ///     content: vec![],
254    ///     model: "".to_string(),
255    ///     stop_reason: "".to_string(),
256    ///     stop_sequence: None,
257    ///     usage: Default::default(),}
258    /// );
259    /// let usage = response.usage();
260    /// println!("Input tokens: {}", usage.input_tokens);
261    /// println!("Output tokens: {}", usage.output_tokens);
262    /// ```
263    pub fn usage(&self) -> CommonUsage {
264        match self {
265            ResponseMessage::Anthropic(response) => CommonUsage {
266                input_tokens: response.usage.input_tokens,
267                output_tokens: response.usage.output_tokens,
268            },
269            ResponseMessage::OpenAI(response) => CommonUsage {
270                input_tokens: response.usage.prompt_tokens,
271                output_tokens: response.usage.completion_tokens,
272            },
273        }
274    }
275}
276
277impl fmt::Display for ResponseMessage {
278    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
279        match self {
280            ResponseMessage::Anthropic(response) => {
281                write!(
282                    f,
283                    "ResponseMessage {{ id: {}, role: {}, content: {:?} }}",
284                    response.id, response.role, response.content
285                )
286            }
287            ResponseMessage::OpenAI(response) => {
288                write!(
289                    f,
290                    "ResponseMessage {{ id: {}, object: {}, model: {}, choices: {:?} }}",
291                    response.id, response.object, response.model, response.choices
292                )
293            }
294        }
295    }
296}
297
298
299/// Tokens represent the underlying cost to llm systems.
300#[derive(Serialize, Deserialize, Debug, Default)]
301pub struct AnthropicUsage {
302    pub input_tokens: usize,
303    pub output_tokens: usize,
304}
305
306#[derive(Serialize, Deserialize, Debug, Default)]
307pub struct CommonUsage {
308    pub input_tokens: usize,
309    pub output_tokens: usize,
310}
311
312#[derive(Serialize, Deserialize, Debug)]
313pub struct OpenAIChoice {
314    pub index: usize,
315    pub message: OpenAIMessage,
316    pub finish_reason: String,
317}
318
319#[derive(Serialize, Deserialize, Debug)]
320pub struct OpenAIMessage {
321    pub role: String,
322    pub content: Option<String>,
323    pub tool_calls: Option<Vec<OpenAIToolCall>>,
324}
325
326#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
327pub struct ToolResponse {
328    pub id: String,
329    pub name: String,
330    pub input: serde_json::Value,
331}
332
333
334#[derive(Serialize, Deserialize, Debug)]
335pub struct OpenAIToolCall {
336    pub id: String,
337    #[serde(rename = "type")]
338    pub call_type: String,
339    pub function: OpenAIFunction,
340}
341
342#[derive(Serialize, Deserialize, Debug)]
343pub struct OpenAIFunction {
344    pub name: String,
345    pub arguments: String,
346}
347
348#[cfg(test)]
349
350mod tests {
351    use super::*;
352    use serde_json::json;
353    use crate::response::{AnthropicContentBlock, AnthropicResponse};
354
355    #[test]
356    fn test_anthropic_response_deserialization() {
357        let json_response = json!({
358            "id": "msg_01KGgxCr7Lm9gi1kfaZWWJUs",
359            "type": "message",
360            "role": "assistant",
361            "model": "claude-3-haiku-20240307",
362            "content": [
363                {
364                    "type": "tool_use",
365                    "id": "toolu_01RQ6pzGpMxBBCirxUcSBokz",
366                    "name": "get_weather",
367                    "input": {
368                        "location": "San Francisco, CA",
369                        "unit": "celsius"
370                    }
371                }
372            ],
373            "stop_reason": "tool_use",
374            "stop_sequence": null,
375            "usage": {
376                "input_tokens": 406,
377                "output_tokens": 73
378            }
379        });
380
381        let response: AnthropicResponse = serde_json::from_value(json_response).unwrap();
382
383        assert_eq!(response.id, "msg_01KGgxCr7Lm9gi1kfaZWWJUs");
384        assert_eq!(response.role, "assistant");
385        assert_eq!(response.model, "claude-3-haiku-20240307");
386        assert_eq!(response.stop_reason, "tool_use");
387        assert_eq!(response.stop_sequence, None);
388        assert_eq!(response.usage.input_tokens, 406);
389        assert_eq!(response.usage.output_tokens, 73);
390
391        assert_eq!(response.content.len(), 1);
392        if let AnthropicContentBlock::ToolUse { id, name, input, .. } = &response.content[0] {
393            assert_eq!(id, "toolu_01RQ6pzGpMxBBCirxUcSBokz");
394            assert_eq!(name, "get_weather");
395            assert_eq!(input["location"], "San Francisco, CA");
396            assert_eq!(input["unit"], "celsius");
397        } else {
398            panic!("Expected ToolUse content block");
399        }
400    }
401
402    #[test]
403    fn test_anthropic_response_text_content() {
404        let json_response = json!({
405            "id": "msg_text_example",
406            "type": "message",
407            "role": "assistant",
408            "model": "claude-3-haiku-20240307",
409            "content": [
410                {
411                    "type": "text",
412                    "text": "This is a text response."
413                }
414            ],
415            "stop_reason": "end_turn",
416            "stop_sequence": null,
417            "usage": {
418                "input_tokens": 10,
419                "output_tokens": 20
420            }
421        });
422
423        let response: AnthropicResponse = serde_json::from_value(json_response).unwrap();
424
425        assert_eq!(response.id, "msg_text_example");
426        assert_eq!(response.role, "assistant");
427        assert_eq!(response.model, "claude-3-haiku-20240307");
428        assert_eq!(response.stop_reason, "end_turn");
429
430        assert_eq!(response.content.len(), 1);
431        if let AnthropicContentBlock::Text { text, .. } = &response.content[0] {
432            assert_eq!(text, "This is a text response.");
433        } else {
434            panic!("Expected Text content block");
435        }
436    }
437
438    #[test]
439    fn test_anthropic_response_mixed_content() {
440        let json_response = json!({
441            "id": "msg_mixed_example",
442            "type": "message",
443            "role": "assistant",
444            "model": "claude-3-haiku-20240307",
445            "content": [
446                {
447                    "type": "text",
448                    "text": "Here's the weather information:"
449                },
450                {
451                    "type": "tool_use",
452                    "id": "toolu_mixed_example",
453                    "name": "get_weather",
454                    "input": {
455                        "location": "New York, NY",
456                        "unit": "fahrenheit"
457                    }
458                }
459            ],
460            "stop_reason": "end_turn",
461            "stop_sequence": null,
462            "usage": {
463                "input_tokens": 50,
464                "output_tokens": 60
465            }
466        });
467
468        let response: AnthropicResponse = serde_json::from_value(json_response).unwrap();
469
470        assert_eq!(response.content.len(), 2);
471
472        match &response.content[0] {
473            AnthropicContentBlock::Text { text, .. } => {
474                assert_eq!(text, "Here's the weather information:");
475            },
476            _ => panic!("Expected Text content block"),
477        }
478
479        match &response.content[1] {
480            AnthropicContentBlock::ToolUse { id, name, input, .. } => {
481                assert_eq!(id, "toolu_mixed_example");
482                assert_eq!(name, "get_weather");
483                assert_eq!(input["location"], "New York, NY");
484                assert_eq!(input["unit"], "fahrenheit");
485            },
486            _ => panic!("Expected ToolUse content block"),
487        }
488    }
489
490    #[test]
491    fn test_openai_response_deserialization() {
492        let json_response = json!({
493            "id": "chatcmpl-9p5LSmflVqlG0Gk6ryp14XHKbNah8",
494            "object": "chat.completion",
495            "created": 1721962302,
496            "model": "gpt-4o-2024-05-13",
497            "choices": [
498                {
499                    "index": 0,
500                    "message": {
501                        "role": "assistant",
502                        "content": null,
503                        "tool_calls": [
504                            {
505                                "id": "call_5dENonKES2CcyWt6yGAXPDtz",
506                                "type": "function",
507                                "function": {
508                                    "name": "get_weather",
509                                    "arguments": "{\"location\":\"San Francisco, CA\"}"
510                                }
511                            }
512                        ]
513                    },
514                    "logprobs": null,
515                    "finish_reason": "tool_calls"
516                }
517            ],
518            "usage": {
519                "prompt_tokens": 106,
520                "completion_tokens": 17,
521                "total_tokens": 123
522            },
523            "system_fingerprint": "fp_400f27fa1f"
524        });
525
526        let response: OpenAIResponse = serde_json::from_value(json_response).unwrap();
527
528        assert_eq!(response.id, "chatcmpl-9p5LSmflVqlG0Gk6ryp14XHKbNah8");
529        assert_eq!(response.object, "chat.completion");
530        assert_eq!(response.created, 1721962302);
531        assert_eq!(response.model, "gpt-4o-2024-05-13");
532
533        assert_eq!(response.choices.len(), 1);
534        let choice = &response.choices[0];
535        assert_eq!(choice.index, 0);
536        assert_eq!(choice.finish_reason, "tool_calls");
537
538        let message = &choice.message;
539        assert_eq!(message.role, "assistant");
540        assert_eq!(message.content, None);
541
542        assert!(message.tool_calls.is_some());
543        let tool_calls = message.tool_calls.as_ref().unwrap();
544        assert_eq!(tool_calls.len(), 1);
545        let tool_call = &tool_calls[0];
546        assert_eq!(tool_call.id, "call_5dENonKES2CcyWt6yGAXPDtz");
547        assert_eq!(tool_call.call_type, "function");
548        assert_eq!(tool_call.function.name, "get_weather");
549        assert_eq!(tool_call.function.arguments, "{\"location\":\"San Francisco, CA\"}");
550
551        assert_eq!(response.usage.prompt_tokens, 106);
552        assert_eq!(response.usage.completion_tokens, 17);
553        assert_eq!(response.usage.total_tokens, 123);
554
555    }
556
557    #[test]
558    fn test_openai_response_tool_calls() {
559        let json_response = json!({
560            "id": "chatcmpl-9p5LSmflVqlG0Gk6ryp14XHKbNah8",
561            "object": "chat.completion",
562            "created": 1721962302,
563            "model": "gpt-4o-2024-05-13",
564            "choices": [
565                {
566                    "index": 0,
567                    "message": {
568                        "role": "assistant",
569                        "content": null,
570                        "tool_calls": [
571                            {
572                                "id": "call_5dENonKES2CcyWt6yGAXPDtz",
573                                "type": "function",
574                                "function": {
575                                    "name": "get_weather",
576                                    "arguments": "{\"location\":\"San Francisco, CA\"}"
577                                }
578                            }
579                        ]
580                    },
581                    "finish_reason": "tool_calls"
582                }
583            ],
584            "usage": {
585                "prompt_tokens": 106,
586                "completion_tokens": 17,
587                "total_tokens": 123
588            }
589        });
590
591        let response: OpenAIResponse = serde_json::from_value(json_response).unwrap();
592        let response_message = ResponseMessage::OpenAI(response);
593
594        if let Some(tools) = response_message.tools() {
595            assert_eq!(tools.len(), 1);
596            assert_eq!(tools[0].name, "get_weather");
597            assert_eq!(tools[0].id, "call_5dENonKES2CcyWt6yGAXPDtz");
598
599            let input: serde_json::Value = serde_json::from_str(&tools[0].input.to_string()).unwrap();
600            assert_eq!(input["location"], "San Francisco, CA");
601        } else {
602            panic!("Expected tool calls, but none were found");
603        }
604
605        assert_eq!(response_message.stop_reason(), "tool_calls");
606    }
607
608    #[test]
609    fn test_openai_response_no_tool_calls() {
610        let json_response = json!({
611            "id": "chatcmpl-123",
612            "object": "chat.completion",
613            "created": 1721962302,
614            "model": "gpt-4o-2024-05-13",
615            "choices": [
616                {
617                    "index": 0,
618                    "message": {
619                        "role": "assistant",
620                        "content": "This is a regular response without tool calls."
621                    },
622                    "finish_reason": "stop"
623                }
624            ],
625            "usage": {
626                "prompt_tokens": 10,
627                "completion_tokens": 10,
628                "total_tokens": 20
629            }
630        });
631
632        let response: OpenAIResponse = serde_json::from_value(json_response).unwrap();
633        let response_message = ResponseMessage::OpenAI(response);
634        assert_eq!(response_message.tools(), None);
635        assert_eq!(response_message.stop_reason(), "stop");
636        assert_eq!(response_message.first_message(), "This is a regular response without tool calls.");
637    }
638}
llm_bridge/response.rs

llm_bridge/
response.rs
