get/save llm keys

This commit is contained in:
Krishav Raj Singh 2025-07-12 13:52:33 +05:30
parent 5566dc5519
commit ecb077e12f
5 changed files with 89 additions and 2 deletions

View File

@ -183,6 +183,9 @@ api_router.include_router(triggers_api.router)
from pipedream import api as pipedream_api from pipedream import api as pipedream_api
api_router.include_router(pipedream_api.router) api_router.include_router(pipedream_api.router)
from local_llm import api as local_llm_api
api_router.include_router(local_llm_api.router)
@api_router.get("/health") @api_router.get("/health")
async def health_check(): async def health_check():
logger.info("Health check endpoint called") logger.info("Health check endpoint called")

30
backend/local_llm/api.py Normal file
View File

@ -0,0 +1,30 @@
from fastapi import APIRouter
from utils.local_api_keys import save_local_api_keys, get_local_api_keys
from utils.config import config, EnvMode
from fastapi import HTTPException
from typing import Dict
from utils.constants import PROVIDERS
router = APIRouter(tags=["local-llm-keys"])
@router.get("/local-llm-keys")
def get_llm_keys() -> Dict[str, str]:
if config.ENV_MODE != EnvMode.LOCAL:
raise HTTPException(status_code=403, detail="API key management only available in local mode")
providers = [f"{provider}_API_KEY" for provider in PROVIDERS]
llm_keys = get_local_api_keys(providers)
return llm_keys
@router.post("/local-llm-keys")
def save_local_llm_keys(request: Dict[str, str]) -> Dict[str, str]:
if config.ENV_MODE != EnvMode.LOCAL:
raise HTTPException(status_code=403, detail="API key management only available in local mode")
print(f"Saving local LLM keys: {request}")
key_saved = save_local_api_keys(request)
if key_saved:
return {"message": "API keys saved successfully"}
else:
raise HTTPException(status_code=500, detail="Failed to save API keys")

View File

@ -19,6 +19,7 @@ import litellm
from litellm.files.main import ModelResponse from litellm.files.main import ModelResponse
from utils.logger import logger from utils.logger import logger
from utils.config import config from utils.config import config
from utils.constants import PROVIDERS
# litellm.set_verbose=True # litellm.set_verbose=True
litellm.modify_params=True litellm.modify_params=True
@ -38,8 +39,7 @@ class LLMRetryError(LLMError):
def setup_api_keys() -> None: def setup_api_keys() -> None:
"""Set up API keys from environment variables.""" """Set up API keys from environment variables."""
providers = ['OPENAI', 'ANTHROPIC', 'GROQ', 'OPENROUTER', 'XAI'] for provider in PROVIDERS:
for provider in providers:
key = getattr(config, f'{provider}_API_KEY') key = getattr(config, f'{provider}_API_KEY')
if key: if key:
logger.debug(f"API key set for provider: {provider}") logger.debug(f"API key set for provider: {provider}")

View File

@ -163,3 +163,5 @@ MODEL_ACCESS_TIERS = {
"tier_125_800": PAID_TIER_MODELS, "tier_125_800": PAID_TIER_MODELS,
"tier_200_1000": PAID_TIER_MODELS, "tier_200_1000": PAID_TIER_MODELS,
} }
PROVIDERS = ['OPENAI', 'ANTHROPIC', 'GROQ', 'OPENROUTER', 'XAI']

View File

@ -0,0 +1,52 @@
"""
Local API key management for LLMs.
This module provides functionality to manage API keys in local mode
by reading and writing to the .env file.
"""
import os
from typing import Dict, Optional, List
from utils.logger import logger
from utils.config import config, EnvMode
from dotenv import load_dotenv, set_key, find_dotenv
from utils.constants import PROVIDERS
def get_local_api_keys(providers: List[str]) -> Dict[str, str]:
"""Get API keys from .env file in local mode."""
if config.ENV_MODE != EnvMode.LOCAL:
return {}
try:
# Load current env vars
load_dotenv()
return {provider: os.getenv(provider) or "" for provider in providers}
except Exception as e:
logger.error(f"Failed to get local API keys: {e}")
return {}
def save_local_api_keys(api_keys: Dict[str, str]) -> bool:
"""Save API keys to .env file in local mode."""
if config.ENV_MODE != EnvMode.LOCAL:
return False
try:
# Find .env file
env_path = find_dotenv()
if not env_path:
logger.error("Could not find .env file")
return False
# Update each API key
for key, value in api_keys.items():
if value: # Only set if value is not empty
set_key(env_path, key, value)
logger.info(f"Updated {key} in .env file")
return True
except Exception as e:
logger.error(f"Failed to save local API keys: {e}")
return False