Swapping models mid-conversation with LlamaIndex

Swapping models mid-conversation with LlamaIndex lets you route requests to the best provider for each turn — cheaper models for simple replies, stronger models for reasoning, or local models when privacy matters. This guide builds a production-ready conversation manager that preserves context across swaps, handles provider failures gracefully, and exposes routing controls your application can drive at runtime.

Step 1: Set up the environment and dependencies

Start with a clean virtual environment. You need LlamaIndex core, the OpenAI and Anthropic integrations, and a way to manage environment variables.

python -m venv .venv
source .venv/bin/activate
pip install llama-index llama-index-llms-openai llama-index-llms-anthropic python-dotenv

Create a .env file with your API keys. If you use n4n.ai as a unified gateway, you only need one key and point the OpenAI-compatible client at the gateway endpoint.

# .env
OPENAI_API_KEY=sk-...
ANTHROPIC_API_KEY=sk-ant-...
# Optional: unified gateway
N4N_API_KEY=n4n-...
N4N_BASE_URL=https://api.n4n.ai/v1

Step 2: Configure multiple LLM providers

Define a registry that holds initialized LLM instances. LlamaIndex’s Settings object controls global defaults, but for per-turn swapping you want explicit instances you can pass to the chat engine.

# llm_registry.py
import os
from dataclasses import dataclass
from typing import Dict, Optional
from llama_index.llms.openai import OpenAI
from llama_index.llms.anthropic import Anthropic
from llama_index.core.llms import LLM
from dotenv import load_dotenv

load_dotenv()

@dataclass
class ModelSpec:
    name: str
    provider: str
    model_id: str
    max_tokens: int = 4096
    temperature: float = 0.7

class LLMRegistry:
    def __init__(self):
        self._models: Dict[str, LLM] = {}
        self._specs: Dict[str, ModelSpec] = {}
        self._register_defaults()

    def _register_defaults(self):
        # OpenAI models
        openai_key = os.getenv("OPENAI_API_KEY")
        if openai_key:
            self.register(ModelSpec(
                name="gpt-4o",
                provider="openai",
                model_id="gpt-4o",
                max_tokens=4096,
            ))
            self.register(ModelSpec(
                name="gpt-4o-mini",
                provider="openai",
                model_id="gpt-4o-mini",
                max_tokens=16384,
            ))

        # Anthropic models
        anthropic_key = os.getenv("ANTHROPIC_API_KEY")
        if anthropic_key:
            self.register(ModelSpec(
                name="claude-3-5-sonnet",
                provider="anthropic",
                model_id="claude-3-5-sonnet-20241022",
                max_tokens=8192,
            ))
            self.register(ModelSpec(
                name="claude-3-haiku",
                provider="anthropic",
                model_id="claude-3-haiku-20240307",
                max_tokens=4096,
            ))

        # Unified gateway (n4n.ai) — single key, 240+ models
        n4n_key = os.getenv("N4N_API_KEY")
        n4n_base = os.getenv("N4N_BASE_URL", "https://api.n4n.ai/v1")
        if n4n_key:
            self.register(ModelSpec(
                name="gpt-4o-via-gateway",
                provider="openai",
                model_id="gpt-4o",
                max_tokens=4096,
            ), base_url=n4n_base, api_key=n4n_key)
            self.register(ModelSpec(
                name="claude-3-5-sonnet-via-gateway",
                provider="anthropic",
                model_id="claude-3-5-sonnet-20241022",
                max_tokens=8192,
            ), base_url=n4n_base, api_key=n4n_key)

    def register(self, spec: ModelSpec, base_url: Optional[str] = None, api_key: Optional[str] = None):
        if spec.provider == "openai":
            llm = OpenAI(
                model=spec.model_id,
                api_key=api_key or os.getenv("OPENAI_API_KEY"),
                base_url=base_url,
                max_tokens=spec.max_tokens,
                temperature=spec.temperature,
            )
        elif spec.provider == "anthropic":
            llm = Anthropic(
                model=spec.model_id,
                api_key=api_key or os.getenv("ANTHROPIC_API_KEY"),
                max_tokens=spec.max_tokens,
                temperature=spec.temperature,
            )
        else:
            raise ValueError(f"Unknown provider: {spec.provider}")

        self._models[spec.name] = llm
        self._specs[spec.name] = spec

    def get(self, name: str) -> LLM:
        if name not in self._models:
            raise KeyError(f"Model '{name}' not registered. Available: {list(self._models.keys())}")
        return self._models[name]

    def list_models(self) -> Dict[str, ModelSpec]:
        return self._specs.copy()

