aster/providers/

openai.rs

use super::api_client::{ApiClient, AuthMethod};
use super::base::{ConfigKey, ModelInfo, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::embedding::{EmbeddingCapable, EmbeddingRequest, EmbeddingResponse};
use super::errors::ProviderError;
use super::formats::openai::{create_request, get_usage, response_to_message};
use super::formats::openai_responses::{
    create_responses_request, get_responses_usage, responses_api_to_message,
    responses_api_to_streaming_message, ResponsesApiResponse,
};
use super::retry::ProviderRetry;
use super::utils::{
    get_model, handle_response_openai_compat, handle_status_openai_compat, stream_openai_compat,
    ImageFormat,
};
use crate::config::declarative_providers::DeclarativeProviderConfig;
use crate::conversation::message::Message;
use anyhow::Result;
use async_stream::try_stream;
use async_trait::async_trait;
use futures::{StreamExt, TryStreamExt};
use reqwest::StatusCode;
use serde_json::Value;
use std::collections::HashMap;
use std::io;
use tokio::pin;
use tokio_util::codec::{FramedRead, LinesCodec};
use tokio_util::io::StreamReader;

use crate::model::ModelConfig;
use crate::providers::base::MessageStream;
use crate::providers::utils::RequestLog;
use rmcp::model::Tool;

pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o";
pub const OPEN_AI_DEFAULT_FAST_MODEL: &str = "gpt-4o-mini";
pub const OPEN_AI_KNOWN_MODELS: &[(&str, usize)] = &[
    ("gpt-4o", 128_000),
    ("gpt-4o-mini", 128_000),
    ("gpt-4.1", 128_000),
    ("gpt-4.1-mini", 128_000),
    ("o1", 200_000),
    ("o3", 200_000),
    ("gpt-3.5-turbo", 16_385),
    ("gpt-4-turbo", 128_000),
    ("o4-mini", 128_000),
    ("gpt-5.1-codex", 400_000),
    ("gpt-5-codex", 400_000),
];

pub const OPEN_AI_DOC_URL: &str = "https://platform.openai.com/docs/models";

#[derive(Debug, serde::Serialize)]
pub struct OpenAiProvider {
    #[serde(skip)]
    api_client: ApiClient,
    base_path: String,
    organization: Option<String>,
    project: Option<String>,
    model: ModelConfig,
    custom_headers: Option<HashMap<String, String>>,
    supports_streaming: bool,
    name: String,
}

impl OpenAiProvider {
    pub async fn from_env(model: ModelConfig) -> Result<Self> {
        let model = model.with_fast(OPEN_AI_DEFAULT_FAST_MODEL.to_string());

        let config = crate::config::Config::global();
        let secrets = config.get_secrets("OPENAI_API_KEY", &["OPENAI_CUSTOM_HEADERS"])?;
        let api_key = secrets.get("OPENAI_API_KEY").unwrap().clone();
        let host: String = config
            .get_param("OPENAI_HOST")
            .unwrap_or_else(|_| "https://api.openai.com".to_string());
        let base_path: String = config
            .get_param("OPENAI_BASE_PATH")
            .unwrap_or_else(|_| "v1/chat/completions".to_string());
        let organization: Option<String> = config.get_param("OPENAI_ORGANIZATION").ok();
        let project: Option<String> = config.get_param("OPENAI_PROJECT").ok();
        let custom_headers: Option<HashMap<String, String>> = secrets
            .get("OPENAI_CUSTOM_HEADERS")
            .cloned()
            .map(parse_custom_headers);
        let timeout_secs: u64 = config.get_param("OPENAI_TIMEOUT").unwrap_or(600);

        let auth = AuthMethod::BearerToken(api_key);
        let mut api_client =
            ApiClient::with_timeout(host, auth, std::time::Duration::from_secs(timeout_secs))?;

        if let Some(org) = &organization {
            api_client = api_client.with_header("OpenAI-Organization", org)?;
        }

        if let Some(project) = &project {
            api_client = api_client.with_header("OpenAI-Project", project)?;
        }

        if let Some(headers) = &custom_headers {
            let mut header_map = reqwest::header::HeaderMap::new();
            for (key, value) in headers {
                let header_name = reqwest::header::HeaderName::from_bytes(key.as_bytes())?;
                let header_value = reqwest::header::HeaderValue::from_str(value)?;
                header_map.insert(header_name, header_value);
            }
            api_client = api_client.with_headers(header_map)?;
        }

        Ok(Self {
            api_client,
            base_path,
            organization,
            project,
            model,
            custom_headers,
            supports_streaming: true,
            name: Self::metadata().name,
        })
    }

    #[doc(hidden)]
    pub fn new(api_client: ApiClient, model: ModelConfig) -> Self {
        Self {
            api_client,
            base_path: "v1/chat/completions".to_string(),
            organization: None,
            project: None,
            model,
            custom_headers: None,
            supports_streaming: true,
            name: Self::metadata().name,
        }
    }

    pub fn from_custom_config(
        model: ModelConfig,
        config: DeclarativeProviderConfig,
    ) -> Result<Self> {
        let global_config = crate::config::Config::global();
        let api_key: String = global_config
            .get_secret(&config.api_key_env)
            .map_err(|_e| anyhow::anyhow!("Missing API key: {}", config.api_key_env))?;

        let url = url::Url::parse(&config.base_url)
            .map_err(|e| anyhow::anyhow!("Invalid base URL '{}': {}", config.base_url, e))?;

        let host = if let Some(port) = url.port() {
            format!(
                "{}://{}:{}",
                url.scheme(),
                url.host_str().unwrap_or(""),
                port
            )
        } else {
            format!("{}://{}", url.scheme(), url.host_str().unwrap_or(""))
        };
        let base_path = url.path().trim_start_matches('/').to_string();
        let base_path = if base_path.is_empty() {
            "v1/chat/completions".to_string()
        } else {
            base_path
        };

        let timeout_secs = config.timeout_seconds.unwrap_or(600);
        let auth = AuthMethod::BearerToken(api_key);
        let mut api_client =
            ApiClient::with_timeout(host, auth, std::time::Duration::from_secs(timeout_secs))?;

        // Add custom headers if present
        if let Some(headers) = &config.headers {
            let mut header_map = reqwest::header::HeaderMap::new();
            for (key, value) in headers {
                let header_name = reqwest::header::HeaderName::from_bytes(key.as_bytes())?;
                let header_value = reqwest::header::HeaderValue::from_str(value)?;
                header_map.insert(header_name, header_value);
            }
            api_client = api_client.with_headers(header_map)?;
        }

        Ok(Self {
            api_client,
            base_path,
            organization: None,
            project: None,
            model,
            custom_headers: config.headers,
            supports_streaming: config.supports_streaming.unwrap_or(true),
            name: config.name.clone(),
        })
    }

    fn uses_responses_api(model_name: &str) -> bool {
        model_name.starts_with("gpt-5-codex") || model_name.starts_with("gpt-5.1-codex")
    }

    async fn post(&self, payload: &Value) -> Result<Value, ProviderError> {
        let response = self
            .api_client
            .response_post(&self.base_path, payload)
            .await?;
        handle_response_openai_compat(response).await
    }

    async fn post_responses(&self, payload: &Value) -> Result<Value, ProviderError> {
        let response = self
            .api_client
            .response_post("v1/responses", payload)
            .await?;
        handle_response_openai_compat(response).await
    }
}

#[async_trait]
impl Provider for OpenAiProvider {
    fn metadata() -> ProviderMetadata {
        let models = OPEN_AI_KNOWN_MODELS
            .iter()
            .map(|(name, limit)| ModelInfo::new(*name, *limit))
            .collect();
        ProviderMetadata::with_models(
            "openai",
            "OpenAI",
            "GPT-4 and other OpenAI models, including OpenAI compatible ones",
            OPEN_AI_DEFAULT_MODEL,
            models,
            OPEN_AI_DOC_URL,
            vec![
                ConfigKey::new("OPENAI_API_KEY", true, true, None),
                ConfigKey::new("OPENAI_HOST", true, false, Some("https://api.openai.com")),
                ConfigKey::new("OPENAI_BASE_PATH", true, false, Some("v1/chat/completions")),
                ConfigKey::new("OPENAI_ORGANIZATION", false, false, None),
                ConfigKey::new("OPENAI_PROJECT", false, false, None),
                ConfigKey::new("OPENAI_CUSTOM_HEADERS", false, true, None),
                ConfigKey::new("OPENAI_TIMEOUT", false, false, Some("600")),
            ],
        )
    }

    fn get_name(&self) -> &str {
        &self.name
    }

    fn get_model_config(&self) -> ModelConfig {
        self.model.clone()
    }

    #[tracing::instrument(
        skip(self, model_config, system, messages, tools),
        fields(model_config, input, output, input_tokens, output_tokens, total_tokens)
    )]
    async fn complete_with_model(
        &self,
        model_config: &ModelConfig,
        system: &str,
        messages: &[Message],
        tools: &[Tool],
    ) -> Result<(Message, ProviderUsage), ProviderError> {
        if Self::uses_responses_api(&model_config.model_name) {
            let payload = create_responses_request(model_config, system, messages, tools)?;
            let mut log = RequestLog::start(&self.model, &payload)?;

            let json_response = self
                .with_retry(|| async {
                    let payload_clone = payload.clone();
                    self.post_responses(&payload_clone).await
                })
                .await
                .inspect_err(|e| {
                    let _ = log.error(e);
                })?;

            let responses_api_response: ResponsesApiResponse =
                serde_json::from_value(json_response.clone()).map_err(|e| {
                    ProviderError::ExecutionError(format!(
                        "Failed to parse responses API response: {}",
                        e
                    ))
                })?;

            let message = responses_api_to_message(&responses_api_response)?;
            let usage = get_responses_usage(&responses_api_response);
            let model = responses_api_response.model.clone();

            log.write(&json_response, Some(&usage))?;
            Ok((message, ProviderUsage::new(model, usage)))
        } else {
            let payload = create_request(
                model_config,
                system,
                messages,
                tools,
                &ImageFormat::OpenAi,
                false,
            )?;

            let mut log = RequestLog::start(&self.model, &payload)?;
            let json_response = self
                .with_retry(|| async {
                    let payload_clone = payload.clone();
                    self.post(&payload_clone).await
                })
                .await
                .inspect_err(|e| {
                    let _ = log.error(e);
                })?;

            let message = response_to_message(&json_response)?;
            let usage = json_response
                .get("usage")
                .map(get_usage)
                .unwrap_or_else(|| {
                    tracing::debug!("Failed to get usage data");
                    Usage::default()
                });

            let model = get_model(&json_response);
            log.write(&json_response, Some(&usage))?;
            Ok((message, ProviderUsage::new(model, usage)))
        }
    }

    async fn fetch_supported_models(&self) -> Result<Option<Vec<String>>, ProviderError> {
        let models_path = self.base_path.replace("v1/chat/completions", "v1/models");
        let response = self.api_client.response_get(&models_path).await?;
        let json = handle_response_openai_compat(response).await?;
        if let Some(err_obj) = json.get("error") {
            let msg = err_obj
                .get("message")
                .and_then(|v| v.as_str())
                .unwrap_or("unknown error");
            return Err(ProviderError::Authentication(msg.to_string()));
        }

        let data = json.get("data").and_then(|v| v.as_array()).ok_or_else(|| {
            ProviderError::UsageError("Missing data field in JSON response".into())
        })?;
        let mut models: Vec<String> = data
            .iter()
            .filter_map(|m| m.get("id").and_then(|v| v.as_str()).map(str::to_string))
            .collect();
        models.sort();
        Ok(Some(models))
    }

    fn supports_embeddings(&self) -> bool {
        true
    }

    async fn create_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>, ProviderError> {
        EmbeddingCapable::create_embeddings(self, texts)
            .await
            .map_err(|e| ProviderError::ExecutionError(e.to_string()))
    }

    fn supports_streaming(&self) -> bool {
        self.supports_streaming
    }

    async fn stream(
        &self,
        system: &str,
        messages: &[Message],
        tools: &[Tool],
    ) -> Result<MessageStream, ProviderError> {
        if Self::uses_responses_api(&self.model.model_name) {
            let mut payload = create_responses_request(&self.model, system, messages, tools)?;
            payload["stream"] = serde_json::Value::Bool(true);

            let mut log = RequestLog::start(&self.model, &payload)?;

            let response = self
                .with_retry(|| async {
                    let payload_clone = payload.clone();
                    let resp = self
                        .api_client
                        .response_post("v1/responses", &payload_clone)
                        .await?;
                    handle_status_openai_compat(resp).await
                })
                .await
                .inspect_err(|e| {
                    let _ = log.error(e);
                })?;

            let stream = response.bytes_stream().map_err(io::Error::other);

            Ok(Box::pin(try_stream! {
                let stream_reader = StreamReader::new(stream);
                let framed = FramedRead::new(stream_reader, LinesCodec::new()).map_err(anyhow::Error::from);

                let message_stream = responses_api_to_streaming_message(framed);
                pin!(message_stream);
                while let Some(message) = message_stream.next().await {
                    let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)))?;
                    log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?;
                    yield (message, usage);
                }
            }))
        } else {
            let payload = create_request(
                &self.model,
                system,
                messages,
                tools,
                &ImageFormat::OpenAi,
                true,
            )?;
            let mut log = RequestLog::start(&self.model, &payload)?;

            let response = self
                .with_retry(|| async {
                    let resp = self
                        .api_client
                        .response_post(&self.base_path, &payload)
                        .await?;
                    handle_status_openai_compat(resp).await
                })
                .await
                .inspect_err(|e| {
                    let _ = log.error(e);
                })?;

            stream_openai_compat(response, log)
        }
    }
}

