Upload 2 files

inference.py

@@ -19,12 +19,19 @@ from typing import Any, Dict, List, Optional
 
 import requests
 from openai import OpenAI
+try:
+    from dotenv import load_dotenv
+except Exception:  # pragma: no cover - optional dependency guard
+    load_dotenv = None
+
+if load_dotenv:
+    load_dotenv()
 
 # Configuration from environment
 MAX_STEPS = int(os.getenv("MAX_STEPS", "10"))
-API_BASE_URL = os.getenv("API_BASE_URL", "https://api.groq.com/openai/v1")
-MODEL_NAME = os.getenv("MODEL_NAME", "llama-3.3-70b-versatile")
-HF_TOKEN = os.getenv("HF_TOKEN", "")
+API_BASE_URL = os.getenv("API_BASE_URL", "<your-active-endpoint>")
+MODEL_NAME = os.getenv("MODEL_NAME", "<your-active-model>")
+HF_TOKEN = os.getenv("HF_TOKEN")
 API_KEY = os.getenv("GROQ_API_KEY", "") or HF_TOKEN or os.getenv("OPENAI_API_KEY", "")
 
 

server/main.py

@@ -12,6 +12,11 @@ from difflib import get_close_matches
 from fastapi import FastAPI, HTTPException, UploadFile, File, Form
 from fastapi.responses import HTMLResponse, JSONResponse
 
