debug

libs/langchain-mongodb/langchain_mongodb/agent_toolkit/database.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import json
-import re
 from datetime import date, datetime
 from typing import Any, Dict, Iterable, List, Optional, Union
 
@@ -205,23 +204,8 @@ def _parse_command(self, command: str) -> Any:
         if command.endswith("]"):
             command += ")"
         agg_command = command[command.index("[") : -1]
-        tokens = re.split("([\[{,:}\]])", agg_command)
-        result = ""
-        patt = re.compile('[-\d"]')
-        markers = set("{,:}[]")
-        for token in tokens:
-            if not token:
-                continue
-            if token in markers:
-                result += token
-            elif token.startswith("'"):
-                result += f'"{token[1:-1]}"'
-            elif re.match(patt, token[0]):
-                result += token
-            else:
-                result += f'"{token}"'
         try:
-            return json.loads(result)
+            return json.loads(agg_command)
         except Exception as e:
             raise ValueError(f"Cannot execute command {command}") from e
 

libs/langchain-mongodb/langchain_mongodb/agent_toolkit/prompt.py

@@ -2,7 +2,7 @@
 
 
 MONGODB_AGENT_SYSTEM_PROMPT = """You are an agent designed to interact with a MongoDB database.
-Given an input question, create a syntactically correct MongoDB aggregation pipeline query to run, then look at the results of the query and return the answer.
+Given an input question, create a syntactically correct MongoDB query to run, then look at the results of the query and return the answer.
 Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most {top_k} results.
 You can order the results by a relevant field to return the most interesting examples in the database.
 Never query for all the fields from a specific collection, only ask for the relevant fields given the question.
@@ -13,7 +13,13 @@
 
 DO NOT make any update, insert, or delete operations.
 
-Only pass the contents of the aggregation, no surrounding syntax.
+The query MUST include the collection name and the contents of the aggregation pipeline.
+
+An example query looks like:
+
+```python
+db.Invoice.aggregate([ {{ "$group": {{ _id: "$BillingCountry", "totalSpent": {{ "$sum": "$Total" }} }} }}, {{ "$sort": {{ "totalSpent": -1 }} }}, {{ "$limit": 5 }} ])
+```
 
 To start you should ALWAYS look at the collections in the database to see what you can query.
 Do NOT skip this step.
@@ -35,6 +41,7 @@
 - Missing content in the aggegregation pipeline
 - Improperly quoting identifiers
 - Improperly quoting operators
+- The content in the aggregation pipeline is not valid JSON
 
 If there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.
 

libs/langchain-mongodb/tests/integration_tests/test_mongodb_database.py

@@ -153,18 +153,3 @@ def test_database_run(db: MongoDBDatabase) -> None:
     del docs[0]["_id"]
     del user["_id"]
     assert docs[0] == user
-
-
-CMDS = [
-    """db.Invoice.aggregate([ { $group: { _id: "$BillingCountry", totalSpent: { $sum: "$Total" } } }, { $sort: { totalSpent: -1 } }, { $limit: 5 } ])""",
-    """db.Invoice.aggregate([{ $group: { _id: '$BillingCountry', totalSpent: { $sum: '$Total' } } }, { $sort: { totalSpent: -1 } }, { $limit: 5 }])""",
-    """db.Invoice.aggregate([{$group: {_id: "$BillingCountry", totalSpent: {$sum: "$Total"}}}, {$sort: {totalSpent: -1}}, {$limit: 5}])""",
-    """db.Invoice.aggregate([ { "$group": { "_id": "$BillingCountry", "totalSpent": { "$sum": "$Total" } } }, { "$sort": { "totalSpent": -1 } }, { "$limit": 5 } ])""",
-    """db.InvoiceLine.aggregate([ { $group: { _id: "$InvoiceId", totalSpent: { $sum: { $multiply: [ "$UnitPrice", "$Quantity" ] }} }}, { $lookup: { from: "Customer", localField: "_id", foreignField: "CustomerId", as: "customerInfo" }}, { $unwind: "$customerInfo" }, { $group: { _id: "$customerInfo.Country", totalSpent: { $sum: "$totalSpent" } }}, { $sort: { totalSpent: -1 }}, { $limit: 5 } ])""",
-    """db.Invoice.aggregate([ { $group: { _id: "$BillingCountry", totalSpent: { $sum: "$Total" } } }, { $sort: { totalSpent: -1 } }, { $limit: 5 } ])""",
-]
-
-
-def test_database_parse_command(db: MongoDBDatabase) -> None:
-    for cmd in CMDS:
-        db._parse_command(cmd)