agnt5-sdk-core 0.3.9

use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

use anyhow::anyhow;
use async_trait::async_trait;
use futures::{Stream, StreamExt};
use serde_json::Value;

use crate::error::{Result as SdkResult, SdkError};

#[async_trait]
pub trait LanguageModel: Send + Sync {
    async fn generate(&self, request: GenerateRequest) -> SdkResult<GenerateResponse>;
    async fn stream(&self, request: StreamRequest) -> SdkResult<StreamHandle>;
}

pub async fn generate<M>(model: &M, request: GenerateRequest) -> SdkResult<GenerateResponse>
where
    M: LanguageModel,
{
    model.generate(request).await
}

pub async fn stream<M>(model: &M, request: StreamRequest) -> SdkResult<StreamHandle>
where
    M: LanguageModel,
{
    model.stream(request).await
}

#[derive(Clone, Debug)]
pub struct GenerateRequest {
    pub model: String,
    pub prompt_ref: Option<PromptRef>,
    pub system_prompt: Option<String>,
    pub messages: Vec<Message>,
    pub tools: Vec<ToolDefinition>,
    pub tool_choice: Option<ToolChoice>,
    pub user_id: Option<String>,
    pub config: GenerationConfig,
    /// Previous response ID for conversation continuation (OpenAI Responses API)
    /// When set, the API will use this to continue a multi-turn conversation
    /// with tool calls without needing to resend the full message history.
    pub previous_response_id: Option<String>,
    /// OpenTelemetry context for trace propagation across async boundaries
    /// This is used internally to ensure LM spans are children of the calling function span
    #[doc(hidden)]
    pub otel_context: Option<opentelemetry::Context>,
}

impl GenerateRequest {
    pub fn new(model: impl Into<String>) -> Self {
        Self {
            model: model.into(),
            prompt_ref: None,
            system_prompt: None,
            messages: Vec::new(),
            tools: Vec::new(),
            tool_choice: None,
            user_id: None,
            config: GenerationConfig::default(),
            previous_response_id: None,
            otel_context: None,
        }
    }

    pub fn prompt_ref(mut self, prompt_ref: PromptRef) -> Self {
        self.prompt_ref = Some(prompt_ref);
        self
    }

    pub fn previous_response_id(mut self, id: impl Into<String>) -> Self {
        self.previous_response_id = Some(id.into());
        self
    }

    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = Some(prompt.into());
        self
    }

    pub fn message(mut self, message: Message) -> Self {
        self.messages.push(message);
        self
    }

    pub fn system_message(mut self, content: impl Into<String>) -> Self {
        self.messages.push(Message::system(content));
        self
    }

    pub fn user_message(mut self, content: impl Into<String>) -> Self {
        self.messages.push(Message::user(content));
        self
    }

    pub fn assistant_message(mut self, content: impl Into<String>) -> Self {
        self.messages.push(Message::assistant(content));
        self
    }

    pub fn add_tool(mut self, tool: ToolDefinition) -> Self {
        self.tools.push(tool);
        self
    }

    pub fn tools(mut self, tools: Vec<ToolDefinition>) -> Self {
        self.tools = tools;
        self
    }

    pub fn tool_choice(mut self, choice: Option<ToolChoice>) -> Self {
        self.tool_choice = choice;
        self
    }

    pub fn user_id(mut self, value: impl Into<String>) -> Self {
        self.user_id = Some(value.into());
        self
    }

    pub fn configure<F>(mut self, configure: F) -> Self
    where
        F: FnOnce(&mut GenerationConfig),
    {
        configure(&mut self.config);
        self
    }

    pub fn with_config(mut self, config: GenerationConfig) -> Self {
        self.config = config;
        self
    }

    pub fn response_format(mut self, format: ResponseFormat) -> Self {
        self.config.response_format = format;
        self
    }
}

pub type StreamRequest = GenerateRequest;

#[derive(Clone, Debug, PartialEq)]
pub struct PromptRef {
    pub id: String,
    pub version: Option<String>,
    pub environment_id: Option<String>,
    pub environment_ref: Option<String>,
    pub variables: serde_json::Map<String, Value>,
}

