Support pluggable LLM backends per analyzer (Bedrock, SageMaker, custom endpoints)

[rbpotter-aws]

Description of Need

As a BADGERS user, I want to configure which LLM backend and model is used for each analyzer task, including non-Bedrock options like SageMaker endpoints or self-hosted container models, So that I can use smaller, fit-for-purpose models that match my cost, latency, and capability requirements rather than being constrained to Bedrock Claude models for every task.

Context

Currently the system is tightly coupled to Amazon Bedrock as the sole LLM provider. Customers are running specialized or fine-tuned models on SageMaker endpoints, ECS/EKS containers, or want to use different Bedrock models (e.g. Haiku, Mistral, Llama) depending on the analysis task. A document classification task doesn't need the same model horsepower as a complex infographic analysis.

Example Use Cases

• Use a fine-tuned SageMaker model for domain-specific document classification
• Use a lightweight Bedrock model (Haiku, Llama) for simple extraction tasks
• Use a self-hosted model behind an API endpoint for air-gapped or compliance-restricted environments
• Keep Sonnet/Opus for complex multi-step analysis where quality matters most

Acceptance Criteria

• Introduce an LLM provider abstraction that supports multiple backends behind a common interface
• Support at minimum: Bedrock (existing), SageMaker endpoints, and generic HTTP/REST endpoints
• Users can configure the provider and model per analyzer in the analyzer config (e.g. analyzer_config.json or equivalent)
• Each provider handles its own auth, request formatting, and response parsing
• A sensible default (current Bedrock behavior) is used when no override is provided
• Existing deployments continue to work with zero configuration changes
• Response format normalization so downstream processing is provider-agnostic

Technical Considerations

• The bedrock_client.py in the foundation layer is the current integration point that would need to be abstracted
• SageMaker endpoints use invoke_endpoint with potentially different payload formats
• Custom endpoints will need configurable auth (IAM, API key, none) and payload templates
• Prompt formatting may differ across models (chat vs. completion, system prompt support, etc.)

Out of Scope (future consideration)

• Automatic model routing/selection based on task complexity
• Model performance benchmarking framework
• Fallback chains (try model A, fall back to model B)

[rbpotter-aws]

Design Proposal: Pluggable LLM Backends for BADGERS Analyzers