Step 3: Build a conversation manager that preserves history

The core challenge when you swap models mid-conversation with LlamaIndex is keeping the message history intact while changing the underlying LLM. LlamaIndex’s ChatEngine holds a reference to an LLM, but you can replace it by creating a new engine with the same memory.

# conversation.py
from typing import List, Optional, AsyncGenerator
from dataclasses import dataclass, field
from llama_index.core.llms import ChatMessage, MessageRole, LLM
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.chat_engine import SimpleChatEngine
from llama_index.core.settings import Settings

@dataclass
class ConversationTurn:
    role: MessageRole
    content: str
    model_used: str
    tokens_in: int = 0
    tokens_out: int = 0

@dataclass
class ConversationState:
    history: List[ConversationTurn] = field(default_factory=list)
    current_model: str = "gpt-4o-mini"
    system_prompt: str = "You are a helpful assistant."

    def to_chat_messages(self) -> List[ChatMessage]:
        msgs = [ChatMessage(role=MessageRole.SYSTEM, content=self.system_prompt)]
        for turn in self.history:
            msgs.append(ChatMessage(role=turn.role, content=turn.content))
        return msgs

class ConversationManager:
    def __init__(self, registry: 'LLMRegistry', initial_model: str = "gpt-4o-mini"):
        self.registry = registry
        self.state = ConversationState(current_model=initial_model)
        self._engine: Optional[SimpleChatEngine] = None
        self._build_engine()

    def _build_engine(self):
        llm = self.registry.get(self.state.current_model)
        memory = ChatMemoryBuffer.from_defaults(token_limit=llm.metadata.context_window - 1024)
        # Pre-populate memory with existing history
        for turn in self.state.history:
            memory.put(ChatMessage(role=turn.role, content=turn.content))
        self._engine = SimpleChatEngine.from_defaults(
            llm=llm,
            memory=memory,
            system_prompt=self.state.system_prompt,
        )

    def swap_model(self, new_model_name: str) -> bool:
        """Swap the active model, preserving full conversation history."""
        if new_model_name == self.state.current_model:
            return False
        if new_model_name not in self.registry.list_models():
            raise ValueError(f"Model '{new_model_name}' not available")

        self.state.current_model = new_model_name
        self._build_engine()
        return True

    def chat(self, user_message: str) -> str:
        """Send a message using the current model, record the turn."""
        response = self._engine.chat(user_message)
        # LlamaIndex response object has raw token counts if available
        tokens_in = getattr(response, 'raw', {}).get('usage', {}).get('prompt_tokens', 0)
        tokens_out = getattr(response, 'raw', {}).get('usage', {}).get('completion_tokens', 0)

        self.state.history.append(ConversationTurn(
            role=MessageRole.USER,
            content=user_message,
            model_used=self.state.current_model,
            tokens_in=tokens_in,
        ))
        self.state.history.append(ConversationTurn(
            role=MessageRole.ASSISTANT,
            content=str(response),
            model_used=self.state.current_model,
            tokens_out=tokens_out,
        ))
        return str(response)

    async def achat(self, user_message: str) -> str:
        response = await self._engine.achat(user_message)
        tokens_in = getattr(response, 'raw', {}).get('usage', {}).get('prompt_tokens', 0)
        tokens_out = getattr(response, 'raw', {}).get('usage', {}).get('completion_tokens', 0)

        self.state.history.append(ConversationTurn(
            role=MessageRole.USER,
            content=user_message,
            model_used=self.state.current_model,
            tokens_in=tokens_in,
        ))
        self.state.history.append(ConversationTurn(
            role=MessageRole.ASSISTANT,
            content=str(response),
            model_used=self.state.current_model,
            tokens_out=tokens_out,
        ))
        return str(response)

    def stream_chat(self, user_message: str) -> AsyncGenerator[str, None]:
        """Stream tokens from the current model."""
        streaming_response = self._engine.stream_chat(user_message)
        full_response = ""
        for token in streaming_response.response_gen:
            full_response += token
            yield token

        # Record after stream completes
        tokens_in = getattr(streaming_response, 'raw', {}).get('usage', {}).get('prompt_tokens', 0)
        tokens_out = getattr(streaming_response, 'raw', {}).get('usage', {}).get('completion_tokens', 0)
        self.state.history.append(ConversationTurn(
            role=MessageRole.USER,
            content=user_message,
            model_used=self.state.current_model,
            tokens_in=tokens_in,
        ))
        self.state.history.append(ConversationTurn(
            role=MessageRole.ASSISTANT,
            content=full_response,
            model_used=self.state.current_model,
            tokens_out=tokens_out,
        ))

