Merge pull request #37 from ChingEnLin/feat/cross-collection-lookup-prompt

Cosmos-safe $lookup prompting for cross-collection query generation

Author: ChingEnLin

backend/routes/query.py

@@ -23,6 +23,7 @@
     execute_mongo_query,
     transform_mongo_result,
     get_database_schema_summary,
+    SCHEMA_FETCH_FAILED,
 )
 from models.analyze import AnalyzeRequest, AnalyzeResponse
 from services.analyze_service import analyze_query_result
@@ -36,9 +37,12 @@
     DeleteResult,
 )
 import ast
+import logging
 import re
 from google import genai
 
+logger = logging.getLogger(__name__)
+
 router = APIRouter()
 
 
@@ -101,6 +105,26 @@ def nl2query(prompt: QueryPrompt = Body(...), authorization: str = Header(...)):
 
     collections = [col.name for col in prompt.db_context.collections]
 
+    relationship_context = ""
+    if (
+        len(collections) > 1
+        and schema_summary
+        and schema_summary != SCHEMA_FETCH_FAILED
+    ):
+        try:
+            rels = generate_schema_relationships(schema_summary, model=prompt.model)
+            if rels.relationships:
+                relationship_context = "\n".join(
+                    f"- {r.source_collection}.{r.source_field} -> "
+                    f"{r.target_collection}.{r.target_field} "
+                    f"(confidence={r.confidence:.2f}) — {r.description}"
+                    for r in rels.relationships
+                )
+        except Exception as e:
+            logger.warning(
+                "Relationship inference failed; continuing without it: %s", e
+            )
+
     return run_query_generator(
         user_input=prompt.user_input,
         database=prompt.db_context.name,
@@ -110,6 +134,7 @@ def nl2query(prompt: QueryPrompt = Body(...), authorization: str = Header(...)):
         connection_string=connection_string,
         max_iterations=prompt.max_iterations,
         model=prompt.model,
+        relationship_context=relationship_context,
     )
 
 

backend/services/mongo_service.py

@@ -1,4 +1,6 @@
 from bson import ObjectId
