ai-transform-runtime 0.1.0

use crate::error::TransformError;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::env;
use std::time::Duration;
use tokio::time::timeout;

// Minimalistic type structure for sending OpenAI requests
#[derive(Serialize, Deserialize)]
struct Message {
    content: String,
}

#[derive(Serialize, Deserialize)]
struct Choice {
    message: Message,
}

#[derive(Serialize, Deserialize)]
struct OpenAIResponse {
    choices: Vec<Choice>,
}

pub struct OpenAIClient {
    client: Client,
    model: String,
    base_url: String,
    api_key: String,
}

fn extract_json_content(content: &str) -> String {
    let trimmed = content.trim();

    // Check for ```json code blocks first
    if let Some(json_start) = trimmed.find("```json") {
        if let Some(json_end) = trimmed[json_start..].find("```") {
            let start_pos = json_start + 7; // "```json".len()
            let end_pos = json_start + json_end;
            if start_pos < end_pos {
                return trimmed[start_pos..end_pos].trim().to_string();
            }
        }
    }

    // Check for generic ``` code blocks
    if let Some(start) = trimmed.find("```") {
        if let Some(end) = trimmed[start + 3..].find("```") {
            let actual_end = start + 3 + end;
            let content_start = start + 3;

            let content_slice = &trimmed[content_start..actual_end];
            if let Some(newline_pos) = content_slice.find('\n') {
                return content_slice[newline_pos + 1..].trim().to_string();
            }
            return content_slice.trim().to_string();
        }
    }

    trimmed.to_string()
}

impl OpenAIClient {
    pub fn new() -> Result<Self, TransformError> {
        let model = env::var("OPENAI_MODEL").unwrap_or("gpt-4o".to_string());
        let base_url =
            env::var("OPENAI_BASE_URL").unwrap_or("https://api.openai.com/v1".to_string());
        let api_key = env::var("OPENAI_API_KEY")
            .map_err(|_| TransformError::EnvVarError("OPENAI_API_KEY".to_string()))?;

        Ok(Self {
            client: Client::new(),
            model,
            base_url,
            api_key,
        })
    }

    pub async fn transform_json(
        &self,
        source_json: &str,
        source_type: &str,
        target_type: &str,
        source_schema: &str,
        target_schema: &str,
    ) -> Result<String, TransformError> {
        let prompt = format!(
            "Transform the following JSON data from type '{source_type}' to type '{target_type}'. \
            The source type has the following field structure (example): {source_schema}\n\
            The target type should have the following field structure (example): {target_schema}\n\n\
            Transform the source JSON to match the target structure. Return only the transformed JSON, no explanations or additional text.\n\n\
            Source JSON to transform:\n{source_json}"
        );

        let payload = json!({
            "model": self.model,
            "messages": [
                {
                    "role": "system",
                    "content": "You are a JSON transformation assistant. Transform the input JSON to match the target type structure. Return only valid JSON."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ]
        });

        let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));

        let response = timeout(
            Duration::from_secs(30),
            self.client
                .post(&url)
                .bearer_auth(&self.api_key)
                .header("Content-Type", "application/json")
                .json(&payload)
                .send(),
        )
        .await
        .map_err(|_| TransformError::TimeoutError)?
        .map_err(TransformError::RequestError)?;

        if response.status().is_success() {
            let openai_response: OpenAIResponse = response
                .json()
                .await
                .map_err(TransformError::RequestError)?;

            let content = openai_response
                .choices
                .first()
                .map(|choice| &choice.message.content)
                .ok_or(TransformError::InvalidResponseFormat)?;

            Ok(extract_json_content(content))
        } else {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            Err(TransformError::ApiError { status, body })
        }
    }
}