Step 4: Implement routing logic for automatic model selection

Hardcoding swaps works for demos. Production systems need routing rules — complexity-based, cost-based, or directive-based. Here’s a policy engine that inspects the user message and conversation state to pick a model.

# routing.py
from typing import Optional, Callable
from dataclasses import dataclass
from enum import Enum
import re

class RoutingStrategy(Enum):
    MANUAL = "manual"
    COMPLEXITY = "complexity"
    COST_OPTIMIZED = "cost_optimized"
    DIRECTIVE = "directive"

@dataclass
class RoutingPolicy:
    strategy: RoutingStrategy = RoutingStrategy.COMPLEXITY
    # Model tiers: (model_name, max_complexity_score)
    tiers: list = None

    def __post_init__(self):
        if self.tiers is None:
            self.tiers = [
                ("gpt-4o-mini", 0.3),
                ("claude-3-haiku", 0.5),
                ("gpt-4o", 0.8),
                ("claude-3-5-sonnet", 1.0),
            ]

class ComplexityScorer:
    """Heuristic complexity scoring for routing decisions."""
    CODE_PATTERNS = [
        r'\b(def|class|function|import|from)\b',
        r'[{}[\]();]',
        r'\b(async|await|Promise|async/await)\b',
    ]
    REASONING_PATTERNS = [
        r'\b(analyze|compare|evaluate|synthesize|reason|step by step)\b',
        r'\b(why|how|explain|justify|prove)\b',
    ]

    def score(self, text: str) -> float:
        score = 0.0
        text_lower = text.lower()

        # Length factor
        score += min(len(text) / 2000, 0.3)

        # Code detection
        for pattern in self.CODE_PATTERNS:
            if re.search(pattern, text):
                score += 0.2
                break

        # Reasoning detection
        for pattern in self.REASONING_PATTERNS:
            if re.search(pattern, text_lower):
                score += 0.3
                break

        # Multi-turn context (would need conversation history)
        # Placeholder for context-aware scoring
        return min(score, 1.0)

