windmill/backend/windmill-ai/src/query_builder.rs

use std::borrow::Cow;

use async_trait::async_trait;
use windmill_common::{client::AuthedClient, error::Error, utils::strip_json_nul};
use windmill_types::s3::S3Object;

use crate::ai_types::OpenAIToolCall;
use crate::proxy::{ProxyBuildArgs, ProxyRequest};
use crate::types::*;

/// Only explicit input-context rejections warrant dropping conversation history.
/// Rate limits, output-token limits and generic validation errors must propagate.
pub fn is_context_length_error(message: &str) -> bool {
    let message = message.to_ascii_lowercase();
    message.contains("context_length_exceeded")
        || message.contains("maximum context length")
        || message.contains("prompt is too long")
        || message.contains("exceed context limit")
        || message.contains("input is too long for requested model")
        || (message.contains("input token count") && message.contains("exceeds the maximum"))
        || message.contains("too many input tokens")
}

#[cfg(test)]
mod context_error_tests {
    use super::is_context_length_error;

    #[test]
    fn recognizes_context_rejections_without_retrying_other_provider_errors() {
        for message in [
            r#"{"error":{"code":"context_length_exceeded"}}"#,
            "This model's maximum context length is 128000 tokens",
            "prompt is too long: 210000 tokens > 200000 maximum",
            "input length and `max_tokens` exceed context limit: 199000 + 8192 > 200000",
            "The input token count (10000) exceeds the maximum number of tokens allowed (8192)",
            "ValidationException: Input is too long for requested model.",
            "ValidationException: Too many input tokens",
        ] {
            assert!(is_context_length_error(message), "{message}");
        }
        for message in [
            "Rate limit exceeded: tokens per minute",
            "max_tokens exceeds the maximum output tokens",
            "Invalid tool schema",
            "Additional properties are not allowed: stream_options",
            "Request body too large",
            "Internal server error",
        ] {
            assert!(!is_context_length_error(message), "{message}");
        }
    }
}

/// Arguments for building an AI request
pub struct BuildRequestArgs<'a> {
    pub messages: &'a [OpenAIMessage],
    pub tools: Option<&'a [ToolDef]>,
    pub model: &'a str,
    pub temperature: Option<f32>,
    /// Provider-native reasoning effort token (e.g. `low`, `high`, `none`).
    /// Each `build_request` maps it onto the provider's thinking config.
    pub reasoning_effort: Option<&'a str>,
    pub max_tokens: Option<u32>,
    pub output_schema: Option<&'a OpenAPISchema>,
    pub output_type: &'a OutputType,
    pub system_prompt: Option<&'a str>,
    pub user_message: &'a str,
    pub attachments: Option<&'a [S3Object]>,
    pub has_websearch: bool,
    /// Routing hint for the provider's prompt cache. Requests sharing a prompt prefix
    /// must reuse one key to land on the same cache, so it is derived from what pins
    /// the prefix (the step), never from the request. `None` retries a key the
    /// endpoint rejected.
    pub prompt_cache_key: Option<&'a str>,
    /// Ask for a summary of the model's reasoning where the provider streams one.
    /// `false` once the provider refused summaries to these credentials.
    pub reasoning_summary: bool,
}

/// Response from AI provider
pub enum ParsedResponse {
    Text {
        content: Option<String>,
        /// The thinking the model streamed before the answer, when it emitted any.
        reasoning: Option<String>,
        tool_calls: Vec<OpenAIToolCall>,
        events_str: Option<String>,
        annotations: Vec<UrlCitation>,
        used_websearch: bool,
        usage: Option<TokenUsage>,
    },
    Image {
        base64_data: String,
    },
}

impl ParsedResponse {
    /// Drops every U+0000 the model emitted. Postgres stores neither a NUL in `text` nor its
    /// `\u0000` escape in `jsonb`, and this output is written to both: tool job args, the flow
    /// status, agent memory, conversation rows.
    pub fn without_nul(mut self) -> Self {
        fn strip(s: &mut String) {
            if s.contains('\0') {
                s.retain(|c| c != '\0');
            }
        }
        if let ParsedResponse::Text { content, reasoning, tool_calls, annotations, .. } = &mut self
        {
            for text in [content, reasoning].into_iter().flatten() {
                strip(text);
            }
            for annotation in annotations {
                strip(&mut annotation.url);
                annotation.title.as_mut().into_iter().for_each(strip);
            }
            for tool_call in tool_calls {
                strip(&mut tool_call.id);
                strip(&mut tool_call.function.name);
                // Arguments are serialized JSON, where a NUL is spelled as its escape.
                strip(&mut tool_call.function.arguments);
                if let Cow::Owned(stripped) = strip_json_nul(&tool_call.function.arguments) {
                    tool_call.function.arguments = stripped;
                }
            }
        }
        self
    }
}

/// Trait for streaming AI events to a sink (e.g., database persistence).
/// Implemented by the worker's StreamEventProcessor.
#[async_trait]
pub trait StreamEventSink: Send + Sync {
    async fn send(&self, event: StreamingEvent, events_str: &mut String) -> Result<(), Error>;
}

/// Trait for building provider-specific AI requests
#[async_trait]
pub trait QueryBuilder: Send + Sync {
    /// Check if this provider supports tools with the given output type
    fn supports_tools_with_output_type(&self, output_type: &OutputType) -> bool;

    /// Build the request body for the provider
    async fn build_request(
        &self,
        args: &BuildRequestArgs<'_>,
        client: &AuthedClient,
        workspace_id: &str,
    ) -> Result<String, Error>;

    /// Build the request body without usage tracking (for retry on incompatible providers)
    /// Default implementation just calls build_request (most providers don't need this)
    async fn build_request_without_usage(
        &self,
        args: &BuildRequestArgs<'_>,
        client: &AuthedClient,
        workspace_id: &str,
    ) -> Result<String, Error> {
        self.build_request(args, client, workspace_id).await
    }

    /// Whether this provider supports retry without usage tracking
    /// Only OtherQueryBuilder (OpenAI-compatible providers) needs this
    fn supports_retry_without_usage(&self) -> bool {
        false
    }

    /// Whether a rejected request may be retried on the `/chat/completions` surface of
    /// the same resource. Azure serves the Responses API for only a subset of models and
    /// regions, and an Azure model name is a user-chosen deployment name, so such a
    /// resource can only find out by being asked.
    fn supports_chat_completions_fallback(&self, _base_url: &str) -> bool {
        false
    }

    /// Build a provider-specific request from an OpenAI-compatible proxy request.
    fn build_proxy_request(&self, args: &ProxyBuildArgs<'_>) -> Result<ProxyRequest, Error> {
        Err(Error::BadRequest(format!(
            "Proxy request building is not supported for provider {:?}",
            args.credentials.provider
        )))
    }

    /// Parse the image response from the provider
    async fn parse_image_response(
        &self,
        response: reqwest::Response,
    ) -> Result<ParsedResponse, Error>;

    /// Parse streaming response from the provider
    async fn parse_streaming_response(
        &self,
        response: reqwest::Response,
        stream_event_sink: Box<dyn StreamEventSink>,
    ) -> Result<ParsedResponse, Error>;

    /// Get the API endpoint for this provider
    fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String;

    /// Get the authentication headers for this provider
    fn get_auth_headers(
        &self,
        api_key: &str,
        base_url: &str,
        output_type: &OutputType,
    ) -> Vec<(&'static str, String)>;
}