Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
27 changes: 27 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 5 additions & 0 deletions backend/models/data_documents.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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
Expand Down
17 changes: 16 additions & 1 deletion backend/models/schemas.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from pydantic import BaseModel
from pydantic import BaseModel, Field


class AccountDetailsRequest(BaseModel):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
3 changes: 3 additions & 0 deletions backend/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@ psycopg2-binary
requests
pymongo
msal
langgraph
langchain-google-genai
langchain-core

# Testing dependencies
pytest
Expand Down
30 changes: 30 additions & 0 deletions backend/routes/data_documents.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from models.data_documents import (
DataDocumentsRequest,
DataDocumentsResponse,
DataDocumentsQueryResponse,
FindByIdRequest,
FindByIdResponse,
UpdateDocumentRequest,
Expand All @@ -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,
Expand Down Expand Up @@ -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(...)
Expand Down
134 changes: 126 additions & 8 deletions backend/routes/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand All @@ -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 "):
Expand All @@ -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,
)


Expand Down Expand Up @@ -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)


Expand All @@ -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,
)
Loading