+try:
+    from dotenv import load_dotenv
+except Exception:  # pragma: no cover - optional dependency guard
+    load_dotenv = None
+
 from .environment import (
     SQLQualityEnvironment,
     analyze_uploaded_database,
@@ -31,6 +36,9 @@ from .models import (
     StepResponse,
 )
 
+if load_dotenv:
+    load_dotenv()
+
 app = FastAPI(
     title="SQL Quality Audit Agent Environment",
     description="OpenEnv-compatible environment for training AI agents on SQL data quality tasks",
@@ -79,6 +87,77 @@ def schema_table_map(schema: list[dict]) -> dict[str, dict]:
     return {table["table_name"]: table for table in schema}
 
 
+def _first_non_empty(*values: str) -> str:
+    """Return the first non-empty stripped string."""
+    for value in values:
+        candidate = str(value or "").strip()
+        if candidate:
+            return candidate
+    return ""
+
+
+def _resolve_llm_runtime() -> dict[str, Any]:
+    """Resolve OpenAI-compatible runtime config from environment variables."""
+    groq_key = _first_non_empty(os.getenv("GROQ_API_KEY", ""))
+    if groq_key:
+        return {
+            "configured": True,
+            "provider": "groq",
+            "api_key": groq_key,
+            "base_url": _first_non_empty(
+                os.getenv("GROQ_BASE_URL", ""),
+                os.getenv("API_BASE_URL", ""),
+                "https://api.groq.com/openai/v1",
+            ),
+            "model": _first_non_empty(
+                os.getenv("GROQ_MODEL", ""),
+                os.getenv("MODEL_NAME", ""),
+                "llama-3.3-70b-versatile",
+            ),
+            "missing": [],
+        }
+
+    hf_token = _first_non_empty(os.getenv("HF_TOKEN", ""))
+    openai_key = _first_non_empty(os.getenv("OPENAI_API_KEY", ""), os.getenv("API_KEY", ""))
+    provider = "huggingface" if hf_token else ("openai" if openai_key else "none")
+    api_key = hf_token or openai_key
+
+    if provider == "openai":
+        base_url = _first_non_empty(
+            os.getenv("OPENAI_BASE_URL", ""),
+            os.getenv("API_BASE_URL", ""),
+            "https://api.openai.com/v1",
+        )
+    else:
+        base_url = _first_non_empty(os.getenv("API_BASE_URL", ""), os.getenv("OPENAI_BASE_URL", ""))
+
+    model_name = _first_non_empty(os.getenv("MODEL_NAME", ""), os.getenv("OPENAI_MODEL", ""))
+
+    missing: list[str] = []
+    if not api_key:
+        missing.append("HF_TOKEN or OPENAI_API_KEY")
+    if not base_url:
+        missing.append("API_BASE_URL")
+    if not model_name:
+        missing.append("MODEL_NAME")
+
+    return {
+        "configured": bool(api_key and base_url and model_name),
+        "provider": provider,
+        "api_key": api_key,
+        "base_url": base_url,
+        "model": model_name,
+        "missing": missing,
+    }
+
+
+def _format_llm_missing_message(missing: list[str]) -> str:
+    """Build a user-friendly missing-config message."""
+    if not missing:
+        return "LLM configuration incomplete."
+    return f"Missing LLM configuration: {', '.join(missing)}"
+
+
 def load_static_html() -> str:
     """Load the static HTML interface for the App section."""
     # Look for static/index.html in various locations
@@ -122,34 +201,21 @@ async def health() -> dict:
     Returns:
         dict: Status response with server and LLM connection status
     """
-    # Check LLM availability - MOCK MODE ENABLED by default for offline operation
-    llm_available = False
-    mock_mode = True  # Default to mock mode so app works without HF_TOKEN
-    
-    conn = None
-    try:
-        import os
-        hf_token = os.getenv("HF_TOKEN", "")
-        api_base_url = os.getenv("API_BASE_URL", "")
-        model_name = os.getenv("MODEL_NAME", "")
-        
-        # Only require LLM if explicitly configured
-        if hf_token and api_base_url and model_name:
-            llm_available = True
-            mock_mode = False
-        else:
-            # Mock mode: app works without LLM for basic SQL operations
-            llm_available = True  # Report as available so UI doesn't hang
-            mock_mode = True
-    except Exception:
-        pass
-    
+    runtime = _resolve_llm_runtime()
+    llm_available = bool(runtime["configured"])
+    mock_mode = not llm_available
+    if llm_available:
+        message = "LLM Connected"
+    else:
+        message = _format_llm_missing_message(runtime["missing"])
+
     return {
         "status": "ok",
         "llm": "mock_mode" if mock_mode else "connected",
         "llm_available": llm_available,
         "mock_mode": mock_mode,
-        "message": "Mock Mode: LLM features disabled. Basic SQL query and task functionality available." if mock_mode else "LLM Connected"
+        "provider": runtime["provider"],
+        "message": message,
     }
 
 
@@ -163,50 +229,41 @@ async def status() -> dict:
     Returns:
         dict: Status response with llm_connected boolean
     """
-    llm_connected = False  # Default to disconnected until proven otherwise
+    runtime = _resolve_llm_runtime()
+    mock_mode = not bool(runtime["configured"])
+    llm_connected = bool(runtime["configured"])
     llm_error = None
-    mock_mode = False  # Default to real mode, determine based on env vars
-    
-    try:
-        import os
-        hf_token = os.getenv("HF_TOKEN", "")
-        api_base_url = os.getenv("API_BASE_URL", "")
-        model_name = os.getenv("MODEL_NAME", "")
-        
-        if hf_token and api_base_url and model_name:
-            # Real LLM mode - attempt connection test
-            mock_mode = False
-            try:
-                import requests
-                headers = {"Authorization": f"Bearer {hf_token}"}
-                # Simple model info request (lightweight)
-                response = requests.get(
-                    f"{api_base_url}/models",
-                    headers=headers,
-                    timeout=5
-                )
-                llm_connected = response.status_code == 200
-            except Exception as e:
-                llm_error = str(e)
-                # Fallback to env check only
-                llm_connected = True  # Assume connected if env vars present
-        else:
-            # No LLM configured - report as disconnected
-            mock_mode = True
-            llm_connected = False
-            llm_error = "HF_TOKEN, API_BASE_URL, and MODEL_NAME required for LLM"
-    except Exception as e:
-        llm_error = str(e)
-        # In case of error, report as disconnected
-        llm_connected = False
-        mock_mode = True
-    
+
+    if llm_connected:
+        try:
+            import requests
+
+            headers = {"Authorization": f"Bearer {runtime['api_key']}"}
+            probe_url = f"{runtime['base_url'].rstrip('/')}/models"
+            response = requests.get(probe_url, headers=headers, timeout=5)
+            if response.status_code == 200:
+                llm_connected = True
+            elif response.status_code in {401, 403}:
+                llm_connected = False
+                llm_error = f"LLM auth failed (HTTP {response.status_code})"
+            else:
+                # Some providers do not expose /models; keep configured-as-connected.
+                llm_connected = True
+                llm_error = f"LLM heartbeat warning (HTTP {response.status_code})"
+        except Exception as e:
+            # Keep configured-as-connected to avoid false negatives from transient network errors.
+            llm_connected = True
+            llm_error = f"LLM heartbeat skipped: {e}"
+    else:
+        llm_error = _format_llm_missing_message(runtime["missing"])
+
     return {
         "status": "online",
         "llm": "mock_mode" if mock_mode else ("connected" if llm_connected else "disconnected"),
         "llm_connected": llm_connected,
         "llm_error": llm_error,
-        "mock_mode": mock_mode
+        "mock_mode": mock_mode,
+        "provider": runtime["provider"],
     }
 
 
@@ -1631,33 +1688,19 @@ async def custom_db_status() -> dict:
 
 def _get_ai_client():
     """Get an OpenAI-compatible client for AI insights.
-    
-    Priority:
-    1. GROQ_API_KEY → Groq API (fast, reliable)
-    2. HF_TOKEN + API_BASE_URL → HuggingFace Inference
-    3. None → fallback to local heuristics
     """
-    groq_key = os.getenv("GROQ_API_KEY", "")
-    if groq_key:
+    runtime = _resolve_llm_runtime()
+    if runtime["configured"]:
         try:
             from openai import OpenAI
+
             return OpenAI(
-                api_key=groq_key,
-                base_url="https://api.groq.com/openai/v1"
-            ), os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
+                api_key=runtime["api_key"],
+                base_url=runtime["base_url"],
+            ), runtime["model"]
         except Exception:
             pass
-    
-    hf_token = os.getenv("HF_TOKEN", "")
-    api_base = os.getenv("API_BASE_URL", "")
-    model_name = os.getenv("MODEL_NAME", "")
-    if hf_token and api_base and model_name:
-        try:
-            from openai import OpenAI
-            return OpenAI(api_key=hf_token, base_url=api_base), model_name
-        except Exception:
-            pass
-    
+
     return None, None