Harnessing Natural Language for MongoDB Queries With Google Gemini

1
import functions_framework
2
import os
3
import json
4
from pymongo import MongoClient
5
from bson import ObjectId
6
import traceback
7
from datetime import datetime
8
9
10
def connect_to_mongodb():
11
    client = MongoClient("<MongoDB Connection String>")
12
    return client
13
14
15
def success_response(body):
16
    return {
17
        'statusCode': '200',
18
        'body': json.dumps(body, cls=DateTimeEncoder),
19
        'headers': {
20
            'Content-Type': 'application/json',
21
        },
22
    }
23
24
25
def error_response(err):
26
    error_message = str(err)
27
    return {
28
        'statusCode': '400',
29
        'body': error_message,
30
        'headers': {
31
            'Content-Type': 'application/json',
32
        },
33
    }
34
35
36
# Used to convert datetime object(s) to string
37
class DateTimeEncoder(json.JSONEncoder):
38
    def default(self, o):
39
        if isinstance(o, datetime):
40
            return o.isoformat()
41
        return super().default(o)
42
43
44
@functions_framework.http
45
def mongodb_crud(request):
46
    client = connect_to_mongodb()
47
    payload = request.get_json(silent=True)
48
    db, coll = payload['database'], payload['collection']
49
    request_args = request.args
50
    op = request.path
51
    try:
52
        if op == "/findOne":
53
            filter_op = payload['filter'] if 'filter' in payload else {}
54
            projection = payload['projection'] if 'projection' in payload else {}
55
            result = {"document": client[db][coll].find_one(filter_op, projection)}
56
            if result['document'] is not None:
57
                if isinstance(result['document']['_id'], ObjectId):
58
                    result['document']['_id'] = str(result['document']['_id'])
59
60
        elif op == "/find":
61
            agg_query = []
62
63
            if 'filter' in payload and payload['filter'] != {}:
64
                agg_query.append({"$match": payload['filter']})
65
66
            if "sort" in payload and payload['sort'] != {}:
67
                agg_query.append({"$sort": payload['sort']})
68
69
            if "skip" in payload:
70
                agg_query.append({"$skip": payload['skip']})
71
72
            if 'limit' in payload:
73
                agg_query.append({"$limit": payload['limit']})
74
75
            if "projection" in payload and payload['projection'] != {}:
76
                agg_query.append({"$project": payload['projection']})
77
78
            result = {"documents": list(client[db][coll].aggregate(agg_query))}
79
            for obj in result['documents']:
80
                if isinstance(obj['_id'], ObjectId):
81
                    obj['_id'] = str(obj['_id'])
82
83
        elif op == "/insertOne":
84
            if "document" not in payload or payload['document'] == {}:
85
                return error_response("Send a document to insert")
86
            insert_op = client[db][coll].insert_one(payload['document'])
87
            result = {"insertedId": str(insert_op.inserted_id)}
88
89
        elif op == "/insertMany":
90
            if "documents" not in payload or payload['documents'] == {}:
91
                return error_response("Send a document to insert")
92
            insert_op = client[db][coll].insert_many(payload['documents'])
93
            result = {"insertedIds": [str(_id) for _id in insert_op.inserted_ids]}
94
95
        elif op in ["/updateOne", "/updateMany"]:
96
            payload['upsert'] = payload['upsert'] if 'upsert' in payload else False
97
            if "_id" in payload['filter']:
98
                payload['filter']['_id'] = ObjectId(payload['filter']['_id'])
99
            if op == "/updateOne":
100
                update_op = client[db][coll].update_one(payload['filter'], payload['update'], upsert=payload['upsert'])
101
            else:
102
                update_op = client[db][coll].update_many(payload['filter'], payload['update'], upsert=payload['upsert'])
103
            result = {"matchedCount": update_op.matched_count, "modifiedCount": update_op.modified_count}
104
105
        elif op in ["/deleteOne", "/deleteMany"]:
106
            payload['filter'] = payload['filter'] if 'filter' in payload else {}
107
            if "_id" in payload['filter']:
108
                payload['filter']['_id'] = ObjectId(payload['filter']['_id'])
109
            if op == "/deleteOne":
110
                result = {"deletedCount": client[db][coll].delete_one(payload['filter']).deleted_count}
111
            else:
112
                result = {"deletedCount": client[db][coll].delete_many(payload['filter']).deleted_count}
113
114
        elif op == "/aggregate":
115
            if "pipeline" not in payload or payload['pipeline'] == []:
116
                return error_response("Send a pipeline")
117
            docs = list(client[db][coll].aggregate(payload['pipeline']))
118
            for obj in docs:
119
                if isinstance(obj['_id'], ObjectId):
120
                    obj['_id'] = str(obj['_id'])
121
            result = {"documents": docs}
122
123
        else:
124
            return error_response("Not a valid operation")
125
126
        return success_response(result)
127
128
    except Exception as e:
129
        print(traceback.format_exc())
130
        return error_response(e)
131
132
    finally:
133
        if client:
134
            client.close()

1
functions-framework==3.*
2
pymongo
3
pymongo[srv]

1
from google.colab import auth
2
auth.authenticate_user("GCP project id")
3
!gcloud config set project {"GCP project id"}

1
!pip install --force-reinstall --quiet google_cloud_aiplatform
2
# This is for printing the Vertex AI service account.
3
!pip install --force-reinstall --quiet langchain==0.0.298
4
5
!pip install --upgrade google-auth
6
!pip install bigframes==0.26.0

