diff --git a/.gitignore b/.gitignore index a4d1ab8..dbd4c2c 100644 --- a/.gitignore +++ b/.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 \ No newline at end of file +kustomize/overlays/*/secret-ghcr.yaml + +qdrant_storage +litellm.env \ No newline at end of file diff --git a/backend/.env.example b/backend/.env.example index bd0e0fc..07db179 100644 --- a/backend/.env.example +++ b/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 diff --git a/backend/app/controllers/llms_controller.py b/backend/app/controllers/llms_controller.py index fe271f5..9a3b6d7 100644 --- a/backend/app/controllers/llms_controller.py +++ b/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)}" ) diff --git a/backend/app/controllers/users_controller.py b/backend/app/controllers/users_controller.py index ad41e9f..2154d25 100644 --- a/backend/app/controllers/users_controller.py +++ b/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) diff --git a/backend/app/core/config.py b/backend/app/core/config.py index c4731b5..3ed43bc 100644 --- a/backend/app/core/config.py +++ b/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() diff --git a/backend/app/core/litellm.py b/backend/app/core/litellm.py new file mode 100644 index 0000000..afe40ea --- /dev/null +++ b/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, + ) diff --git a/backend/app/core/singleton.py b/backend/app/core/singleton.py index 85c4716..8d2e845 100644 --- a/backend/app/core/singleton.py +++ b/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(): diff --git a/backend/app/routes/llms_router.py b/backend/app/routes/llms_router.py index b563761..4643969 100644 --- a/backend/app/routes/llms_router.py +++ b/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} diff --git a/backend/app/routes/users_router.py b/backend/app/routes/users_router.py index ce5fd40..a7cc7e3 100644 --- a/backend/app/routes/users_router.py +++ b/backend/app/routes/users_router.py @@ -1,14 +1,16 @@ 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), ): @@ -16,4 +18,15 @@ async def sync_current_user( endpoint to Sync authenticated Auth0 user into Supabase. """ - return await sync_current_user_controller(payload) \ No newline at end of file + 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) \ No newline at end of file diff --git a/backend/app/schemas/llms_schema.py b/backend/app/schemas/llms_schema.py new file mode 100644 index 0000000..49f8492 --- /dev/null +++ b/backend/app/schemas/llms_schema.py @@ -0,0 +1,20 @@ +from typing import Literal + +from pydantic import BaseModel, Field + +LLMModel = Literal["openai-model", "local-model"] + + +class ChatRequest(BaseModel): + conversation_history: list[dict] = Field(..., min_length=1) + model: LLMModel = "openai-model" + + +class ChatResponse(BaseModel): + response: str + bmkeys: list = Field(default_factory=list) + model_used: str + + +class AnalysisResponse(BaseModel): + response: str diff --git a/backend/app/schemas/users_schema.py b/backend/app/schemas/users_schema.py new file mode 100644 index 0000000..94355cc --- /dev/null +++ b/backend/app/schemas/users_schema.py @@ -0,0 +1,28 @@ +from datetime import datetime +from typing import Optional + +from pydantic import BaseModel, Field + + +class User(BaseModel): + user_id: str + auth0_sub: str + email: Optional[str] = None + name: Optional[str] = None + role: str + token_limit: int + tokens_used: int + created_at: datetime + last_login: datetime + litellm_virtual_key: Optional[str] = Field(default=None, exclude=True) + + +class SyncUserResponse(BaseModel): + status: str + user: Optional[User] = None + + +class UserBudgetResponse(BaseModel): + spend: float + max_budget: Optional[float] = None + remaining_budget: Optional[float] = None diff --git a/backend/app/services/knowledge_base_service.py b/backend/app/services/knowledge_base_service.py index 86f61d9..e354d82 100644 --- a/backend/app/services/knowledge_base_service.py +++ b/backend/app/services/knowledge_base_service.py @@ -3,7 +3,7 @@ from markitdown import MarkItDown from typing import List, Dict, Any, Optional from app.core.config import settings -from app.core.singleton import get_openai_client, get_qdrant_client +from app.core.singleton import get_embeddings_client, get_qdrant_client from langchain.text_splitter import RecursiveCharacterTextSplitter from app.services.qdrant_service import ( create_qdrant_collection, @@ -13,10 +13,10 @@ ) from langfuse import observe -openai_client = get_openai_client() +embeddings_client = get_embeddings_client() qdrant_client = get_qdrant_client() markitdown_client = MarkItDown( - llm_client=openai_client, model=settings.AZURE_DEPLOYMENT_NAME + llm_client=embeddings_client, model=settings.AZURE_DEPLOYMENT_NAME ) KB_COLLECTION_NAME = settings.QDRANT_COLLECTION_NAME @@ -53,7 +53,7 @@ def embed_text(text: str): Args: text (str): The text to embed. """ - response = openai_client.embeddings.create( + response = embeddings_client.embeddings.create( input=text, model=settings.AZURE_EMBEDDING_DEPLOYMENT_NAME ) return response.data[0].embedding diff --git a/backend/app/services/litellm_service.py b/backend/app/services/litellm_service.py new file mode 100644 index 0000000..8d74967 --- /dev/null +++ b/backend/app/services/litellm_service.py @@ -0,0 +1,100 @@ +import httpx +from supabase import Client + +from app.core.config import settings +from app.core.logger import get_logger + +logger = get_logger("litellm_service") + + +async def provision_user(auth0_sub: str, email: str) -> str: + """ + Create a user in LiteLLM and return the virtual key it generates. + + Args: + auth0_sub (str): The Auth0 subject claim, used as LiteLLM's user_id. + email (str): The user's email, stored on the LiteLLM user record. + + Returns: + str: The virtual key ("sk-...") LiteLLM generated for this user. + """ + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post( + f"{settings.LITELLM_URL}/user/new", + headers={"Authorization": f"Bearer {settings.LITELLM_MASTER_KEY}"}, + json={ + "user_id": auth0_sub, + "user_email": email, + "max_budget": float(settings.DEFAULT_USER_BUDGET), + "budget_duration": settings.DEFAULT_BUDGET_DURATION, + }, + ) + response.raise_for_status() + data = response.json() + + logger.info(f"Provisioned LiteLLM virtual key for user {auth0_sub}") + return data["key"] + + +async def get_or_create_virtual_key(auth0_sub: str, email: str, supabase: Client) -> str: + """ + Return the user's existing LiteLLM virtual key, provisioning a new one if + none is stored in Supabase yet. + + Args: + auth0_sub (str): The Auth0 subject claim, used as LiteLLM's user_id. + email (str): The user's email, passed to provision_user on first login. + supabase (Client): Supabase client used to look up / persist the virtual key. + + Returns: + str: The user's LiteLLM virtual key. + """ + response = ( + supabase.table("users") + .select("litellm_virtual_key") + .eq("auth0_sub", auth0_sub) + .limit(1) + .execute() + ) + existing_key = response.data[0].get("litellm_virtual_key") if response.data else None + if existing_key: + return existing_key + + virtual_key = await provision_user(auth0_sub, email) + + supabase.table("users").upsert( + {"auth0_sub": auth0_sub, "litellm_virtual_key": virtual_key}, + on_conflict="auth0_sub", + ).execute() + + return virtual_key + + +async def get_user_budget_info(auth0_sub: str) -> dict: + """ + Fetch spend and budget details for a user from LiteLLM. + + Args: + auth0_sub (str): The Auth0 subject claim, used as LiteLLM's user_id. + + Returns: + dict: spend, max_budget, and remaining_budget for the user. + """ + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.get( + f"{settings.LITELLM_URL}/user/info", + headers={"Authorization": f"Bearer {settings.LITELLM_MASTER_KEY}"}, + params={"user_id": auth0_sub}, + ) + response.raise_for_status() + user_info = response.json()["user_info"] + + spend = user_info["spend"] + max_budget = user_info["max_budget"] + remaining_budget = max_budget - spend if max_budget is not None else None + + return { + "spend": spend, + "max_budget": max_budget, + "remaining_budget": remaining_budget, + } diff --git a/backend/app/services/llms_service.py b/backend/app/services/llms_service.py index 5d5e91d..a0d8173 100644 --- a/backend/app/services/llms_service.py +++ b/backend/app/services/llms_service.py @@ -12,21 +12,68 @@ from app.utils.system_prompt import SYSTEM_PROMPT from app.schemas.vcelldb_schema import BiomodelRequestParams -from app.core.singleton import get_openai_client +from app.core.litellm import get_litellm_client from app.core.config import settings import json from app.core.logger import get_logger logger = get_logger("llm_service") -client = get_openai_client() +LOCAL_MODEL = "local-model" -async def get_llm_response(system_prompt: str, user_prompt: str): + +def _key_for_model(virtual_key: str, model: str) -> str: + """ + A user's budget is enforced across every model on their virtual key, so + retrying a budget-exceeded request against local-model with that same + key would just fail again. Use the unconstrained master key instead. + """ + if model == LOCAL_MODEL and settings.LITELLM_MASTER_KEY: + return settings.LITELLM_MASTER_KEY + return virtual_key + + +def _is_budget_error(error: Exception) -> bool: + status_code = getattr(error, "status_code", None) + error_text = str(error).lower() + return status_code in {400, 402, 429} and any( + marker in error_text for marker in ("budget", "quota", "limit", "exceeded") + ) + + +async def _create_chat_completion(virtual_key: str, model: str, **kwargs): + """ + Call the requested model; on a budget-exceeded error, silently retry + against local-model using the master key. + """ + client = get_litellm_client(_key_for_model(virtual_key, model)) + try: + return await client.chat.completions.create(model=model, **kwargs) + except Exception as error: + if model != LOCAL_MODEL and _is_budget_error(error): + logger.info( + f"Budget limit reached for {model}; falling back to {LOCAL_MODEL}" + ) + local_client = get_litellm_client(_key_for_model(virtual_key, LOCAL_MODEL)) + return await local_client.chat.completions.create( + model=LOCAL_MODEL, **kwargs + ) + raise + + +async def get_llm_response( + system_prompt: str, + user_prompt: str, + virtual_key: str, + model: str, +): """ Helper function to get a response from the LLM. args: system_prompt (str): The system prompt to guide the LLM. user_prompt (str): The user's query or request. + virtual_key (str): The caller's LiteLLM virtual key. + model (str): The LiteLLM model alias to use. returns: str: The response from the LLM. """ @@ -35,16 +82,20 @@ async def get_llm_response(system_prompt: str, user_prompt: str): {"role": "user", "content": user_prompt}, ] - response = client.chat.completions.create( - name="GET_LLM_RESPONSE", - model=settings.AZURE_DEPLOYMENT_NAME, + response = await _create_chat_completion( + virtual_key, + model, messages=messages, ) return response.choices[0].message.content -async def get_response_with_tools(conversation_history: list[dict]): +async def get_response_with_tools( + conversation_history: list[dict], + virtual_key: str, + model: str, +) -> tuple[str, list, str]: messages = [ { "role": "system", @@ -58,68 +109,72 @@ async def get_response_with_tools(conversation_history: list[dict]): logger.info(f"User prompt: {user_prompt}") - response = client.chat.completions.create( - name="GET_RESPONSE_WITH_TOOLS::RETRIEVE_TOOLS", - model=settings.AZURE_DEPLOYMENT_NAME, + response = await _create_chat_completion( + virtual_key, + model, messages=messages, tools=tools, tool_choice="auto", ) # Handle the tool calls - tool_calls = response.choices[0].message.tool_calls + response_message = response.choices[0].message + tool_calls = response_message.tool_calls - messages.append(response.choices[0].message) + messages.append(response_message.model_dump(exclude_none=True)) bmkeys = [] - if tool_calls: - for tool_call in tool_calls: - # Extract the function name and arguments - name = tool_call.function.name - args = json.loads(tool_call.function.arguments) + if not tool_calls: + final_response = response_message.content or "" + logger.info(f"LLM Response: {final_response}") + return final_response, bmkeys, response.model - logger.info(f"Tool Call: {name} with args: {args}") + for tool_call in tool_calls: + # Extract the function name and arguments + name = tool_call.function.name + args = json.loads(tool_call.function.arguments) - # Execute the tool function - result = await execute_tool(name, args) + logger.info(f"Tool Call: {name} with args: {args}") - logger.info(f"Tool Result: {str(result)[:500]}") + # Execute the tool function + result = await execute_tool(name, args) - # Extract bmkeys only if result is a dictionary and contains the expected key - if isinstance(result, dict): - bmkeys = result.get("unique_model_keys (bmkey)", []) + logger.info(f"Tool Result: {str(result)[:500]}") - # Send the result back to the model - messages.append( - {"role": "tool", "tool_call_id": tool_call.id, "content": str(result)} - ) + # Extract bmkeys only if result is a dictionary and contains the expected key + if isinstance(result, dict): + bmkeys = result.get("unique_model_keys (bmkey)", []) + + # Send the result back to the model + messages.append( + {"role": "tool", "tool_call_id": tool_call.id, "content": str(result)} + ) logger.info(str(messages)) # Send back the final response incorporating the tool result - completion = client.chat.completions.create( - name="GET_RESPONSE_WITH_TOOLS::PROCESS_TOOL_RESULTS", - model=settings.AZURE_DEPLOYMENT_NAME, + completion = await _create_chat_completion( + virtual_key, + model, messages=messages, - metadata={ - "tool_calls": tool_calls, - }, ) final_response = completion.choices[0].message.content logger.info(f"LLM Response: {final_response}") - return final_response, bmkeys + return final_response, bmkeys, completion.model -async def analyse_vcml(biomodel_id: str): +async def analyse_vcml(biomodel_id: str, virtual_key: str, model: str): """ Analyze VCML content for a given biomodel. args: biomodel_id (str): The ID of the biomodel to analyze. + virtual_key (str): The caller's LiteLLM virtual key. + model (str): The LiteLLM model alias to use. returns: str: The VCML analysis response. """ @@ -133,7 +188,9 @@ async def analyse_vcml(biomodel_id: str): ) vcml_system_prompt = "You are a VCell BioModel Assistant, designed to help users understand and interact with biological models in VCell. Your task is to provide human-readable, concise responses based on the given VCML." vcml_prompt = f"Analyze the following VCML content for Biomodel {biomodel_id}: {str(vcml)}" - vcml_analysis = await get_llm_response(vcml_system_prompt, vcml_prompt) + vcml_analysis = await get_llm_response( + vcml_system_prompt, vcml_prompt, virtual_key, model + ) return vcml_analysis except Exception as e: logger.error( @@ -142,13 +199,17 @@ async def analyse_vcml(biomodel_id: str): return f"An error occurred during VCML analysis: {str(e)}" -async def analyse_biomodel(biomodel_id: str, user_prompt: str): +async def analyse_biomodel( + biomodel_id: str, user_prompt: str, virtual_key: str, model: str +): """ Analyze user query with biomodel context. args: biomodel_id (str): The ID of the biomodel to analyze. user_prompt (str): The user's query or request. + virtual_key (str): The caller's LiteLLM virtual key. + model (str): The LiteLLM model alias to use. returns: str: The AI analysis response. """ @@ -172,7 +233,7 @@ async def analyse_biomodel(biomodel_id: str, user_prompt: str): # Analyze the user prompt with added biomodel context system_prompt = "You are a VCell BioModel Assistant, designed to help users understand and interact with biological models in VCell. Your task is to provide human-readable, accurate responses based on the given data. Give a response to the user's query, considering the provided biomodel information." user_analysis_response = await get_llm_response( - system_prompt, enhanced_user_prompt + system_prompt, enhanced_user_prompt, virtual_key, model ) return user_analysis_response except Exception as e: @@ -180,12 +241,14 @@ async def analyse_biomodel(biomodel_id: str, user_prompt: str): return f"An error occurred during AI analysis: {str(e)}" -async def analyse_diagram(biomodel_id: str): +async def analyse_diagram(biomodel_id: str, virtual_key: str, model: str): """ Analyze diagram for a given biomodel. args: biomodel_id (str): The ID of the biomodel to analyze. + virtual_key (str): The caller's LiteLLM virtual key. + model (str): The LiteLLM model alias to use. returns: str: The diagram analysis response. """ @@ -218,9 +281,9 @@ async def analyse_diagram(biomodel_id: str): {"type": "text", "text": diagram_analysis_prompt}, {"type": "image_url", "image_url": {"url": diagram_url}}, ] - response = client.chat.completions.create( - name="ANALYSE_DIAGRAM", - model=settings.AZURE_DEPLOYMENT_NAME, + response = await _create_chat_completion( + virtual_key, + model, messages=[ { "role": "user", diff --git a/backend/app/services/users_service.py b/backend/app/services/users_service.py index aa6e92c..e5c26de 100644 --- a/backend/app/services/users_service.py +++ b/backend/app/services/users_service.py @@ -1,6 +1,7 @@ from datetime import datetime, timezone from app.core.singleton import get_supabase_client +from app.services.litellm_service import get_or_create_virtual_key def sync_auth0_user(payload: dict) -> dict | None: @@ -32,3 +33,20 @@ def sync_auth0_user(payload: dict) -> dict | None: ) return response.data[0] if response.data else None + + +async def sync_current_user(payload: dict) -> dict: + """ + Ensure the user has a LiteLLM virtual key, then sync them into Supabase. + """ + auth0_sub = payload["sub"] + + supabase = get_supabase_client() + await get_or_create_virtual_key(auth0_sub, payload.get("email") or "", supabase) + + user = sync_auth0_user(payload) + + return { + "status": "success", + "user": user, + } diff --git a/docker-compose.yml b/docker-compose.yml index 958ec35..7683c92 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -10,6 +10,19 @@ services: - qdrant-data:/qdrant/storage restart: unless-stopped + ### LiteLLM Proxy (OpenAI-compatible LLM gateway) + litellm: + image: ghcr.io/berriai/litellm:main-stable + container_name: litellm-vcellai + ports: + - "4000:4000" + volumes: + - ./litellm_config.yaml:/app/config.yaml + env_file: + - ./litellm.env + command: --config /app/config.yaml + restart: unless-stopped + ### Backend (FastAPI with Poetry) backend: build: ./backend diff --git a/frontend/components/ChatBox.tsx b/frontend/components/ChatBox.tsx index e861870..3dacb07 100644 --- a/frontend/components/ChatBox.tsx +++ b/frontend/components/ChatBox.tsx @@ -3,15 +3,25 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { ScrollArea } from "@/components/ui/scroll-area"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; import { MarkdownRenderer } from "@/components/markdown-renderer"; import { MessageSquare, Send, Bot, User, Loader2 } from "lucide-react"; import { getAccessToken } from "@auth0/nextjs-auth0/client"; +type ModelId = "openai-model" | "local-model"; + interface Message { id: string; role: "user" | "assistant"; content: string; timestamp: Date; + modelUsed?: string; } interface QuickAction { @@ -78,6 +88,7 @@ export const ChatBox: React.FC = ({ ); const [inputMessage, setInputMessage] = useState(""); const [isLoading, setIsLoading] = useState(false); + const [selectedModel, setSelectedModel] = useState("openai-model"); const messagesEndRef = useRef(null); const inputRef = useRef(null); @@ -194,6 +205,7 @@ export const ChatBox: React.FC = ({ role: msg.role, content: msg.content, })), + model: selectedModel, }), }, ); @@ -210,6 +222,7 @@ export const ChatBox: React.FC = ({ role: "assistant", content: formattedResponse, timestamp: new Date(), + modelUsed: data.model_used, }; setMessages((prev) => [...prev, assistantMessage]); } catch (error) { @@ -283,6 +296,12 @@ export const ChatBox: React.FC = ({ ) : ( )} + {message.role === "assistant" && + message.modelUsed === "local-model" && ( +
+ Responded via Local LLM +
+ )} @@ -319,6 +338,22 @@ export const ChatBox: React.FC = ({
+ { + if (value === null) { + return "Unlimited"; + } + + if (Math.abs(value) > 0 && Math.abs(value) < 1) { + return `$${Number(value.toFixed(4)).toString()}`; + } + + return new Intl.NumberFormat("en-US", { + style: "currency", + currency: "USD", + maximumFractionDigits: 2, + }).format(value); +}; + export function AppSidebar() { const pathname = usePathname(); const { state } = useSidebar(); + const { user, isLoading: isUserLoading } = useUser(); + const [budget, setBudget] = useState(null); const isCollapsed = state === "collapsed"; + useEffect(() => { + if (isUserLoading || !user) { + return; + } + + const controller = new AbortController(); + + const fetchBudget = async () => { + try { + const token = await getAccessToken(); + const apiUrl = process.env.NEXT_PUBLIC_API_URL; + + if (!apiUrl) { + throw new Error("NEXT_PUBLIC_API_URL is not configured"); + } + + const response = await fetch(`${apiUrl}/users/me/budget`, { + headers: { + Authorization: `Bearer ${token}`, + accept: "application/json", + }, + signal: controller.signal, + }); + + if (!response.ok) { + throw new Error(`Budget request failed with status ${response.status}`); + } + + setBudget((await response.json()) as BudgetInfo); + } catch (error) { + if (!controller.signal.aborted) { + console.error("Failed to fetch budget", error); + } + } + }; + + fetchBudget(); + + return () => { + controller.abort(); + }; + }, [isUserLoading, user?.sub]); + if (pathname == "/") { return null; } + const usagePercent = + budget?.max_budget && budget.max_budget > 0 + ? Math.min((budget.spend / budget.max_budget) * 100, 100) + : 0; + const remainingText = !budget + ? "Loading..." + : `${formatBudget(budget.remaining_budget)} remaining`; + const budgetSummary = !budget + ? "Loading..." + : budget.max_budget === null + ? `${formatBudget(budget.spend)} spent, unlimited budget` + : `${formatBudget(budget.spend)} spent of ${formatBudget(budget.max_budget)} budget`; + const budgetCollapsedText = !budget ? "--" : formatBudget(budget.spend); + return ( @@ -244,25 +327,27 @@ export function AppSidebar() {
- {/* Usage Progress Bar */} + {/* Budget Progress Bar */}
{!isCollapsed && (
- Monthly Token Limit - 600K remaining + Budget + {remainingText}
)}
{isCollapsed ? ( -
0.4M
+
+ {budgetCollapsedText} +
) : (
- 400K of 1M tokens used this month + {budgetSummary}
)}
diff --git a/litellm.env.example b/litellm.env.example new file mode 100644 index 0000000..5e3d230 --- /dev/null +++ b/litellm.env.example @@ -0,0 +1,23 @@ +# LiteLLM Proxy +# Generate a strong value for local/prod use, for example: +# openssl rand -hex 32 +LITELLM_MASTER_KEY=sk-litellm-change-me + +# Hosted OpenAI provider for the "openai-model" LiteLLM alias. +OPENAI_API_KEY=sk-your-openai-api-key + +# Local provider for the "local-model" LiteLLM alias. +# For Ollama, keep the model prefixed with "ollama/". +LOCAL_LLM_MODEL=ollama/llama3.1:8b +LOCAL_LLM_API_BASE=http://host.docker.internal:11434 + +# Langfuse tracing callback (litellm_settings.success_callback in litellm_config.yaml). +LANGFUSE_PUBLIC_KEY=pk-lf-your-public-key +LANGFUSE_SECRET_KEY=sk-lf-your-secret-key +LANGFUSE_HOST=https://cloud.langfuse.com + +# Required: Postgres connection LiteLLM uses to persist users, virtual keys, +# and spend/budget tracking (general_settings.database_url in litellm_config.yaml). +# Without this, virtual keys and budgets do not survive a container restart. +# Use Supabase's direct connection or session pooler URL (port 5432) with SSL. +DATABASE_URL=postgresql://postgres:your-db-password@db.your-project-ref.supabase.co:5432/postgres?sslmode=require diff --git a/litellm_config.yaml b/litellm_config.yaml new file mode 100644 index 0000000..c4d0581 --- /dev/null +++ b/litellm_config.yaml @@ -0,0 +1,19 @@ +model_list: + - model_name: openai-model + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + + - model_name: local-model + litellm_params: + model: os.environ/LOCAL_LLM_MODEL + api_base: os.environ/LOCAL_LLM_API_BASE + +litellm_settings: + success_callback: ["langfuse"] + fallbacks: + - openai-model: ["local-model"] + +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL