Building a Custom LLM Provider for FlowSquad in 45 Minutes

What Is FlowSquad's LLM Provider Pattern?

FlowSquad is built around a pluggable LLM provider pattern — a Python abstract base class that all LLM integrations implement. Every feature in the platform — Epic/Story/Task decomposition, HLD/LLD design generation, test case generation, code analysis, and AskProduct — talks to LLMs through this single interface. Swapping providers is a config change, not a code refactor.

The design pays dividends in three ways. First, it enables true Multi-LLM support: teams can route different task types to different models (e.g., Claude for code review, GPT-4 for story writing, Ollama for local dev). Second, the LLM cost tracker hooks into the interface at the base class level, so every provider's token consumption is automatically tracked — you never write cost-tracking code in a provider. Third, the pattern enables the admin UI's provider health dashboard: every provider exposes a health_check() method that the platform polls every 60 seconds.

What Interface Does a Provider Class Need to Implement?

The base class lives at flowsquad/providers/base.py. It is a standard Python ABC with 4 abstract methods and several optional hooks. The 4 required methods map to the 4 LLM interaction types that FlowSquad uses across its entire feature set.

# flowsquad/providers/base.py — the abstract base every provider implements
from abc import ABC, abstractmethod
from typing import AsyncIterator, Optional
from dataclasses import dataclass


@dataclass
class CompletionResponse:
    content: str
    model: str
    prompt_tokens: int
    completion_tokens: int
    latency_ms: float


class BaseLLMProvider(ABC):
    """
    Abstract base class for all FlowSquad LLM providers.
    Subclass this and implement the 4 abstract methods.
    Register your subclass in providers.yaml.
    """

    provider_name: str = "base"   # override in subclass

    @abstractmethod
    async def complete(
        self,
        prompt: str,
        system: Optional[str] = None,
        temperature: float = 0.3,
        max_tokens: int = 4096,
        **kwargs
    ) -> CompletionResponse:
        """Single-turn completion. Used by story decomposition, design gen."""

    @abstractmethod
    async def embed(self, text: str) -> list[float]:
        """Return embedding vector. Used by AskProduct / VectorDB indexing."""

    @abstractmethod
    async def stream(
        self,
        prompt: str,
        system: Optional[str] = None,
        **kwargs
    ) -> AsyncIterator[str]:
        """Token-streaming completion. Used by the live generation UI."""

    @abstractmethod
    async def health_check(self) -> dict:
        """
        Return {"status": "ok"|"degraded"|"down", "latency_ms": float}.
        Called every 60s by the provider health dashboard.
        """

How Do You Create the Azure OpenAI Provider Class?

Create a new file at flowsquad/providers/azure_openai.py. The Azure OpenAI SDK is API-compatible with the OpenAI SDK but uses a different base URL, API version, and deployment name instead of model name. These 3 differences are the entire substance of the custom provider.

# flowsquad/providers/azure_openai.py
import time
import httpx
from typing import AsyncIterator, Optional
from openai import AsyncAzureOpenAI
from flowsquad.providers.base import BaseLLMProvider, CompletionResponse
from flowsquad.config import settings
from flowsquad.telemetry import cost_tracker   # automatic cost tracking hook


