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.

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,
    raw_body: &[u8],
    path: &str,
) -> Result<ProxyRequest, Error>;

Where ProxyRequest contains the transformed body, endpoint URL, and auth headers:

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

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 → ProviderWithResource
2. Create QueryBuilder via create_query_builder
3. Call query_builder.build_proxy_request(body, path) → 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 (or refactor to produce ProviderWithResource)
• 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.