toast_api/

server.rs

use axum::{
    extract::State,
    http::StatusCode,
    response::IntoResponse,
    routing::{get, post},
    Json, Router,
};
use clap::Parser;
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;
use std::sync::Arc;
use tower_http::cors::{Any, CorsLayer};
use tracing::{info, warn};
use uuid::Uuid;

use crate::api::{Claude, Session};

/// Command line arguments for server configuration
#[derive(Parser, Debug, Clone)]
#[clap(author, version, about = "OpenAI-compatible Claude API server")]
pub struct Args {
    /// Port to listen on - public field
    #[clap(long, short, default_value = "3000")]
    pub port: u16,

    /// Enable debug logging
    #[clap(long)]
    pub debug: bool,

    /// Default Claude model to use
    #[clap(long, default_value = "claude-3-7-sonnet-latest")]
    pub model: String,

    /// Host address to bind to
    #[clap(long, default_value = "0.0.0.0")]
    pub host: String,
}

/// Server state shared across all requests
#[derive(Clone)]
pub struct AppState {
    pub claude: Arc<Claude>,
    // Map to track conversations: Vec<Message> -> (conversation_id, Vec<Message>)
    pub conversations: Arc<DashMap<String, (String, Vec<Message>)>>,
}

/// Base response structure for OpenAI-compatible API
#[derive(Serialize, Debug)]
pub struct OpenAIResponse {
    id: String,
    object: &'static str,
    created: u64,
    model: String,
    system_fingerprint: String,
    choices: Vec<Choice>,
    usage: Usage,
}

/// Token usage information
#[derive(Serialize, Debug)]
pub struct Usage {
    prompt_tokens: u32,
    completion_tokens: u32,
    total_tokens: u32,
}

/// Chat message structure (OpenAI compatible)
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
pub struct Message {
    role: String,
    content: ContentValue,
}

/// Content value can be string, array of strings, or structured content
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
#[serde(untagged)]
pub enum ContentValue {
    String(String),
    Array(Vec<String>),
    Object(ContentObject),
    ContentBlocks(Vec<ContentBlock>),
}

/// Content object with type and other fields
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
pub struct ContentObject {
    #[serde(default)]
    pub text: Option<String>,
    #[serde(default)]
    pub r#type: Option<String>,
    // Allow for other fields
    #[serde(flatten)]
    pub other: std::collections::HashMap<String, serde_json::Value>,
}

/// Structured content block
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
pub struct ContentBlock {
    pub r#type: String,
    #[serde(default)]
    pub text: Option<String>,
    // Allow for other fields
    #[serde(flatten)]
    pub other: std::collections::HashMap<String, serde_json::Value>,
}

