feat(backends): pluggable model backend abstraction layer with router, fallback chain, and yaml config (Issue #128)

[DZDasherKTB]

Overview

This PR implements a production-grade pluggable backend abstraction layer for sugar-ai, making the entire inference pipeline model-agnostic. Switching from Qwen2 to Groq to Gemini to a local Ollama instance is now a one-line change in sugar_ai.yaml with zero code modifications. The existing RAG pipeline, LangChain chains, prompts, FAISS retrieval, and all routes are completely untouched.

Closes #128

---

The Problem

The existing app/ai.py hardcodes HuggingFace Transformers as the only inference backend. Every model change requires code modifications. Schools and individuals with Groq credits, Google Workspace accounts, or local Ollama setups have no way to use their own resources. The change-model endpoint only switches between HuggingFace models and cannot switch
inference providers entirely.

---

Architecture

sugar_ai.yaml
|
SugarAIConfig.load()
|
BackendRouter.from_config()
|
├── primary: HuggingFaceBackend (Qwen2, SmolLM2, Gemma, Phi)
├── fallback 1: OpenAICompatBackend (Groq)
├── fallback 2: OpenAICompatBackend (Gemini)
└── fallback 3: OpenAICompatBackend (Ollama, local)
|
_RouterCallable (wraps router as HuggingFace pipeline)
|
RAGAgent.model (existing LangChain chains unchanged)

---

Files Added

app/backends/
init.py -> public API exports
base.py -> ModelBackend ABC, all shared data structures
config.py -> SugarAIConfig yaml loader with env var support
router.py -> BackendRouter + BackendRegistry
huggingface.py -> HuggingFace Transformers backend
llamacpp.py -> llama-cpp-python GGUF backend (existing path wrapped)
onnx.py -> ONNX Runtime backend (zero C++ compilation, ARM-ready)
openai_compat.py -> OpenAI-compatible API backend
sugar_ai.yaml ->configuration file with all providers documented

Files Modified

app/ai.py -> wired router into RAGAgent via _RouterCallable
requirements.txt -> added pyyaml, openai
.example.env -> added backend API key placeholders

---

How It Works

1. ModelBackend abstract base class

Every backend implements the same interface:

class ModelBackend(ABC):
    def ask(self, question, history, config) -> BackendResponse: ...
    def stream(self, question, history, config) -> Iterator[str]: ...
    def is_available(self) -> bool: ...
    def capabilities(self) -> BackendCapabilities: ...
    def count_tokens(self, text) -> int: ...
    def truncate_history(self, history, max_tokens) -> list[Message]: ...

The rest of the application never imports a specific backend class directly.
Everything goes through BackendRouter.

2. BackendRouter with fallback chain

router = BackendRouter(
    primary=HuggingFaceBackend(config),
    fallbacks=[GroqBackend(config), GeminiBackend(config), OllamaBackend(config)]
)
response = router.ask("What is photosynthesis?")

If the primary backend fails or is unavailable, the router silently tries the next one in the fallback chain. No manual error handling needed anywhere in application code. This is especially important for schools running on unreliable connectivity where a cloud API might time out.

3. _RouterCallable: the key integration trick

The existing LangChain chains use | self.model | expecting a HuggingFace pipeline callable. Rather than rewriting every chain, _RouterCallable wraps the router so it can be called exactly like a pipeline:

first_chain = (
    chain_input
    | self.prompt
    | combine_messages
    | self.model          # this is now _RouterCallable(router)
    | extract_answer_from_output
)

The wrapper receives a string or message list and returns output in HuggingFace pipeline format. Every LangChain chain, every prompt template, every extraction function is untouched.

4. sugar_ai.yaml configuration

primary_backend:
  type: huggingface
  model_name: Qwen/Qwen2-1.5B-Instruct
  device: cpu

fallback_backends:
  - type: openai_compat
    provider: groq
    api_key: ${GROQ_API_KEY}
    model_name: llama-3.1-8b-instant

  - type: openai_compat
    provider: gemini
    api_key: ${GEMINI_API_KEY}
    model_name: gemini-2.0-flash

  - type: openai_compat
    provider: ollama
    model_name: llama3

The ${VAR_NAME} syntax interpolates environment variables at load time so API keys never appear in the config file. The yaml file itself is safe to commit.

5. Environment variable overrides

Every yaml value can be overridden via environment variables without touching the config file:

SUGAR_AI_BACKEND_TYPE=openai_compat
SUGAR_AI_MODEL_NAME=gemini-2.0-flash
SUGAR_AI_API_KEY=your_key

This makes Docker and CI deployments clean with no yaml file needed at all if env vars are set.

6. Token budget enforcement

The router automatically truncates conversation history when it exceeds max_history_tokens before sending to the backend. Oldest turns are dropped first, system messages are always preserved. This prevents context window overflows silently in long sessions.

7. Backends implemented (7 total)

Backend	Models	Local	Notes
huggingface	Qwen2, SmolLM2, Gemma, Phi	Yes	existing path, hot-swap supported
llamacpp	any .gguf file	Yes	existing GGUF path, wrapped
onnx	Phi-3-mini-onnx, any .onnx	Yes	zero C++ compilation, ARM-ready
openai_compat (openai)	gpt-4o, gpt-4o-mini	No	standard OpenAI
openai_compat (groq)	llama-3.1, mixtral	No	very fast, free tier
openai_compat (gemini)	gemini-2.0-flash	No	free tier for schools
openai_compat (ollama)	any ollama model	Yes	local, no API key needed

Anthropic Claude and Together AI are also supported via the same openai_compat backend with different provider and base_url values.

8. set_model() now actually works

The existing set_model() only worked for HuggingFace models. The updated version checks if the primary backend supports hot-swap via load_model() and uses it if available. If not, it re-initialises the router with the new model name. Either way the switch works correctly.

9. backend_health() endpoint ready

A new backend_health() method on RAGAgent returns the availability and capabilities of all configured backends. This can be exposed as a monitoring route:

agent.backend_health()
# {
#   "primary": {"name": "huggingface", "available": true, "capabilities": {...}},
#   "fallbacks": [{"name": "openai_compat", "available": true}, ...]
# }

---

What This Enables Going Forward

With this abstraction in place, future additions require zero changes to ai.py or the router:

• vLLM backend for high-throughput local inference on schools with GPUs
• Bedrock backend for institutions with AWS credits
• Custom fine-tuned Sugar-specific model backend
• Streaming endpoint via FastAPI, the router already supports stream()

---

Testing

Manually tested the pipeline end to end with the HuggingFace backend on the default Qwen2-1.5B-Instruct model, verifying that RAG retrieval, chat completion, custom prompt generation, and model switching all work correctly through the new abstraction layer. Existing behaviour is preserved exactly.

Config loading, env var interpolation, fallback chain, token truncation, history normalization, and all seven backend is_available() checks have been verified to work correctly.

---

How to Switch Backends

No code changes needed. Just update sugar_ai.yaml:

# Switch to Groq (fast, free tier)
primary_backend:
  type: openai_compat
  provider: groq
  api_key: ${GROQ_API_KEY}
  model_name: llama-3.1-8b-instant

# Switch to local Ollama
primary_backend:
  type: openai_compat
  provider: ollama
  model_name: llama3

# Switch to ONNX (XO hardware, no C++ compilation)
primary_backend:
  type: onnx
  model_path: ./models/phi3-onnx/model.onnx

Or entirely via environment variables with no yaml file:

SUGAR_AI_BACKEND_TYPE=openai_compat
SUGAR_AI_MODEL_NAME=gemini-2.0-flash
SUGAR_AI_API_KEY=your_gemini_key

---

Author: Dashpreet Singh | dashpreetsinghhanda@gmail.com
IIT Jammu, B.Tech CSE 2024-2028