impl PromptRef {
    pub fn new(id: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            version: None,
            environment_id: None,
            environment_ref: None,
            variables: serde_json::Map::new(),
        }
    }

    pub fn version(mut self, version: impl Into<String>) -> Self {
        self.version = Some(version.into());
        self
    }

    pub fn environment_id(mut self, environment_id: impl Into<String>) -> Self {
        self.environment_id = Some(environment_id.into());
        self
    }

    pub fn environment_ref(mut self, environment_ref: impl Into<String>) -> Self {
        self.environment_ref = Some(environment_ref.into());
        self
    }

    pub fn variable(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
        self.variables.insert(key.into(), value.into());
        self
    }
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct PromptCacheConfig {
    /// Enable provider-native prompt caching for providers that require an
    /// explicit request signal. Providers with automatic prompt caching may
    /// ignore this while still reporting cache usage.
    pub enabled: bool,
    /// Optional cache lifetime hint. Anthropic currently supports "1h".
    pub ttl: Option<String>,
    /// Optional provider-side routing key. OpenAI Responses maps this to
    /// `prompt_cache_key`.
    pub key: Option<String>,
    /// Optional retention hint. OpenAI Responses maps this to
    /// `prompt_cache_retention`.
    pub retention: Option<String>,
    /// Explicit reusable context-cache resource name. Gemini maps this to
    /// `cachedContent`.
    pub resource: Option<String>,
}

impl PromptCacheConfig {
    pub fn enabled() -> Self {
        Self {
            enabled: true,
            ..Default::default()
        }
    }

    pub fn ttl(mut self, ttl: impl Into<String>) -> Self {
        self.ttl = Some(ttl.into());
        self.enabled = true;
        self
    }

    pub fn key(mut self, key: impl Into<String>) -> Self {
        self.key = Some(key.into());
        self.enabled = true;
        self
    }

    pub fn retention(mut self, retention: impl Into<String>) -> Self {
        self.retention = Some(retention.into());
        self.enabled = true;
        self
    }

    pub fn resource(mut self, resource: impl Into<String>) -> Self {
        self.resource = Some(resource.into());
        self.enabled = true;
        self
    }
}

#[derive(Clone, Debug, Default)]
pub struct GenerationConfig {
    pub temperature: Option<f32>,
    pub top_p: Option<f32>,
    pub max_output_tokens: Option<u32>,
    pub response_format: ResponseFormat,
    pub prompt_cache: Option<PromptCacheConfig>,
    /// Provider-neutral reasoning effort.
    /// Maps to OpenAI reasoning effort and Gemini 3 thinking level.
    pub reasoning_effort: Option<ReasoningEffort>,
    /// Output modalities (text, audio, image)
    /// Only supported by OpenAI Responses API
    pub modalities: Option<Vec<Modality>>,
    /// Built-in tools (web search, code interpreter, file search)
    /// Only supported by OpenAI Responses API
    pub built_in_tools: Vec<BuiltInTool>,
    /// Per-request timeout override. When set, overrides the provider-level timeout
    /// for this specific request.
    pub timeout: Option<Duration>,
}

impl GenerationConfig {
    pub fn temperature(mut self, value: f32) -> Self {
        self.temperature = Some(value);
        self
    }

    pub fn top_p(mut self, value: f32) -> Self {
        self.top_p = Some(value);
        self
    }

    pub fn max_output_tokens(mut self, value: u32) -> Self {
        self.max_output_tokens = Some(value);
        self
    }

    pub fn response_format(mut self, format: ResponseFormat) -> Self {
        self.response_format = format;
        self
    }

    pub fn prompt_cache(mut self, config: PromptCacheConfig) -> Self {
        self.prompt_cache = Some(config);
        self
    }

    pub fn enable_prompt_cache(mut self) -> Self {
        self.prompt_cache = Some(PromptCacheConfig::enabled());
        self
    }

    pub fn reasoning_effort(mut self, effort: ReasoningEffort) -> Self {
        self.reasoning_effort = Some(effort);
        self
    }

