From ecb077e12f0219299b8289c25355b129cd7817ff Mon Sep 17 00:00:00 2001
From: Krishav Raj Singh <krishavrajsingh@gmail.com>
Date: Sat, 12 Jul 2025 13:52:33 +0530
Subject: [PATCH] get/save llm keys

---
 backend/api.py                  |  3 ++
 backend/local_llm/api.py        | 30 +++++++++++++++++++
 backend/services/llm.py         |  4 +--
 backend/utils/constants.py      |  2 ++
 backend/utils/local_api_keys.py | 52 +++++++++++++++++++++++++++++++++
 5 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 backend/local_llm/api.py
 create mode 100644 backend/utils/local_api_keys.py

diff --git a/backend/api.py b/backend/api.py
index b452658f..f39ab14d 100644
--- a/backend/api.py
+++ b/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")
diff --git a/backend/local_llm/api.py b/backend/local_llm/api.py
new file mode 100644
index 00000000..4f89baa0
--- /dev/null
+++ b/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")
\ No newline at end of file
diff --git a/backend/services/llm.py b/backend/services/llm.py
index 6187e247..24c270b3 100644
--- a/backend/services/llm.py
+++ b/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}")
diff --git a/backend/utils/constants.py b/backend/utils/constants.py
index 9d85684f..45c90295 100644
--- a/backend/utils/constants.py
+++ b/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']
\ No newline at end of file
diff --git a/backend/utils/local_api_keys.py b/backend/utils/local_api_keys.py
new file mode 100644
index 00000000..0379390d
--- /dev/null
+++ b/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 
\ No newline at end of file