feat: enable custom api key and database upload via frontend

Author: ridash2005

agent/retriever.py

@@ -30,42 +30,48 @@
 _collection: Optional[chromadb.Collection] = None
 
 
-def _get_collection() -> chromadb.Collection:
-    """Lazily initialise and cache the ChromaDB collection."""
-    global _client, _collection
-    if _collection is not None:
+_last_api_key: Optional[str] = None
+
+def _get_collection(api_key: str | None = None) -> chromadb.Collection:
+    """Lazily initialise and cache the ChromaDB collection. Reinitialize if api_key changes."""
+    global _client, _collection, _last_api_key
+    
+    current_key = api_key or os.getenv("OPENAI_API_KEY", "")
+    
+    if _collection is not None and _last_api_key == current_key:
         return _collection
 
     persist_dir = os.getenv("CHROMA_PERSIST_DIR", "./chroma_store")
     embedding_model = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
-    openai_api_key = os.getenv("OPENAI_API_KEY", "")
 
     _client = chromadb.PersistentClient(path=persist_dir)
     embedding_fn = OpenAIEmbeddingFunction(
-        api_key=openai_api_key,
+        api_key=current_key,
         model_name=embedding_model,
     )
     _collection = _client.get_collection(
         name=COLLECTION_NAME,
         embedding_function=embedding_fn,
     )
+    _last_api_key = current_key
     return _collection
 
 
-def get_relevant_schema(query: str, k: int = 3) -> str:
+def get_relevant_schema(query: str, k: int = 3, api_key: str | None = None) -> str:
     """
     Embed *query*, similarity-search ChromaDB, and return a formatted
     table+column block suitable for injection into an LLM prompt.
 
     Args:
         query: The user's natural-language question.
         k:     Number of most relevant tables to retrieve.
+        api_key: Dynamic OpenAI API Key.
 
     Returns:
         A newline-delimited schema block string.
     """
     try:
-        collection = _get_collection()
+        collection = _get_collection(api_key)
         collection_count = collection.count()
         if collection_count == 0:
             logger.debug("Schema collection is empty; returning empty context for query: %s", query)

agent/sql_chain.py

@@ -49,14 +49,21 @@
 _llm: ChatOpenAI | None = None
 
 
-def _get_llm() -> ChatOpenAI:
-    global _llm
-    if _llm is None:
-        _llm = ChatOpenAI(
-            model=os.getenv("OPENAI_MODEL", "gpt-4o"),
-            temperature=0,
-            api_key=os.getenv("OPENAI_API_KEY", ""),
-        )
+_last_llm_key: str | None = None
+
+def _get_llm(api_key: str | None = None) -> ChatOpenAI:
+    global _llm, _last_llm_key
+    current_key = api_key or os.getenv("OPENAI_API_KEY", "")
+    
+    if _llm is not None and _last_llm_key == current_key:
+        return _llm
+        
+    _llm = ChatOpenAI(
+        model=os.getenv("OPENAI_MODEL", "gpt-4o"),
+        temperature=0,
+        api_key=current_key,
+    )
+    _last_llm_key = current_key
     return _llm
 
 SYSTEM_PROMPT = """You are an expert SQL analyst for the Olist Brazilian E-Commerce database.
@@ -193,7 +200,7 @@ def _execute_sql(sql: str) -> list[dict[str, Any]]:
         return [dict(zip(columns, row)) for row in rows]
 
 
-async def run_query(question: str) -> dict[str, Any]:
+async def run_query(question: str, api_key: str | None = None) -> dict[str, Any]:
     """
     Full LCEL pipeline: natural-language question → SQL → executed results.
 
@@ -213,7 +220,7 @@ async def run_query(question: str) -> dict[str, Any]:
 
     try:
         # Step 1: Retrieve relevant schema snippets via RAG
-        schema_context = get_relevant_schema(question, k=3)
+        schema_context = get_relevant_schema(question, k=3, api_key=api_key)
 
         # Step 2: Load few-shot examples
         few_shot = _load_few_shot_examples()
@@ -227,7 +234,7 @@ async def run_query(question: str) -> dict[str, Any]:
         )
 
         # Step 4: Get cached LLM instance (temperature=0 for deterministic SQL)
-        llm = _get_llm()
+        llm = _get_llm(api_key=api_key)
 
         # Step 5: Build LCEL chain: prompt | llm | str parser
         chain = prompt | llm | StrOutputParser()

api/main.py

@@ -10,7 +10,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 
-from api.routes import health, query, schema
+from api.routes import health, query, schema, upload
 
 load_dotenv()
 
@@ -53,3 +53,4 @@ async def global_exception_handler(request: Request, exc: Exception) -> JSONResp
 app.include_router(query.router)
 app.include_router(schema.router)
 app.include_router(health.router)
+app.include_router(upload.router)

api/routes/query.py

@@ -6,7 +6,7 @@
 import logging
 from typing import Any
 
-from fastapi import APIRouter, HTTPException
+from fastapi import APIRouter, HTTPException, Header
 from pydantic import BaseModel, Field
 
 from agent.hitl_guard import check_sql
@@ -15,11 +15,9 @@
 logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api", tags=["query"])
 
-
 class QueryRequest(BaseModel):
     question: str = Field(..., min_length=1, description="Natural-language question to convert to SQL")
 