class ModelRouter:
    def __init__(self, policy: RoutingPolicy, registry: 'LLMRegistry'):
        self.policy = policy
        self.registry = registry
        self.scorer = ComplexityScorer()

    def select_model(self, user_message: str, conversation_state: 'ConversationState',
                     directive: Optional[str] = None) -> str:
        """Return the model name to use for the next turn."""
        # Explicit directive wins (e.g., from UI or API header)
        if directive and directive in self.registry.list_models():
            return directive

        if self.policy.strategy == RoutingStrategy.MANUAL:
            return conversation_state.current_model

        if self.policy.strategy == RoutingStrategy.DIRECTIVE:
            # Check for inline directives like "@model:gpt-4o"
            match = re.search(r'@model:(\w+)', user_message)
            if match and match.group(1) in self.registry.list_models():
                return match.group(1)
            return conversation_state.current_model

        if self.policy.strategy == RoutingStrategy.COMPLEXITY:
            complexity = self.scorer.score(user_message)
            for model_name, max_complexity in self.policy.tiers:
                if complexity <= max_complexity and model_name in self.registry.list_models():
                    return model_name
            return self.policy.tiers[-1][0]

        if self.policy.strategy == RoutingStrategy.COST_OPTIMIZED:
            # Prefer cheapest model that meets minimum capability
            # This requires a cost map — simplified here
            cheap_models = ["gpt-4o-mini", "claude-3-haiku"]
            for m in cheap_models:
                if m in self.registry.list_models():
                    return m
            return conversation_state.current_model

        return conversation_state.current_model

Step 5: Wire it together with a CLI for testing

A minimal CLI demonstrates the full flow: chat, swap models mid-conversation with LlamaIndex, and verify context carries over.

# main.py
import asyncio
import sys
from llm_registry import LLMRegistry
from conversation import ConversationManager
from routing import ModelRouter, RoutingPolicy, RoutingStrategy

async def main():
    registry = LLMRegistry()
    print("Available models:")
    for name, spec in registry.list_models().items():
        print(f"  {name} ({spec.provider}/{spec.model_id})")

    # Start with a cheap model
    manager = ConversationManager(registry, initial_model="gpt-4o-mini")
    router = ModelRouter(RoutingPolicy(strategy=RoutingStrategy.COMPLEXITY), registry)

    print("\nCommands:")
    print("  /swap <model>     - Switch model manually")
    print("  /policy <strategy> - Change routing: manual, complexity, cost, directive")
    print("  /history          - Show conversation history")
    print("  /quit             - Exit")
    print()

    while True:
        try:
            user_input = input(f"[{manager.state.current_model}] You: ").strip()
        except (EOFError, KeyboardInterrupt):
            break

        if not user_input:
            continue

        if user_input.startswith("/swap "):
            model = user_input[6:].strip()
            try:
                manager.swap_model(model)
                print(f"Swapped to {model}")
            except ValueError as e:
                print(f"Error: {e}")
            continue

        if user_input.startswith("/policy "):
            strategy_name = user_input[8:].strip().upper()
            try:
                router.policy.strategy = RoutingStrategy[strategy_name]
                print(f"Routing policy: {strategy_name}")
            except KeyError:
                print(f"Unknown strategy. Options: {[s.name for s in RoutingStrategy]}")
            continue

        if user_input == "/history":
            for turn in manager.state.history:
                print(f"  [{turn.model_used}] {turn.role.value}: {turn.content[:80]}...")
            continue

        if user_input == "/quit":
            break

        # Determine model for this turn
        selected_model = router.select_model(user_input, manager.state)
        if selected_model != manager.state.current_model:
            manager.swap_model(selected_model)
            print(f"[Auto-routed to {selected_model}]")

        # Stream response
        print(f"[{manager.state.current_model}] Assistant: ", end="", flush=True)
        async for token in manager.stream_chat(user_input):
            print(token, end="", flush=True)
        print()

if __name__ == "__main__":
    asyncio.run(main())

Run it:

python main.py

Step 6: Add provider fallback for resilience

When a provider returns 429 or 5xx, you want automatic fallback without losing the conversation. Wrap the chat call with a retry/fallback chain.

# fallback.py
import asyncio
from typing import List, Optional
from llama_index.core.llms import LLM
from llama_index.core.base.llms.types import ChatResponse, ChatResponseGen
from llama_index.core.llms import ChatMessage

