windmill/docs/windmill-ai-refactor-plan.md

Refactor Plan: windmill-ai Crate

Context

AI provider logic is currently split across three crates with duplicate code:

• windmill-common — base types (ai_types, ai_providers, ai_google, ai_bedrock, ai_cache)
• windmill-api — chat proxy (ai.rs, google.rs, bedrock.rs) with its own request building for Google/Bedrock, plus AIRequestConfig::prepare_request for auth/URL handling
• windmill-worker — agent execution (ai/ module) with QueryBuilder trait, SSE parsers, provider implementations

The goal: a single windmill-ai crate with all AI provider logic. Both the API proxy and worker agent use QueryBuilder for every provider — no more duplicate logic.

Dependency Direction

windmill-ai  →  windmill-common (for DB, Error, AgentAction, AuthedClient, etc.)
             →  windmill-types (for S3Object)
             →  windmill-parser (for Typ, used in OpenAPISchema)

windmill-api    →  windmill-ai
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.

---

Step 1: Create windmill-ai crate, move base types from windmill-common ✅

Create backend/windmill-ai/Cargo.toml and backend/windmill-ai/src/lib.rs.

Move from windmill-common/src/ to windmill-ai/src/:

• ai_types.rs — OpenAI-compatible message types
• ai_providers.rs — AIProvider enum, AIPlatform, base URLs, ProviderConfig
• ai_google.rs — Gemini types and OpenAI↔Gemini conversion
• ai_bedrock.rs — Bedrock SDK wrapper (feature-gated on bedrock)
• ai_cache.rs — instance AI config revision tracking

Update all imports (windmill_common::ai_* → windmill_ai::ai_*).

---

Step 2: Move worker AI types to windmill-ai ✅

Move from windmill-worker/src/ai/types.rs to windmill-ai/src/types.rs:

