LlamaIndex: Add MCP server example

Author: amotl

topic/machine-learning/llama-index/demo_mcp.py

@@ -0,0 +1,82 @@
+"""
+Use an LLM to query a database in human language via MCP.
+Example code using LlamaIndex with vanilla Open AI and Azure Open AI.
+
+https://github.com/run-llama/llama_index/tree/main/llama-index-integrations/tools/llama-index-tools-mcp
+
+## Start CrateDB MCP Server
+```
+export CRATEDB_CLUSTER_URL="http://localhost:4200/"
+cratedb-mcp serve --transport=streamable-http
+```
+
+## Usage
+```
+source env.standalone
+export OPENAI_API_KEY=sk-XJZ7pfog5Gp8Kus8D--invalid--0CJ5lyAKSefZLaV1Y9S1
+python demo_mcp.py
+```
+"""
+import asyncio
+from cratedb_about.instruction import Instructions
+
+from dotenv import load_dotenv
+from llama_index.core.agent.workflow import FunctionAgent
+from llama_index.llms.openai import OpenAI
+from llama_index.tools.mcp import BasicMCPClient, McpToolSpec
+
+from boot import configure_llm
+
+
+class Agent:
+
+    async def get_tools(self):
+        # Connect to the CrateDB MCP server using `streamable-http` transport.
+        mcp_client = BasicMCPClient("http://127.0.0.1:8000/mcp/")
+        mcp_tool_spec = McpToolSpec(
+            client=mcp_client,
+            # Optional: Filter the tools by name
+            # allowed_tools=["tool1", "tool2"],
+            # Optional: Include resources in the tool list
+            # include_resources=True,
+        )
+        return await mcp_tool_spec.to_tool_list_async()
+
+    async def get_agent(self):
+        return FunctionAgent(
+            name="Agent",
+            description="CrateDB text-to-SQL agent",
+            llm=OpenAI(model="gpt-4o"),
+            tools=await self.get_tools(),
+            system_prompt=Instructions.full(),
+        )
+
+    async def aquery(self, query):
+        return await (await self.get_agent()).run(query)
+
+    def query(self, query):
+        return asyncio.run(self.aquery(query))
+
+
+def main():
+    """
+    Use an LLM to query a database in human language.
+    """
+
+    # Configure application.
+    load_dotenv()
+    configure_llm()
+
+    # Use an agent that uses the CrateDB MCP server.
+    agent = Agent()
+
+    # Invoke an inquiry.
+    print("Running query")
+    QUERY_STR = "What is the average value for sensor 1?"
+    answer = agent.query(QUERY_STR)
+    print("Query was:", QUERY_STR)
+    print("Answer was:", answer)
+
+
+if __name__ == "__main__":
+    main()

topic/machine-learning/llama-index/requirements.txt

@@ -1,7 +1,9 @@
+cratedb-about @ git+https://github.com/crate/about.git@instructions
 langchain-openai<0.4
 llama-index-embeddings-langchain<0.4
 llama-index-embeddings-openai<0.4
 llama-index-llms-azure-openai<0.4
 llama-index-llms-openai<0.5
+llama-index-tools-mcp<0.3
 python-dotenv
 sqlalchemy-cratedb

topic/machine-learning/llama-index/test.py

@@ -38,3 +38,21 @@ def test_nlsql(cratedb, capsys):
     # Verify the outcome.
     out = capsys.readouterr().out
     assert "Answer was: The average value for sensor 1 is approximately 17.03." in out
+
+
+def test_mcp(cratedb, capsys):
+    """
+    Execute `demo_mcp.py` and verify outcome.
+    """
+
+    # Load the standalone configuration also for software testing.
+    # On CI, `OPENAI_API_KEY` will need to be supplied externally.
+    load_dotenv("env.standalone")
+
+    # Invoke the workload, in-process.
+    from demo_mcp import main
+    main()
+
+    # Verify the outcome.
+    out = capsys.readouterr().out
+    assert "Answer was: The average value for sensor 1 is approximately 17.03." in out