impl Message {
    // Helper to get content as a string regardless of format
    pub fn content_as_string(&self) -> String {
        match &self.content {
            ContentValue::String(content) => content.clone(),
            ContentValue::Array(content) => content.join(" "),
            ContentValue::Object(obj) => {
                if let Some(text) = &obj.text {
                    text.clone()
                } else {
                    // Try to convert to JSON string if no text field
                    serde_json::to_string(obj).unwrap_or_default()
                }
            }
            ContentValue::ContentBlocks(blocks) => blocks
                .iter()
                .filter_map(|block| block.text.as_ref())
                .cloned()
                .collect::<Vec<_>>()
                .join(" "),
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct MessageInner {
    role: String,
    content: String,
}

impl From<MessageInner> for Message {
    fn from(inner: MessageInner) -> Self {
        Self {
            role: inner.role,
            content: ContentValue::String(inner.content),
        }
    }
}

/// Deserializer for message content that supports both strings and arrays
/// Custom serialization/deserialization for message content
/// Supports either a single string or an array of strings
mod content_string_or_array {
    use serde::{self, Deserialize, Deserializer, Serializer};

    pub fn _serialize<S>(content: &str, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(content)
    }

    pub fn _deserialize<'de, D>(deserializer: D) -> Result<String, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(untagged)]
        enum StringOrArray {
            String(()),
            Array(()),
        }

        let value = StringOrArray::deserialize(deserializer)?;
        match value {
            StringOrArray::String(_) => Ok(String::new()),
            StringOrArray::Array(_) => Ok(String::new()),
        }
    }
}

/// Chat completion request (OpenAI compatible)
#[derive(Deserialize, Serialize, Debug)]
pub struct ChatCompletionRequest {
    model: String,
    messages: Vec<Message>,
    #[serde(default)]
    user: Option<String>,  // Optional user identifier for thread association
}

/// Chat completion response choices (OpenAI compatible)
#[derive(Serialize, Debug)]
pub struct Choice {
    index: i32,
    message: Message,
    logprobs: Option<serde_json::Value>,
    finish_reason: String,
}

/// Helper function to get model from request or default to Sonnet
fn get_model(model_name: &str) -> &'static str {
    match model_name {
        "gpt-3.5-turbo" => crate::config::HAIKU_MODEL,
        "gpt-4" | "gpt-4-turbo" => crate::config::SONNET_MODEL,
        "gpt-4o" | "gpt-4.1" => crate::config::OPUS_MODEL,
        _ => crate::config::SONNET_MODEL,
    }
}

/// Normalize roles for Claude API
fn _normalize_role(role: &str) -> String {
    match role {
        "assistant" => "assistant".to_string(),
        "user" => "human".to_string(),
        "system" | "developer" => "system".to_string(),
        _ => "human".to_string(), // Default to human for unknown roles
    }
}

/// Convert messages to a single Claude prompt
fn _messages_to_prompt(messages: &[Message]) -> String {
    let mut prompt = String::new();

    for msg in messages {
        let role = _normalize_role(&msg.role);
        let content = msg.content_as_string();
        match role.as_str() {
            "system" => prompt.push_str(&format!("System: {content}\n\n")),
            "human" => prompt.push_str(&format!("Human: {content}\n\n")),
            "assistant" => prompt.push_str(&format!("Assistant: {content}\n\n")),
            _ => prompt.push_str(&format!("{role}: {content}\n\n")),
        }
    }

    prompt
}

/// Generate a consistent conversation key based on message context
fn generate_conversation_key(messages: &[Message]) -> String {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    
    // Create a consistent hash based on the conversation context (all messages except the last user message)
    let context_messages = if messages.len() > 1 {
        &messages[..messages.len() - 1]
    } else {
        &[]
    };
    
    let mut hasher = DefaultHasher::new();
    for msg in context_messages {
        msg.role.hash(&mut hasher);
        msg.content_as_string().hash(&mut hasher);
    }
    
    format!("ctx:{:016x}", hasher.finish())
}

/// Find a matching conversation by context hash
/// Returns the key of the matching conversation if found
fn find_matching_conversation(conversations: &DashMap<String, (String, Vec<Message>)>, messages: &[Message]) -> Option<String> {
    if messages.is_empty() {
        return None;
    }

    // Generate the context key for the current request
    let context_key = generate_conversation_key(messages);
    
    // Check if we have an exact context match
    if conversations.contains_key(&context_key) {
        return Some(context_key);
    }
    
    // If no exact match, try to find a conversation that could be extended
    // Look for conversations where our context is a prefix of their context
    let our_context = if messages.len() > 1 {
        &messages[..messages.len() - 1]
    } else {
        &[]
    };
    
    for item in conversations.iter() {
        let (_, stored_history) = item.value();
        
        // Check if our context is a prefix of the stored history
        if stored_history.len() >= our_context.len() {
            let prefix_matches = our_context.iter().zip(stored_history.iter())
                .all(|(a, b)| a.role == b.role && a.content_as_string() == b.content_as_string());
            
            if prefix_matches {
                return Some(item.key().clone());
            }
        }
    }
    
    None
}

/// Chat completion API handler
async fn chat_completion_handler(
    State(state): State<AppState>,
    Json(request): Json<ChatCompletionRequest>,
) -> impl IntoResponse {
    // Debug the raw client request
    // dbg!(&request);

    // Log when array content is used
    // Log message content types
    for (i, msg) in request.messages.iter().enumerate() {
        match &msg.content {
            ContentValue::Array(_) => info!("Message at index {} has array content", i),
            ContentValue::Object(_) => info!("Message at index {} has object content", i),
            ContentValue::ContentBlocks(_) => info!("Message at index {} has content blocks", i),
            _ => {} // String content is the default, no need to log
        }
    }

    let model_id = get_model(&request.model);

    if request.messages.is_empty() {
        return (
            StatusCode::BAD_REQUEST,
            Json(serde_json::json!({
                "error": "No messages provided in the request"
            })),
        );
    }

    let messages = request.messages;

    // Generate a request identifier from the user field or a new UUID if not provided
    // We don't directly use this anymore, but keep it for future use
    let _request_id = request.user.clone().unwrap_or_else(|| Uuid::new_v4().to_string());
    
    // The actual user query is the last message
    let user_query = messages.last().unwrap().clone();

    let claude = state.claude;
    let conversations = state.conversations;

    // Try to find existing conversation by matching message history
    let mut conversation_id: Option<String> = None;
    let mut message_history: Vec<Message> = Vec::new();
    
    info!("Looking for matching conversation by message history");
    info!("Available conversations: {}", conversations.len());

    // Find matching conversation based on message history
    let matching_key = find_matching_conversation(&conversations, &messages);

    if let Some(ref key) = matching_key {
        if let Some(existing) = conversations.get(key) {
            conversation_id = Some(existing.0.clone());
            message_history = existing.1.clone();
            info!("Found conversation with matching history, key: '{}'", key);
        }
    } else {
        info!("No existing conversation with matching history found");
    }

    // If no matching conversation, create a new one
    if conversation_id.is_none() {
        match claude.create_chat().await {
            Ok(new_id) => {
                conversation_id = Some(new_id.clone());
                message_history = messages[..messages.len() - 1].to_vec();

                // Send system prompt if present and include it in history
                if let Some(system_msg) = messages
                    .iter()
                    .find(|m| m.role == "system" || m.role == "developer")
                {
                    if let Err(e) = claude
                        .send_message(&new_id, &system_msg.content_as_string(), &[])
                        .await
                    {
                        warn!("Failed to send system prompt: {}", e);
                    }
                }

                // Generate context-based key for new conversation
                let context_key = generate_conversation_key(&messages);
                conversations.insert(
                    context_key,
                    (new_id.clone(), message_history.clone()),
                );
            }
            Err(e) => {
                let error_response = serde_json::json!({
                    "error": format!("Failed to create conversation: {}", e)
                });
                dbg!(&error_response);
                return (StatusCode::INTERNAL_SERVER_ERROR, Json(error_response));
            }
        }
    }

    // Send the user message and get response
    let completion = match claude
        .send_message(
            &conversation_id.clone().unwrap(),
            &user_query.content_as_string(),
            &[],
        )
        .await
    {
        Ok(response) => response,
        Err(e) => {
            let error_response = serde_json::json!({
                "error": format!("Failed to get completion: {}", e)
            });
            dbg!(&error_response);
            return (StatusCode::INTERNAL_SERVER_ERROR, Json(error_response));
        }
    };

    // Update the message history with the new assistant response
    let new_message = Message {
        role: "assistant".to_string(),
        content: ContentValue::String(completion.clone()),
    };

    // Update message history and store it
    let mut updated_history = message_history.clone();
    updated_history.push(user_query.clone());
    updated_history.push(new_message);

    let conv_id = conversation_id.unwrap();
    
    // Generate context-based key for the updated conversation
    let mut extended_messages = messages.clone();
    extended_messages.push(Message {
        role: "assistant".to_string(),
        content: ContentValue::String(completion.clone()),
    });
    let updated_key = generate_conversation_key(&extended_messages);
    
    // Remove old entry if we had a matching key and it's different from the new key
    if let Some(ref old_key) = matching_key {
        if old_key != &updated_key
            && conversations.remove(old_key).is_some() {
                info!("Removed old conversation entry with key: '{}'", old_key);
            }
    }

    // Store with the updated context key
    conversations.insert(
        updated_key.clone(),
        (conv_id.clone(), updated_history.clone()),
    );
    
    // Log the conversation storage
    info!(
        "Stored conversation with ID: {} under context key '{}'",
        conv_id,
        updated_key
    );

    // Create OpenAI-compatible response
    let response = OpenAIResponse {
        id: Uuid::new_v4().to_string(),
        object: "chat.completion",
        created: chrono::Utc::now().timestamp() as u64,
        model: model_id.to_string(),
        system_fingerprint: "fp_toast".to_string(),
        choices: vec![Choice {
            index: 0,
            message: Message {
                role: "assistant".to_string(),
                content: ContentValue::String(completion.clone()),
            },
            logprobs: None,
            finish_reason: "stop".to_string(),
        }],
        usage: Usage {
            prompt_tokens: user_query.content_as_string().chars().count() as u32 / 4, // Rough estimate
            completion_tokens: completion.chars().count() as u32 / 4, // Rough estimate
            total_tokens: (user_query.content_as_string().chars().count()
                + completion.chars().count()) as u32
                / 4,
        },
    };

    // Debug the response being sent to the client
    // dbg!(&response);

    (
        StatusCode::OK,
        Json(serde_json::to_value(response).unwrap()),
    )
}

/// Health check endpoint
async fn health_check() -> impl IntoResponse {
    (StatusCode::OK, "OK")
}

/// Start the Axum server
pub async fn run() -> anyhow::Result<()> {
    // For backward compatibility, parse arguments directly when run() is called
    run_with_args(Args::parse()).await
}

/// Start the Axum server with provided arguments
pub async fn run_with_args(args: Args) -> anyhow::Result<()> {
    // Use provided arguments instead of parsing them

    // Initialize tracing
    if args.debug {
        // Full logging in debug mode
        tracing_subscriber::fmt::init();
    } else {
        // Only show warnings and errors in normal mode
        tracing_subscriber::fmt()
            .with_max_level(tracing::Level::WARN)
            .init();
    }

    // Load config
    let config_dir = dirs::config_dir()
        .ok_or_else(|| anyhow::anyhow!("Could not determine config directory"))?
        .join("toast");

    let cookie_path = config_dir.join("cookie");
    let org_id_path = config_dir.join("org_id");

    // Check if config directory exists
    if !config_dir.exists() {
        return Err(anyhow::anyhow!(
            "Configuration directory does not exist at {:?}",
            config_dir
        ));
    }

    // Check and load cookie
    let cookie = if cookie_path.exists() {
        std::fs::read_to_string(&cookie_path)?.trim().to_string()
    } else {
        return Err(anyhow::anyhow!(
            "Cookie file not found at {:?}",
            cookie_path
        ));
    };

    // Check and load org_id
    let org_id = if org_id_path.exists() {
        std::fs::read_to_string(&org_id_path)?.trim().to_string()
    } else {
        // Try to extract org_id from cookie
        if let Some(extracted_org_id) = crate::utils::extract_org_id_from_cookie(&cookie) {
            // Save the extracted org_id to the file for future use
            std::fs::write(&org_id_path, &extracted_org_id)?;
            info!(
                "Extracted organization ID from cookie and saved to {:?}",
                org_id_path
            );
            extracted_org_id
        } else {
            return Err(anyhow::anyhow!(
                "Organization ID file not found at {:?} and couldn't extract it from cookie",
                org_id_path
            ));
        }
    };

    let user_agent =
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:137.0) Gecko/20100101 Firefox/137.0"
            .to_string();

