ralph/providers/

mod.rs

pub mod deepseek;
mod openai;

use crate::errors::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};

// ── Shared message types ─────────────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
    Tool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub role: Role,
    pub content: MessageContent,
    /// Tool call id — only set on Role::Tool result messages.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
    /// Name — only set on Role::Tool result messages.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    /// DeepSeek thinking content — must be echoed back on subsequent API calls.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_content: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessageContent {
    Text(String),
    Parts(Vec<ContentPart>),
}

impl MessageContent {
    pub fn as_text(&self) -> String {
        match self {
            MessageContent::Text(s) => s.clone(),
            MessageContent::Parts(parts) => parts
                .iter()
                .filter_map(|p| match p {
                    ContentPart::Text { text } => Some(text.as_str()),
                    _ => None,
                })
                .collect::<Vec<_>>()
                .join("\n"),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentPart {
    Text {
        text: String,
    },
    ToolUse {
        id: String,
        name: String,
        input: serde_json::Value,
    },
    ToolResult {
        tool_use_id: String,
        content: String,
    },
}

impl Message {
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: Role::System,
            content: MessageContent::Text(content.into()),
            tool_call_id: None,
            name: None,
            reasoning_content: None,
        }
    }

    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: Role::User,
            content: MessageContent::Text(content.into()),
            tool_call_id: None,
            name: None,
            reasoning_content: None,
        }
    }

    pub fn assistant(content: impl Into<String>) -> Self {
        Self {
            role: Role::Assistant,
            content: MessageContent::Text(content.into()),
            tool_call_id: None,
            name: None,
            reasoning_content: None,
        }
    }

    pub fn tool_result(
        tool_call_id: impl Into<String>,
        name: impl Into<String>,
        result: impl Into<String>,
    ) -> Self {
        Self {
            role: Role::Tool,
            content: MessageContent::Text(result.into()),
            tool_call_id: Some(tool_call_id.into()),
            name: Some(name.into()),
            reasoning_content: None,
        }
    }
}

// ── Tool definitions ─────────────────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDef {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

// ── Tool calls returned by the LLM ──────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub id: String,
    pub name: String,
    pub arguments: serde_json::Value,
}

// ── LLM response ─────────────────────────────────────────────────────────────

#[derive(Debug, Clone)]
pub struct LlmResponse {
    /// Text the model produced (may be empty if only tool calls were made).
    pub text: Option<String>,
    /// Tool calls the model wants to make.
    pub tool_calls: Vec<ToolCall>,
    /// Total tokens used (input + output + reasoning) — kept for compat.
    pub tokens_used: u64,
    /// Input (prompt) tokens.
    pub input_tokens: u64,
    /// Output (completion) tokens.
    pub output_tokens: u64,
    /// Reasoning/thinking tokens (DeepSeek extended thinking).
    pub reasoning_tokens: u64,
    /// Raw reasoning content text (DeepSeek thinking mode) — must be echoed back.
    pub reasoning_content: Option<String>,
    /// Stop reason.
    pub stop_reason: StopReason,
}

#[derive(Debug, Clone, PartialEq)]
pub enum StopReason {
    EndTurn,
    ToolUse,
    MaxTokens,
    Stop,
}

// ── Provider trait ────────────────────────────────────────────────────────────

#[async_trait]
pub trait LlmProvider: Send + Sync {
    /// Send messages + tool definitions, receive a response.
    async fn chat(&self, messages: &[Message], tools: &[ToolDef]) -> Result<LlmResponse>;

    /// Stream text tokens via `token_tx` while accumulating the full response.
    /// The default falls back to `chat()` and delivers the entire text as one chunk.
    async fn chat_streaming(
        &self,
        messages: &[Message],
        tools: &[ToolDef],
        token_tx: &tokio::sync::mpsc::UnboundedSender<String>,
    ) -> Result<LlmResponse> {
        let resp = self.chat(messages, tools).await?;
        if let Some(ref text) = resp.text {
            if !text.is_empty() {
                let _ = token_tx.send(text.clone());
            }
        }
        Ok(resp)
    }

    /// Whether this provider supports true token-by-token streaming.
    /// Providers that override `chat_streaming` with SSE should return true.
    fn supports_streaming(&self) -> bool {
        false
    }

    /// Human-readable name for this provider.
    fn name(&self) -> &str;

    /// Context window size in tokens.
    fn context_window(&self) -> u64;

    /// Default model name.
    fn default_model(&self) -> &str;
}

// ── Retry helper ─────────────────────────────────────────────────────────────

pub async fn with_retry<F, Fut>(
    provider_name: &str,
    max_attempts: u32,
    mut f: F,
) -> Result<LlmResponse>
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<Output = Result<LlmResponse>>,
{
    use crate::errors::RalphError;
    let mut delay = std::time::Duration::from_millis(500);
    for attempt in 1..=max_attempts {
        match f().await {
            Ok(r) => return Ok(r),
            Err(RalphError::LlmRateLimit { .. }) | Err(RalphError::LlmApi { .. })
                if attempt < max_attempts =>
            {
                tokio::time::sleep(delay).await;
                delay *= 2;
            }
            Err(e) => return Err(e),
        }
    }
    Err(RalphError::LlmRateLimit {
        provider: provider_name.to_string(),
        attempts: max_attempts,
    })
}