Merge pull request #28 from ChingEnLin/dev

Release

Author: ChingEnLin

.github/workflows/ci.yml

@@ -50,6 +50,8 @@ jobs:
         
     - name: Run tests with pytest
       working-directory: ./backend
+      env:
+        GEMINI_API_KEY: "dummy-key-for-tests"
       run: |
         PYTHONPATH=. pytest --cov=. --cov-report=term-missing --cov-report=xml --verbose
         

README.md

@@ -128,6 +128,33 @@ QueryPal follows a secure **Backend-for-Frontend (BFF)** pattern with enterprise
 └─────────────────────┘
 ```
 
+### Autonomous Query Generation with ReAct Agent
+
+QueryPal employs a powerful LangGraph-based ReAct (Reasoning and Acting) agent to autonomously generate, sandbox-test, and evaluate MongoDB queries. This ensures unparalleled accuracy while maintaining strict safety boundaries for write operations.
+
+```mermaid
+graph TD
+    A[User Intent] -->|Natural Language| B(Generate Query)
+    B --> C{Write Operation?}
+    C -- Yes --> D[Return Query to User]
+    D --> E((Manual Review & Run))
+    E --> F[Execute Write]
+    F --> G[Log to Audit Database]
+    F --> H[Evaluate with AI]
+    
+    C -- No (Read Only) --> I(Sandbox Execution)
+    I --> J{Evaluate Result}
+    J -- Success --> K[Return Final Query & Data]
+    J -- Failure/Error --> B
+    
+    subgraph LangGraph Agent Loop
+        B
+        C
+        I
+        J
+    end
+```
+
 **Security Features:**
 - ✅ **Zero-Trust Architecture**: No secrets stored in frontend
 - 🔐 **Token-Based Authentication**: MSAL with automatic token refresh

backend/models/data_documents.py

@@ -6,6 +6,7 @@ class DataDocumentsFilter(BaseModel):
     key: str
     value: Any = ""  # Can be str, int, etc. depending on the field type
     operator: Optional[str] = "equals"  # 'equals', 'exists', 'not_exists'
+    type: Optional[str] = "string"
 
 
 class DataDocumentsRequest(BaseModel):
@@ -18,6 +19,10 @@ class DataDocumentsRequest(BaseModel):
     filters: Optional[List[DataDocumentsFilter]] = None
 
 
+class DataDocumentsQueryResponse(BaseModel):
+    query_code: str
+
+
 class DataDocumentsResponse(BaseModel):
     documents: List[Dict[str, Any]]
     currentPage: int

backend/models/schemas.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class AccountDetailsRequest(BaseModel):
@@ -36,6 +36,9 @@ class QueryPrompt(BaseModel):
     intermediate_context: object | None = (
         None  # Optional intermediate context for complex queries
     )
+    max_iterations: int = Field(
+        default=3, ge=1, le=10
+    )  # Server-enforced agent iteration cap
 
 
 class GeneratedCode(BaseModel):
@@ -74,3 +77,15 @@ class Relationship(BaseModel):
 
 class SchemaRelationshipsResponse(BaseModel):
     relationships: list[Relationship]
+
+
+class EvaluateWriteRequest(BaseModel):
+    user_intent: str
+    query_code: str
+    write_result: dict
+    account_id: str
+    database_name: str
+
+
+class EvaluateWriteResponse(BaseModel):
+    evaluation: str

backend/requirements.txt

@@ -8,6 +8,9 @@ psycopg2-binary
 requests
 pymongo
 msal
+langgraph
+langchain-google-genai
+langchain-core
 
 # Testing dependencies
 pytest

backend/routes/data_documents.py

@@ -4,6 +4,7 @@
 from models.data_documents import (
     DataDocumentsRequest,
     DataDocumentsResponse,
+    DataDocumentsQueryResponse,
     FindByIdRequest,
     FindByIdResponse,
     UpdateDocumentRequest,
@@ -17,6 +18,8 @@
 from services.data_documents_service import (
     find_document_by_id,
     fetch_documents,
+    build_mongo_query,
+    generate_mongo_query_string,
     update_document,
     get_single_document,
     insert_document,
@@ -50,6 +53,33 @@ def get_documents(
     )
 
 
+@router.post("/documents/query_code", response_model=DataDocumentsQueryResponse)
+def get_documents_query_code(
+    body: DataDocumentsRequest = Body(...), authorization: str = Header(...)
+):
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid token format")
+    user_token = authorization.replace("Bearer ", "")
+    access_token = exchange_token_obo(user_token)
+    connection_string = get_connection_string(body.account_id, access_token)
+
+    # Needs to initialize MongoClient briefly to access find_one if 'all' keys are used in the algorithm
+    from pymongo import MongoClient
+
+    client = MongoClient(connection_string)
+    db = client[body.database_name]
+    collection = db[body.collection_name]
+
+    query = build_mongo_query(
+        collection=collection,
+        filter=body.filter.model_dump() if body.filter else None,
+        filters=[f.model_dump() for f in body.filters] if body.filters else None,
+    )
+
+    query_str = generate_mongo_query_string(body.collection_name, query)
+    return DataDocumentsQueryResponse(query_code=query_str)
+
+
 @router.put("/documents", response_model=dict)
 def put_update_document(
     body: UpdateDocumentRequest = Body(...), authorization: str = Header(...)

backend/routes/query.py

@@ -10,9 +10,11 @@
     DebugSuggestionResponse,
     SchemaRelationshipsRequest,
     SchemaRelationshipsResponse,
+    EvaluateWriteRequest,
+    EvaluateWriteResponse,
 )
+from services.react_agent_service import run_query_generator
 from services.gemini_service import (
-    generate_query_from_prompt,
     generate_suggestion_from_query_error,
     generate_schema_relationships,
 )
@@ -23,10 +25,46 @@
 )
 from models.analyze import AnalyzeRequest, AnalyzeResponse
 from services.analyze_service import analyze_query_result
+from services.evaluate_write_service import evaluate_write_result
+from services.data_documents_service import log_write_operation
+from services.user_queries_service import get_user_id_from_token
+from pymongo.results import (
+    UpdateResult,
+    InsertOneResult,
+    InsertManyResult,
+    DeleteResult,
+)
+import ast
+import re
 
 router = APIRouter()
 
 
+def extract_collection_name(query_str: str) -> str:
+    """Extract collection name from a PyMongo query string using AST."""
+    try:
+        tree = ast.parse(query_str)
+        for node in ast.walk(tree):
+            # Match: db["collection"] or db['collection']
+            if (
+                isinstance(node, ast.Subscript)
+                and isinstance(node.value, ast.Name)
+                and node.value.id == "db"
+            ):
+                if isinstance(node.slice, ast.Constant):
+                    return str(node.slice.value)
+            # Match: db.collection
+            if (
+                isinstance(node, ast.Attribute)
+                and isinstance(node.value, ast.Name)
+                and node.value.id == "db"
+            ):
+                return node.attr
+    except SyntaxError:
+        pass
+    return "unknown"
+
+
 @router.post("/nl2query")
 def nl2query(prompt: QueryPrompt = Body(...), authorization: str = Header(...)):
     if not authorization.startswith("Bearer "):
@@ -51,18 +89,24 @@ def nl2query(prompt: QueryPrompt = Body(...), authorization: str = Header(...)):
             schema_summary = get_database_schema_summary(
                 prompt.account_id, prompt.db_context.name, access_token
             )
+        connection_string = get_connection_string(prompt.account_id, access_token)
     except Exception as e:
-        print(f"Error fetching schema context: {e}")
+        print(
+            f"[ERROR] Failed to fetch schema/connection for account {prompt.account_id}: {e}"
+        )
         schema_summary = "Could not fetch schema summary."
+        connection_string = ""
 
     collections = [col.name for col in prompt.db_context.collections]
-    return generate_query_from_prompt(
-        prompt.user_input,
-        collections,
-        prompt.db_context.name,
-        collection_context=None,
+
+    return run_query_generator(
+        user_input=prompt.user_input,
+        database=prompt.db_context.name,
+        collections=collections,
+        schema_context=schema_summary,
         intermediate_context=prompt.intermediate_context,
-        all_collections_schema=schema_summary,
+        connection_string=connection_string,
+        max_iterations=prompt.max_iterations,
     )
 
 
@@ -106,6 +150,50 @@ def execute(query: ExecuteInput = Body(...), authorization: str = Header(...)):
             status_code=500,
             detail=f"MongoDB query error: {result['error']} ({result['exception_type']})",
         )
+
+    # LOG WRITE OPERATIONS
+    if isinstance(
+        result, (UpdateResult, InsertOneResult, InsertManyResult, DeleteResult)
+    ):
+        try:
+            user_email = get_user_id_from_token(user_token)
+
+            operation = "query_generator"
+            if isinstance(result, UpdateResult):
+                operation = "update"
+            elif isinstance(result, (InsertOneResult, InsertManyResult)):
+                operation = "insert"
+            elif isinstance(result, DeleteResult):
+                operation = "delete"
+
+            collection = extract_collection_name(query.query)
+
+            match = re.search(r"//([^:@]+)", connection_string)
+            account_name = match.group(1) if match else "unknown"
+            account_database = f"{account_name}.{query.database_name}"
+
+            transformed_res = transform_mongo_result(result)
+            query_info = {
+                "query": query.query,
+                "source": "query_generator",
+                "result": transformed_res,
+            }
+
+            before_data = query_info if operation == "delete" else None
+            after_data = query_info if operation in ["update", "insert"] else None
+
+            log_write_operation(
+                user_email=user_email,
+                operation=operation,
+                database_name=account_database,
+                collection_name=collection,
+                document_id="query_generator",
+                before_data=before_data,
+                after_data=after_data,
+            )
+        except Exception as e:
+            print(f"Error logging query generator write: {e}")
+
     return transform_mongo_result(result)
 
 
@@ -123,3 +211,33 @@ def analyze(body: AnalyzeRequest = Body(...)):
     Sends a query result to the AI for analysis and visualization suggestions.
     """
     return analyze_query_result(body.query_result)
+
+
+@router.post("/evaluate-write", response_model=EvaluateWriteResponse)
+def evaluate_write(
+    body: EvaluateWriteRequest = Body(...), authorization: str = Header(...)
+):
+    """
+    Evaluates the result of a write operation against the user's original intent.
+    Requires authentication via Bearer token.
+    """
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid token format")
+
+    user_token = authorization.replace("Bearer ", "")
+    access_token = exchange_token_obo(user_token)
+    try:
+        connection_string = get_connection_string(body.account_id, access_token)
+    except Exception as e:
+        print(
+            f"[ERROR] Failed to fetch connection string for account {body.account_id}: {e}"
+        )
+        connection_string = ""
+
+    return evaluate_write_result(
+        user_intent=body.user_intent,
+        query_code=body.query_code,
+        write_result=body.write_result,
+        connection_string=connection_string,
+        database_name=body.database_name,
+    )