1use super::llm::{LLMProvider,LLMInterface, OutputFormat};
2use reqwest::header::{HeaderMap, HeaderValue, CONTENT_TYPE, AUTHORIZATION};
3use serde_json::{json, Value};
4use std::env;
5use crate::utils::llm::FromLLMResponse;
6use tokio::time::Duration;
7use crate::utils::lib::*;
8use anyhow::{Context, Result};
9use std::path::PathBuf;
10use reqwest::multipart::{Form, Part};
11use log::{info, debug, error};
12
13
14pub struct OpenAI {
15 model: String,
16 temperature: f32,
17 max_tokens: u32,
18 max_retries: u32,
19 delay: Duration
20}
21
22impl OpenAI {
23 pub fn new(model: String, temperature: f32, max_tokens: u32) -> Self {
24 Self { model, temperature, max_tokens, max_retries: 3, delay: Duration::from_secs(1) }
25 }
26
27 pub fn with_retries(mut self, max_retries: u32) -> Self {
28 self.max_retries = max_retries;
29 self
30 }
31
32 pub fn with_delay(mut self, delay: Duration) -> Self {
33 self.delay = delay;
34 self
35 }
36}
37
38impl LLMProvider for OpenAI {
39 fn generate_headers(&self) -> Result<HeaderMap> {
40 let mut headers = HeaderMap::new();
41 let api_key = env::var("OPENAI_API_KEY").context("OPENAI_API_KEY must be set")?;
42 headers.insert(AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {}", api_key))
43 .context("Failed to create Authorization header")?);
44 headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
45 Ok(headers)
46 }
47
48 fn generate_request_body(&self, sys_prompt: &str, user_prompt: &str, output_format: &OutputFormat) -> Result<Value> {
49 let messages = vec![
50 json!({"role": "system", "content": sys_prompt}),
51 json!({"role": "user", "content": user_prompt}),
52 ];
53
54 let mut body = json!({
55 "model": self.model,
56 "messages": messages,
57 "temperature": self.temperature,
58 "max_tokens": self.max_tokens,
59 });
60
61 match output_format {
62 OutputFormat::String => {},
63 OutputFormat::Json => {
64 body["response_format"] = json!({"type": "json_object"});
65 },
66 OutputFormat::StrictJson(schema) => {
67 body["response_format"] = schema.clone();
68 },
69 }
70
71 Ok(body)
72 }
73
74
75}
76
77impl LLMInterface for OpenAI {
78 async fn send_request<T: FromLLMResponse + Send + Sync>(&self, sys_prompt: &str, user_prompt: &str) -> Result<T> {
79 retry(self.max_retries, self.delay, || async {
80 let client = reqwest::Client::new();
81 let headers = self.generate_headers()?;
82 let output_format = T::output_format();
83 let body = self.generate_request_body(sys_prompt, user_prompt, &output_format)?;
84 let response = client.post("https://api.openai.com/v1/chat/completions")
85 .headers(headers)
86 .json(&body)
87 .send()
88 .await
89 .context("Failed to send request to OpenAI API")?;
90
91 if response.status().is_success() {
92 let response_value = response.json::<Value>().await
93 .context("Failed to parse OpenAI API response as JSON")?;
94 let content = response_value["choices"][0]["message"]["content"]
95 .as_str()
96 .context("Failed to extract content from OpenAI API response")?
97 .to_string();
98 info!("OpenAI API request successful");
99 debug!("Response: {:?}", response_value);
100 T::from_llm_response(content)
101 } else {
102 let error_text = response.text().await
103 .context("Failed to get error text from OpenAI API")?;
104 error!("OpenAI API request failed: {}", error_text);
105 anyhow::bail!("OpenAI API request failed: {}", error_text)
106 }
107 }).await
108 }
109
110 async fn upload_file(&self, file_path: PathBuf) -> Result<Value> {
111 let api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY must be set");
112 let client = reqwest::Client::new();
113
114 let purpose = "fine-tune";
115
116 let content = tokio::fs::read(&file_path)
117 .await
118 .context("Failed to read file")?;
119
120 let file_name = file_path.file_name()
121 .and_then(|n| n.to_str())
122 .context("Failed to get file name")?
123 .to_string();
124
125 let part = Part::bytes(content)
126 .file_name(file_name)
127 .mime_str("application/json")
128 .context("Failed to set MIME type")?;
129
130 let form = Form::new()
131 .part("file", part)
132 .text("purpose", purpose.to_string());
133
134 let response = client.post("https://api.openai.com/v1/files")
135 .header(AUTHORIZATION, format!("Bearer {}", api_key))
136 .multipart(form)
137 .send()
138 .await
139 .context("Failed to send file upload request")?;
140
141 if response.status().is_success() {
142 let response_json: Value = response.json().await
143 .context("Failed to parse upload response as JSON")?;
144 info!("File uploaded successfully");
145 debug!("Response: {:?}", response_json);
146 Ok(response_json)
147 } else {
148 let error_text = response.text().await
149 .context("Failed to get error text from upload response")?;
150 error!("File upload failed: {}", error_text);
151 Err(anyhow::anyhow!("File upload failed: {}", error_text))
152 }
153 }
154
155 async fn create_fine_tuning_job(&self, training_file: &str) -> Result<Value> {
156 let api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY must be set");
157 let client = reqwest::Client::new();
158
159 let body = json!({
160 "training_file": training_file,
161 "model": self.model,
162 });
163
164 let response = client.post("https://api.openai.com/v1/fine_tuning/jobs")
165 .header(AUTHORIZATION, format!("Bearer {}", api_key))
166 .header(CONTENT_TYPE, "application/json")
167 .json(&body)
168 .send()
169 .await
170 .context("Failed to send fine-tuning job request")?;
171
172 if response.status().is_success() {
173 let response_json: Value = response.json().await
174 .context("Failed to parse fine-tuning job response as JSON")?;
175 info!("Fine-tuning job created successfully");
176 debug!("Response: {:?}", response_json);
177 Ok(response_json)
178 } else {
179 let error_text = response.text().await
180 .context("Failed to get error text from fine-tuning job response")?;
181 error!("Fine-tuning job creation failed: {}", error_text);
182 Err(anyhow::anyhow!("Fine-tuning job creation failed: {}", error_text))
183 }
184 }
185
186 async fn train(
187 &self,
188 file_path: PathBuf
189 ) -> Result<()> {
190 let upload_response = self.upload_file(file_path).await
191 .context("Failed to upload training file")?;
192
193 let training_file_id = upload_response["id"].as_str()
194 .context("Failed to get training file ID from upload response")?;
195
196 let fine_tuning_response = self.create_fine_tuning_job(
197 training_file_id,
198 ).await
199 .context("Failed to create fine-tuning job")?;
200
201 let status = fine_tuning_response["status"].as_str()
202 .context("Failed to get status from fine-tuning job response")?
203 .to_string();
204
205 println!("Fine-tuning job created. Status: {}", status);
206 debug!("Response: {:?}", fine_tuning_response);
207 Ok(())
208 }
209
210
211
212
213
214}
215
216