Building a MongoDB Query Chatbot with Azure OpenAI Agent

In this article, I’ll Walk you through how I’ve built a chatbot using Azure OpenAI models and agents to interact with MongoDB and convert user queries into MongoDB queries. This setup leverages the power of large language models (LLMs) to interpret user input and generate a valid MongoDB query in response.

Prerequisites

Before diving into the code, here’s what you need:

  • Azure OpenAI Subscription: You’ll need an Azure OpenAI subscription to use the LLM models, including GPT-4 and others.
  • MongoDB: A MongoDB instance to store your data and perform queries on.
  • Streamlit: For building the user interface of the chatbot.

Overview of the System

The main components of this system include:

  • Azure OpenAI API: Used to process user input and convert it into a MongoDB query.
  • MongoDB: Where the chatbot queries data stored in collections.
  • Streamlit: For providing an interface where users can enter queries and see results.

Step-by-Step Implementation

  • Setting Up Environment Variables
    First, we load all necessary credentials and settings from environment variables. This includes your OpenAI endpoint, deployment name, and MongoDB connection string.
import os import dotenv 
dotenv.load_dotenv()  
endpoint = os.getenv("ENDPOINT_URL", "https://aoai-moti-lab.openai.azure.com/")
deployment = os.getenv("DEPLOYMENT_NAME", "gpt-4o-mini") 
subscription_key = os.getenv("AZURE_OPENAI_API_KEY", "")
  • Creating the User Interface
    I use Streamlit to create a simple interface where users can input search queries and interact with the chatbot.
import streamlit as st 
st.title("Azure OpenAI Mongo Agent Search")
query = st.text_input("Enter search query:")
  • MongoDB Setup
    We connect to MongoDB using pymongo, which allows us to query data and handle the results. The script assumes an existing MongoDB collection, like invoices.
import pymongo  
CONNECTION_STRING = os.getenv("MONGO_CONNECTION_STRING") 
DB_NAME = os.getenv("DB_NAME") 
client = pymongo.MongoClient(CONNECTION_STRING) 
db = client[DB_NAME] 
ms_invoice_collection = db["invoices"]
  • Creating the Chatbot Agent
    With the Azure AIProjectClient, I create a new agent that processes user queries. The agent uses the CodeInterpreterTool, which is a pre-built tool in Azure AI to assist in interpreting code and running relevant queries.
from azure.ai.projects import AIProjectClient 
from azure.ai.projects.models import CodeInterpreterTool 
from azure.identity import DefaultAzureCredential  

code_interpreter = CodeInterpreterTool() 
project_client = AIProjectClient.from_connection_string(credential=DefaultAzureCredential(), 
      conn_str=os.environ["PROJECT_CONNECTION_STRING"] )

In this code, the mongo-agent is created with a set of instructions to convert a user’s query into a pymongo query.

Processing the Query
The user’s query is processed by the agent, which generates Python code for querying MongoDB and outputs it in the form of a .py file. This file is then executed to retrieve the results.

mongo_agent = project_client.agents.create_agent(
  model="gpt-4o-2",
  name="mongo-agent",
  instructions=f"""
                   You are a helpful assistant that can convert a user's search query into a pymongo query     
                   to search for a document in the {ms_invoice_collection} collection.   
                """,     
  tools=code_interpreter.definitions,
  tool_resources=code_interpreter.resources, )
  • Running the Generated Code
    After the query is processed, the generated code is saved as a Python file (main.py), which is then executed to retrieve the result.
code_path.write_text(code) os.system(f"python3.10 {file_path._str}")
  • Returning the Results
    The results of the MongoDB query are saved in a result.json file, which is then read by the script and displayed to the user.
with open("result.json", "r") as f:
     result = f.read()
  • Chatbot Output
    Finally, after processing the data, the chatbot generates a human-readable summary of the results, making it easy for the user to understand the output.
completion_result = completion.choices[0].message.content st.write(completion_result)
  • completion_result = completion.choices[0].message.content st.write(completion_result)

Key Benefits

Easy Integration: The system can be easily expanded or integrated with other tools, such as different databases or services.

Human-like Interaction: By leveraging Azure OpenAI, the chatbot interprets user queries in a human-like manner and generates appropriate MongoDB queries.

Automation: The process automates the task of converting user queries into MongoDB queries, saving time and reducing errors.

Written by Moti Malka, Head of Cloud Solutions, UBTECH.