    let session = Session {
        cookie,
        user_agent,
        organization_id: org_id,
    };

    // Use the configurable model from arguments or default to Sonnet
    let model_str = crate::config::SONNET_MODEL;

    let claude = Arc::new(Claude::new(session.clone(), model_str)?);

    // Create app state
    let app_state = AppState {
        claude,
        conversations: Arc::new(DashMap::new()),
    };

    // Set up CORS
    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods(Any)
        .allow_headers(Any);

    // Create router
    let app = Router::new()
        .route("/health", get(health_check))
        .route("/v1/chat/completions", post(chat_completion_handler))
        .route("/", post(chat_completion_handler))
        .route("/chat/completions", post(chat_completion_handler))
        .layer(cors)
        .with_state(app_state);

    // Parse host address
    let host_parts: Vec<u8> = args
        .host
        .split('.')
        .map(|s| s.parse::<u8>().unwrap_or(0))
        .collect();

    let host_addr = if host_parts.len() == 4 {
        [host_parts[0], host_parts[1], host_parts[2], host_parts[3]]
    } else {
        [0, 0, 0, 0]
    };

    // Run server
    let addr = SocketAddr::from((host_addr, args.port));
    info!("Server listening on {}", addr);

    // Print usage example
    info!("Example request:");
    info!(
        "curl -X POST http://localhost:{}{} \\",
        args.port, "/v1/chat/completions  # Also available at / and /chat/completions"
    );
    info!("  -H \"Content-Type: application/json\" \\");
    info!("  -d '{{");
    info!("    \"model\": \"gpt-4.1\",");
    info!("    \"messages\": [");
    info!("      {{");
    info!("        \"role\": \"developer\",");
    info!("        \"content\": \"You are a helpful assistant.\"");
    info!("      }},");
    info!("      {{");
    info!("        \"role\": \"user\",");
    info!("        \"content\": \"Hello! (Simple string content)\"");
    info!("      }},");
    info!("      {{");
    info!("        \"role\": \"user\",");
    info!("        \"content\": [\"Array of\", \"string content\"]");
    info!("      }},");
    info!("      {{");
    info!("        \"role\": \"user\",");
    info!("        \"content\": {{\"type\": \"text\", \"text\": \"Object with text field\"}}");
    info!("      }},");
    info!("      {{");
    info!("        \"role\": \"user\",");
    info!("        \"content\": [");
    info!("          {{\"type\": \"text\", \"text\": \"Array of content blocks\"}},");
    info!("          {{\"type\": \"text\", \"text\": \"with multiple entries\"}}");
    info!("        ]");
    info!("      }}");
    info!("    ]");
    info!("  }}'");

    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;

    Ok(())
}