class FallbackLLM(LLM):
    """LLM wrapper that tries a primary model, then falls back through a chain."""
    def __init__(self, primary: LLM, fallbacks: List[LLM]):
        self.primary = primary
        self.fallbacks = fallbacks
        self._current = primary

    @property
    def metadata(self):
        return self._current.metadata

    def chat(self, messages: List[ChatMessage], **kwargs) -> ChatResponse:
        return self._execute_with_fallback(lambda m: m.chat(messages, **kwargs))

    async def achat(self, messages: List[ChatMessage], **kwargs) -> ChatResponse:
        return await self._aexecute_with_fallback(lambda m: m.achat(messages, **kwargs))

    def stream_chat(self, messages: List[ChatMessage], **kwargs) -> ChatResponseGen:
        return self._execute_with_fallback(lambda m: m.stream_chat(messages, **kwargs))

    async def astream_chat(self, messages: List[ChatMessage], **kwargs) -> ChatResponseGen:
        return await self._aexecute_with_fallback(lambda m: m.astream_chat(messages, **kwargs))

    def _execute_with_fallback(self, fn):
        last_error = None
        models = [self.primary] + self.fallbacks
        for model in models:
            try:
                self._current = model
                return fn(model)
            except Exception as e:
                last_error = e
                # Check if error is retryable (rate limit, server error)
                if not self._is_retryable(e):
                    raise
                continue
        raise last_error or RuntimeError("All models failed")

    async def _aexecute_with_fallback(self, fn):
        last_error = None
        models = [self.primary] + self.fallbacks
        for model in models:
            try:
                self._current = model
                return await fn(model)
            except Exception as e:
                last_error = e
                if not self._is_retryable(e):
                    raise
                await asyncio.sleep(0.5)  # Brief backoff
                continue
        raise last_error or RuntimeError("All models failed")

    def _is_retryable(self, error: Exception) -> bool:
        error_str = str(error).lower()
        retryable_indicators = [
            "rate limit", "429", "503", "502", "504",
            "timeout", "connection", "unavailable", "overloaded"
        ]
        return any(indicator in error_str for indicator in retryable_indicators)

Integrate it into the registry:

# In llm_registry.py, add to LLMRegistry class:
def get_with_fallback(self, primary_name: str, fallback_names: List[str]) -> FallbackLLM:
    primary = self.get(primary_name)
    fallbacks = [self.get(name) for name in fallback_names if name in self._models]
    return FallbackLLM(primary, fallbacks)

Then in ConversationManager._build_engine, use registry.get_with_fallback("gpt-4o", ["claude-3-5-sonnet", "gpt-4o-mini"]) instead of registry.get().

Step 7: Verify the implementation

Run these checks to confirm swapping works correctly.

7.1 Context preservation test

# test_context.py
import asyncio
from llm_registry import LLMRegistry
from conversation import ConversationManager

async def test_context_preservation():
    registry = LLMRegistry()
    manager = ConversationManager(registry, initial_model="gpt-4o-mini")

    # Turn 1: Establish context
    resp1 = await manager.achat("My name is Alex and I work on distributed systems.")
    print(f"Turn 1: {resp1}")

    # Turn 2: Swap to a different provider
    manager.swap_model("claude-3-haiku")
    resp2 = await manager.achat("What's my name?")
    print(f"Turn 2 (after swap): {resp2}")

    # Turn 3: Swap again, ask for details
    manager.swap_model("gpt-4o")
    resp3 = await manager.achat("What do I work on?")
    print(f"Turn 3 (after second swap): {resp3}")

    # Verify history has all turns with correct model attribution
    assert len(manager.state.history) == 6  # 3 user + 3 assistant
    assert manager.state.history[0].model_used == "gpt-4o-mini"
    assert manager.state.history[2].model_used == "claude-3-haiku"
    assert manager.state.history[4].model_used == "gpt-4o"

    print("\n✓ Context preserved across 3 model swaps")
    print(f"Models used in sequence: {[t.model_used for t in manager.state.history[::2]]}")

asyncio.run(test_context_preservation())

Run it:

python test_context.py