-
 class QueryResponse(BaseModel):
     sql: str
     results: list[dict[str, Any]]
@@ -28,26 +26,26 @@ class QueryResponse(BaseModel):
     approval_reason: str = ""
     latency_ms: int
 
-
 class ApproveRequest(BaseModel):
     sql: str = Field(..., min_length=1)
     approved: bool
 
-
 class ApproveResponse(BaseModel):
     executed: bool
     results: list[dict[str, Any]]
     message: str
 
-
 @router.post("/query", response_model=QueryResponse)
-async def query_endpoint(request: QueryRequest) -> QueryResponse:
+async def query_endpoint(
+    request: QueryRequest,
+    x_openai_key: str | None = Header(None)
+) -> QueryResponse:
     """
     Convert a natural-language question to SQL and execute it.
     Write operations require explicit approval via /api/approve.
     """
     try:
-        result = await run_query(request.question)
+        result = await run_query(request.question, api_key=x_openai_key)
         return QueryResponse(**result)
     except Exception as exc:
         logger.error("Query endpoint error: %s", exc, exc_info=True)

api/routes/upload.py

@@ -0,0 +1,67 @@
+import os
+import sqlite3
+import shutil
+import logging
+from fastapi import APIRouter, UploadFile, File, Header, HTTPException
+from model.database import set_engine_url
+from agent.semantic_layer import SEMANTIC_SCHEMA
+from agent.build_index import build_index
+
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+@router.post("/api/upload")
+async def upload_database(
+    file: UploadFile = File(...),
+    x_openai_key: str | None = Header(None)
+):
+    try:
+        # Require API Key to build the Chroma index
+        api_key = x_openai_key or os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            raise HTTPException(status_code=400, detail="OpenAI API Key is required to index the database.")
+            
+        # Temporarily set the env var so build_index uses it
+        os.environ["OPENAI_API_KEY"] = api_key
+
+        upload_dir = os.path.join(os.getcwd(), "data")
+        os.makedirs(upload_dir, exist_ok=True)
+        file_path = os.path.join(upload_dir, "uploaded.db")
+        
+        with open(file_path, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+            
+        # Update database engine
+        set_engine_url(f"sqlite:///{file_path}")
+        
+        # Connect and extract tables
+        conn = sqlite3.connect(file_path)
+        cursor = conn.cursor()
+        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+        tables = cursor.fetchall()
+        
+        new_schema = []
+        for table in tables:
+            table_name = table[0]
+            cursor.execute(f"PRAGMA table_info({table_name});")
+            columns = cursor.fetchall()
+            
+            new_schema.append({
+                "table_name": table_name,
+                "description": f"Custom uploaded table {table_name}",
+                "columns": [{"name": col[1], "description": f"Column {col[1]} of type {col[2]}"} for col in columns]
+            })
+            
+        conn.close()
+        
+        # Update global semantic schema
+        SEMANTIC_SCHEMA.clear()
+        SEMANTIC_SCHEMA.extend(new_schema)
+        
+        # Build ChromaDB index
+        build_index()
+        
+        return {"success": True, "message": f"Successfully loaded and indexed {len(new_schema)} tables."}
+    except Exception as e:
+        logger.error(f"Error processing upload: {e}")
+        raise HTTPException(status_code=500, detail=str(e))

frontend/src/api.ts

@@ -1,8 +1,18 @@
 import axios from 'axios'
 
 const savedUrl = localStorage.getItem('qm_api_url') || '/'
+const savedKey = localStorage.getItem('qm_openai_key') || ''
+
 const api = axios.create({ baseURL: savedUrl })
 
+api.interceptors.request.use((config) => {
+  const key = localStorage.getItem('qm_openai_key')
+  if (key && config.headers) {
+    config.headers['x-openai-key'] = key
+  }
+  return config
+})
+
 export function setApiUrl(url: string) {
   localStorage.setItem('qm_api_url', url)
   api.defaults.baseURL = url
@@ -12,6 +22,14 @@ export function getApiUrl(): string {
   return localStorage.getItem('qm_api_url') || ''
 }
 
+export function setApiKey(key: string) {
+  localStorage.setItem('qm_openai_key', key)
+}
+
+export function getApiKey(): string {
+  return localStorage.getItem('qm_openai_key') || ''
+}
+
 export interface QueryResponse {
   sql: string
   results: Record<string, unknown>[]
@@ -38,6 +56,29 @@ export interface SchemaTable {
   columns: SchemaColumn[]
 }
 
+export async function approveQuery(sql: string): Promise<QueryResponse> {
+  const { data } = await api.post<QueryResponse>('/api/approve', { sql })
+  return data
+}
+
+export async function uploadDatabase(file: File, onProgress?: (percent: number) => void): Promise<{ success: boolean; message: string }> {
+  const formData = new FormData()
+  formData.append('file', file)
+  
+  const { data } = await api.post('/api/upload', formData, {
+    headers: {
+      'Content-Type': 'multipart/form-data'
+    },
+    onUploadProgress: (progressEvent) => {
+      if (progressEvent.total && onProgress) {
+        const percentCompleted = Math.round((progressEvent.loaded * 100) / progressEvent.total)
+        onProgress(percentCompleted)
+      }
+    }
+  })
+  return data
+}
+
 export async function postQuery(question: string): Promise<QueryResponse> {
   const { data } = await api.post<QueryResponse>('/api/query', { question })
   return data