Integrate LiteLLM: multi-model chat, per-user budgets, and automatic fallback

.gitignore

@@ -19,4 +19,7 @@ kustomize/overlays/*/secrets.dat
 # per-overlay by secrets.sh; not committed in this repo).
 kustomize/overlays/*/secret-backend.yaml
 kustomize/overlays/*/secret-frontend.yaml
-kustomize/overlays/*/secret-ghcr.yaml
+kustomize/overlays/*/secret-ghcr.yaml
+
+qdrant_storage
+litellm.env

backend/.env.example

@@ -26,3 +26,9 @@ AUTH0_AUDIENCE=your_auth0_audience
 # Supabase Configuration
 SUPABASE_URL=your_supabase_url
 SUPABASE_SERVICE_ROLE_KEY=your_supabase_service_role_key
+
+# LiteLLM Proxy Configuration
+LITELLM_URL=http://litellm:4000
+LITELLM_MASTER_KEY=sk-litellm-change-me
+DEFAULT_USER_BUDGET=10.00
+DEFAULT_BUDGET_DURATION=30d

backend/app/controllers/llms_controller.py

@@ -1,4 +1,8 @@
 from fastapi import HTTPException
+from supabase import Client
+
+from app.core.singleton import get_supabase_client
+from app.services.litellm_service import get_or_create_virtual_key
 from app.services.llms_service import (
     get_response_with_tools,
     analyse_biomodel,
@@ -7,70 +11,114 @@
 )
 
 
-async def get_llm_response(conversation_history: list[dict]) -> tuple[str, list]:
+async def _get_virtual_key(payload: dict, supabase: Client) -> str:
+    auth0_sub = payload.get("sub")
+    if not auth0_sub:
+        raise HTTPException(status_code=401, detail="Missing Auth0 subject claim")
+
+    return await get_or_create_virtual_key(
+        auth0_sub=auth0_sub,
+        email=payload.get("email") or "",
+        supabase=supabase,
+    )
+
+
+async def get_llm_response(
+    conversation_history: list[dict],
+    model: str,
+    payload: dict,
+) -> tuple[str, list, str]:
     """
     Controller function to interact with the LLM service.
     Args:
         conversation_history (list[dict]): The conversation history containing user prompts and responses.
+        model (str): The LiteLLM model alias to use.
+        payload (dict): The verified Auth0 token payload for the caller.
     Returns:
-        tuple[str, list]: A tuple containing the final response and bmkeys list.
+        tuple[str, list, str]: The final response, bmkeys list, and model actually used.
     """
     try:
-        result, bmkeys = await get_response_with_tools(conversation_history)
-        return result, bmkeys
+        supabase = get_supabase_client()
+        virtual_key = await _get_virtual_key(payload, supabase)
+        result, bmkeys, model_used = await get_response_with_tools(
+            conversation_history, virtual_key, model
+        )
+        return result, bmkeys, model_used
     except Exception as e:
+        if isinstance(e, HTTPException):
+            raise e
         raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
 
 
-async def analyse_vcml_controller(biomodel_id: str) -> str:
+async def analyse_vcml_controller(biomodel_id: str, model: str, payload: dict) -> str:
     """
     Controller function to analyze VCML content for a given biomodel.
     Args:
         biomodel_id (str): The ID of the biomodel to analyze.
+        model (str): The LiteLLM model alias to use.
+        payload (dict): The verified Auth0 token payload for the caller.
     Returns:
         str: The VCML analysis response.
     """
     try:
-        result = await analyse_vcml(biomodel_id)
+        supabase = get_supabase_client()
+        virtual_key = await _get_virtual_key(payload, supabase)
+        result = await analyse_vcml(biomodel_id, virtual_key, model)
         return result
     except Exception as e:
+        if isinstance(e, HTTPException):
+            raise e
         raise HTTPException(
             status_code=500,
             detail=f"Error analyzing VCML for biomodel {biomodel_id}: {str(e)}",
         )
 
 
-async def analyse_diagram_controller(biomodel_id: str) -> str:
+async def analyse_diagram_controller(biomodel_id: str, model: str, payload: dict) -> str:
     """
     Controller function to analyze diagram for a given biomodel.
     Args:
         biomodel_id (str): The ID of the biomodel to analyze.
+        model (str): The LiteLLM model alias to use.
+        payload (dict): The verified Auth0 token payload for the caller.
     Returns:
         str: The diagram analysis response.
     """
     try:
-        result = await analyse_diagram(biomodel_id)
+        supabase = get_supabase_client()
+        virtual_key = await _get_virtual_key(payload, supabase)
+        result = await analyse_diagram(biomodel_id, virtual_key, model)
         return result
     except Exception as e:
+        if isinstance(e, HTTPException):
+            raise e
         raise HTTPException(
             status_code=500,
             detail=f"Error analyzing diagram for biomodel {biomodel_id}: {str(e)}",
         )
 
 
-async def analyse_biomodel_controller(biomodel_id: str, user_prompt: str) -> str:
+async def analyse_biomodel_controller(
+    biomodel_id: str, user_prompt: str, model: str, payload: dict
+) -> str:
     """
     Controller function to analyze a biomodel using the LLM service.
     Args:
         biomodel_id (str): The ID of the biomodel to be analyzed.
         user_prompt (str): The query or input provided by the user.
+        model (str): The LiteLLM model alias to use.
+        payload (dict): The verified Auth0 token payload for the caller.
     Returns:
         str: The analysis result from the LLM service.
     """
     try:
-        result = await analyse_biomodel(biomodel_id, user_prompt)
+        supabase = get_supabase_client()
+        virtual_key = await _get_virtual_key(payload, supabase)
+        result = await analyse_biomodel(biomodel_id, user_prompt, virtual_key, model)
         return result
     except Exception as e:
+        if isinstance(e, HTTPException):
+            raise e
         raise HTTPException(
             status_code=500, detail=f"Error analyzing biomodel {biomodel_id}: {str(e)}"
         )

backend/app/controllers/users_controller.py

@@ -1,6 +1,7 @@
 from fastapi import HTTPException
 
-from app.services.users_service import sync_auth0_user
+from app.services.litellm_service import get_user_budget_info
+from app.services.users_service import sync_current_user
 
 
 async def sync_current_user_controller(
@@ -18,7 +19,20 @@ async def sync_current_user_controller(
             detail="Missing Auth0 subject claim",
         )
 
-    return {
-        "status": "success",
-        "user": sync_auth0_user(payload),
-    }
+    return await sync_current_user(payload)
+
+
+async def get_current_user_budget_controller(payload: dict) -> dict:
+    """
+    Fetch the authenticated user's LiteLLM spend and budget.
+    """
+
+    auth0_sub = payload.get("sub")
+
+    if not auth0_sub:
+        raise HTTPException(
+            status_code=400,
+            detail="Missing Auth0 subject claim",
+        )
+
+    return await get_user_budget_info(auth0_sub)

backend/app/core/config.py

@@ -1,5 +1,6 @@
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from typing import Optional
+from decimal import Decimal
 
 
 class Settings(BaseSettings):
@@ -39,5 +40,11 @@ class Settings(BaseSettings):
     SUPABASE_URL: Optional[str] = None
     SUPABASE_SERVICE_ROLE_KEY: Optional[str] = None
 
+    # LiteLLM Proxy Config
+    LITELLM_URL: str = "http://litellm:4000"
+    LITELLM_MASTER_KEY: Optional[str] = None
+    DEFAULT_USER_BUDGET: Decimal = Decimal("10.00")
+    DEFAULT_BUDGET_DURATION: str = "30d"
+
 
 settings = Settings()

backend/app/core/litellm.py

@@ -0,0 +1,10 @@
+from openai import AsyncOpenAI
+
+from app.core.config import settings
+
+
+def get_litellm_client(virtual_key: str) -> AsyncOpenAI:
+    return AsyncOpenAI(
+        api_key=virtual_key,
+        base_url=settings.LITELLM_URL,
+    )

backend/app/core/singleton.py

@@ -1,38 +1,37 @@
-from langfuse.openai import AzureOpenAI, OpenAI
+from openai import AzureOpenAI, OpenAI
 from qdrant_client import QdrantClient
 from app.core.config import settings
 from supabase import Client, create_client
 
-openai_client = None
+embeddings_client = None
 qdrant_client = None
 supabase_client = None
 
 
-# OpenAI
-def connect_openai():
-    global openai_client
-    if openai_client is None:
+# Embeddings / document extraction client
+def connect_embeddings_client():
+    global embeddings_client
+    if embeddings_client is None:
         if settings.PROVIDER == "azure":
-            openai_client = AzureOpenAI(
+            embeddings_client = AzureOpenAI(
                 api_key=settings.AZURE_API_KEY,
                 api_version=settings.AZURE_API_VERSION,
                 azure_endpoint=settings.AZURE_ENDPOINT,
             )
         else:
             ## THIS IS FOR LOCAL LLM ONLY
-            openai_client = OpenAI(
+            embeddings_client = OpenAI(
                 api_key=settings.AZURE_API_KEY,
                 base_url=settings.AZURE_ENDPOINT,
                 project=None,
                 organization=None,
             )
-    return openai_client
+    return embeddings_client
 
 
-def get_openai_client():
-    connect_openai()
-    openai = openai_client
-    return openai
+def get_embeddings_client():
+    connect_embeddings_client()
+    return embeddings_client
 
 
 def connect_qdrant():

backend/app/routes/llms_router.py

@@ -7,33 +7,37 @@
     analyse_diagram_controller,
 )
 from app.core.auth import verify_auth0_token
+from app.schemas.llms_schema import AnalysisResponse, ChatRequest, ChatResponse, LLMModel
 
 router = APIRouter()
 
 
-@router.post("/query")
+@router.post("/query", response_model=ChatResponse)
 async def query_llm(
-    conversation_history: dict,
-    _payload: dict = Depends(verify_auth0_token),
+    request: ChatRequest,
+    payload: dict = Depends(verify_auth0_token),
 ):
     """
     Endpoint to query the LLM and execute the necessary tools.
     Args:
-        conversation_history (dict): The conversation history containing user prompts and responses.
+        request (ChatRequest): The conversation history and model choice.
     Returns:
         dict: The final response after processing the prompt with the tools.
     """
-    result, bmkeys = await get_llm_response(
-        conversation_history.get("conversation_history", [])
+    result, bmkeys, model_used = await get_llm_response(
+        request.conversation_history,
+        request.model,
+        payload,
     )
-    return {"response": result, "bmkeys": bmkeys}
+    return {"response": result, "bmkeys": bmkeys, "model_used": model_used}
 
 
-@router.post("/analyse/{biomodel_id}")
+@router.post("/analyse/{biomodel_id}", response_model=AnalysisResponse)
 async def analyse_biomodel(
     biomodel_id: str,
     user_prompt: str,
-    _payload: dict = Depends(verify_auth0_token),
+    model: LLMModel = "openai-model",
+    payload: dict = Depends(verify_auth0_token),
 ):
     """
     Endpoint to analyze a biomodel using the LLM service.
@@ -43,14 +47,15 @@ async def analyse_biomodel(
     Returns:
         dict: The analysis result from the LLM service.
     """
-    result = await analyse_biomodel_controller(biomodel_id, user_prompt)
+    result = await analyse_biomodel_controller(biomodel_id, user_prompt, model, payload)
     return {"response": result}
 
 
-@router.post("/analyse/{biomodel_id}/vcml")
+@router.post("/analyse/{biomodel_id}/vcml", response_model=AnalysisResponse)
 async def analyse_vcml(
     biomodel_id: str,
-    _payload: dict = Depends(verify_auth0_token),
+    model: LLMModel = "openai-model",
+    payload: dict = Depends(verify_auth0_token),
 ):
     """
     Endpoint to analyze VCML content for a given biomodel.
@@ -59,14 +64,15 @@ async def analyse_vcml(
     Returns:
         dict: The VCML analysis response.
     """
-    result = await analyse_vcml_controller(biomodel_id)
+    result = await analyse_vcml_controller(biomodel_id, model, payload)
     return {"response": result}
 
 
-@router.post("/analyse/{biomodel_id}/diagram")
+@router.post("/analyse/{biomodel_id}/diagram", response_model=AnalysisResponse)
 async def analyse_diagram(
     biomodel_id: str,
-    _payload: dict = Depends(verify_auth0_token),
+    model: LLMModel = "openai-model",
+    payload: dict = Depends(verify_auth0_token),
 ):
     """
     Endpoint to analyze diagram for a given biomodel.
@@ -75,5 +81,5 @@ async def analyse_diagram(
     Returns:
         dict: The diagram analysis response.
     """
-    result = await analyse_diagram_controller(biomodel_id)
+    result = await analyse_diagram_controller(biomodel_id, model, payload)
     return {"response": result}

backend/app/routes/users_router.py

@@ -1,19 +1,32 @@
 from fastapi import APIRouter, Depends
 
 from app.controllers.users_controller import (
+    get_current_user_budget_controller,
     sync_current_user_controller,
 )
 from app.core.auth import verify_auth0_token
+from app.schemas.users_schema import SyncUserResponse, UserBudgetResponse
 
 router = APIRouter()
 
 
-@router.post("/users/me", response_model=dict)
+@router.post("/users/me", response_model=SyncUserResponse)
 async def sync_current_user(
     payload: dict = Depends(verify_auth0_token),
 ):
     """
     endpoint to Sync authenticated Auth0 user into Supabase.
     """
 
-    return await sync_current_user_controller(payload)
+    return await sync_current_user_controller(payload)
+
+
+@router.get("/users/me/budget", response_model=UserBudgetResponse)
+async def get_current_user_budget(
+    payload: dict = Depends(verify_auth0_token),
+):
+    """
+    Endpoint to retrieve authenticated user's LiteLLM spend and budget.
+    """
+
+    return await get_current_user_budget_controller(payload)