Issue: [#11 — Support pluggable LLM backends per analyzer](#11)
Status: Proposal
Date: 2026-02-23
Author: Randall Potter rbpotter-aws

---

Summary

This proposal outlines an approach to implementing pluggable LLM backends in BADGERS that preserves the system's core architectural advantage — prompt-orchestrated extraction — while enabling cost optimization, latency reduction, and deployment flexibility through backend-agnostic model routing.

The key design principle: the inference backend is a deployment concern, not an extraction logic concern. The classifier, prompt templates, and structured output schemas remain the authoritative extraction pipeline regardless of which model serves the inference.

---

Architectural Context

BADGERS currently operates on a prompt-orchestrated pattern:

Document → Classifier (identifies document type)
         → Prompt Template Selection (task-specific extraction logic)
         → Foundation Model Inference (Bedrock API)
         → Structured Output

The extraction intelligence lives in the prompt layer and classifier, not in model weights. This is a deliberate architectural choice with specific trade-offs worth understanding before introducing pluggable backends.

Why This Matters for Backend Abstraction

When we introduce pluggable backends, the abstraction boundary needs to sit between the prompt layer and the inference call. The analyzer's prompt template produces a fully-formed request; the backend provider is responsible only for executing that request against a model and returning a normalized response. This keeps extraction logic decoupled from infrastructure decisions.

Analyzer
  ├── Classifier          (document routing — unchanged)
  ├── Prompt Template     (extraction logic — unchanged)
  └── InferenceProvider   (new abstraction)
        ├── BedrockProvider        (existing behavior)
        ├── SageMakerProvider      (new)
        └── HTTPEndpointProvider   (new)

---

Proposed Provider Interface

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional

@dataclass
class InferenceRequest:
    """Backend-agnostic inference request."""
    messages: list[dict]            # Prompt in messages format
    system_prompt: Optional[str]    # System prompt if supported
    max_tokens: int
    temperature: float
    images: list[bytes]             # Vision inputs (base64-encoded)
    output_schema: Optional[dict]   # Structured output schema if supported

@dataclass
class InferenceResponse:
    """Normalized response from any backend."""
    content: str                    # Extracted text/JSON response
    input_tokens: int               # For cost tracking
    output_tokens: int
    model_id: str                   # Actual model that served the request
    latency_ms: float
    provider: str                   # "bedrock", "sagemaker", "http"

class InferenceProvider(ABC):
    """Base class for all LLM backend providers."""

    @abstractmethod
    async def invoke(self, request: InferenceRequest) -> InferenceResponse:
        """Execute an inference request and return a normalized response."""
        ...

    @abstractmethod
    def supports_vision(self) -> bool:
        """Whether this provider/model supports image inputs."""
        ...

    @abstractmethod
    def supports_structured_output(self) -> bool:
        """Whether this provider/model supports schema-constrained output."""
        ...

This interface intentionally does not include prompt formatting logic. Each provider is responsible for translating InferenceRequest into its backend-specific payload format (Bedrock Converse API, SageMaker invoke_endpoint, raw HTTP POST, etc.), but the content of the prompt is determined upstream by the analyzer.

---

Per-Analyzer Configuration

Extend the existing analyzer configuration to include an optional provider override:

{
  "analyzer_name": "table_extraction",
  "prompt_template": "table_extraction_v3",
  "inference": {
    "provider": "bedrock",
    "model_id": "us.amazon.nova-lite-v1:0",
    "temperature": 0.1,
    "max_tokens": 4096
  }
}

{
  "analyzer_name": "document_classifier",
  "prompt_template": "classifier_v2",
  "inference": {
    "provider": "sagemaker",
    "endpoint_name": "badgers-classifier-endpoint",
    "temperature": 0.0,
    "max_tokens": 512
  }
}

{
  "analyzer_name": "clinical_form_extraction",
  "prompt_template": "clinical_extraction_v1",
  "inference": {
    "provider": "http",
    "endpoint_url": "https://internal-llm.corp.example.com/v1/chat/completions",
    "auth_type": "api_key",
    "auth_secret_arn": "arn:aws:secretsmanager:us-east-1:...",
    "temperature": 0.1,
    "max_tokens": 4096
  }
}

When inference is omitted, the system falls back to the current default Bedrock behavior. This satisfies the acceptance criteria that existing deployments require zero configuration changes.

---

Technical Considerations by Provider

Bedrock (existing, refactored)

Refactor the current bedrock_client.py into a BedrockProvider that implements the interface. This is largely a structural change — the existing Converse API integration maps cleanly to InferenceRequest.

Bedrock's advantage here is transparent scaling, no infrastructure management, and access to the full model catalog (Claude, Nova, Llama, Mistral) through a single API surface. For most analyzers, particularly those processing moderate document volumes or requiring high-capability models for complex extraction, Bedrock remains the recommended default.

SageMaker Endpoints

SageMaker endpoints use invoke_endpoint with model-specific payload formats. The SageMakerProvider will need to handle payload serialization per model framework (HuggingFace TGI, vLLM, Ollama containers, etc.). This is the primary integration point for fine-tuned or custom models.

Considerations:

• Payload format variability. Different model serving frameworks expect different request schemas. The provider should support configurable payload templates or a small set of known framework presets (tgi, vllm, ollama).
• Vision support. Not all SageMaker-hosted models support multimodal inputs. The provider must validate supports_vision() and surface clear errors when an analyzer requiring vision is routed to a text-only endpoint.
• Concurrency. SageMaker endpoint scaling policies directly affect throughput. Analyzers routed to SageMaker should be aware of endpoint capacity limits.

Generic HTTP Endpoints

Covers self-hosted models behind load balancers, ECS/EKS deployments, or third-party APIs. The provider supports configurable authentication (IAM SigV4, API key header, bearer token, none) and expects an OpenAI-compatible chat completions interface by default, with configurable request/response mapping for non-standard APIs.

This provider enables air-gapped and compliance-restricted deployments where inference must remain within specific network boundaries.

---

Cost and Operational Trade-off Guidance

A pluggable backend system enables per-task cost optimization, but the decision of which backend to use for a given analyzer should be informed by workload characteristics. The following guidance helps users make appropriate choices:

When Bedrock is the right choice

• The analyzer processes moderate document volume where per-token API pricing is cheaper than dedicated infrastructure.
• The analyzer requires high-capability models (e.g., Sonnet, Nova Pro) for complex multi-element extraction.
• The extraction schema is evolving and rapid prompt iteration is more important than marginal inference cost savings.
• The team does not have dedicated ML infrastructure expertise.

When SageMaker / self-hosted is the right choice

• A fine-tuned model exists that demonstrably outperforms prompted general models for a specific document type.
• Document volume for a specific analyzer is high enough that fixed infrastructure cost is lower than cumulative per-token cost.
• Data residency or compliance requirements mandate that inference occurs within a specific VPC or network boundary.
• Latency requirements are strict and API round-trip overhead is unacceptable.

When the decision isn't clear

For document types where a fine-tuned model is being considered, a useful validation pattern is to use the existing prompt-orchestrated pipeline to generate labeled extraction data, then fine-tune a smaller model on those outputs and compare accuracy on a held-out test set. This treats the prompt pipeline as a data generation source for model distillation — the prompt system produces the ground truth, and the fine-tuned model compresses it for efficient serving. If the fine-tuned model meets accuracy thresholds, the analyzer config is updated to route to the new endpoint. If accuracy drifts over time, the prompt pipeline regenerates training data.

This approach preserves the classifier and prompt templates as the system of record while allowing selective optimization through fine-tuning where it is empirically justified.

---

Implementation Plan

Phase 1: Provider Abstraction (non-breaking)

• Define InferenceProvider interface and InferenceRequest/InferenceResponse data classes.
• Refactor bedrock_client.py into BedrockProvider implementing the interface.
• Introduce a ProviderFactory that reads analyzer config and instantiates the appropriate provider.
• All existing analyzers continue to use BedrockProvider by default. No behavior change.

Phase 2: SageMaker Provider

• Implement SageMakerProvider with support for TGI and vLLM payload formats.
• Add endpoint health checking and capacity-aware routing.
• Document SageMaker endpoint deployment patterns for common model serving frameworks.

Phase 3: HTTP Endpoint Provider

• Implement HTTPEndpointProvider with configurable auth and payload mapping.
• Support OpenAI-compatible chat completions format as the default.
• Add configurable request/response transformation for non-standard APIs.

Phase 4: Observability and Cost Tracking

• Extend the existing inference profiles (cost tracking) stack to capture per-provider, per-analyzer metrics.
• Surface provider latency, token usage, and error rates through the existing monitoring infrastructure.

---

Items Intentionally Deferred

Per the original issue, the following are out of scope for this proposal but are natural extensions of the pluggable backend architecture:

• Automatic model routing based on task complexity. The provider abstraction makes this possible in the future — a routing layer could examine document characteristics and select a backend dynamically. However, static per-analyzer configuration is the right starting point.
• Fallback chains. (Try model A, fall back to model B.) The provider interface supports this as a composition pattern (a FallbackProvider wrapping multiple providers), but it introduces retry semantics and timeout configuration that warrant separate design work.
• Model performance benchmarking framework. The normalized InferenceResponse with latency and token counts provides the raw data for benchmarking. A structured evaluation framework comparing provider/model combinations per analyzer is a valuable follow-on.

---

Discussion

This proposal is open for community feedback. Key questions to consider:

1. Payload format handling for SageMaker. Should we support a small set of known framework presets (tgi, vllm, ollama) or a fully configurable payload template system? Presets are simpler; templates are more flexible.
2. Vision input normalization. Different backends handle image inputs differently (base64 in messages, separate multipart fields, pre-signed URLs). What level of normalization should the provider interface enforce?
3. Configuration validation. Should the system validate at deployment time that a configured provider/model supports the capabilities required by the analyzer (vision, structured output, context window size)?

Feedback welcome via comments on [#11](#11).

[hshahrokni2]

The challenge with pluggable LLM backends is that the abstraction boundary tends to sit in the wrong place. Most implementations put it at the HTTP client layer, which means you still end up with provider-specific logic leaking into prompt formatting, response parsing, and error handling. A cleaner split is to define the interface at the task semantics level — document classification, extraction, multi-step analysis — so that each task type encapsulates the formatting and parsing needed for whatever backend it's talking to. That also makes it natural to assign different backends per task type, which maps exactly to your use case of using lighter models for simple extraction vs. heavier ones for complex analysis.

At Komilion we've built exactly this kind of routing layer — we route requests across 400+ models based on task complexity, and in our benchmarks the "simple classification" tier costs about 60x less per request than sending everything to a frontier model, with measurably equivalent accuracy on those tasks. The savings stack up fast once you have high-volume extraction pipelines.

If it's useful context, the routing API is OpenAI-compatible — happy to share config examples showing how we set up per-task model assignments.