get/save llm keys

backend/api.py

@@ -183,6 +183,9 @@ api_router.include_router(triggers_api.router)
 from pipedream import api as pipedream_api
 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")
 async def health_check():
     logger.info("Health check endpoint called")

backend/local_llm/api.py

@@ -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")

backend/services/llm.py

@@ -19,6 +19,7 @@ import litellm
 from litellm.files.main import ModelResponse
 from utils.logger import logger
 from utils.config import config
+from utils.constants import PROVIDERS
 
 # litellm.set_verbose=True
 litellm.modify_params=True
@@ -38,8 +39,7 @@ class LLMRetryError(LLMError):
 
 def setup_api_keys() -> None:
     """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')
         if key:
             logger.debug(f"API key set for provider: {provider}")

backend/utils/constants.py

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

backend/utils/local_api_keys.py

@@ -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