feat: add BYOK support and dynamic LLM provider selection via headers

Author: soneylegal

src/app/api/v1/chat.py

@@ -1,6 +1,6 @@
 import logging
 
-from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi import APIRouter, Depends, Header, HTTPException, Request
 from langgraph.graph.state import CompiledStateGraph
 
 from src.app.core.dependencies import get_agent_workflow
@@ -18,6 +18,9 @@ async def chat_endpoint(
     request: Request,
     payload: ChatRequest,
     workflow: CompiledStateGraph = Depends(get_agent_workflow),
+    authorization: str | None = Header(default=None),
+    x_api_key: str | None = Header(default=None),
+    x_provider: str | None = Header(default=None),
 ):
     """
     Submit a question to the Vortex agentic workflow.
@@ -26,6 +29,25 @@ async def chat_endpoint(
     - The RAG pipeline (retrieve → grade → generate) for technical questions.
     - A direct LLM response for general queries.
     """
+    provider = x_provider.lower() if x_provider else None
+    if provider and provider not in ["gemini", "anthropic", "ollama"]:
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                f"Unsupported LLM provider: {x_provider}. "
+                "Supported providers: gemini, anthropic, ollama"
+            ),
+        )
+
+    api_key = None
+    if authorization:
+        if authorization.startswith("Bearer "):
+            api_key = authorization[7:]
+        else:
+            api_key = authorization
+    elif x_api_key:
+        api_key = x_api_key
+
     try:
         initial_state = {
             "question": payload.query,
@@ -34,6 +56,8 @@ async def chat_endpoint(
             "steps": ["received_query"],
             "route": "",
             "retry_count": 0,
+            "api_key": api_key,
+            "provider": provider,
         }
 
         result = await workflow.ainvoke(initial_state)

src/app/core/llm.py

@@ -5,17 +5,20 @@
 from .config import settings
 
 
-@functools.lru_cache(maxsize=1)
-def get_llm() -> BaseChatModel:
+@functools.lru_cache(maxsize=32)
+def get_llm(api_key: str | None = None, provider: str | None = None) -> BaseChatModel:
     """
-    Returns the appropriate LLM instance based on configuration.
+    Returns the appropriate LLM instance based on configuration or dynamic inputs.
     Supports Gemini (free), Ollama (local), and Anthropic (paid).
 
     The instance is cached so that all graph nodes share a single LLM client
-    per process, avoiding redundant re-initialization on every node call.
+    per process/key, avoiding redundant re-initialization on every node call.
     """
-    if settings.llm_provider == "gemini":
-        if not settings.gemini_api_key:
+    active_provider = provider or settings.llm_provider
+
+    if active_provider == "gemini":
+        active_key = api_key or settings.gemini_api_key
+        if not active_key:
             raise ValueError(
                 "GEMINI_API_KEY must be set to use Gemini. "
                 "Get a free key at https://aistudio.google.com/"
@@ -24,11 +27,11 @@ def get_llm() -> BaseChatModel:
 
         return ChatGoogleGenerativeAI(
             model=settings.gemini_model,
-            google_api_key=settings.gemini_api_key,
+            google_api_key=active_key,
             temperature=0.0,
         )
 
-    elif settings.llm_provider == "ollama":
+    elif active_provider == "ollama":
         try:
             from langchain_ollama import ChatOllama
         except ImportError as exc:
@@ -42,8 +45,9 @@ def get_llm() -> BaseChatModel:
             temperature=0.0,
         )
 
-    elif settings.llm_provider == "anthropic":
-        if not settings.anthropic_api_key:
+    elif active_provider == "anthropic":
+        active_key = api_key or settings.anthropic_api_key
+        if not active_key:
             raise ValueError("ANTHROPIC_API_KEY must be set to use Anthropic.")
         try:
             from langchain_anthropic import ChatAnthropic
@@ -54,9 +58,9 @@ def get_llm() -> BaseChatModel:
             ) from exc
         return ChatAnthropic(
             model=settings.anthropic_model,
-            api_key=settings.anthropic_api_key,
+            api_key=active_key,
             temperature=0.0,
         )
 
     else:
-        raise ValueError(f"Unsupported LLM provider: {settings.llm_provider}")
+        raise ValueError(f"Unsupported LLM provider: {active_provider}")

src/app/graph/nodes.py

@@ -49,7 +49,7 @@ def grade_documents_node(state: AgentState) -> dict[str, Any]:
     steps = list(state.get("steps", []))
     steps.append("grade_documents")
 
