fix(critical): LLM gateway actually passes user API keys to litellm calls

backend/app/services/llm_gateway.py

@@ -1,7 +1,6 @@
 """
-LLM Gateway - Multi-provider with retry, JSON validation, and proper key config.
+LLM Gateway - Multi-provider with retry, JSON validation, user key injection.
 """
-import hashlib
 import json
 import logging
 import time
@@ -9,20 +8,16 @@ from enum import Enum
 from typing import Any
 
 import litellm
-from pydantic import ValidationError
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
 
 from app.core.config import settings
 from app.core.encryption import decrypt_api_key
 from app.models.llm import LLMRequest
-from app.models.resume import AIArtifact
 from app.models.user import User, UserAPIKey
 
 logger = logging.getLogger(__name__)
 
-# Disable litellm telemetry
 litellm.telemetry = False
 litellm.drop_params = True
 
@@ -46,52 +41,46 @@ TASK_MODEL_PREFERENCES = {
 }
 
 TASK_PROMPTS = {
-    TaskType.RESUME_REVIEW: """You are an expert ATS resume reviewer. Analyze and return ONLY valid JSON with these fields:
+    TaskType.RESUME_REVIEW: """You are an expert ATS resume reviewer. Return ONLY valid JSON:
 {"score": int 0-100, "strengths": ["..."], "weaknesses": ["..."], "ats_issues": ["..."], "missing_keywords": ["..."], "rewrite_suggestions": [{"section": "...", "original": "...", "suggestion": "..."}], "overall_feedback": "2-3 sentences"}""",
-    TaskType.JOB_MATCH: """You are a job matching expert. Compare resume vs job and return ONLY valid JSON:
+    TaskType.JOB_MATCH: """You are a job matching expert. Return ONLY valid JSON:
 {"match_score": int 0-100, "requirements_met": ["..."], "requirements_missing": ["..."], "skills_gap": ["..."], "strengths_for_role": ["..."], "risks": ["..."], "recommendation": "string"}""",
     TaskType.COVER_LETTER: """Write a tailored cover letter. Return ONLY valid JSON:
-{"cover_letter": "full text", "key_points_addressed": ["..."]}
-Rules: Never fabricate. Under 400 words. Reference real experience only.""",
-    TaskType.BULLET_REWRITE: """Rewrite resume bullets for impact. Return ONLY valid JSON:
+{"cover_letter": "full text under 400 words", "key_points_addressed": ["..."]}
+Never fabricate experiences.""",
+    TaskType.BULLET_REWRITE: """Rewrite resume bullets. Return ONLY valid JSON:
 {"rewrites": [{"original": "...", "rewritten": "...", "improvement_notes": "..."}]}
-Rules: Preserve all facts. Never add fabricated metrics.""",
+Never add fabricated metrics.""",
 }
 
 
-def _configure_provider_keys():
-    """Configure all LLM provider API keys from settings."""
+def _configure_global_keys():
+    """Set system-level API keys from env vars."""
+    import os
     if settings.OPENAI_API_KEY:
         litellm.openai_key = settings.OPENAI_API_KEY
     if settings.ANTHROPIC_API_KEY:
         litellm.anthropic_key = settings.ANTHROPIC_API_KEY
     if settings.GOOGLE_API_KEY:
-        import os
-        os.environ["GEMINI_API_KEY"] = settings.GOOGLE_API_KEY
+        os.environ.setdefault("GEMINI_API_KEY", settings.GOOGLE_API_KEY)
     if settings.MISTRAL_API_KEY:
-        import os
-        os.environ["MISTRAL_API_KEY"] = settings.MISTRAL_API_KEY
+        os.environ.setdefault("MISTRAL_API_KEY", settings.MISTRAL_API_KEY)
     if settings.DEEPSEEK_API_KEY:
-        import os
-        os.environ["DEEPSEEK_API_KEY"] = settings.DEEPSEEK_API_KEY
+        os.environ.setdefault("DEEPSEEK_API_KEY", settings.DEEPSEEK_API_KEY)
     if settings.XAI_API_KEY:
-        import os
-        os.environ["XAI_API_KEY"] = settings.XAI_API_KEY
+        os.environ.setdefault("XAI_API_KEY", settings.XAI_API_KEY)
     if settings.HUGGINGFACE_API_KEY:
-        import os
-        os.environ["HUGGINGFACE_API_KEY"] = settings.HUGGINGFACE_API_KEY
-
+        os.environ.setdefault("HUGGINGFACE_API_KEY", settings.HUGGINGFACE_API_KEY)
     litellm.ollama_api_base = settings.OLLAMA_BASE_URL
 
 
-# Configure on import
-_configure_provider_keys()
+_configure_global_keys()
 
 
-def _repair_json(text: str) -> str:
-    """Attempt to repair malformed JSON from LLMs."""
+def _repair_json(text: str) -> dict:
+    """Parse JSON with repair: strip fences, find object, retry."""
+    import re
     text = text.strip()
-    # Remove markdown code fences
     if text.startswith("```json"):
         text = text[7:]
     if text.startswith("```"):
@@ -100,44 +89,34 @@ def _repair_json(text: str) -> str:
         text = text[:-3]
     text = text.strip()
 
-    # Try to find JSON object boundaries
-    start = text.find("{")
-    end = text.rfind("}") + 1
-    if start >= 0 and end > start:
-        text = text[start:end]
-
-    return text
-
-
-def _parse_json_response(content: str) -> dict:
-    """Parse JSON from LLM response with repair attempts."""
-    # Attempt 1: Direct parse
+    # Direct parse
     try:
-        return json.loads(content)
+        return json.loads(text)
     except json.JSONDecodeError:
         pass
 
-    # Attempt 2: Repair and parse
-    repaired = _repair_json(content)
-    try:
-        return json.loads(repaired)
-    except json.JSONDecodeError:
-        pass
+    # Find JSON object
+    start = text.find("{")
+    end = text.rfind("}") + 1
+    if start >= 0 and end > start:
+        try:
+            return json.loads(text[start:end])
+        except json.JSONDecodeError:
+            pass
 
-    # Attempt 3: Try to extract any JSON object
-    import re
-    json_match = re.search(r'\{[\s\S]*\}', content)
-    if json_match:
+    # Regex extract
+    match = re.search(r'\{[\s\S]*\}', text)
+    if match:
         try:
-            return json.loads(json_match.group())
+            return json.loads(match.group())
         except json.JSONDecodeError:
             pass
 
-    raise ValueError(f"Could not parse JSON from LLM response: {content[:200]}")
+    raise ValueError(f"Cannot parse JSON: {text[:200]}")
 
 
 class LLMGateway:
-    """Multi-provider LLM gateway with retry, validation, and audit."""
+    """Multi-provider LLM gateway with user key support, retry, and audit."""
 
     def __init__(self, user: User, db: AsyncSession):
         self.user = user
@@ -158,7 +137,6 @@ class LLMGateway:
         return None
 
     def _resolve_model(self, task_type: TaskType, provider: str | None, model: str | None) -> str:
-        """Resolve which model to use."""
         if model:
             return f"{provider}/{model}" if provider else model
 
@@ -170,7 +148,6 @@ class LLMGateway:
                     if m.startswith(f"{provider}/"):
                         return m
 
-        # Use user preferences
         if self.user.llm_preferences:
             pref_provider = self.user.llm_preferences.get("default_provider")
             if pref_provider:
@@ -179,18 +156,17 @@ class LLMGateway:
                         if m.startswith(f"{pref_provider}/"):
                             return m
 
-        # Default: Ollama in dev, first available in prod
         if settings.APP_ENV == "development":
             return "ollama/llama3.2"
 
         for m in prefs.get("quality", []) + prefs.get("budget", []):
             p = m.split("/")[0]
-            if self._has_key(p):
+            if self._has_system_key(p):
                 return m
 
         return "ollama/llama3.2"
 
-    def _has_key(self, provider: str) -> bool:
+    def _has_system_key(self, provider: str) -> bool:
         key_map = {
             "openai": settings.OPENAI_API_KEY,
             "anthropic": settings.ANTHROPIC_API_KEY,
@@ -198,72 +174,79 @@ class LLMGateway:
             "mistral": settings.MISTRAL_API_KEY,
             "deepseek": settings.DEEPSEEK_API_KEY,
             "xai": settings.XAI_API_KEY,
+            "huggingface": settings.HUGGINGFACE_API_KEY,
         }
         return bool(key_map.get(provider, ""))
 
     async def execute(self, task_type: TaskType, input_data: dict[str, Any], provider: str | None = None, model: str | None = None) -> dict[str, Any]:
-        """Execute an AI task with retry and JSON validation."""
+        """Execute AI task with user key injection, retry, and JSON validation."""
         resolved_model = self._resolve_model(task_type, provider, model)
         system_prompt = TASK_PROMPTS.get(task_type, "Return valid JSON.")
         user_message = json.dumps(input_data, indent=2)
 
+        # Resolve API key: user key > system key
+        provider_name = resolved_model.split("/")[0]
+        api_key = await self._get_user_api_key(provider_name)
+        # If no user key, litellm will use system-level keys configured globally
+
         start_time = time.time()
-        last_error = None
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_message},
+        ]
 
-        # Retry up to 3 times for JSON parsing failures
         for attempt in range(3):
             try:
-                response = await litellm.acompletion(
-                    model=resolved_model,
-                    messages=[
-                        {"role": "system", "content": system_prompt},
-                        {"role": "user", "content": user_message},
-                    ],
-                    response_format={"type": "json_object"},
-                    temperature=0.3,
-                    max_tokens=4096,
-                )
+                # Build kwargs — pass api_key if user has one
+                kwargs: dict[str, Any] = {
+                    "model": resolved_model,
+                    "messages": messages,
+                    "response_format": {"type": "json_object"},
+                    "temperature": 0.3,
+                    "max_tokens": 4096,
+                }
+                if api_key:
+                    kwargs["api_key"] = api_key
+
+                response = await litellm.acompletion(**kwargs)
 
                 content = response.choices[0].message.content
-                result = _parse_json_response(content)
+                result = _repair_json(content)
 
                 latency_ms = int((time.time() - start_time) * 1000)
-
-                # Log success
                 await self._log_request(
-                    provider=resolved_model.split("/")[0], model=resolved_model,
+                    provider=provider_name, model=resolved_model,
                     task_type=task_type.value,
                     tokens_input=response.usage.prompt_tokens if response.usage else None,
                     tokens_output=response.usage.completion_tokens if response.usage else None,
                     latency_ms=latency_ms, success=True,
                 )
-
                 return result
 
             except (ValueError, json.JSONDecodeError) as e:
-                last_error = e
-                logger.warning(f"JSON parse failed (attempt {attempt+1}/3) for {task_type}: {e}")
+                logger.warning(f"JSON parse fail (attempt {attempt+1}/3): {e}")
                 if attempt < 2:
-                    continue  # Retry
+                    # Add corrective message for retry
+                    messages.append({"role": "assistant", "content": content if 'content' in dir() else ""})
+                    messages.append({"role": "user", "content": "Your previous response was not valid JSON. Return ONLY a valid JSON object matching the schema above. No markdown, no explanation."})
+                    continue
+
             except Exception as e:
-                last_error = e
                 latency_ms = int((time.time() - start_time) * 1000)
                 await self._log_request(
-                    provider=resolved_model.split("/")[0], model=resolved_model,
+                    provider=provider_name, model=resolved_model,
                     task_type=task_type.value, latency_ms=latency_ms,
                     success=False, error_message=str(e)[:500],
                 )
                 raise
 
-        # All retries exhausted
-        raise ValueError(f"Failed to get valid JSON after 3 attempts: {last_error}")
+        raise ValueError(f"Failed to get valid JSON after 3 attempts for {task_type}")
 
     async def _log_request(self, **kwargs):
         log = LLMRequest(user_id=self.user.id, **kwargs)
         self.db.add(log)
 
     async def list_available_providers(self) -> list[dict]:
-        """List all configured providers."""
         return [
             {"name": "ollama", "display_name": "Ollama (Local)", "is_local": True, "is_available": True, "models": [{"id": "llama3.2", "name": "Llama 3.2"}]},
             {"name": "openai", "display_name": "OpenAI", "is_local": False, "is_available": bool(settings.OPENAI_API_KEY), "models": [{"id": "gpt-4o", "name": "GPT-4o"}, {"id": "gpt-4o-mini", "name": "GPT-4o Mini"}]},