Merge pull request #35 from ChingEnLin/worktree-multiple_model

feat: add dynamic model selection for LLM queries

Author: ChingEnLin

backend/models/analyze.py

@@ -4,6 +4,7 @@
 
 class AnalyzeRequest(BaseModel):
     query_result: List[Dict[str, Any]]
+    model: str = "gemini-2.5-flash"
 
 
 class AnalyzeResponse(BaseModel):

backend/models/schemas.py

@@ -39,6 +39,7 @@ class QueryPrompt(BaseModel):
     max_iterations: int = Field(
         default=3, ge=1, le=10
     )  # Server-enforced agent iteration cap
+    model: str = "gemini-2.5-flash"
 
 
 class GeneratedCode(BaseModel):
@@ -54,6 +55,7 @@ class ExecuteInput(BaseModel):
 class DebugQueryRequest(BaseModel):
     query: str
     error_message: str
+    model: str = "gemini-2.5-flash"
 
 
 class DebugSuggestionResponse(BaseModel):
@@ -64,6 +66,7 @@ class SchemaRelationshipsRequest(BaseModel):
     account_id: str
     database_name: str
     collection_names: list[str]
+    model: str = "gemini-2.5-flash"
 
 
 class Relationship(BaseModel):
@@ -85,6 +88,7 @@ class EvaluateWriteRequest(BaseModel):
     write_result: dict
     account_id: str
     database_name: str
+    model: str = "gemini-2.5-flash"
 
 
 class EvaluateWriteResponse(BaseModel):

backend/routes/audit.py

@@ -10,6 +10,7 @@
 
 class AuditQueryRequest(BaseModel):
     question: str
+    model: str = "gemini-2.5-flash"
 
 
 class AuditQueryResponse(BaseModel):
@@ -35,7 +36,7 @@ def query_audit_log(
     if not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Invalid token format")
 
-    response = process_audit_question(body.question)
+    response = process_audit_question(body.question, model=body.model)
     return AuditQueryResponse(**response)
 
 

backend/routes/query.py

@@ -1,4 +1,5 @@
 from fastapi import APIRouter, Header, Body, HTTPException
+from typing import List
 from services.azure_auth import exchange_token_obo
 from services.azure_cosmos_resources import (
     get_connection_string,
@@ -36,6 +37,7 @@
 )
 import ast
 import re
+from google import genai
 
 router = APIRouter()
 
@@ -107,6 +109,7 @@ def nl2query(prompt: QueryPrompt = Body(...), authorization: str = Header(...)):
         intermediate_context=prompt.intermediate_context,
         connection_string=connection_string,
         max_iterations=prompt.max_iterations,
+        model=prompt.model,
     )
 
 
@@ -129,7 +132,7 @@ def infer_relationships(
             collection_filter=request.collection_names,
         )
 
-        return generate_schema_relationships(schema_summary)
+        return generate_schema_relationships(schema_summary, model=request.model)
 
     except Exception as e:
         print(f"Error inferring relationships: {e}")
@@ -202,15 +205,17 @@ def debug(body: DebugQueryRequest = Body(...)):
     """
     Sends a failed query and error message to Gemini for debugging suggestion.
     """
-    return generate_suggestion_from_query_error(body.query, body.error_message)
+    return generate_suggestion_from_query_error(
+        body.query, body.error_message, model=body.model
+    )
 
 
 @router.post("/analyze", response_model=AnalyzeResponse)
 def analyze(body: AnalyzeRequest = Body(...)):
     """
     Sends a query result to the AI for analysis and visualization suggestions.
     """
-    return analyze_query_result(body.query_result)
+    return analyze_query_result(body.query_result, model=body.model)
 
 
 @router.post("/evaluate-write", response_model=EvaluateWriteResponse)
@@ -240,4 +245,35 @@ def evaluate_write(
         write_result=body.write_result,
         connection_string=connection_string,
         database_name=body.database_name,
+        model=body.model,
     )
