From 27acbbf3d595f16b3e73b81011c4ae4df989c2a7 Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Mon, 11 May 2026 12:01:45 +0200 Subject: [PATCH] refactor: move ai sse plumbing to windmill-ai (#9059) * docs: refine windmill ai refactor plan * refactor: move ai sse plumbing to windmill-ai * refactor: remove ai re-export shims * fix: update ee ai memory ref * chore: update ee-repo-ref to d3bc7fa85195b46b7a38d43c2f806520bf8b5454 This commit updates the EE repository reference after PR #560 was merged in windmill-ee-private. Previous ee-repo-ref: ff35bf7cc198e13884b33654e1d6dbd8a8b314d3 New ee-repo-ref: d3bc7fa85195b46b7a38d43c2f806520bf8b5454 Automated by sync-ee-ref workflow. --------- Co-authored-by: windmill-internal-app[bot] --- backend/Cargo.lock | 2 + backend/ee-repo-ref.txt | 2 +- backend/windmill-ai/Cargo.toml | 2 + backend/windmill-ai/src/lib.rs | 2 + .../src/ai => windmill-ai/src}/sse.rs | 19 +++-- backend/windmill-ai/src/utils.rs | 56 ++++++++++++++ backend/windmill-api/src/ai.rs | 29 +------ .../windmill-worker/src/ai/image_handler.rs | 3 +- backend/windmill-worker/src/ai/mod.rs | 2 - .../src/ai/providers/anthropic.rs | 11 +-- .../src/ai/providers/bedrock.rs | 17 ++-- .../src/ai/providers/google_ai.rs | 18 ++--- .../src/ai/providers/openai.rs | 12 +-- .../src/ai/providers/openrouter.rs | 11 ++- .../windmill-worker/src/ai/providers/other.rs | 10 +-- .../windmill-worker/src/ai/query_builder.rs | 19 ++--- backend/windmill-worker/src/ai/tools.rs | 6 +- backend/windmill-worker/src/ai/types.rs | 2 - backend/windmill-worker/src/ai/utils.rs | 29 +------ backend/windmill-worker/src/ai_executor.rs | 49 +++--------- backend/windmill-worker/src/memory_common.rs | 2 +- backend/windmill-worker/src/memory_oss.rs | 4 +- docs/windmill-ai-refactor-plan.md | 77 +++++++++++++++++-- 23 files changed, 209 insertions(+), 175 deletions(-) rename backend/{windmill-worker/src/ai => windmill-ai/src}/sse.rs (99%) create mode 100644 backend/windmill-ai/src/utils.rs delete mode 100644 backend/windmill-worker/src/ai/types.rs diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 6952682391..842799d5f9 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -16114,12 +16114,14 @@ dependencies = [ "aws-sdk-bedrockruntime", "aws-smithy-types", "base64 0.22.1", + "eventsource-stream", "lazy_static", "reqwest 0.13.1", "serde", "serde_json", "sqlx", "tokio", + "tokio-stream", "tracing", "uuid", "windmill-common", diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 401a8213f7..2e43e667c1 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -c8d100d74b8de6bd26fc973d5edbd8853d54dd8b +d3bc7fa85195b46b7a38d43c2f806520bf8b5454 diff --git a/backend/windmill-ai/Cargo.toml b/backend/windmill-ai/Cargo.toml index 69cc0f181b..3cfea846b9 100644 --- a/backend/windmill-ai/Cargo.toml +++ b/backend/windmill-ai/Cargo.toml @@ -21,6 +21,7 @@ windmill-mcp = { workspace = true, optional = true } async-trait.workspace = true base64.workspace = true +eventsource-stream.workspace = true reqwest.workspace = true serde.workspace = true serde_json.workspace = true @@ -29,6 +30,7 @@ uuid.workspace = true lazy_static.workspace = true tracing.workspace = true tokio.workspace = true +tokio-stream.workspace = true # Bedrock (optional) aws-config = { workspace = true, optional = true } diff --git a/backend/windmill-ai/src/lib.rs b/backend/windmill-ai/src/lib.rs index dfa34e7d43..8b786cb741 100644 --- a/backend/windmill-ai/src/lib.rs +++ b/backend/windmill-ai/src/lib.rs @@ -5,4 +5,6 @@ pub mod ai_google; pub mod ai_providers; pub mod ai_types; pub mod query_builder; +pub mod sse; pub mod types; +pub mod utils; diff --git a/backend/windmill-worker/src/ai/sse.rs b/backend/windmill-ai/src/sse.rs similarity index 99% rename from backend/windmill-worker/src/ai/sse.rs rename to backend/windmill-ai/src/sse.rs index 2cc03b148e..35a4e43acc 100644 --- a/backend/windmill-worker/src/ai/sse.rs +++ b/backend/windmill-ai/src/sse.rs @@ -3,17 +3,15 @@ use std::collections::HashMap; use eventsource_stream::Eventsource; use reqwest::Response; use serde::Deserialize; -use serde_json; use tokio_stream::StreamExt; -use windmill_ai::{ - ai_google::{parse_gemini_sse_event, GeminiUsageMetadata}, - ai_types::{ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall}, -}; use windmill_common::{error::Error, utils::rd_string}; -use crate::ai::{ +use crate::{ + ai_google::{parse_gemini_sse_event, GeminiUsageMetadata}, + ai_types::UrlCitation, + ai_types::{ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall}, query_builder::StreamEventSink, - types::{StreamingEvent, UrlCitation}, + types::StreamingEvent, }; #[derive(Deserialize)] @@ -64,6 +62,7 @@ lazy_static::lazy_static! { .parse::() .unwrap_or(false); } +#[allow(async_fn_in_trait)] pub trait SSEParser { async fn parse_event_data(&mut self, data: &str) -> Result<(), Error>; @@ -459,11 +458,11 @@ impl SSEParser for AnthropicSSEParser { // Gemini SSE Parser // ============================================================================ -/// Accumulates Gemini streaming events and converts them into the worker's -/// internal [`OpenAIToolCall`] / [`StreamingEvent`] representation. +/// Accumulates Gemini streaming events and converts them into the shared +/// [`OpenAIToolCall`] / [`StreamingEvent`] representation. /// /// The actual SSE parsing is delegated to [`parse_gemini_sse_event`] from -/// `windmill_common::ai_google` so the logic can be shared with the API proxy. +/// `windmill_ai::ai_google` so the logic can be shared with the API proxy. pub struct GeminiSSEParser { pub accumulated_content: String, pub accumulated_tool_calls: HashMap, diff --git a/backend/windmill-ai/src/utils.rs b/backend/windmill-ai/src/utils.rs new file mode 100644 index 0000000000..b401282df1 --- /dev/null +++ b/backend/windmill-ai/src/utils.rs @@ -0,0 +1,56 @@ +use crate::{ + ai_providers::AIProvider, + ai_types::{ContentPart, OpenAIContent}, +}; + +lazy_static::lazy_static! { + /// Parse AI_HTTP_HEADERS environment variable into a vector of (header_name, header_value) tuples + /// Format: "header1: value1, header2: value2" + pub static ref AI_HTTP_HEADERS: Vec<(String, String)> = { + std::env::var("AI_HTTP_HEADERS") + .ok() + .map(|headers_str| { + headers_str + .split(',') + .filter_map(|header| { + let parts: Vec<&str> = header.splitn(2, ':').collect(); + if parts.len() == 2 { + let name = parts[0].trim().to_string(); + let value = parts[1].trim().to_string(); + if !name.is_empty() && !value.is_empty() { + Some((name, value)) + } else { + None + } + } else { + None + } + }) + .collect() + }) + .unwrap_or_default() + }; +} + +/// AWS Bedrock do not handle structured output query param, so we use a tool for structured output. Same for every Claude models. +pub fn should_use_structured_output_tool(provider: &AIProvider, model: &str) -> bool { + model.contains("claude") || provider == &AIProvider::AWSBedrock +} + +/// Extract text content from OpenAIContent, joining parts with space if multiple +pub fn extract_text_content(content: &OpenAIContent) -> String { + match content { + OpenAIContent::Text(text) => text.clone(), + OpenAIContent::Parts(parts) => parts + .iter() + .filter_map(|p| { + if let ContentPart::Text { text } = p { + Some(text.as_str()) + } else { + None + } + }) + .collect::>() + .join(""), + } +} diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index a2c4ceb811..8b415997c2 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -15,11 +15,12 @@ use serde::{Deserialize, Serialize}; use serde_json::{json, value::RawValue}; use std::collections::HashMap; use std::time::Duration; -use windmill_audit::{audit_oss::audit_log, ActionKind}; use windmill_ai::ai_cache::current_instance_ai_config_revision; use windmill_ai::ai_providers::{ empty_string_as_none, AIPlatform, AIProvider, ProviderConfig, ProviderModel, }; +use windmill_ai::utils::AI_HTTP_HEADERS; +use windmill_audit::{audit_oss::audit_log, ActionKind}; use windmill_common::db::UserDB; use windmill_common::error::{to_anyhow, Error, Result}; use windmill_common::utils::configure_client; @@ -101,32 +102,6 @@ lazy_static::lazy_static! { pub static ref AI_REQUEST_CACHE: Cache<(String, AIProvider), ExpiringAIRequestConfig> = Cache::new(500); - /// Parse AI_HTTP_HEADERS environment variable into a vector of (header_name, header_value) tuples - /// Format: "header1: value1, header2: value2" - static ref AI_HTTP_HEADERS: Vec<(String, String)> = { - std::env::var("AI_HTTP_HEADERS") - .ok() - .map(|headers_str| { - headers_str - .split(',') - .filter_map(|header| { - let parts: Vec<&str> = header.splitn(2, ':').collect(); - if parts.len() == 2 { - let name = parts[0].trim().to_string(); - let value = parts[1].trim().to_string(); - if !name.is_empty() && !value.is_empty() { - Some((name, value)) - } else { - None - } - } else { - None - } - }) - .collect() - }) - .unwrap_or_default() - }; } pub(crate) fn invalidate_ai_request_cache_for_workspace(workspace_id: &str) { diff --git a/backend/windmill-worker/src/ai/image_handler.rs b/backend/windmill-worker/src/ai/image_handler.rs index 4fc0bcb4db..089d73dc42 100644 --- a/backend/windmill-worker/src/ai/image_handler.rs +++ b/backend/windmill-worker/src/ai/image_handler.rs @@ -1,12 +1,11 @@ use base64::Engine; use futures; use ulid; +use windmill_ai::types::*; use windmill_common::{client::AuthedClient, error::Error}; use windmill_queue::MiniPulledJob; use windmill_types::s3::S3Object; -use crate::ai::types::*; - /// Upload image to S3 and return S3Object pub async fn upload_image_to_s3( base64_image: &str, diff --git a/backend/windmill-worker/src/ai/mod.rs b/backend/windmill-worker/src/ai/mod.rs index 4ad67b9e6e..2423fca372 100644 --- a/backend/windmill-worker/src/ai/mod.rs +++ b/backend/windmill-worker/src/ai/mod.rs @@ -4,7 +4,5 @@ pub mod image_handler; pub mod providers; pub mod query_builder; -pub mod sse; pub mod tools; -pub mod types; pub mod utils; diff --git a/backend/windmill-worker/src/ai/providers/anthropic.rs b/backend/windmill-worker/src/ai/providers/anthropic.rs index 2bbfd83768..f2c310923f 100644 --- a/backend/windmill-worker/src/ai/providers/anthropic.rs +++ b/backend/windmill-worker/src/ai/providers/anthropic.rs @@ -1,16 +1,17 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; -use windmill_ai::{ai_google::parse_data_url, ai_providers::AIProvider}; -use windmill_common::{client::AuthedClient, error::Error}; - -use crate::ai::{ - image_handler::prepare_messages_for_api, +use windmill_ai::{ + ai_google::parse_data_url, + ai_providers::AIProvider, query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink}, sse::{AnthropicSSEParser, SSEParser}, types::*, utils::{extract_text_content, should_use_structured_output_tool}, }; +use windmill_common::{client::AuthedClient, error::Error}; + +use crate::ai::image_handler::prepare_messages_for_api; /// Anthropic API version for standard API const ANTHROPIC_VERSION_STANDARD: &str = "2023-06-01"; diff --git a/backend/windmill-worker/src/ai/providers/bedrock.rs b/backend/windmill-worker/src/ai/providers/bedrock.rs index ce87d21a44..ec499f5fcd 100644 --- a/backend/windmill-worker/src/ai/providers/bedrock.rs +++ b/backend/windmill-worker/src/ai/providers/bedrock.rs @@ -6,25 +6,22 @@ //! - Stream event parsing //! - Helper utilities -use crate::ai::{ - image_handler::prepare_messages_for_api, - query_builder::{ParsedResponse, StreamEventSink}, - types::StreamingEvent, - types::TokenUsage, - types::{OpenAIMessage, ToolDef}, -}; +use crate::ai::image_handler::prepare_messages_for_api; use std::collections::HashMap; +use windmill_ai::{ + query_builder::{ParsedResponse, StreamEventSink}, + types::{OpenAIMessage, StreamingEvent, TokenUsage, ToolDef}, +}; use windmill_common::{client::AuthedClient, error::Error}; -// Re-export from shared module for use by other parts of the worker +// Import shared Bedrock helpers for worker-specific orchestration. use windmill_ai::ai_bedrock::{ bedrock_model_supports_prompt_caching, bedrock_stream_event_is_block_stop, bedrock_stream_event_to_text, bedrock_stream_event_to_tool_delta, bedrock_stream_event_to_tool_start, build_tool_config, create_inference_config, format_bedrock_error, openai_messages_to_bedrock, streaming_tool_calls_to_openai, - StreamingToolCall, + BedrockClient, StreamingToolCall, }; -pub use windmill_ai::ai_bedrock::{check_env_credentials, BedrockClient}; // ============================================================================ // Query Builder (Worker-specific orchestration) diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs index 6098cb22d8..2a00011de7 100644 --- a/backend/windmill-worker/src/ai/providers/google_ai.rs +++ b/backend/windmill-worker/src/ai/providers/google_ai.rs @@ -1,17 +1,17 @@ use async_trait::async_trait; -use windmill_ai::ai_google::{ - openai_messages_to_gemini, openai_tools_to_gemini, GeminiGenerationConfig, GeminiImageContent, - GeminiImageRequest, GeminiImageResponse, GeminiInlineData, GeminiPart, GeminiPredictContent, - GeminiTextRequest, GeminiTool, -}; -use windmill_common::{client::AuthedClient, error::Error}; - -use crate::ai::{ - image_handler::{download_and_encode_s3_image, prepare_messages_for_api}, +use windmill_ai::{ + ai_google::{ + openai_messages_to_gemini, openai_tools_to_gemini, GeminiGenerationConfig, + GeminiImageContent, GeminiImageRequest, GeminiImageResponse, GeminiInlineData, GeminiPart, + GeminiPredictContent, GeminiTextRequest, GeminiTool, + }, query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink}, sse::{GeminiSSEParser, SSEParser}, types::*, }; +use windmill_common::{client::AuthedClient, error::Error}; + +use crate::ai::image_handler::{download_and_encode_s3_image, prepare_messages_for_api}; // ============================================================================ // Query Builder Implementation diff --git a/backend/windmill-worker/src/ai/providers/openai.rs b/backend/windmill-worker/src/ai/providers/openai.rs index 52feb2eca7..691f8fd5e3 100644 --- a/backend/windmill-worker/src/ai/providers/openai.rs +++ b/backend/windmill-worker/src/ai/providers/openai.rs @@ -1,17 +1,17 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; -use windmill_ai::ai_providers::AIProvider; -use windmill_ai::ai_types::OpenAIToolCall; -use windmill_common::{client::AuthedClient, error::Error}; - -use crate::ai::{ - image_handler::{prepare_messages_for_api, s3_object_to_content_part}, +use windmill_ai::{ + ai_providers::AIProvider, + ai_types::OpenAIToolCall, query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink}, sse::{OpenAIResponsesSSEParser, SSEParser}, types::*, utils::extract_text_content, }; +use windmill_common::{client::AuthedClient, error::Error}; + +use crate::ai::image_handler::{prepare_messages_for_api, s3_object_to_content_part}; // Responses API structures #[derive(Deserialize)] diff --git a/backend/windmill-worker/src/ai/providers/openrouter.rs b/backend/windmill-worker/src/ai/providers/openrouter.rs index 68ad1922d3..9bfc47ed8f 100644 --- a/backend/windmill-worker/src/ai/providers/openrouter.rs +++ b/backend/windmill-worker/src/ai/providers/openrouter.rs @@ -1,15 +1,14 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json; -use windmill_ai::ai_providers::AIProvider; -use windmill_common::{client::AuthedClient, error::Error}; - -use crate::ai::{ - image_handler::prepare_messages_for_api, - providers::other::OtherQueryBuilder, +use windmill_ai::{ + ai_providers::AIProvider, query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink}, types::*, }; +use windmill_common::{client::AuthedClient, error::Error}; + +use crate::ai::{image_handler::prepare_messages_for_api, providers::other::OtherQueryBuilder}; // OpenRouter-specific types #[derive(Serialize)] diff --git a/backend/windmill-worker/src/ai/providers/other.rs b/backend/windmill-worker/src/ai/providers/other.rs index f8ee60f287..7982dc739d 100644 --- a/backend/windmill-worker/src/ai/providers/other.rs +++ b/backend/windmill-worker/src/ai/providers/other.rs @@ -1,16 +1,16 @@ use async_trait::async_trait; use serde::Serialize; use serde_json; -use windmill_ai::ai_providers::AIProvider; -use windmill_common::{client::AuthedClient, error::Error}; - -use crate::ai::{ - image_handler::prepare_messages_for_api, +use windmill_ai::{ + ai_providers::AIProvider, query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink}, sse::{OpenAISSEParser, SSEParser}, types::*, utils::should_use_structured_output_tool, }; +use windmill_common::{client::AuthedClient, error::Error}; + +use crate::ai::image_handler::prepare_messages_for_api; #[derive(Serialize, Debug, Clone)] #[serde(rename_all = "lowercase")] diff --git a/backend/windmill-worker/src/ai/query_builder.rs b/backend/windmill-worker/src/ai/query_builder.rs index d74c45c098..606aa8f9db 100644 --- a/backend/windmill-worker/src/ai/query_builder.rs +++ b/backend/windmill-worker/src/ai/query_builder.rs @@ -1,24 +1,19 @@ use async_trait::async_trait; +use windmill_ai::{ + query_builder::{QueryBuilder, StreamEventSink}, + types::*, +}; use windmill_common::{error::Error, worker::Connection}; use windmill_queue::MiniPulledJob; use crate::{ - ai::{ - providers::{ - anthropic::AnthropicQueryBuilder, google_ai::GoogleAIQueryBuilder, - openai::OpenAIQueryBuilder, openrouter::OpenRouterQueryBuilder, - other::OtherQueryBuilder, - }, - types::*, + ai::providers::{ + anthropic::AnthropicQueryBuilder, google_ai::GoogleAIQueryBuilder, + openai::OpenAIQueryBuilder, openrouter::OpenRouterQueryBuilder, other::OtherQueryBuilder, }, job_logger::append_result_stream, }; -// Re-export from windmill_ai -pub use windmill_ai::query_builder::{ - BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink, -}; - /// Factory function to create the appropriate query builder for a provider pub fn create_query_builder(provider: &ProviderWithResource) -> Box { use windmill_ai::ai_providers::AIProvider; diff --git a/backend/windmill-worker/src/ai/tools.rs b/backend/windmill-worker/src/ai/tools.rs index c92fc2a94e..fcc9cdf3c9 100644 --- a/backend/windmill-worker/src/ai/tools.rs +++ b/backend/windmill-worker/src/ai/tools.rs @@ -1,6 +1,4 @@ -use crate::ai::query_builder::{StreamEventProcessor, StreamEventSink}; -use crate::ai::types::McpToolSource; -use crate::ai::types::*; +use crate::ai::query_builder::StreamEventProcessor; use crate::ai::utils::{ add_message_to_conversation, execute_mcp_tool, get_step_name_from_flow, is_completed_input_transform, update_flow_status_module_with_actions, @@ -20,7 +18,7 @@ use mappable_rc::Marc; use serde_json::value::RawValue; use std::{collections::HashMap, sync::Arc}; use uuid::Uuid; -use windmill_ai::ai_types::OpenAIToolCall; +use windmill_ai::{ai_types::OpenAIToolCall, query_builder::StreamEventSink, types::*}; use windmill_common::jobs::JobPayload; #[cfg(feature = "mcp")] diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs deleted file mode 100644 index 57619f0dd9..0000000000 --- a/backend/windmill-worker/src/ai/types.rs +++ /dev/null @@ -1,2 +0,0 @@ -// Re-export all types from windmill_ai::types -pub use windmill_ai::types::*; diff --git a/backend/windmill-worker/src/ai/utils.rs b/backend/windmill-worker/src/ai/utils.rs index f1e8fdd22f..74e75ef0a5 100644 --- a/backend/windmill-worker/src/ai/utils.rs +++ b/backend/windmill-worker/src/ai/utils.rs @@ -1,5 +1,3 @@ -pub use crate::ai::types::McpToolSource; -use crate::ai::types::ToolDef; use anyhow::Context; use serde_json::value::RawValue; use sqlx::types::Json; @@ -8,7 +6,7 @@ use std::{ sync::Arc, }; use uuid::Uuid; -use windmill_ai::ai_providers::AIProvider; +use windmill_ai::types::*; use windmill_common::flows::FlowModuleValue; use windmill_common::{ db::DB, @@ -24,7 +22,7 @@ use windmill_common::{ use windmill_mcp::{McpClient, McpResource, McpTool}; use windmill_queue::{flow_status::get_step_of_flow_status, MiniPulledJob}; -use crate::{ai::types::*, parse_sig_of_lang}; +use crate::parse_sig_of_lang; pub fn parse_raw_script_schema( content: &str, @@ -323,11 +321,6 @@ pub fn get_step_name_from_flow( ) } -/// AWS Bedrock do not handle structured output query param, so we use a tool for structured output. Same for every Claude models. -pub fn should_use_structured_output_tool(provider: &AIProvider, model: &str) -> bool { - model.contains("claude") || provider == &AIProvider::AWSBedrock -} - /// Cleanup MCP clients by gracefully shutting down connections #[cfg(feature = "mcp")] pub async fn cleanup_mcp_clients(mcp_clients: HashMap>) { @@ -713,21 +706,3 @@ pub fn any_tool_needs_previous_result(tools: &[Tool]) -> bool { false }) } - -/// Extract text content from OpenAIContent, joining parts with space if multiple -pub fn extract_text_content(content: &OpenAIContent) -> String { - match content { - OpenAIContent::Text(text) => text.clone(), - OpenAIContent::Parts(parts) => parts - .iter() - .filter_map(|p| { - if let ContentPart::Text { text } = p { - Some(text.as_str()) - } else { - None - } - }) - .collect::>() - .join(""), - } -} diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 03153df0a4..4dfeee3571 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1,12 +1,10 @@ -#[cfg(feature = "bedrock")] -use crate::ai::providers::bedrock::check_env_credentials; use crate::ai::tools::{execute_tool_calls, ToolAbortHandles, ToolExecutionContext}; use crate::ai::utils::{ add_message_to_conversation, any_tool_needs_previous_result, cleanup_mcp_clients, filter_schema_by_input_transforms, find_unique_tool_name, get_flow_context, get_flow_job_runnable_and_raw_flow, get_step_name_from_flow, load_mcp_tools, - parse_raw_script_schema, should_use_structured_output_tool, - update_flow_status_module_with_actions, update_flow_status_module_with_actions_success, + parse_raw_script_schema, update_flow_status_module_with_actions, + update_flow_status_module_with_actions_success, }; use crate::memory_oss::{read_from_memory, write_to_memory}; use crate::worker_flow::{get_previous_job_result, get_transform_context}; @@ -15,12 +13,19 @@ use regex::Regex; use serde_json::value::RawValue; use std::{collections::HashMap, sync::Arc}; use uuid::Uuid; +#[cfg(feature = "bedrock")] +use windmill_ai::ai_bedrock::check_env_credentials; #[cfg(feature = "mcp")] use windmill_mcp::McpClient; #[cfg(not(feature = "mcp"))] use crate::ai::tools::McpClientStub as McpClient; -use windmill_ai::ai_providers::AIProvider; +use windmill_ai::{ + ai_providers::AIProvider, + query_builder::{BuildRequestArgs, ParsedResponse}, + types::*, + utils::{should_use_structured_output_tool, AI_HTTP_HEADERS}, +}; use windmill_common::{ cache, client::AuthedClient, @@ -40,10 +45,7 @@ use windmill_queue::{cancel_single_job, CanceledBy, MiniPulledJob}; use crate::{ ai::{ image_handler::upload_image_to_s3, - query_builder::{ - create_query_builder, BuildRequestArgs, ParsedResponse, StreamEventProcessor, - }, - types::*, + query_builder::{create_query_builder, StreamEventProcessor}, }, common::{build_args_map, resolve_job_timeout, OccupancyMetrics, StreamNotifier}, handle_child::{run_future_with_polling_update_job_poller_graceful, GracefulPollOutcome}, @@ -52,33 +54,6 @@ use crate::{ lazy_static::lazy_static! { static ref TOOL_NAME_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_]+$").unwrap(); - /// Parse AI_HTTP_HEADERS environment variable into a vector of (header_name, header_value) tuples - /// Format: "header1: value1, header2: value2" - static ref AI_HTTP_HEADERS: Vec<(String, String)> = { - std::env::var("AI_HTTP_HEADERS") - .ok() - .map(|headers_str| { - headers_str - .split(',') - .filter_map(|header| { - let parts: Vec<&str> = header.splitn(2, ':').collect(); - if parts.len() == 2 { - let name = parts[0].trim().to_string(); - let value = parts[1].trim().to_string(); - if !name.is_empty() && !value.is_empty() { - Some((name, value)) - } else { - None - } - } else { - None - } - }) - .collect() - }) - .unwrap_or_default() - }; - static ref AI_AGENT_TOOL_SCHEMA: Box = to_raw_value(&serde_json::json!({ "type": "object", "properties": { @@ -791,7 +766,7 @@ pub async fn run_agent( let mut actions = vec![]; let mut content = None; - let mut final_usage: Option = None; + let mut final_usage: Option = None; // Check if this provider supports tools with the current output type let supports_tools = query_builder.supports_tools_with_output_type(output_type); diff --git a/backend/windmill-worker/src/memory_common.rs b/backend/windmill-worker/src/memory_common.rs index 440199702f..e024acea69 100644 --- a/backend/windmill-worker/src/memory_common.rs +++ b/backend/windmill-worker/src/memory_common.rs @@ -1,5 +1,5 @@ -use crate::ai::types::OpenAIMessage; use uuid::Uuid; +use windmill_ai::types::OpenAIMessage; use windmill_common::{db::DB, error::Error}; pub const MAX_MEMORY_SIZE_BYTES: usize = 100_000; // 100KB per memory entry in database diff --git a/backend/windmill-worker/src/memory_oss.rs b/backend/windmill-worker/src/memory_oss.rs index 74771a79e8..dcdab14e2a 100644 --- a/backend/windmill-worker/src/memory_oss.rs +++ b/backend/windmill-worker/src/memory_oss.rs @@ -3,7 +3,9 @@ pub use crate::memory_ee::*; #[cfg(not(all(feature = "private", feature = "enterprise")))] -use {crate::ai::types::OpenAIMessage, crate::memory_common, uuid::Uuid, windmill_common::db::DB}; +use { + crate::memory_common, uuid::Uuid, windmill_ai::types::OpenAIMessage, windmill_common::db::DB, +}; /// Read AI agent memory from storage /// In OSS: always reads from database diff --git a/docs/windmill-ai-refactor-plan.md b/docs/windmill-ai-refactor-plan.md index 81ab5d7b86..78b82ddcc6 100644 --- a/docs/windmill-ai-refactor-plan.md +++ b/docs/windmill-ai-refactor-plan.md @@ -23,6 +23,58 @@ windmill-worker → windmill-ai windmill-common does **NOT** re-export from windmill-ai (would be circular). All consumers update imports. +## Reviewer Note: Keep the Next PR Small + +The first merged PR established the crate boundary; it did not yet remove the duplicated API-vs-worker provider paths. The remaining work should stay split by dependency risk, not by the final desired module layout. + +Do not jump directly from the current state to provider moves, proxy unification, and credential unification in one PR. The riskiest part is the API proxy because it combines request transformation, endpoint selection, auth headers, custom headers, OAuth user injection, Azure URL handling, Anthropic Vertex handling, Bedrock SDK calls, and SSE keepalive behavior. + +Pull the shared plumbing forward before moving provider implementations: +- Move tiny shared utilities first, including `AI_HTTP_HEADERS`, `extract_text_content`, and `should_use_structured_output_tool`. +- Move SSE parsers next, using the existing `StreamEventSink` abstraction, and update callers to import from `windmill_ai` directly. +- Leave provider implementations, image upload/download handling, API proxy changes, and credential unification out of that PR. + +Avoid adding modules whose only purpose is to re-export moved code. Direct imports from `windmill_ai` make ownership and dependency direction clearer at each call site. + +Also do not make `build_proxy_request(raw_body, path)` too narrow. The proxy path needs method, incoming headers, resolved credentials, base URL/platform, organization/user fields, custom headers, and Bedrock/Azure/Vertex-specific context. Introduce a structured `ProxyBuildArgs`/`ProviderCredentials` shape before deleting `AIRequestConfig::prepare_request`, `google.rs`, or `bedrock.rs`. + +## Next Phase PR: Shared Plumbing Only + +Goal: make `windmill-ai` own the provider-independent helper code that later provider moves will need, without changing API proxy behavior or agent request behavior. + +Suggested PR title: `refactor(ai): move shared SSE plumbing into windmill-ai`. + +Scope: +- Add `windmill-ai/src/utils.rs`. +- Move the duplicated `AI_HTTP_HEADERS` parsing into `windmill_ai::utils` with identical parsing behavior. +- Move `extract_text_content` and `should_use_structured_output_tool` from `windmill-worker/src/ai/utils.rs` to `windmill_ai::utils`. +- Move `windmill-worker/src/ai/sse.rs` to `windmill-ai/src/sse.rs`. +- Delete `windmill-worker/src/ai/sse.rs` and update callers to import parser types from `windmill_ai::sse`. +- Update callers of moved utility functions to import from `windmill_ai::utils` directly. +- Add the minimal new `windmill-ai` dependencies required by `sse.rs` (`eventsource-stream`, `tokio-stream`) and avoid adding worker/queue dependencies. + +Out of scope: +- Do not move provider implementations. +- Do not move `image_handler`. +- Do not change `QueryBuilder` method signatures. +- Do not add `build_proxy_request`. +- Do not change API proxy routing, request preparation, credential resolution, audit logging, cache behavior, or Bedrock/Google special cases. +- Do not remove `windmill-api/src/google.rs`, `windmill-api/src/bedrock.rs`, or `AIRequestConfig::prepare_request`. + +Implementation checklist: +1. Add `utils.rs` to `windmill-ai` and export it from `lib.rs`. +2. Move `AI_HTTP_HEADERS` exactly once, then update `windmill-api/src/ai.rs` and `windmill-worker/src/ai_executor.rs` to import it. +3. Move the two provider-independent helper functions into `windmill_ai::utils`; leave worker-specific flow/MCP/conversation utilities in `windmill-worker/src/ai/utils.rs`. +4. Move `sse.rs` into `windmill-ai`, change imports from `crate::ai::{query_builder, types}` to `crate::{query_builder, types}`, and keep behavior unchanged. +5. Remove worker `ai/sse.rs` and update provider imports to use `windmill_ai::sse` directly. +6. Run focused grep checks for duplicate `AI_HTTP_HEADERS`, old local helper definitions, and accidental `windmill_queue`/worker dependencies from `windmill-ai`. +7. Validate with `cargo check -p windmill-ai`, `cargo check -p windmill-worker`, and `cargo check -p windmill-api`. For `bedrock` builds, also check the existing bedrock feature path. + +Review expectations: +- The diff should be mostly moved code and import updates. +- The behavior should be byte-for-byte equivalent where practical. +- Tests are only needed if helper behavior changes. For a pure move, existing backend checks plus manual AI streaming verification are enough. + ## Step-by-Step Plan Each step produces a compiling, working backend. @@ -131,18 +183,27 @@ This is the key unification step. Add a new method to the `QueryBuilder` trait: /// Used by the API chat proxy. Handles format conversion for non-OpenAI providers. fn build_proxy_request( &self, - raw_body: &[u8], - path: &str, + args: &ProxyBuildArgs<'_>, ) -> Result; ``` -Where `ProxyRequest` contains the transformed body, endpoint URL, and auth headers: +Where `ProxyBuildArgs` carries the API proxy context that provider implementations need: +```rust +pub struct ProxyBuildArgs<'a> { + pub method: http::Method, + pub path: &'a str, + pub headers: &'a http::HeaderMap, + pub body: &'a [u8], + pub credentials: &'a ProviderCredentials, +} +``` + +And `ProxyRequest` contains the transformed request: ```rust pub struct ProxyRequest { pub url: String, pub body: Vec, - pub auth_headers: Vec<(String, String)>, - pub is_sse: bool, + pub headers: Vec<(String, String)>, } ``` @@ -153,9 +214,9 @@ pub struct ProxyRequest { - **Bedrock**: Convert OpenAI format → Bedrock SDK calls. Replaces `windmill-api/src/bedrock.rs`. **Refactor API proxy** (`windmill-api/src/ai.rs`): -1. Parse provider from headers, resolve credentials → `ProviderWithResource` +1. Parse provider from headers, resolve credentials → `ProviderCredentials` 2. Create `QueryBuilder` via `create_query_builder` -3. Call `query_builder.build_proxy_request(body, path)` → `ProxyRequest` +3. Call `query_builder.build_proxy_request(&proxy_args)` → `ProxyRequest` 4. Send the request, return response with SSE keepalive injection **Remove** from windmill-api: @@ -166,7 +227,7 @@ pub struct ProxyRequest { - `supports_native_fim`, `transform_fim_to_chat_completions` — moved to windmill-ai **Keep** in API: -- `AIRequestConfig::new` credential resolution (or refactor to produce `ProviderWithResource`) +- `AIRequestConfig::new` credential resolution until it is refactored to produce `ProviderCredentials` - HTTP routes, audit logging, request caching - `inject_keepalives`, `is_sse_response` helpers - `AIConfig`, `ExpiringAIRequestConfig` caching types