class AzureOpenAIProvider(BaseLLMProvider):
    """
    FlowSquad LLM provider for Azure OpenAI Service.

    Reads configuration from environment variables (injected via K8s Secret):
      AZURE_OPENAI_ENDPOINT   — e.g. https://my-resource.openai.azure.com/
      AZURE_OPENAI_API_KEY    — your Azure OpenAI API key
      AZURE_OPENAI_DEPLOYMENT — deployment name (e.g. gpt-4-prod)
      AZURE_OPENAI_API_VERSION — e.g. 2024-02-01
      AZURE_EMBED_DEPLOYMENT  — embedding deployment name (e.g. ada-002-prod)
    """

    provider_name = "azure_openai"

    def __init__(self):
        self._client = AsyncAzureOpenAI(
            azure_endpoint=settings.AZURE_OPENAI_ENDPOINT,
            api_key=settings.AZURE_OPENAI_API_KEY,
            api_version=settings.AZURE_OPENAI_API_VERSION,
        )
        self._deployment = settings.AZURE_OPENAI_DEPLOYMENT
        self._embed_deployment = settings.AZURE_EMBED_DEPLOYMENT

    async def complete(
        self,
        prompt: str,
        system: Optional[str] = None,
        temperature: float = 0.3,
        max_tokens: int = 4096,
        **kwargs
    ) -> CompletionResponse:
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        t0 = time.monotonic()
        response = await self._client.chat.completions.create(
            model=self._deployment,   # Azure uses deployment name here
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        latency_ms = (time.monotonic() - t0) * 1000

        result = CompletionResponse(
            content=response.choices[0].message.content,
            model=self._deployment,
            prompt_tokens=response.usage.prompt_tokens,
            completion_tokens=response.usage.completion_tokens,
            latency_ms=latency_ms,
        )
        # The cost_tracker hook fires automatically — no manual call needed
        await cost_tracker.record(self.provider_name, result)
        return result

    async def embed(self, text: str) -> list[float]:
        response = await self._client.embeddings.create(
            model=self._embed_deployment,
            input=text
        )
        return response.data[0].embedding

    async def stream(
        self,
        prompt: str,
        system: Optional[str] = None,
        **kwargs
    ) -> AsyncIterator[str]:
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        async with await self._client.chat.completions.create(
            model=self._deployment,
            messages=messages,
            stream=True,
            **kwargs
        ) as stream:
            async for chunk in stream:
                delta = chunk.choices[0].delta.content
                if delta:
                    yield delta

    async def health_check(self) -> dict:
        try:
            t0 = time.monotonic()
            await self._client.models.list()
            latency_ms = (time.monotonic() - t0) * 1000
            return {"status": "ok", "latency_ms": latency_ms}
        except Exception as e:
            return {"status": "down", "error": str(e)}

How Do You Register the Provider with FlowSquad?

Provider registration uses providers.yaml — a config file that maps provider names to Python class paths. The pluggable Python class pattern means FlowSquad discovers providers at startup through Python's importlib; no changes to core code are required.

# config/providers.yaml — add these 3 lines
providers:
  openai:
    class: flowsquad.providers.openai_provider.OpenAIProvider
  anthropic:
    class: flowsquad.providers.anthropic_provider.AnthropicProvider
  ollama:
    class: flowsquad.providers.ollama_provider.OllamaProvider

  # ↓ Add your new provider here
  azure_openai:
    class: flowsquad.providers.azure_openai.AzureOpenAIProvider
    default_model: gpt-4-prod   # shown in the admin UI
    embed_model: ada-002-prod

After updating the config, apply it without restarting the full application using the CLI:

# Reload providers at runtime (no downtime)
flowsquad reload-providers --config config/providers.yaml

# Verify the new provider appears in the admin panel
flowsquad providers list
# Output:
#   openai        ✔ healthy (142ms)
#   anthropic     ✔ healthy (198ms)
#   ollama        ✔ healthy (23ms)
#   azure_openai  ✔ healthy (167ms)  ← new

Once registered, the provider appears in the Multi-LLM support dropdown in every project's settings page. Any user with the product admin RBAC role can set it as the default for their project.

How Do You Test Your New Provider?

FlowSquad ships a provider test harness at tests/providers/. Run the suite against the new provider to validate all 4 interface methods before deploying. The full test suite runs in under 12 seconds against a live Azure endpoint.

# tests/providers/test_azure_openai.py
import pytest
import asyncio
from flowsquad.providers.azure_openai import AzureOpenAIProvider


@pytest.fixture
def provider():
    return AzureOpenAIProvider()


@pytest.mark.asyncio
async def test_complete_returns_content(provider):
    response = await provider.complete(
        prompt="Write one sentence about software quality.",
        system="You are a concise technical writer.",
        max_tokens=100
    )
    assert response.content, "Expected non-empty response"
    assert response.prompt_tokens > 0
    assert response.latency_ms > 0


@pytest.mark.asyncio
async def test_embed_returns_1536_dims(provider):
    vector = await provider.embed("FlowSquad SDLC platform")
    assert len(vector) == 1536     # ada-002 dimension
    assert all(isinstance(x, float) for x in vector)


@pytest.mark.asyncio
async def test_stream_yields_chunks(provider):
    chunks = []
    async for chunk in provider.stream("Count to three."):
        chunks.append(chunk)
    assert len(chunks) > 3      # should stream multiple tokens


@pytest.mark.asyncio
async def test_health_check_ok(provider):
    result = await provider.health_check()
    assert result["status"] == "ok"
    assert result["latency_ms"] < 2000   # SLA: health check under 2s


# Run: pytest tests/providers/test_azure_openai.py -v

For integration testing against FlowSquad features end-to-end, set the project's LLM provider to azure_openai in the admin panel, then run the Epic/Story/Task decomposition flow on a test project. Verify that the LLM cost tracker records tokens against the azure_openai provider name — this confirms the cost_tracker.record() hook fired correctly in complete().