|
| 1 | +/* |
| 2 | + * Parseable Server (C) 2022 - 2023 Parseable, Inc. |
| 3 | + * |
| 4 | + * This program is free software: you can redistribute it and/or modify |
| 5 | + * it under the terms of the GNU Affero General Public License as |
| 6 | + * published by the Free Software Foundation, either version 3 of the |
| 7 | + * License, or (at your option) any later version. |
| 8 | + * |
| 9 | + * This program is distributed in the hope that it will be useful, |
| 10 | + * but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 11 | + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 12 | + * GNU Affero General Public License for more details. |
| 13 | + * |
| 14 | + * You should have received a copy of the GNU Affero General Public License |
| 15 | + * along with this program. If not, see <http://www.gnu.org/licenses/>. |
| 16 | + * |
| 17 | + */ |
| 18 | + |
| 19 | +use actix_web::{http::header::ContentType, web, HttpResponse, Result}; |
| 20 | +use http::{header, StatusCode}; |
| 21 | +use itertools::Itertools; |
| 22 | +use reqwest; |
| 23 | +use serde_json::{json, Value}; |
| 24 | + |
| 25 | +use crate::{ |
| 26 | + metadata::{error::stream_info::MetadataError, STREAM_INFO}, |
| 27 | + option::CONFIG, |
| 28 | +}; |
| 29 | + |
| 30 | +const OPEN_AI_URL: &str = "https://api.openai.com/v1/chat/completions"; |
| 31 | + |
| 32 | +// Deserialize types for OpenAI Response |
| 33 | +#[derive(serde::Deserialize, Debug)] |
| 34 | +struct ResponseData { |
| 35 | + choices: Vec<Choice>, |
| 36 | +} |
| 37 | + |
| 38 | +#[derive(serde::Deserialize, Debug)] |
| 39 | +struct Choice { |
| 40 | + message: Message, |
| 41 | +} |
| 42 | + |
| 43 | +#[derive(serde::Deserialize, Debug)] |
| 44 | +struct Message { |
| 45 | + content: String, |
| 46 | +} |
| 47 | + |
| 48 | +// Request body |
| 49 | +#[derive(serde::Deserialize, Debug)] |
| 50 | +pub struct AiPrompt { |
| 51 | + prompt: String, |
| 52 | + stream: String, |
| 53 | +} |
| 54 | + |
| 55 | +// Temperory type |
| 56 | +#[derive(Debug, serde::Serialize)] |
| 57 | +struct Field { |
| 58 | + name: String, |
| 59 | + data_type: String, |
| 60 | +} |
| 61 | + |
| 62 | +impl From<&arrow_schema::Field> for Field { |
| 63 | + fn from(field: &arrow_schema::Field) -> Self { |
| 64 | + Self { |
| 65 | + name: field.name().clone(), |
| 66 | + data_type: field.data_type().to_string(), |
| 67 | + } |
| 68 | + } |
| 69 | +} |
| 70 | + |
| 71 | +fn build_prompt(stream: &str, prompt: &str, schema_json: &str) -> String { |
| 72 | + format!( |
| 73 | + r#"I have a table called {}. |
| 74 | +It has the columns:\n{} |
| 75 | +Based on this, generate valid SQL for the query: "{}" |
| 76 | +Generate only SQL as output. Also add comments in SQL syntax to explain your actions. |
| 77 | +Don't output anything else. |
| 78 | +If it is not possible to generate valid SQL, output an SQL comment saying so."#, |
| 79 | + stream, schema_json, prompt |
| 80 | + ) |
| 81 | +} |
| 82 | + |
| 83 | +fn build_request_body(ai_prompt: String) -> impl serde::Serialize { |
| 84 | + json!({ |
| 85 | + "model": "gpt-3.5-turbo", |
| 86 | + "messages": [{ "role": "user", "content": ai_prompt}], |
| 87 | + "temperature": 0.6, |
| 88 | + }) |
| 89 | +} |
| 90 | + |
| 91 | +pub async fn make_llm_request(body: web::Json<AiPrompt>) -> Result<HttpResponse, LLMError> { |
| 92 | + let api_key = match &CONFIG.parseable.open_ai_key { |
| 93 | + Some(api_key) if api_key.len() > 3 => api_key, |
| 94 | + _ => return Err(LLMError::InvalidAPIKey), |
| 95 | + }; |
| 96 | + |
| 97 | + let stream_name = &body.stream; |
| 98 | + let schema = STREAM_INFO.schema(stream_name)?; |
| 99 | + let filtered_schema = schema |
| 100 | + .all_fields() |
| 101 | + .into_iter() |
| 102 | + .map(Field::from) |
| 103 | + .collect_vec(); |
| 104 | + |
| 105 | + let schema_json = |
| 106 | + serde_json::to_string(&filtered_schema).expect("always converted to valid json"); |
| 107 | + |
| 108 | + let prompt = build_prompt(stream_name, &body.prompt, &schema_json); |
| 109 | + let body = build_request_body(prompt); |
| 110 | + |
| 111 | + let client = reqwest::Client::new(); |
| 112 | + let response = client |
| 113 | + .post(OPEN_AI_URL) |
| 114 | + .header(header::CONTENT_TYPE, "application/json") |
| 115 | + .bearer_auth(api_key) |
| 116 | + .json(&body) |
| 117 | + .send() |
| 118 | + .await?; |
| 119 | + |
| 120 | + if response.status().is_success() { |
| 121 | + let body: ResponseData = response |
| 122 | + .json() |
| 123 | + .await |
| 124 | + .expect("OpenAI response is always the same"); |
| 125 | + Ok(HttpResponse::Ok() |
| 126 | + .content_type("application/json") |
| 127 | + .json(&body.choices[0].message.content)) |
| 128 | + } else { |
| 129 | + let body: Value = response.json().await?; |
| 130 | + let message = body |
| 131 | + .as_object() |
| 132 | + .and_then(|body| body.get("error")) |
| 133 | + .and_then(|error| error.as_object()) |
| 134 | + .and_then(|error| error.get("message")) |
| 135 | + .map(|message| message.to_string()) |
| 136 | + .unwrap_or_else(|| "Error from OpenAI".to_string()); |
| 137 | + |
| 138 | + Err(LLMError::APIError(message)) |
| 139 | + } |
| 140 | +} |
| 141 | + |
| 142 | +pub async fn is_llm_active(_body: web::Json<AiPrompt>) -> HttpResponse { |
| 143 | + let is_active = matches!(&CONFIG.parseable.open_ai_key, Some(api_key) if api_key.len() > 3); |
| 144 | + HttpResponse::Ok() |
| 145 | + .content_type("application/json") |
| 146 | + .json(json!({"is_active": is_active})) |
| 147 | +} |
| 148 | + |
| 149 | +#[derive(Debug, thiserror::Error)] |
| 150 | +pub enum LLMError { |
| 151 | + #[error("Either OpenAI key was not provided or was invalid")] |
| 152 | + InvalidAPIKey, |
| 153 | + #[error("Failed to call OpenAI endpoint: {0}")] |
| 154 | + FailedRequest(#[from] reqwest::Error), |
| 155 | + #[error("{0}")] |
| 156 | + APIError(String), |
| 157 | + #[error("{0}")] |
| 158 | + StreamDoesNotExist(#[from] MetadataError), |
| 159 | +} |
| 160 | + |
| 161 | +impl actix_web::ResponseError for LLMError { |
| 162 | + fn status_code(&self) -> http::StatusCode { |
| 163 | + match self { |
| 164 | + Self::InvalidAPIKey => StatusCode::INTERNAL_SERVER_ERROR, |
| 165 | + Self::FailedRequest(_) => StatusCode::INTERNAL_SERVER_ERROR, |
| 166 | + Self::APIError(_) => StatusCode::INTERNAL_SERVER_ERROR, |
| 167 | + Self::StreamDoesNotExist(_) => StatusCode::INTERNAL_SERVER_ERROR, |
| 168 | + } |
| 169 | + } |
| 170 | + |
| 171 | + fn error_response(&self) -> actix_web::HttpResponse<actix_web::body::BoxBody> { |
| 172 | + actix_web::HttpResponse::build(self.status_code()) |
| 173 | + .insert_header(ContentType::plaintext()) |
| 174 | + .body(self.to_string()) |
| 175 | + } |
| 176 | +} |
0 commit comments