-    llm = get_llm()
+    llm = get_llm(api_key=state.get("api_key"), provider=state.get("provider"))
     structured_llm_grader = llm.with_structured_output(Grade)
 
     system = (
@@ -91,7 +91,7 @@ def generate_node(state: AgentState) -> dict[str, Any]:
     steps = list(state.get("steps", []))
     steps.append("generate_answer")
 
-    llm = get_llm()
+    llm = get_llm(api_key=state.get("api_key"), provider=state.get("provider"))
     prompt = ChatPromptTemplate.from_template(
         "You are an expert IT infrastructure support "
         "assistant for the Cortex and Sentinel platforms. "
@@ -126,7 +126,7 @@ def direct_response_node(state: AgentState) -> dict[str, Any]:
     steps = list(state.get("steps", []))
     steps.append("direct_response")
 
-    llm = get_llm()
+    llm = get_llm(api_key=state.get("api_key"), provider=state.get("provider"))
     prompt = ChatPromptTemplate.from_template(
         "You are a helpful IT support assistant. Answer the following question "
         "directly and concisely.\n\n"
@@ -157,7 +157,7 @@ def rewrite_query_node(state: AgentState) -> dict[str, Any]:
     steps = list(state.get("steps", []))
     steps.append("rewrite_query")
 
-    llm = get_llm()
+    llm = get_llm(api_key=state.get("api_key"), provider=state.get("provider"))
     system = (
         "You are a question re-writer that converts an input question to a better "
         "version optimized for vector store retrieval. Reason about the underlying "

src/app/graph/router.py

@@ -38,7 +38,7 @@ def router_node(state: AgentState) -> dict[str, Any]:
     steps = list(state.get("steps", []))
     steps.append("router")
 
-    llm = get_llm()
+    llm = get_llm(api_key=state.get("api_key"), provider=state.get("provider"))
     structured_llm = llm.with_structured_output(RouteDecision)
 
     system = (

src/app/graph/state.py

@@ -1,4 +1,4 @@
-from typing import TypedDict
+from typing import NotRequired, TypedDict
 
 from langchain_core.documents import Document
 
@@ -14,6 +14,8 @@ class AgentState(TypedDict):
         steps:       Ordered list of node names executed (for observability).
         route:       Router decision — "retrieve" or "direct" (for API response).
         retry_count: Number of query-rewrite retries performed (loop protection).
+        api_key:     Optional dynamic API key for the LLM provider.
+        provider:    Optional dynamic LLM provider (gemini, anthropic, ollama).
     """
 
     question: str
@@ -22,3 +24,5 @@ class AgentState(TypedDict):
     steps: list[str]
     route: str
     retry_count: int
+    api_key: NotRequired[str | None]
+    provider: NotRequired[str | None]

tests/test_auth.py

@@ -0,0 +1,97 @@
+from unittest.mock import AsyncMock
+
+from fastapi.testclient import TestClient
+
+from src.app.core.dependencies import get_agent_workflow
+from src.app.main import app
+
+client = TestClient(app)
+
+
+def test_chat_byok_headers():
+    """
+    Test that Authorization and X-Provider headers are correctly
+    extracted and passed to the workflow state.
+    """
+    mock_workflow = AsyncMock()
+    mock_workflow.ainvoke.return_value = {
+        "question": "test question",
+        "generation": "mock response",
+        "documents": [],
+        "steps": ["received_query"],
+        "route": "direct",
+        "retry_count": 0,
+    }
+
+    app.dependency_overrides[get_agent_workflow] = lambda: mock_workflow
+
+    # 1. Test Bearer token in Authorization header and X-Provider header
+    headers = {
+        "Authorization": "Bearer my-custom-gemini-key",
+        "X-Provider": "gemini",
+    }
+    response = client.post(
+        "/api/v1/chat",
+        json={"query": "test question"},
+        headers=headers,
+    )
+
+    assert response.status_code == 200
+    assert mock_workflow.ainvoke.call_count == 1
+    call_state = mock_workflow.ainvoke.call_args[0][0]
+    assert call_state["api_key"] == "my-custom-gemini-key"
+    assert call_state["provider"] == "gemini"
+
+    # 2. Test X-API-Key header and different casing for X-Provider
+    mock_workflow.reset_mock()
+    headers = {
+        "X-API-Key": "my-custom-anthropic-key",
+        "X-Provider": "Anthropic",
+    }
+    response = client.post(
+        "/api/v1/chat",
+        json={"query": "test question"},
+        headers=headers,
+    )
+
+    assert response.status_code == 200
+    assert mock_workflow.ainvoke.call_count == 1
+    call_state = mock_workflow.ainvoke.call_args[0][0]
+    assert call_state["api_key"] == "my-custom-anthropic-key"
+    assert call_state["provider"] == "anthropic"
+
+    # 3. Test when no headers are provided (should pass None)
+    mock_workflow.reset_mock()
+    response = client.post(
+        "/api/v1/chat",
+        json={"query": "test question"},
+    )
+
+    assert response.status_code == 200
+    assert mock_workflow.ainvoke.call_count == 1
+    call_state = mock_workflow.ainvoke.call_args[0][0]
+    assert call_state["api_key"] is None
+    assert call_state["provider"] is None
+
+    app.dependency_overrides.clear()
+
+
+def test_chat_invalid_provider():
+    """Test that an unsupported provider returns a 400 Bad Request error."""
+    mock_workflow = AsyncMock()
+    app.dependency_overrides[get_agent_workflow] = lambda: mock_workflow
+
+    headers = {
+        "X-Provider": "unsupported-llm-brand",
+    }
+    response = client.post(
+        "/api/v1/chat",
+        json={"query": "test question"},
+        headers=headers,
+    )
+
+    assert response.status_code == 400
+    assert "Unsupported LLM provider" in response.json()["detail"]
+    assert mock_workflow.ainvoke.call_count == 0
+
+    app.dependency_overrides.clear()