lm_studio_api/chat/
response.rs

1use crate::prelude::*;
2use super::{ Choice, StreamChoice, Usage, Role, Message, Embedding };
3
4/// Chat response
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct Response {
7    #[serde(skip_serializing_if = "Option::is_none")]
8    pub id: Option<String>,
9    pub object: String,
10    #[serde(skip_serializing_if = "Option::is_none")]
11    pub created: Option<u64>,
12    pub model: String,
13    #[serde(default)]
14    #[serde(skip_serializing_if = "Vec::is_empty")]
15    pub choices: Vec<Choice>,
16    #[serde(default)]
17    #[serde(skip_serializing_if = "Vec::is_empty")]
18    pub data: Vec<Embedding>,
19    pub usage: Usage,
20    #[serde(skip_serializing_if = "Option::is_none")]
21    pub stats: Option<HashMap<String, Value>>,
22    #[serde(rename = "system_fingerprint")]
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub fingerprint: Option<String>,
25}
26
27impl Response {
28    /// Returns response text
29    pub fn text(&self) -> Option<String> {
30        if !self.choices.is_empty() {
31            let text = self.choices[0].text().clone().unwrap();
32            
33            Some(text)
34        } else {
35            None
36        }
37    }
38
39    /// Returns embedding data
40    pub fn data(&self) -> Option<Vec<Vec<f32>>> {
41        if !self.data.is_empty() {
42            let data = self.data.clone()
43                .into_iter()
44                .map(|embed| embed.embedding)
45                .collect::<Vec<_>>();
46
47            Some(data)
48        } else {
49            None
50        }
51    }
52}
53
54
55use futures::StreamExt;
56use tokio_stream::wrappers::UnboundedReceiverStream;
57
58/// Chat stream reader
59#[derive(Debug)]
60pub struct ResponseReader {
61    pub receiver: UnboundedReceiverStream<Result<StreamChoice>>,
62    pub message: Message,
63    pub is_ready: bool,
64    pub context: bool
65}
66
67impl ResponseReader {
68    /// Creates a new stream reader
69    pub fn new(receiver: UnboundedReceiverStream<Result<StreamChoice>>, context: bool) -> Self {
70        Self {
71            receiver,
72            message: Message { role: Role::Assistant, content: str!("") },
73            is_ready: false,
74            context
75        }
76    }
77
78    /// Iters next stream choice
79    pub async fn next(&mut self) -> Option<Result<StreamChoice>> {
80        let result = self.receiver.next().await;
81
82        match result {
83            Some(result) => {
84                match result {
85                    Ok(choice) => {
86                        if let Some(text) = choice.text() {
87                            self.message.content.push_str(&text);
88                        }
89
90                        Some(Ok(choice))
91                    },
92
93                    Err(e) => Some(Err(e))
94                }
95            },
96
97            _ => {
98                self.is_ready = true;
99                None
100            }
101        }
102    }
103}