1  import asyncio
 2  import os
 3  
 4  from langchain.chat_models import init_chat_model
 5  from langchain_mcp_adapters.tools import load_mcp_tools
 6  from langgraph.prebuilt import create_react_agent
 7  from mcp import ClientSession
 8  from mcp.client.sse import sse_client
 9  
10  base_url = os.getenv("OPENAI_API_BASE_URL")
11  model = os.getenv("MODEL_NAME")
12  mcp_server_url = os.getenv("MCP_SERVER_URL")
13  api_key = os.getenv("OPENAI_API_KEY", "does_not_matter")
14  
15  system_prompt = """
16  You are an agent designed to interact with a SQL database.
17  Given an input question, create a syntactically correct {dialect} query to run,
18  then look at the results of the query and return the answer. Unless the user
19  specifies a specific number of examples they wish to obtain, always limit your
20  query to at most {top_k} results.
21  
22  You can order the results by a relevant column to return the most interesting
23  examples in the database. Never query for all the columns from a specific table,
24  only ask for the relevant columns given the question.
25  
26  You MUST double check your query before executing it. If you get an error while
27  executing a query, rewrite the query and try again.
28  
29  DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the
30  database.
31  
32  To start you should ALWAYS look at the tables in the database to see what you
33  can query. Do NOT skip this step.
34  
35  Then you should query the schema of the most relevant tables.
36  
37  For example, for PostgreSQL, you can use the following query to get the tables:
38  
39  SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';
40  
41  And to retrieve all columns of a specific table, you can use:
42  SELECT column_name, data_type FROM information_schema.columns WHERE table_name = 'your_table_name';
43  
44  """.format(
45      dialect=os.getenv("DATABASE_DIALECT"),
46      top_k=5,
47  )
48  
49  
50  async def main():
51      if mcp_server_url is None:
52          raise ValueError("Please set the MCP_SERVER_URL environment variable.")
53  
54      llm = init_chat_model(
55          model, model_provider="openai", api_key=api_key, base_url=base_url
56      )
57      async with sse_client(
58          url=mcp_server_url,
59          timeout=60,
60      ) as (read, write):
61          async with ClientSession(read, write) as session:
62              await session.initialize()
63  
64              tools = await load_mcp_tools(session)
65              print(f"MCP tools loaded: {tools}")
66              agent = create_react_agent(
67                  llm,
68                  tools=tools,
69                  prompt=system_prompt,
70              )
71  
72              question = os.getenv("QUESTION")
73              if not question:
74                  raise ValueError(
75                      "Please set the QUESTION environment variable with your question."
76                  )
77  
78              async for step in agent.astream(
79                  {"messages": [{"role": "user", "content": question}]},
80                  stream_mode="values",
81              ):
82                  step["messages"][-1].pretty_print()
83  
84  
85  asyncio.run(main())