Expected output shows each response answers correctly using context from previous turns, and the history records which model handled each turn.

7.2 Fallback verification

# test_fallback.py
import asyncio
from llm_registry import LLMRegistry
from fallback import FallbackLLM
from llama_index.llms.openai import OpenAI
from llama_index.core.llms import ChatMessage

async def test_fallback_chain():
    registry = LLMRegistry()

    # Create a failing primary (invalid key) and working fallback
    bad_llm = OpenAI(model="gpt-4o", api_key="sk-invalid")
    good_llm = registry.get("gpt-4o-mini")

    fallback_llm = FallbackLLM(bad_llm, [good_llm])

    messages = [ChatMessage(role="user", content="Say 'fallback worked'")]
    response = await fallback_llm.achat(messages)

    assert "fallback worked" in str(response).lower()
    assert fallback_llm._current == good_llm
    print("✓ Fallback chain executed successfully")

asyncio.run(test_fallback_chain())

7.3 Routing policy verification

# test_routing.py
from routing import ModelRouter, RoutingPolicy, RoutingStrategy, ComplexityScorer
from llm_registry import LLMRegistry
from conversation import ConversationState

def test_complexity_routing():
    registry = LLMRegistry()
    policy = RoutingPolicy(strategy=RoutingStrategy.COMPLEXITY)
    router = ModelRouter(policy, registry)
    state = ConversationState(current_model="gpt-4o-mini")

    # Simple query -> cheap model
    simple = "What's 2+2?"
    model = router.select_model(simple, state)
    assert model in ["gpt-4o-mini", "claude-3-haiku"], f"Expected cheap model, got {model}"

    # Complex query -> stronger model
    complex_q = "Analyze the tradeoffs between eventual consistency and strong consistency in distributed databases, then synthesize a recommendation for a financial ledger system."
    model = router.select_model(complex_q, state)
    assert model in ["gpt-4o", "claude-3-5-sonnet"], f"Expected strong model, got {model}"

    print("✓ Complexity routing selects appropriate tier")

def test_directive_routing():
    registry = LLMRegistry()
    policy = RoutingPolicy(strategy=RoutingStrategy.DIRECTIVE)
    router = ModelRouter(policy, registry)
    state = ConversationState(current_model="gpt-4o-mini")

    # Inline directive
    model = router.select_model("Explain quantum computing @model:gpt-4o", state)
    assert model == "gpt-4o"

    print("✓ Directive routing honors @model: tags")

test_complexity_routing()
test_directive_routing()

Step 8: Production considerations

Token budgeting across swaps

Different models have different context windows. When you swap models mid-conversation with LlamaIndex, the ChatMemoryBuffer token limit should respect the new model’s window. The ConversationManager._build_engine method already handles this by recreating the memory with the current model’s context_window.

Observability

Log every swap with metadata:

import logging
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("model_swap")

def log_swap(old_model: str, new_model: str, reason: str, turn: int):
    logger.info(json.dumps({
        "event": "model_swap",
        "from": old_model,
        "to": new_model,
        "reason": reason,  # "manual", "complexity", "fallback", "directive"
        "turn": turn,
        "timestamp": "2024-01-15T10:30:00Z",
    }))

Cost tracking

Accumulate per-model token counts from ConversationTurn records. Multiply by your provider’s per-token rates for real-time cost dashboards.

Streaming with swaps

If a user requests a model swap mid-stream, cancel the current generation, swap the engine, and restart. The ConversationManager preserves history, so the new model sees the full context including the partial assistant response (which you may want to trim or mark as incomplete).

Common pitfalls

Next steps

• Add a cost-aware router that checks real-time pricing from your provider or gateway
• Implement model-specific system prompts (Claude likes XML tags, GPT prefers markdown)
• Build a WebSocket server around ConversationManager for real-time UIs
• Add evaluation hooks to log model quality per task type and auto-tune routing policies

The pattern scales: one conversation object, multiple models, deterministic swaps. Your application decides the policy; the infrastructure executes it.