fn parse_custom_headers(s: String) -> HashMap<String, String> {
    s.split(',')
        .filter_map(|header| {
            let mut parts = header.splitn(2, '=');
            let key = parts.next().map(|s| s.trim().to_string())?;
            let value = parts.next().map(|s| s.trim().to_string())?;
            Some((key, value))
        })
        .collect()
}

#[async_trait]
impl EmbeddingCapable for OpenAiProvider {
    async fn create_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(vec![]);
        }

        let embedding_model = std::env::var("ASTER_EMBEDDING_MODEL")
            .unwrap_or_else(|_| "text-embedding-3-small".to_string());

        let request = EmbeddingRequest {
            input: texts,
            model: embedding_model,
        };

        let response = self
            .with_retry(|| async {
                let request_clone = EmbeddingRequest {
                    input: request.input.clone(),
                    model: request.model.clone(),
                };
                let request_value = serde_json::to_value(request_clone)
                    .map_err(|e| ProviderError::ExecutionError(e.to_string()))?;
                self.api_client
                    .api_post("v1/embeddings", &request_value)
                    .await
                    .map_err(|e| ProviderError::ExecutionError(e.to_string()))
            })
            .await?;

        if response.status != StatusCode::OK {
            let error_text = response
                .payload
                .as_ref()
                .and_then(|p| p.as_str())
                .unwrap_or("Unknown error");
            return Err(anyhow::anyhow!("Embedding API error: {}", error_text));
        }

        let embedding_response: EmbeddingResponse = serde_json::from_value(
            response
                .payload
                .ok_or_else(|| anyhow::anyhow!("Empty response body"))?,
        )?;

        Ok(embedding_response
            .data
            .into_iter()
            .map(|d| d.embedding)
            .collect())
    }
}