+
+
+@router.get("/models", response_model=List[str])
+def list_models():
+    """
+    Returns available Gemini model IDs filtered to generative models.
+    Intentionally unauthenticated — model names are non-sensitive and
+    this endpoint is called on page load before auth completes.
+    """
+    try:
+        client = genai.Client()
+        models = list(client.models.list())
+        model_ids = [
+            m.name.replace("models/", "")
+            for m in models
+            if m.name
+            and "gemini" in m.name.lower()
+            and hasattr(m, "supported_actions")
+            and "generateContent" in (m.supported_actions or [])
+        ]
+        if not model_ids:
+            # Fallback: return all gemini models if supported_actions is absent
+            model_ids = [
+                m.name.replace("models/", "")
+                for m in models
+                if m.name and "gemini" in m.name.lower()
+            ]
+        return sorted(set(model_ids))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to list models: {str(e)}")

backend/services/analyze_service.py

@@ -4,7 +4,9 @@
 from models.analyze import AnalyzeResponse
 
 
-def analyze_query_result(query_result: list[dict]) -> AnalyzeResponse:
+def analyze_query_result(
+    query_result: list[dict], model: str = "gemini-2.5-flash"
+) -> AnalyzeResponse:
     prompt = f"""
 You are a data analyst assistant. Given the following MongoDB query result, provide:
 1. A concise textual insight or summary of the data.
@@ -18,7 +20,7 @@ def analyze_query_result(query_result: list[dict]) -> AnalyzeResponse:
 """
     client = genai.Client()
     response = client.models.generate_content(
-        model="gemini-2.5-flash",
+        model=model,
         contents=prompt,
         config=types.GenerateContentConfig(
             thinking_config=types.ThinkingConfig(thinking_budget=0)

backend/services/audit_service.py

@@ -66,14 +66,16 @@ def get_recent_activity(user_email: str, limit: int = 10) -> List[Dict[str, Any]
         return []
 
 
-def process_audit_question(question: str) -> Dict[str, Any]:
+def process_audit_question(
+    question: str, model: str = "gemini-2.5-flash"
+) -> Dict[str, Any]:
     """
     Orchestrates the process of answering a user's audit question:
     1. Generate SQL from NL question (via Gemini)
     2. Execute SQL
     3. Summarize results (via Gemini)
     """
-    sql_query = generate_audit_sql(question)
+    sql_query = generate_audit_sql(question, model=model)
 
     # If the generator returned an error query or invalid SQL, return it
     if "Error:" in sql_query:
@@ -93,7 +95,9 @@ def process_audit_question(question: str) -> Dict[str, Any]:
             "summary": f"Error executing query: {results[0]['error']}",
         }
 
-    summary_response = summarize_audit_results(question, sql_query, results)
+    summary_response = summarize_audit_results(
+        question, sql_query, results, model=model
+    )
 
     return {
         "sql_query": sql_query,

backend/services/evaluate_write_service.py

@@ -11,6 +11,7 @@ def evaluate_write_result(
     write_result: dict,
     connection_string: str = "",
     database_name: str = "",
+    model: str = "gemini-2.5-flash",
 ) -> EvaluateWriteResponse:
     prompt = f"""
 You are an expert MongoDB database administrator and assistant.
@@ -63,7 +64,7 @@ def query_database(query: str) -> str:
     tools = [query_database] if connection_string else None
 
     response = client.models.generate_content(
-        model="gemini-2.5-flash",
+        model=model,
         contents=prompt,
         config=types.GenerateContentConfig(
             tools=tools, thinking_config=types.ThinkingConfig(thinking_budget=0)
@@ -84,7 +85,7 @@ def query_database(query: str) -> str:
 
                 # Send the tool result back to the model
                 response = client.models.generate_content(
-                    model="gemini-2.5-flash",
+                    model=model,
                     contents=[
                         prompt,
                         response.candidates[0].content,

backend/services/gemini_service.py

@@ -119,6 +119,7 @@ def generate_query_from_prompt(
     collection_context: CollectionContext = None,
     intermediate_context: dict = None,
     all_collections_schema: str = "",
+    model: str = "gemini-2.5-flash",
 ) -> GeneratedCode:
     # Prune intermediate_context to remove image/large data
     safe_intermediate_context = (
@@ -136,7 +137,7 @@ def generate_query_from_prompt(
     )
     client = genai.Client()
     response = client.models.generate_content(
-        model="gemini-2.5-flash",
+        model=model,
         contents=full_prompt,
         config=types.GenerateContentConfig(
             thinking_config=types.ThinkingConfig(thinking_budget=0)  # Disables thinking
@@ -146,14 +147,16 @@ def generate_query_from_prompt(
     return GeneratedCode(generated_code=code)
 
 
-def generate_suggestion_from_query_error(query: str, error_message: str) -> str:
+def generate_suggestion_from_query_error(
+    query: str, error_message: str, model: str = "gemini-2.5-flash"
+) -> str:
     """
     Sends a failed query and error message to Gemini for debugging suggestion.
     """
     full_prompt = PROMPT_TEMPLATE_DEBUG.format(query=query, error_message=error_message)
     client = genai.Client()
     response = client.models.generate_content(
-        model="gemini-2.5-flash",
+        model=model,
         contents=full_prompt,
         config=types.GenerateContentConfig(
             thinking_config=types.ThinkingConfig(thinking_budget=0)  # Disables thinking
@@ -202,11 +205,11 @@ def generate_suggestion_from_query_error(query: str, error_message: str) -> str:
 """
 
 
-def generate_audit_sql(user_input: str) -> str:
+def generate_audit_sql(user_input: str, model: str = "gemini-2.5-flash") -> str:
     full_prompt = PROMPT_TEMPLATE_AUDIT_SQL.format(user_input=user_input)
     client = genai.Client()
     response = client.models.generate_content(
-        model="gemini-2.5-flash",
+        model=model,
         contents=full_prompt,
         config=types.GenerateContentConfig(
             thinking_config=types.ThinkingConfig(thinking_budget=0)
@@ -220,7 +223,7 @@ def generate_audit_sql(user_input: str) -> str:
 
 
 def summarize_audit_results(
-    user_input: str, sql_query: str, results: list
+    user_input: str, sql_query: str, results: list, model: str = "gemini-2.5-flash"
 ) -> AuditSummaryResponse:
     # Truncate results if too large to avoid token limits
     results_str = str(results)[:10000]
@@ -229,7 +232,7 @@ def summarize_audit_results(
     )
     client = genai.Client()
     response = client.models.generate_content(
-        model="gemini-2.5-flash",
+        model=model,
         contents=full_prompt,
         config=types.GenerateContentConfig(
             response_mime_type="application/json",
@@ -281,13 +284,15 @@ def summarize_audit_results(
 """
 
 
-def generate_schema_relationships(schema_summary: str) -> SchemaRelationshipsResponse:
+def generate_schema_relationships(
+    schema_summary: str, model: str = "gemini-2.5-flash"
+) -> SchemaRelationshipsResponse:
     from models.schemas import SchemaRelationshipsResponse
 
     full_prompt = PROMPT_TEMPLATE_RELATIONSHIPS.format(schema_summary=schema_summary)
     client = genai.Client()
     response = client.models.generate_content(
-        model="gemini-2.5-flash",
+        model=model,
         contents=full_prompt,
         config=types.GenerateContentConfig(
             response_mime_type="application/json",

backend/services/react_agent_service.py

@@ -56,6 +56,7 @@ class AgentState(TypedDict):
     schema_context: str
     intermediate_context: dict
     connection_string: str
+    model: str
 
     generated_query: str
     is_write_action: bool
@@ -127,7 +128,7 @@ def generate_query_node(state: AgentState):
 
     try:
         response = client.models.generate_content(
-            model="gemini-2.5-flash",
+            model=state.get("model", "gemini-2.5-flash"),
             contents=prompt,
             config=types.GenerateContentConfig(
                 thinking_config=types.ThinkingConfig(thinking_budget=0)
@@ -210,7 +211,7 @@ def evaluate_node(state: AgentState):
 
     try:
         response = client.models.generate_content(
-            model="gemini-2.5-flash",
+            model=state.get("model", "gemini-2.5-flash"),
             contents=prompt,
             config=types.GenerateContentConfig(
                 response_mime_type="application/json",
@@ -271,6 +272,7 @@ def run_query_generator(
     intermediate_context: dict,
     connection_string: str,
     max_iterations: int = 3,
+    model: str = "gemini-2.5-flash",
 ):
     logger.info(f"Starting ReAct Agent workflow for input: '{user_input}'")
     initial_state = {
@@ -280,6 +282,7 @@ def run_query_generator(
         "schema_context": schema_context,
         "intermediate_context": intermediate_context,
         "connection_string": connection_string,
+        "model": model,
         "iterations": 0,
         "max_iterations": max_iterations,
     }