• ProviderWithResource, ProviderResource — credential types
• TokenUsage — token usage tracking
• OutputType, SchemaType, AdditionalProperties — output configuration
• OpenAPISchema — tool parameter schema (depends on windmill-parser::Typ)
• Tool, Message, ResponseFormat, JsonSchemaFormat — agent types
• StreamingEvent — SSE event enum
• AIAgentArgs, AIAgentArgsRaw, AIAgentResult — agent job args
• Memory — agent memory enum
• S3ObjectWithType — S3 image type
• McpToolSource stub (with same #[cfg(feature = "mcp")] pattern)

Worker ai/types.rs becomes a re-export: pub use windmill_ai::types::*.

---

Step 3: Move QueryBuilder trait, ParsedResponse, and StreamEventSink abstraction to windmill-ai ✅

Move from windmill-worker/src/ai/query_builder.rs to windmill-ai/src/query_builder.rs:

• BuildRequestArgs struct
• ParsedResponse enum
• QueryBuilder trait (with all existing methods)

New StreamEventSink trait in windmill-ai:

#[async_trait]
pub trait StreamEventSink: Send + Sync {
    async fn send(&self, event: StreamingEvent, events_str: &mut String) -> Result<(), Error>;
}

StreamEventSink abstracts the worker's StreamEventProcessor so windmill-ai doesn't depend on windmill-queue or the worker's job logger. The worker's StreamEventProcessor implements StreamEventSink. All provider parse_streaming_response methods and SSE parsers accept Box<dyn StreamEventSink>.

---

Step 4: Move SSE parsers to windmill-ai

Move from windmill-worker/src/ai/sse.rs to windmill-ai/src/sse.rs:

• SSEParser trait
• OpenAISSEParser, AnthropicSSEParser, GeminiSSEParser, OpenAIResponsesSSEParser
• All associated types (delta types, usage types, etc.)

---

Step 5: Move provider implementations to windmill-ai ✅

Move from windmill-worker/src/ai/providers/ to windmill-ai/src/providers/:

• anthropic.rs — AnthropicQueryBuilder
• openai.rs — OpenAIQueryBuilder
• google_ai.rs — GoogleAIQueryBuilder
• bedrock.rs — BedrockQueryBuilder (feature-gated)
• other.rs — OtherQueryBuilder (Mistral, DeepSeek, Groq, TogetherAI, CustomAI)
• openrouter.rs — OpenRouterQueryBuilder
• mod.rs with create_query_builder factory

Move utility functions providers depend on:

• should_use_structured_output_tool (from utils.rs)
• extract_text_content (from utils.rs)

---

Step 6: Move image_handler to windmill-ai ✅

Move from windmill-worker/src/ai/image_handler.rs to windmill-ai/src/image_handler.rs:

• download_and_encode_s3_image — no signature change needed
• prepare_messages_for_api — no signature change needed
• upload_image_to_s3 — refactor: (base64_image, workspace_id, job_id, client) instead of (base64_image, &MiniPulledJob, client) to remove windmill-queue dependency

---

Step 7: Move shared utilities to windmill-ai

Move AI_HTTP_HEADERS lazy_static (currently duplicated in windmill-api/src/ai.rs and windmill-worker/src/ai_executor.rs) to windmill_ai::utils. Both consumers import from windmill-ai.

---

Step 8: Add proxy support to QueryBuilder — API uses QueryBuilder for all providers

This is the key unification step. Add a new method to the QueryBuilder trait:

/// Build a request from a raw OpenAI-format proxy request.
/// Used by the API chat proxy. Handles format conversion for non-OpenAI providers.
fn build_proxy_request(
    &self,
    args: &ProxyBuildArgs<'_>,
) -> Result<ProxyRequest, Error>;

Where ProxyBuildArgs carries the API proxy context that provider implementations need:

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:

pub struct ProxyRequest {
    pub url: String,
    pub body: Vec<u8>,
    pub headers: Vec<(String, String)>,
}

Provider implementations:

• OpenAI-compatible (OpenAI, Mistral, DeepSeek, Groq, TogetherAI, CustomAI, OpenRouter): Minimal transformation — pass body through, build URL and auth headers.
• Anthropic: Handle standard vs Vertex AI. For Vertex: transform body (extract model, add anthropic_version). For standard: pass through with appropriate headers.
• Google AI: Convert OpenAI format → Gemini format (using existing ai_google functions). Replaces windmill-api/src/google.rs.
• 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 → ProviderCredentials
2. Create QueryBuilder via create_query_builder
3. Call query_builder.build_proxy_request(&proxy_args) → ProxyRequest
4. Send the request, return response with SSE keepalive injection

Remove from windmill-api:

• AIRequestConfig::prepare_request — replaced by QueryBuilder::build_proxy_request
• google.rs — replaced by GoogleAIQueryBuilder::build_proxy_request
• bedrock.rs — replaced by BedrockQueryBuilder::build_proxy_request
• transform_anthropic_for_vertex — moved to AnthropicQueryBuilder
• supports_native_fim, transform_fim_to_chat_completions — moved to windmill-ai

Keep in API:

• 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

---

Step 9: Unify credential resolution

Merge AIRequestConfig (API-side) and ProviderWithResource (worker-side) into a single type in windmill-ai.

Both currently carry: api_key, base_url, region, platform, custom_headers, AWS credentials. The API's AIRequestConfig::new resolves credentials from DB (workspace/instance settings). The worker's ProviderWithResource gets credentials from the flow module definition.

Create windmill_ai::ProviderCredentials that both can produce:

pub struct ProviderCredentials {
    pub provider: AIProvider,
    pub api_key: Option<String>,
    pub base_url: String,
    pub platform: AIPlatform,
    pub region: Option<String>,
    pub aws_access_key_id: Option<String>,
    pub aws_secret_access_key: Option<String>,
    pub aws_session_token: Option<String>,
    pub custom_headers: HashMap<String, String>,
    pub organization_id: Option<String>,
    pub enable_1m_context: bool,
}

The create_query_builder factory takes &ProviderCredentials instead of &ProviderWithResource.

---

Final Crate Structure

windmill-ai/src/
├── lib.rs              # re-exports, AI_HTTP_HEADERS
├── ai_types.rs         # OpenAI-compatible message types
├── ai_providers.rs     # AIProvider enum, base URLs, config
├── ai_google.rs        # Gemini types and conversions
├── ai_bedrock.rs       # Bedrock SDK wrapper (feature: bedrock)
├── ai_cache.rs         # Instance AI config revision
├── types.rs            # ProviderCredentials, TokenUsage, Tool, OpenAPISchema, etc.
├── query_builder.rs    # QueryBuilder trait, BuildRequestArgs, ParsedResponse, ProxyRequest, StreamEventSink
├── sse.rs              # SSE parsers (OpenAI, Anthropic, Gemini, Responses)
├── image_handler.rs    # S3 image upload/download
├── utils.rs            # extract_text_content, should_use_structured_output_tool
└── providers/
    ├── mod.rs           # create_query_builder factory
    ├── anthropic.rs     # build_request + build_proxy_request
    ├── openai.rs        # build_request + build_proxy_request
    ├── google_ai.rs     # build_request + build_proxy_request
    ├── bedrock.rs       # build_request + build_proxy_request (feature: bedrock)
    ├── other.rs         # build_request + build_proxy_request
    └── openrouter.rs    # build_request + build_proxy_request

windmill-worker keeps: ai_executor.rs, ai/tools.rs, ai/utils.rs (flow/conversation/MCP logic), StreamEventProcessor (impl of StreamEventSink).

windmill-api keeps: HTTP routes (ai.rs proxy endpoints), audit logging, caching, credential resolution from DB. google.rs and bedrock.rs deleted.