+
+SCHEMA_FETCH_FAILED = "Could not fetch schema summary."
 import pymongo
 from pymongo.results import (
     UpdateResult,
@@ -26,23 +28,23 @@ def execute_mongo_query(connection_string: str, database: str, query: str):
         return query_result
 
 
+def _convert_object_ids(value):
+    # Recursively stringify ObjectIds inside nested dicts/lists so FastAPI can
+    # serialize $lookup-joined documents (which carry nested _id ObjectIds).
+    if isinstance(value, ObjectId):
+        return str(value)
+    if isinstance(value, dict):
+        return {k: _convert_object_ids(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_convert_object_ids(v) for v in value]
+    return value
+
+
 def transform_mongo_result(result):
-    # If result is a list of dicts, convert ObjectIds
     if isinstance(result, list):
-        if result and isinstance(result[0], dict):
-            for doc in result:
-                for k, v in doc.items():
-                    if isinstance(v, ObjectId):
-                        doc[k] = str(v)
-            return result
-        else:
-            # List of primitives (e.g., from distinct)
-            return result
+        return _convert_object_ids(result)
     elif isinstance(result, dict):
-        for k, v in result.items():
-            if isinstance(v, ObjectId):
-                result[k] = str(v)
-        return result
+        return _convert_object_ids(result)
     elif isinstance(result, InsertOneResult):
         return {"inserted_id": str(result.inserted_id)}
     elif isinstance(result, InsertManyResult):
@@ -90,4 +92,4 @@ def get_database_schema_summary(
         return "\n\n".join(summary)
     except Exception as e:
         print(f"Error fetching schema summary: {e}")
-        return "Could not fetch schema summary."
+        return SCHEMA_FETCH_FAILED

backend/services/react_agent_service.py

@@ -54,6 +54,7 @@ class AgentState(TypedDict):
     database: str
     collections: list[str]
     schema_context: str
+    relationship_context: str
     intermediate_context: dict
     connection_string: str
     model: str
@@ -71,25 +72,92 @@ class AgentState(TypedDict):
 # Define LLM clients and prompts
 client = genai.Client()
 
+CROSS_COLLECTION_GUIDANCE = """
+
+Cross-Collection Join Guidance (MULTIPLE collections selected):
+- Target: Azure Cosmos DB (MongoDB API). IMPORTANT — Cosmos does NOT support `let` or `pipeline` inside `$lookup`. Only the plain `localField`/`foreignField` form works. Never emit `let` or a `pipeline` array inside `$lookup` — it will fail with `CommandNotSupported`.
+- Produce ONE aggregate() pipeline. Pick the most-filtering collection as the driver.
+- Use the Inferred Relationships block above to choose join fields. Do NOT invent field names that are not in the schema or relationships.
+- TYPE-MISMATCH HANDLING (this is the #1 reason $lookup returns empty arrays on Cosmos): if one side is an ObjectId and the other is a string, you MUST pre-convert with an `$addFields` stage BEFORE `$lookup`, then $lookup on the converted field. Do not rely on $lookup to coerce types.
+  - String → ObjectId: `{"$addFields": {"<field>_oid": {"$toObjectId": "$<field>"}}}` (or `$map` over an array of strings).
+  - ObjectId → String: `{"$addFields": {"<field>_str": {"$toString": "$<field>"}}}`.
+- After $lookup, $unwind the joined array (use `preserveNullAndEmptyArrays: True` for LEFT-JOIN semantics) before filtering on joined fields, OR filter with `joined.field` dotted notation.
+- Always include a $limit (<=50) unless the user explicitly asks for all rows.
+
+Example A — simple equality join (same field type on both sides):
+```python
+db['orders'].aggregate([
+    {"$match": {"status": "paid"}},
+    {"$lookup": {
+        "from": "users",
+        "localField": "userId",
+        "foreignField": "_id",
+        "as": "user"
+    }},
+    {"$unwind": {"path": "$user", "preserveNullAndEmptyArrays": True}},
+    {"$limit": 20}
+])
+```
+
+Example B — array-of-strings → ObjectId join (Cosmos-safe; pre-convert via $addFields):
+```python
+db['patient-cohort'].aggregate([
+    {"$addFields": {
+        "patient_oids": {"$map": {"input": "$patient_ids", "in": {"$toObjectId": "$$this"}}}
+    }},
+    {"$lookup": {
+        "from": "patient",
+        "localField": "patient_oids",
+        "foreignField": "_id",
+        "as": "patients_info"
+    }},
+    {"$match": {"patients_info.origin_ethnicity": "caucasian"}},
+    {"$limit": 50}
+])
+```
+
+Example C — single ObjectId → string join (Cosmos-safe; pre-convert the other side):
+```python
+db['orders'].aggregate([
+    {"$addFields": {"userId_str": {"$toString": "$userId"}}},
+    {"$lookup": {
+        "from": "users",
+        "localField": "userId_str",
+        "foreignField": "external_id",
+        "as": "user"
+    }},
+    {"$unwind": "$user"},
+    {"$limit": 20}
+])
+```
+"""
+
 GENERATE_PROMPT = """
-You are an expert MongoDB architect. 
+You are an expert MongoDB architect.
 Your task is to generate a PyMongo query based on the user's request.
 
 User Request: {user_input}
 Database: {database}
 Collections: {collections}
 Schema summary: {schema_context}
+Inferred Relationships (foreign keys between collections):
+{relationship_context}
 Intermediate Context (optional): {intermediate_context}
 
-Previous Evaluation Feedback (if this is a retry): 
+Previous Attempt (if this is a retry — DO NOT repeat the same query verbatim):
+{previous_query}
+
+Previous Evaluation Feedback (if this is a retry):
 {evaluation}
 
 Instructions:
-1. Write ONLY the PyMongo query code. 
+0. If a Previous Attempt is shown above, your new query MUST be materially different — change the join form, fields, types, or stages in response to the critique. Producing the same query (or a trivial reformat) is a failure. If the previous $lookup returned empty arrays for every input doc, the cause is almost always a type mismatch on the join key: add an `$addFields` stage BEFORE the `$lookup` that applies `$toObjectId` (string→OID) or `$toString` (OID→string), then $lookup on the converted field. NEVER use `let` or `pipeline` inside `$lookup` — Cosmos rejects both with CommandNotSupported.
+1. Write ONLY the PyMongo query code.
 2. Use variables appropriately (e.g., db['collection_name'].find(...) or db['collection_name'].aggregate(...)).
 3. Do not include markdown formatting or explanations, just return the code, but you can use ```python if you must.
 4. If the user is asking for a visualization, ensure the query retrieves the necessary data format.
 5. You are allowed to use both find() and aggregate() pipelines. If using aggregate(), be mindful of the resulting data size and include $limit stages if applicable.
+{cross_collection_guidance}
 """
 
 EVALUATE_PROMPT = """
@@ -105,6 +173,12 @@ class AgentState(TypedDict):
 Determine if this query successfully answers the user's request based on the code and the result.
 If there is an error in the query result, or if it clearly does not match the intent, it is NOT valid.
 If it is a write action, we cannot test the result, but you should evaluate if the code looks correct for the user's intent.
+If the query uses $lookup, specifically verify (target is Azure Cosmos DB MongoDB API):
+  - The query does NOT use `let` or `pipeline` inside `$lookup` — Cosmos rejects both with `CommandNotSupported`. If you see either, it is NOT valid; recommend pre-converting the type with an `$addFields` stage and then using plain `localField`/`foreignField`.
+  - The `localField`/`foreignField` reference fields that actually exist on each side.
+  - When the joined array is empty for every input doc, suspect a type mismatch (ObjectId vs string). The query is NOT valid; recommend an `$addFields` stage BEFORE `$lookup` that applies `$toObjectId` (string→OID) or `$toString` (OID→string), and a `$map` if the local field is an array of strings.
+  - Do not rationalize an empty result with "maybe no such records exist" when the filter is on a field inside the joined array — empty joined arrays will make the downstream $match drop everything; treat that as a join-correctness failure unless the join itself is clearly populated.
+  - The driving collection is the right one (smallest filtered set first).
 
 Respond in JSON format:
 {{
@@ -117,13 +191,22 @@ class AgentState(TypedDict):
 def generate_query_node(state: AgentState):
     logger.info(f"--- GENERATE NODE (Iteration {state.get('iterations', 0)}) ---")
 
+    is_multi_collection = len(state.get("collections", [])) > 1
+    is_retry = state.get("iterations", 0) > 0
+    previous_query = state.get("generated_query", "") if is_retry else ""
     prompt = GENERATE_PROMPT.format(
         user_input=state["user_input"],
         database=state["database"],
         collections=", ".join(state["collections"]),
         schema_context=state["schema_context"],
+        relationship_context=state.get("relationship_context")
+        or "None (single collection or no inferred relationships).",
         intermediate_context=state.get("intermediate_context", {}),
+        previous_query=previous_query or "None (this is the first attempt).",
         evaluation=state.get("evaluation", "None"),
+        cross_collection_guidance=(
+            CROSS_COLLECTION_GUIDANCE if is_multi_collection else ""
+        ),
     )
 
     try:
@@ -277,13 +360,15 @@ def run_query_generator(
     connection_string: str,
     max_iterations: int = 3,
     model: str = "gemini-2.5-flash",
+    relationship_context: str = "",
 ):
     logger.info(f"Starting ReAct Agent workflow for input: '{user_input}'")
     initial_state = {
         "user_input": user_input,
         "database": database,
         "collections": collections,
         "schema_context": schema_context,
+        "relationship_context": relationship_context,
         "intermediate_context": intermediate_context,
         "connection_string": connection_string,
         "model": model,