ragit_api/
response.rs

1use crate::api_provider::ApiProvider;
2use crate::error::Error;
3
4mod anthropic;
5mod cohere;
6mod google;
7mod openai;
8
9pub use anthropic::AnthropicResponse;
10pub use cohere::CohereResponse;
11pub use google::GoogleResponse;
12pub use openai::OpenAiResponse;
13
14pub trait IntoChatResponse {
15    fn into_chat_response(&self) -> Result<Response, Error>;
16}
17
18pub struct Response {
19    messages: Vec<String>,
20    reasonings: Vec<Option<String>>,
21    output_tokens: usize,
22    prompt_tokens: usize,
23    total_tokens: usize,
24}
25
26impl Response {
27    pub fn dummy(s: String) -> Self {
28        Response {
29            messages: vec![s],
30            reasonings: vec![None],
31            output_tokens: 0,
32            prompt_tokens: 0,
33            total_tokens: 0,
34        }
35    }
36
37    pub fn from_str(s: &str, api_provider: &ApiProvider) -> Result<Self, Error> {
38        api_provider.parse_chat_response(s)?.into_chat_response()
39    }
40
41    pub fn get_output_token_count(&self) -> usize {
42        self.output_tokens
43    }
44
45    pub fn get_prompt_token_count(&self) -> usize {
46        self.prompt_tokens
47    }
48
49    pub fn get_total_token_count(&self) -> usize {
50        self.total_tokens
51    }
52
53    pub fn get_message(&self, index: usize) -> Option<&str> {
54        self.messages.get(index).map(|s| s.as_str())
55    }
56
57    pub fn get_reasoning(&self, index: usize) -> Option<&str> {
58        match self.reasonings.get(index) {
59            Some(s) => s.as_ref().map(|s| s.as_str()),
60            None => None,
61        }
62    }
63}
ragit_api/response.rs

ragit_api/
response.rs
