baked_potato/
mock.rs

1use crate::error::MockError;
2use mockito;
3use potato_agent::agents::provider::{gemini::GenerateContentResponse, openai::OpenAIChatResponse};
4use serde_json;
5
6use pyo3::prelude::*;
7
8pub const OPENAI_CHAT_COMPLETION_RESPONSE: &str =
9    include_str!("assets/openai/openai_chat_completion_response.json");
10
11pub const OPENAI_CHAT_STRUCTURED_RESPONSE: &str =
12    include_str!("assets/openai/chat_completion_structured_response.json");
13
14pub const OPENAI_CHAT_STRUCTURED_SCORE_RESPONSE: &str =
15    include_str!("assets/openai/chat_completion_structured_score_response.json");
16
17pub const OPENAI_CHAT_STRUCTURED_RESPONSE_PARAMS: &str =
18    include_str!("assets/openai/chat_completion_structured_response_params.json");
19
20pub const OPENAI_CHAT_STRUCTURED_TASK_OUTPUT: &str =
21    include_str!("assets/openai/chat_completion_structured_task_output.json");
22
23pub const GEMINI_CHAT_COMPLETION_RESPONSE: &str =
24    include_str!("assets/gemini/chat_completion.json");
25
26pub const GEMINI_CHAT_COMPLETION_RESPONSE_WITH_SCORE: &str =
27    include_str!("assets/gemini/chat_completion_with_score.json");
28
29pub struct LLMApiMock {
30    pub url: String,
31    pub server: mockito::ServerGuard,
32}
33
34impl LLMApiMock {
35    pub fn new() -> Self {
36        let mut server = mockito::Server::new();
37        // load the OpenAI chat completion response
38        let chat_msg_response: OpenAIChatResponse =
39            serde_json::from_str(OPENAI_CHAT_COMPLETION_RESPONSE).unwrap();
40        let chat_structured_response: OpenAIChatResponse =
41            serde_json::from_str(OPENAI_CHAT_STRUCTURED_RESPONSE).unwrap();
42        let chat_structured_score_response: OpenAIChatResponse =
43            serde_json::from_str(OPENAI_CHAT_STRUCTURED_SCORE_RESPONSE).unwrap();
44        let chat_structured_response_params: OpenAIChatResponse =
45            serde_json::from_str(OPENAI_CHAT_STRUCTURED_RESPONSE_PARAMS).unwrap();
46        let chat_structured_task_output: OpenAIChatResponse =
47            serde_json::from_str(OPENAI_CHAT_STRUCTURED_TASK_OUTPUT).unwrap();
48
49        // load the Gemini chat completion response
50        let gemini_chat_response: GenerateContentResponse =
51            serde_json::from_str(GEMINI_CHAT_COMPLETION_RESPONSE).unwrap();
52        let gemini_chat_response_with_score: GenerateContentResponse =
53            serde_json::from_str(GEMINI_CHAT_COMPLETION_RESPONSE_WITH_SCORE).unwrap();
54
55        server
56            .mock("POST", "/chat/completions")
57            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
58                "response_format": {
59                    "type": "json_schema",
60                    "json_schema": {
61                        "name": "Parameters",
62                         "schema": {
63                              "$schema": "https://json-schema.org/draft/2020-12/schema",
64                              "properties": {
65                                  "variable1": {
66                                  "format": "int32",
67                                  "type": "integer"
68                                  },
69                                  "variable2": {
70                                  "format": "int32",
71                                  "type": "integer"
72                                  }
73                              },
74                              "required": [
75                                  "variable1",
76                                  "variable2"
77                              ],
78                              "title": "Parameters",
79                              "type": "object"
80                              },
81                        "strict": true
82                    }
83
84                }
85            })))
86            .expect(usize::MAX)
87            .with_status(200)
88            .with_header("content-type", "application/json")
89            .with_body(serde_json::to_string(&chat_structured_response_params).unwrap())
90            .create();
91
92        server
93            .mock("POST", "/chat/completions")
94            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
95               "response_format": {
96                    "type": "json_schema",
97                    "json_schema": {
98                        "name": "TaskOutput",
99                    }
100                }
101            })))
102            .expect(usize::MAX)
103            .with_status(200)
104            .with_header("content-type", "application/json")
105            .with_body(serde_json::to_string(&chat_structured_task_output).unwrap())
106            .create();
107
108        server
109            .mock("POST", "/chat/completions")
110            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
111                "response_format": {
112                    "type": "json_schema",
113                    "json_schema": {
114                        "name": "Score",
115                    }
116                }
117            })))
118            .expect(usize::MAX)
119            .with_status(200)
120            .with_header("content-type", "application/json")
121            .with_body(serde_json::to_string(&chat_structured_score_response).unwrap())
122            .create();
123
124        server
125            .mock("POST", "/chat/completions")
126            .match_body(mockito::Matcher::Regex(
127                r#".*"name"\s*:\s*"Score".*"#.to_string(),
128            ))
129            .expect(usize::MAX)
130            .with_status(200)
131            .with_header("content-type", "application/json")
132            .with_body(serde_json::to_string(&chat_structured_score_response).unwrap())
133            .create();
134
135        server
136            .mock("POST", "/chat/completions")
137            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
138                "response_format": {
139                    "type": "json_schema"
140                }
141            })))
142            .expect(usize::MAX)
143            .with_status(200)
144            .with_header("content-type", "application/json")
145            .with_body(serde_json::to_string(&chat_structured_response).unwrap())
146            .create();
147
148        // mock the Gemini chat completion response
149        server
150            .mock(
151                "POST",
152                mockito::Matcher::Regex(r".*/.*:generateContent$".to_string()),
153            )
154            .match_header("x-goog-api-key", mockito::Matcher::Any)
155            .match_header("content-type", "application/json")
156            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
157                "contents": [
158                    {
159                        "parts": [
160                            {
161                                "text":  "You are a helpful assistant"
162                            }
163                        ]
164                    }
165                ]
166            })))
167            .expect(usize::MAX) // More specific expectation than usize::MAX
168            .with_status(200)
169            .with_header("content-type", "application/json")
170            .with_body(serde_json::to_string(&gemini_chat_response).unwrap())
171            .create();
172
173        // mock structured response
174        server
175            .mock(
176                "POST",
177                mockito::Matcher::Regex(r".*/.*:generateContent$".to_string()),
178            )
179            .match_header("x-goog-api-key", mockito::Matcher::Any)
180            .match_header("content-type", "application/json")
181            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
182                "generation_config": {
183                    "responseMimeType": "application/json"
184                }
185            })))
186            .expect(usize::MAX)
187            .with_status(200)
188            .with_header("content-type", "application/json")
189            .with_body(serde_json::to_string(&gemini_chat_response_with_score).unwrap())
190            .create();
191
192        // Openai chat completion mock
193        server
194            .mock("POST", "/chat/completions")
195            .expect(usize::MAX)
196            .with_status(200)
197            .with_header("content-type", "application/json")
198            .with_body(serde_json::to_string(&chat_msg_response).unwrap())
199            .create();
200
201        Self {
202            url: server.url(),
203            server,
204        }
205    }
206}
207
208impl Default for LLMApiMock {
209    fn default() -> Self {
210        Self::new()
211    }
212}
213
214#[pyclass]
215#[allow(dead_code)]
216pub struct LLMTestServer {
217    openai_server: Option<LLMApiMock>,
218}
219
220#[pymethods]
221impl LLMTestServer {
222    #[new]
223    pub fn new() -> Self {
224        LLMTestServer {
225            openai_server: None,
226        }
227    }
228
229    pub fn start_mock_server(&mut self) -> Result<(), MockError> {
230        let llm_server = LLMApiMock::new();
231        println!("Mock LLM Server started at {}", llm_server.url);
232        self.openai_server = Some(llm_server);
233        Ok(())
234    }
235
236    pub fn stop_mock_server(&mut self) {
237        if let Some(server) = self.openai_server.take() {
238            drop(server);
239            std::env::remove_var("OPENAI_API_URL");
240            std::env::remove_var("OPENAI_API_KEY");
241        }
242        println!("Mock LLM Server stopped");
243    }
244
245    pub fn set_env_vars_for_client(&self) -> Result<(), MockError> {
246        {
247            std::env::set_var("APP_ENV", "dev_client");
248            std::env::set_var("OPENAI_API_KEY", "test_key");
249            std::env::set_var("GEMINI_API_KEY", "gemini");
250            std::env::set_var(
251                "OPENAI_API_URL",
252                self.openai_server.as_ref().unwrap().url.clone(),
253            );
254            std::env::set_var(
255                "GEMINI_API_URL",
256                self.openai_server.as_ref().unwrap().url.clone(),
257            );
258            Ok(())
259        }
260    }
261
262    pub fn start_server(&mut self) -> Result<(), MockError> {
263        self.cleanup()?;
264
265        println!("Starting Mock GenAI Server...");
266        self.start_mock_server()?;
267        self.set_env_vars_for_client()?;
268
269        // set server env vars
270        std::env::set_var("APP_ENV", "dev_server");
271
272        Ok(())
273    }
274
275    pub fn stop_server(&mut self) -> Result<(), MockError> {
276        self.cleanup()?;
277
278        Ok(())
279    }
280
281    pub fn remove_env_vars_for_client(&self) -> Result<(), MockError> {
282        std::env::remove_var("OPENAI_API_URI");
283        std::env::remove_var("OPENAI_API_KEY");
284        std::env::remove_var("GEMINI_API_KEY");
285        std::env::remove_var("GEMINI_API_URL");
286        Ok(())
287    }
288
289    fn cleanup(&self) -> Result<(), MockError> {
290        // unset env vars
291        self.remove_env_vars_for_client()?;
292
293        Ok(())
294    }
295
296    fn __enter__(mut self_: PyRefMut<Self>) -> Result<PyRefMut<Self>, MockError> {
297        self_.start_server()?;
298
299        Ok(self_)
300    }
301
302    fn __exit__(
303        &mut self,
304        _exc_type: PyObject,
305        _exc_value: PyObject,
306        _traceback: PyObject,
307    ) -> Result<(), MockError> {
308        self.stop_server()
309    }
310}
311
312impl Default for LLMTestServer {
313    fn default() -> Self {
314        Self::new()
315    }
316}
baked_potato/mock.rs

baked_potato/
mock.rs