1
import IPython
2
app = IPython.Application.instance()
3
app.kernel.do_shutdown(True)  # Re-run the Env variable cell again after Kernel restart

1
import os
2
## This is just a sample values please replace accordingly to your project
3
# Setting up the GCP project
4
os.environ['PROJECT_ID'] = 'gcp project id'  # GCP Project ID
5
os.environ['REGION'] =  "us-central1" # Project Region
6
## GCS Bucket location
7
os.environ['STAGING_BUCKET'] =  "gs://vertexai_extensions"
8
## Extension Config
9
os.environ['EXTENSION_DISPLAY_HOME'] =  "MongoDb Vertex API Interpreter"
10
os.environ['EXTENSION_DESCRIPTION'] =  "This extension makes api call to mongodb to do all crud operations"
11
12
13
## OPEN API SPec config
14
os.environ['MANIFEST_NAME'] =  "mdb_crud_interpreter"
15
os.environ['MANIFEST_DESCRIPTION'] =  "This extension makes api call to mongodb to do all crud operations"
16
os.environ['OPENAPI_GCS_URI'] =  "gs://vertexai_extensions/mongodbopenapispec.yaml"
17
18
19
## API KEY secret location
20
os.environ['API_SECRET_LOCATION'] = "projects/787220387490/secrets/mdbapikey/versions/1"
21
22
23
##LLM config
24
os.environ['LLM_MODEL'] = "gemini-1.5-pro"

1
from google.cloud import aiplatformfrom google.cloud.aiplatform.private_preview import llm_extension
2
3
PROJECT_ID = os.environ['PROJECT_ID']
4
REGION = os.environ['REGION']
5
STAGING_BUCKET = os.environ['STAGING_BUCKET']
6
7
aiplatform.init(
8
    project=PROJECT_ID,
9
    location=REGION,
10
    staging_bucket=STAGING_BUCKET,
11
)

1
from google.cloud import aiplatform
2
from vertexai.preview import extensions
3
4
mdb_crud = extensions.Extension.create(
5
         display_name = EXTENSION_DISPLAY_HOME,
6
         # Optional.
7
         description = EXTENSION_DESCRIPTION,
8
         manifest = {
9
             "name": MANIFEST_NAME,
10
             "description": MANIFEST_DESCRIPTION,
11
             "api_spec": {
12
                 "open_api_gcs_uri": (
13
                     OPENAPI_GCS_URI
14
                 ),
15
             },
16
             "authConfig": {
17
                "authType": "OAUTH",
18
                "oauthConfig": {"service_account": "<Service account name here>"}
19
              },
20
         },
21
     )
22
mdb_crud

1
print("Name:", mdb_crud.gca_resource.name)print("Display Name:", mdb_crud.gca_resource.display_name)print("Description:", mdb_crud.gca_resource.description)
2
import pprint
3
pprint.pprint(mdb_crud.operation_schemas())

1
## Please replace accordingly to your project
2
## Operation Ids
3
os.environ['FIND_ONE_OP_ID'] = "findone_mdb"
4
5
6
7
8
## NL Queries
9
os.environ['FIND_ONE_NL_QUERY'] = "Find the release year of the movie 'A Corner in Wheat' from VertexAI-POC cluster, sample_mflix, movies"
10
11
12
## Mongodb Config
13
os.environ['DATA_SOURCE'] = "VertexAI-POC"
14
os.environ['DB_NAME'] = "sample_mflix"
15
os.environ['COLLECTION_NAME'] = "movies"
16
17
18
### Test data setup
19
os.environ['TITLE_FILTER_CLAUSE'] = "A Corner in Wheat"
20
21
from vertexai.preview.generative_models import GenerativeModel, Tool
22
23
fc_chat = GenerativeModel(os.environ['LLM_MODEL']).start_chat()
24
findOneResponse = fc_chat.send_message(os.environ['FIND_ONE_NL_QUERY'],
25
    tools=[Tool.from_dict({        "function_declarations": mdb_crud.operation_schemas()
26
   })],
27
)
28
29
findOneResponse

1
response = mdb_crud.execute(
2
    operation_id = findOneResponse.candidates[0].content.parts[0].function_call.name,
3
    operation_params = findOneResponse.candidates[0].content.parts[0].function_call.args
4
)
5
6
response

1
## This is just a sample values please replace accordingly to your project
2
## Operation Ids
3
os.environ['FIND_MANY_OP_ID'] = "findmany_mdb"
4
5
6
## NL Queries
7
os.environ['FIND_MANY_NL_QUERY'] = "give me movies released in year 1924 from VertexAI-POC cluster, sample_mflix, movies"
8
9
10
11
12
## Mongodb Config
13
os.environ['DATA_SOURCE'] = "VertexAI-POC"
14
os.environ['DB_NAME'] = "sample_mflix"
15
os.environ['COLLECTION_NAME'] = "movies"
16
os.environ['YEAR'] = 1924
17
fc_chat = GenerativeModel(os.environ['LLM_MODEL']).start_chat()
18
findmanyResponse = fc_chat.send_message(os.environ['FIND_MANY_NL_QUERY'],
19
    tools=[Tool.from_dict({        "function_declarations": mdb_crud.operation_schemas()
20
   })],
21
)
22
23
findmanyResponse

1
response = mdb_crud.execute(
2
    operation_id = findmanyResponse.candidates[0].content.parts[0].function_call.name,
3
    operation_params = findmanyResponse.candidates[0].content.parts[0].function_call.args
4
)
5
6
response