    pub fn modalities(mut self, modalities: Vec<Modality>) -> Self {
        self.modalities = Some(modalities);
        self
    }

    pub fn add_built_in_tool(mut self, tool: BuiltInTool) -> Self {
        self.built_in_tools.push(tool);
        self
    }

    pub fn built_in_tools(mut self, tools: Vec<BuiltInTool>) -> Self {
        self.built_in_tools = tools;
        self
    }

    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum ResponseFormat {
    Text,
    Json,
    JsonSchema(JsonSchemaFormat),
}

impl Default for ResponseFormat {
    fn default() -> Self {
        ResponseFormat::Text
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct JsonSchemaFormat {
    pub name: String,
    pub schema: Value,
    pub strict: bool,
}

impl JsonSchemaFormat {
    pub fn new(name: impl Into<String>, schema: Value) -> Self {
        Self {
            name: name.into(),
            schema,
            strict: true,
        }
    }

    pub fn with_strict(mut self, strict: bool) -> Self {
        self.strict = strict;
        self
    }
}

/// Provider-neutral control for how much compute a model uses for reasoning.
#[derive(Clone, Debug, PartialEq)]
pub enum ReasoningEffort {
    /// No reasoning (gpt-6; gpt-5 rejects it)
    None,
    /// Minimal reasoning effort (gpt-5; gpt-6 rejects it)
    Minimal,
    /// Low reasoning effort
    Low,
    /// Medium reasoning effort (balanced)
    Medium,
    /// High reasoning effort (most thorough)
    High,
}

/// Provider-hosted tools enabled via `Agent(built_in_tools=[...])`. Each variant
/// maps differently per provider — OpenAI emits `web_search_preview`, Anthropic
/// emits `web_search_20250305`, Gemini emits `google_search`, etc. Variants
/// that a given provider doesn't host are silently skipped at request build.
#[derive(Clone, Debug, PartialEq)]
pub enum BuiltInTool {
    /// Provider-hosted web search.
    WebSearch,
    /// Python code interpreter (OpenAI Responses API only).
    CodeInterpreter,
    /// File search over uploaded documents (OpenAI Responses API only).
    FileSearch,
    /// Provider-hosted page fetch (Anthropic only today; OpenAI's web_search
    /// already includes fetched content, Gemini has no separate fetch).
    WebFetch,
}

impl BuiltInTool {
    /// Resolve a provider-side tool name back to its `BuiltInTool` variant.
    ///
    /// Recognizes both OpenAI Responses API names (`web_search_preview`,
    /// `code_interpreter`, `file_search`) and the cross-provider bare names
    /// surfaced by Anthropic in `tool_use` blocks (`web_search`, `web_fetch`).
    /// Returns `None` for unknown names; callers decide whether that's an
    /// error or a silent skip.
    pub fn from_provider_name(name: &str) -> Option<Self> {
        match name {
            "web_search_preview" | "web_search" => Some(Self::WebSearch),
            "code_interpreter" => Some(Self::CodeInterpreter),
            "file_search" => Some(Self::FileSearch),
            "web_fetch" => Some(Self::WebFetch),
            _ => None,
        }
    }
}

#[cfg(test)]
mod built_in_tool_tests {
    use super::BuiltInTool;

    #[test]
    fn openai_names_map_correctly() {
        assert_eq!(
            BuiltInTool::from_provider_name("web_search_preview"),
            Some(BuiltInTool::WebSearch)
        );
        assert_eq!(
            BuiltInTool::from_provider_name("code_interpreter"),
            Some(BuiltInTool::CodeInterpreter)
        );
        assert_eq!(
            BuiltInTool::from_provider_name("file_search"),
            Some(BuiltInTool::FileSearch)
        );
    }

    #[test]
    fn anthropic_bare_names_map_correctly() {
        // Anthropic surfaces the bare `name` field of the tool spec, not the type.
        assert_eq!(
            BuiltInTool::from_provider_name("web_search"),
            Some(BuiltInTool::WebSearch)
        );
        assert_eq!(
            BuiltInTool::from_provider_name("web_fetch"),
            Some(BuiltInTool::WebFetch)
        );
    }

    #[test]
    fn web_search_aliases_collapse_to_same_variant() {
        // Both names must map to the same variant so the dispatch loop and the
        // request builder agree about which provider tool is in play.
        assert_eq!(
            BuiltInTool::from_provider_name("web_search_preview"),
            BuiltInTool::from_provider_name("web_search")
        );
    }

    #[test]
    fn unknown_name_returns_none() {
        assert_eq!(BuiltInTool::from_provider_name("unknown_tool"), None);
        assert_eq!(BuiltInTool::from_provider_name(""), None);
        assert_eq!(BuiltInTool::from_provider_name("WebSearch"), None); // case-sensitive
    }
}

/// Output modalities supported by the model.
#[derive(Clone, Debug, PartialEq)]
pub enum Modality {
    /// Text output
    Text,
    /// Audio output
    Audio,
    /// Image output
    Image,
}

#[derive(Clone, Debug)]
pub struct Message {
    pub role: MessageRole,
    pub content: String,
    /// Tool calls made by the assistant (for assistant messages)
    pub tool_calls: Option<Vec<ToolCall>>,
    /// ID of the tool call this message is responding to (for tool result messages)
    pub tool_call_id: Option<String>,
}

impl Message {
    pub fn new(role: MessageRole, content: impl Into<String>) -> Self {
        Self {
            role,
            content: content.into(),
            tool_calls: None,
            tool_call_id: None,
        }
    }

    pub fn system(content: impl Into<String>) -> Self {
        Self::new(MessageRole::System, content)
    }

    pub fn user(content: impl Into<String>) -> Self {
        Self::new(MessageRole::User, content)
    }

    pub fn assistant(content: impl Into<String>) -> Self {
        Self::new(MessageRole::Assistant, content)
    }

    /// Create an assistant message with tool calls
    pub fn assistant_with_tool_calls(
        content: impl Into<String>,
        tool_calls: Vec<ToolCall>,
    ) -> Self {
        Self {
            role: MessageRole::Assistant,
            content: content.into(),
            tool_calls: Some(tool_calls),
            tool_call_id: None,
        }
    }

    /// Create a tool result message
    pub fn tool_result(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            role: MessageRole::User, // Tool results are typically sent as user role
            content: content.into(),
            tool_calls: None,
            tool_call_id: Some(tool_call_id.into()),
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MessageRole {
    System,
    User,
    Assistant,
}

impl MessageRole {
    pub fn as_str(&self) -> &'static str {
        match self {
            MessageRole::System => "system",
            MessageRole::User => "user",
            MessageRole::Assistant => "assistant",
        }
    }
}

#[derive(Clone, Debug)]
pub struct ToolDefinition {
    pub name: String,
    pub description: Option<String>,
    pub parameters: Option<Value>,
    pub strict: Option<bool>,
}

impl ToolDefinition {
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            description: None,
            parameters: None,
            strict: None,
        }
    }

    pub fn description(mut self, text: impl Into<String>) -> Self {
        self.description = Some(text.into());
        self
    }

    pub fn parameters(mut self, parameters: Value) -> Self {
        self.parameters = Some(parameters);
        self
    }

    pub fn strict(mut self, strict: bool) -> Self {
        self.strict = Some(strict);
        self
    }
}

#[derive(Clone, Debug)]
pub enum ToolChoice {
    Auto,
    None,
    Required,
    Tool { name: String },
}

/// Metadata from the HTTP response headers of an LM API call.
#[derive(Clone, Debug, Default)]
pub struct ResponseMetadata {
    pub status_code: Option<u16>,
    pub request_id: Option<String>,
    pub rate_limit_remaining: Option<u32>,
    pub rate_limit_reset: Option<Duration>,
}

#[derive(Clone, Debug)]
pub struct GenerateResponse {
    pub id: String,
    pub model: String,
    pub created: Option<u64>,
    pub text: String,
    pub usage: Option<TokenUsage>,
    pub finish_reason: Option<String>,
    pub tool_calls: Option<Vec<ToolCall>>,
    pub object: Option<Value>,
    pub raw: Option<Value>,
    pub metadata: Option<ResponseMetadata>,
}

#[derive(Clone, Debug)]
pub struct ToolCall {
    pub id: String,
    pub name: String,
    pub arguments: String,
    /// Opaque provider state that must be replayed with this call.
    /// Callers should preserve it unchanged in subsequent assistant messages.
    pub provider_data: Option<Value>,
}

#[derive(Clone, Debug, Default)]
pub struct TokenUsage {
    pub prompt_tokens: Option<u32>,
    pub completion_tokens: Option<u32>,
    pub total_tokens: Option<u32>,
    /// Input tokens served from a prompt cache (cache hits).
    ///
    /// This is a subset of `prompt_tokens`, following the OpenAI usage
    /// convention. Maps to OpenAI `prompt_tokens_details.cached_tokens`,
    /// Anthropic/Bedrock `cache_read_input_tokens`, and Gemini
    /// `cachedContentTokenCount`. `None` when the provider does not report it.
    pub cached_tokens: Option<u32>,
    /// Input tokens written to the prompt cache on this request (cache writes).
    ///
    /// Anthropic-style providers only (`cache_creation_input_tokens`); these
    /// are billed at a premium and are also counted within `prompt_tokens`.
    /// `None` for providers that do not report or charge for cache creation.
    pub cache_creation_tokens: Option<u32>,
}

/// Type of content block in a streaming response.
///
/// Used to distinguish between regular assistant output and thinking/reasoning
/// content from extended thinking models (Claude) or o-series models (OpenAI).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum ContentBlockType {
    /// Regular text content from the assistant
    #[default]
    Text,
    /// Thinking/reasoning content (extended thinking, chain-of-thought)
    Thinking,
}

/// A chunk of streaming response from the language model.
///
/// Streaming responses emit these events in order:
/// 1. `ContentBlockStart` - indicates a new content block is beginning
/// 2. `Delta` - incremental content (may be emitted multiple times)
/// 3. `ContentBlockStop` - indicates the content block has finished
/// 4. `Completed` - final response with full text, usage, and tool calls
///
/// A response may contain multiple content blocks (e.g., thinking then text).
pub enum StreamChunk {
    /// Start of a new content block.
    ContentBlockStart {
        /// Index of this content block (0-indexed)
        index: u32,
        /// Type of content block (text or thinking)
        block_type: ContentBlockType,
    },
    /// Incremental content within a content block.
    Delta {
        /// The content text
        content: String,
        /// Index of the content block this delta belongs to
        index: u32,
        /// Type of content block (text or thinking)
        block_type: ContentBlockType,
    },
    /// End of a content block.
    ContentBlockStop {
        /// Index of the content block that ended
        index: u32,
    },
    /// Stream completed with final response.
    Completed(GenerateResponse),
}

pub struct StreamHandle {
    inner: Pin<Box<dyn Stream<Item = SdkResult<StreamChunk>> + Send>>,
}

impl StreamHandle {
    pub(crate) fn new(inner: Pin<Box<dyn Stream<Item = SdkResult<StreamChunk>> + Send>>) -> Self {
        Self { inner }
    }

    pub async fn collect_text(mut self) -> SdkResult<GenerateResponse> {
        let mut final_response: Option<GenerateResponse> = None;
        while let Some(item) = self.inner.next().await {
            match item? {
                StreamChunk::ContentBlockStart { .. } => {}
                StreamChunk::Delta { .. } => {}
                StreamChunk::ContentBlockStop { .. } => {}
                StreamChunk::Completed(response) => {
                    final_response = Some(response);
                    break;
                }
            }
        }

        final_response
            .ok_or_else(|| SdkError::Other(anyhow!("stream ended without a completion response")))
    }

    pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = SdkResult<StreamChunk>> + Send>> {
        self.inner
    }
}

impl Stream for StreamHandle {
    type Item = SdkResult<StreamChunk>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // Safety: we never move the inner stream after pinning.
        let inner = unsafe { self.map_unchecked_mut(|s| &mut s.inner) };
        inner.poll_next(cx)
    }
}