Consultar MongoDB con lenguaje natural usando LangChain y LangGraph

export OPENAI_API_KEY=<your-openai-api-key>

1
import os
2
from langchain_openai import ChatOpenAI
3
from langgraph.prebuilt import create_react_agent
4
from langchain_mongodb.agent_toolkit import (
5
    MONGODB_AGENT_SYSTEM_PROMPT,
6
    MongoDBDatabase,
7
    MongoDBDatabaseToolkit,
8
)
9
10
MONGODB_URI = '<connection-string>'
11
DB_NAME = 'sample_restaurants'
12
NATURAL_LANGUAGE_QUERY = 'Find all restaurants that serve hamburgers.'
13
14
class NaturalLanguageToMQL:
15
    def __init__(self):
16
        self.llm = ChatOpenAI(model="gpt-4o-mini", timeout=60)
17
        self.system_message = MONGODB_AGENT_SYSTEM_PROMPT.format(top_k=5)
18
        self.db_wrapper = MongoDBDatabase.from_connection_string(
19
                            MONGODB_URI, 
20
                            database=DB_NAME)
21
        self.toolkit = MongoDBDatabaseToolkit(db=self.db_wrapper, llm=self.llm)
22
        self.agent = create_react_agent(
23
                        self.llm, 
24
                        self.toolkit.get_tools(), 
25
                        state_modifier=self.system_message)
26
        self.messages = []
27
28
    def convert_to_mql_and_execute_query(self, query):
29
        # Start the agent with the agent.stream() method
30
        events = self.agent.stream(
31
            {"messages": [("user", query)]},
32
            stream_mode="values",
33
        )
34
        # Add output (events) from the agent to the self.messages list
35
        for event in events:
36
            self.messages.extend(event["messages"])
37
    
38
    def print_results(self):
39
        # Print the the end-user's expected output from 
40
        # the final message produced by the agent.
41
        print(self.messages[-1].content)
42
43
def main():
44
    converter = NaturalLanguageToMQL()
45
    converter.convert_to_mql_and_execute_query(NATURAL_LANGUAGE_QUERY)
46
    converter.print_results()
47
48
if __name__ == '__main__':
49
    main()