diff --git a/.github/workflows/ml-llamaindex.yml b/.github/workflows/ml-llamaindex.yml index 67fa34e9..a5129a57 100644 --- a/.github/workflows/ml-llamaindex.yml +++ b/.github/workflows/ml-llamaindex.yml @@ -40,10 +40,15 @@ jobs: 'ubuntu-latest', ] python-version: [ - '3.8', + '3.10', '3.13', ] - cratedb-version: [ 'nightly' ] + cratedb-version: [ + 'nightly', + ] + cratedb-mcp-version: [ + 'main', + ] services: cratedb: @@ -53,6 +58,15 @@ jobs: - 5432:5432 env: CRATE_HEAP_SIZE: 4g + cratedb-mcp: + image: ghcr.io/crate/cratedb-mcp:${{ matrix.cratedb-mcp-version }} + ports: + - 8000:8000 + env: + CRATEDB_MCP_TRANSPORT: streamable-http + CRATEDB_MCP_HOST: 0.0.0.0 + CRATEDB_MCP_PORT: 8000 + CRATEDB_CLUSTER_URL: http://crate:crate@cratedb:4200/ env: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} diff --git a/topic/machine-learning/llama-index/README.md b/topic/machine-learning/llama-index/README.md index 2c287e8c..316afa42 100644 --- a/topic/machine-learning/llama-index/README.md +++ b/topic/machine-learning/llama-index/README.md @@ -1,12 +1,15 @@ -# Connecting CrateDB Data to an LLM with LlamaIndex and Azure OpenAI +# NL2SQL with LlamaIndex: Querying CrateDB using natural language -This folder contains the codebase for [this tutorial](https://community.cratedb.com/t/how-to-connect-your-cratedb-data-to-llm-with-llamaindex-and-azure-openai/1612) on the CrateDB community forum. You should read the tutorial for instructions on how to set up the components that you need on Azure, and use this README for setting up CrateDB and the Python code. +Connecting CrateDB to an LLM with LlamaIndex and Azure OpenAI, +optionally using MCP. See also the [LlamaIndex Text-to-SQL Guide]. -This has been tested using: +This folder contains the codebase for the tutorial +[How to connect your CrateDB data to LLM with LlamaIndex and Azure OpenAI] +on the CrateDB community forum. -* Python 3.12 -* macOS -* CrateDB 5.8 and higher +You should read the tutorial for instructions on how to set up the components +that you need on Azure, and use this README for setting up CrateDB and the +Python code. ## Database Setup @@ -57,7 +60,7 @@ VALUES Create and activate a virtual environment: -``` +```shell python3 -m venv .venv source .venv/bin/activate ``` @@ -81,7 +84,7 @@ OPENAI_AZURE_ENDPOINT=https:// EMBEDDING_MODEL_INSTANCE= -CRATEDB_SQLALCHEMY_URL="crate://:@:4200/?ssl=true" +CRATEDB_SQLALCHEMY_URL=crate://:@:4200/?ssl=true CRATEDB_TABLE_NAME=time_series_data ``` @@ -89,15 +92,17 @@ Save your changes. ## Run the Code -Run the code like so: +### NLSQL +[LlamaIndex's NLSQLTableQueryEngine] is a natural language SQL table query engine. + +Run the code like so: ```bash -python main.py +python demo_nlsql.py ``` Here's the expected output: - -``` +```text Creating SQLAlchemy engine... Connecting to CrateDB... Creating SQLDatabase instance... @@ -124,4 +129,32 @@ Answer was: The average value for sensor 1 is 17.033333333333335. 'avg(value)' ] } -``` \ No newline at end of file +``` + +### MCP + +Spin up the [CrateDB MCP server], connecting it to CrateDB on localhost. +```bash +export CRATEDB_CLUSTER_URL=http://crate:crate@localhost:4200/ +export CRATEDB_MCP_TRANSPORT=http +uvx cratedb-mcp serve +``` + +Run the code using OpenAI API: +```bash +export OPENAI_API_KEY= +python demo_mcp.py +``` +Expected output: +```text +Running query +Inquiring MCP server +Query was: What is the average value for sensor 1? +Answer was: The average value for sensor 1 is approximately 17.03. +``` + + +[CrateDB MCP server]: https://cratedb.com/docs/guide/integrate/mcp/cratedb-mcp.html +[How to connect your CrateDB data to LLM with LlamaIndex and Azure OpenAI]: https://community.cratedb.com/t/how-to-connect-your-cratedb-data-to-llm-with-llamaindex-and-azure-openai/1612 +[LlamaIndex's NLSQLTableQueryEngine]: https://docs.llamaindex.ai/en/stable/api_reference/query_engine/NL_SQL_table/ +[LlamaIndex Text-to-SQL Guide]: https://docs.llamaindex.ai/en/stable/examples/index_structs/struct_indices/SQLIndexDemo/ diff --git a/topic/machine-learning/llama-index/boot.py b/topic/machine-learning/llama-index/boot.py new file mode 100644 index 00000000..08e894cc --- /dev/null +++ b/topic/machine-learning/llama-index/boot.py @@ -0,0 +1,62 @@ +import os +from typing import Tuple + +import openai +import llama_index.core +from langchain_openai import AzureOpenAIEmbeddings +from langchain_openai import OpenAIEmbeddings +from llama_index.core.base.embeddings.base import BaseEmbedding +from llama_index.core.llms import LLM +from llama_index.llms.azure_openai import AzureOpenAI +from llama_index.llms.openai import OpenAI +from llama_index.embeddings.langchain import LangchainEmbedding + + +MODEL_NAME = "gpt-4o" + + +def configure_llm(debug: bool = False) -> Tuple[LLM, BaseEmbedding]: + """ + Configure LLM. Use either vanilla Open AI, or Azure Open AI. + """ + + openai.api_type = os.getenv("OPENAI_API_TYPE") + openai.azure_endpoint = os.getenv("OPENAI_AZURE_ENDPOINT") + openai.api_version = os.getenv("OPENAI_AZURE_API_VERSION") + openai.api_key = os.getenv("OPENAI_API_KEY") + + # https://docs.llamaindex.ai/en/stable/understanding/tracing_and_debugging/tracing_and_debugging/ + if debug: + llama_index.core.set_global_handler("simple") + + if openai.api_type == "openai": + llm = OpenAI( + model=MODEL_NAME, + temperature=0.0, + api_key=os.getenv("OPENAI_API_KEY"), + ) + elif openai.api_type == "azure": + llm = AzureOpenAI( + model=MODEL_NAME, + temperature=0.0, + engine=os.getenv("LLM_INSTANCE"), + azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"), + api_key = os.getenv("OPENAI_API_KEY"), + api_version = os.getenv("OPENAI_AZURE_API_VERSION"), + ) + else: + raise ValueError(f"Open AI API type not defined or invalid: {openai.api_type}") + + if openai.api_type == "openai": + embed_model = LangchainEmbedding(OpenAIEmbeddings(model=MODEL_NAME)) + elif openai.api_type == "azure": + embed_model = LangchainEmbedding( + AzureOpenAIEmbeddings( + azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"), + model=os.getenv("EMBEDDING_MODEL_INSTANCE") + ) + ) + else: + embed_model = None + + return llm, embed_model diff --git a/topic/machine-learning/llama-index/demo_mcp.py b/topic/machine-learning/llama-index/demo_mcp.py new file mode 100644 index 00000000..dd7b31b3 --- /dev/null +++ b/topic/machine-learning/llama-index/demo_mcp.py @@ -0,0 +1,89 @@ +""" +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=http +``` + +## Usage +``` +source env.standalone +export OPENAI_API_KEY=sk-XJZ7pfog5Gp8Kus8D--invalid--0CJ5lyAKSefZLaV1Y9S1 +python demo_mcp.py +``` +""" +import asyncio +import os + +from cratedb_about.instruction import GeneralInstructions + +from dotenv import load_dotenv +from llama_index.core.agent.workflow import FunctionAgent +from llama_index.core.llms import LLM +from llama_index.tools.mcp import BasicMCPClient, McpToolSpec + +from boot import configure_llm + + +class Agent: + + def __init__(self, llm: LLM): + self.llm = llm + + async def get_tools(self): + # Connect to the CrateDB MCP server using the `http` transport. + mcp_url = os.getenv("CRATEDB_MCP_URL", "http://127.0.0.1:8000/mcp/") + mcp_client = BasicMCPClient(mcp_url) + 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="DemoAgent", + description="CrateDB text-to-SQL agent", + llm=self.llm, + tools=await self.get_tools(), + system_prompt=GeneralInstructions().render(), + ) + + async def aquery(self, query): + return await (await self.get_agent()).run(query) + + def query(self, query): + print("Inquiring MCP server") + return asyncio.run(self.aquery(query)) + + +def main(): + """ + Use an LLM to query a database in human language. + """ + + # Configure application. + load_dotenv() + llm, embed_model = configure_llm() + + # Use an agent that uses the CrateDB MCP server. + agent = Agent(llm) + + # Invoke an inquiry. + print("Running query") + QUERY_STR = os.getenv("DEMO_QUERY", "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() diff --git a/topic/machine-learning/llama-index/demo_nlsql.py b/topic/machine-learning/llama-index/demo_nlsql.py new file mode 100644 index 00000000..043750af --- /dev/null +++ b/topic/machine-learning/llama-index/demo_nlsql.py @@ -0,0 +1,50 @@ +""" +Use an LLM to query a database in human language via NLSQLTableQueryEngine. +Example code using LlamaIndex with vanilla Open AI and Azure Open AI. +""" + +import os +import sqlalchemy as sa + +from dotenv import load_dotenv +from llama_index.core.utilities.sql_wrapper import SQLDatabase +from llama_index.core.query_engine import NLSQLTableQueryEngine + +from boot import configure_llm + + +def main(): + """ + Use an LLM to query a database in human language. + """ + + # Configure application. + load_dotenv() + llm, embed_model = configure_llm() + + # Configure database connection and query engine. + print("Connecting to CrateDB") + engine_crate = sa.create_engine(os.getenv("CRATEDB_SQLALCHEMY_URL")) + engine_crate.connect() + + print("Creating LlamaIndex QueryEngine") + sql_database = SQLDatabase(engine_crate, include_tables=[os.getenv("CRATEDB_TABLE_NAME")]) + query_engine = NLSQLTableQueryEngine( + sql_database=sql_database, + tables=[os.getenv("CRATEDB_TABLE_NAME")], + llm=llm, + embed_model=embed_model, + ) + + # Invoke an inquiry. + print("Running query") + QUERY_STR = os.getenv("DEMO_QUERY", "What is the average value for sensor 1?") + answer = query_engine.query(QUERY_STR) + print(answer.get_formatted_sources()) + print("Query was:", QUERY_STR) + print("Answer was:", answer) + print(answer.metadata) + + +if __name__ == "__main__": + main() diff --git a/topic/machine-learning/llama-index/env.azure b/topic/machine-learning/llama-index/env.azure index f97346ef..e0c5cf77 100644 --- a/topic/machine-learning/llama-index/env.azure +++ b/topic/machine-learning/llama-index/env.azure @@ -1,8 +1,8 @@ -OPENAI_API_KEY=TODO -OPENAI_API_TYPE=azure -OPENAI_AZURE_ENDPOINT=https://TODO.openai.azure.com -OPENAI_AZURE_API_VERSION=2024-08-01-preview -LLM_INSTANCE=TODO -EMBEDDING_MODEL_INSTANCE=TODO -CRATEDB_SQLALCHEMY_URL="crate://USER:PASSWORD@HOST:4200/?ssl=true" -CRATEDB_TABLE_NAME=time_series_data +export OPENAI_API_KEY=TODO +export OPENAI_API_TYPE=azure +export OPENAI_AZURE_ENDPOINT=https://TODO.openai.azure.com +export OPENAI_AZURE_API_VERSION=2024-08-01-preview +export LLM_INSTANCE=TODO +export EMBEDDING_MODEL_INSTANCE=TODO +export CRATEDB_SQLALCHEMY_URL="crate://USER:PASSWORD@HOST:4200/?ssl=true" +export CRATEDB_TABLE_NAME=time_series_data diff --git a/topic/machine-learning/llama-index/env.standalone b/topic/machine-learning/llama-index/env.standalone index 9ad450ef..19d4fa7b 100644 --- a/topic/machine-learning/llama-index/env.standalone +++ b/topic/machine-learning/llama-index/env.standalone @@ -1,4 +1,4 @@ # OPENAI_API_KEY=sk-XJZ7pfog5Gp8Kus8D--invalid--0CJ5lyAKSefZLaV1Y9S1 -OPENAI_API_TYPE=openai -CRATEDB_SQLALCHEMY_URL="crate://crate@localhost:4200/" -CRATEDB_TABLE_NAME=time_series_data +export OPENAI_API_TYPE=openai +export CRATEDB_SQLALCHEMY_URL="crate://crate@localhost:4200/" +export CRATEDB_TABLE_NAME=time_series_data diff --git a/topic/machine-learning/llama-index/main.py b/topic/machine-learning/llama-index/main.py deleted file mode 100644 index 077f67d5..00000000 --- a/topic/machine-learning/llama-index/main.py +++ /dev/null @@ -1,92 +0,0 @@ -""" -Use an LLM to query a database in human language. -Example code using LlamaIndex with vanilla Open AI and Azure Open AI. -""" - -import os -import openai -import sqlalchemy as sa - -from dotenv import load_dotenv -from langchain_openai import AzureOpenAIEmbeddings -from langchain_openai import OpenAIEmbeddings -from llama_index.llms.azure_openai import AzureOpenAI -from llama_index.llms.openai import OpenAI -from llama_index.embeddings.langchain import LangchainEmbedding -from llama_index.core.utilities.sql_wrapper import SQLDatabase -from llama_index.core.query_engine import NLSQLTableQueryEngine -from llama_index.core import Settings - - -def configure_llm(): - """ - Configure LLM. Use either vanilla Open AI, or Azure Open AI. - """ - - openai.api_type = os.getenv("OPENAI_API_TYPE") - openai.azure_endpoint = os.getenv("OPENAI_AZURE_ENDPOINT") - openai.api_version = os.getenv("OPENAI_AZURE_API_VERSION") - openai.api_key = os.getenv("OPENAI_API_KEY") - - if openai.api_type == "openai": - llm = OpenAI( - api_key=os.getenv("OPENAI_API_KEY"), - temperature=0.0 - ) - elif openai.api_type == "azure": - llm = AzureOpenAI( - engine=os.getenv("LLM_INSTANCE"), - azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"), - api_key = os.getenv("OPENAI_API_KEY"), - api_version = os.getenv("OPENAI_AZURE_API_VERSION"), - temperature=0.0 - ) - else: - raise ValueError(f"Open AI API type not defined or invalid: {openai.api_type}") - - Settings.llm = llm - if openai.api_type == "openai": - Settings.embed_model = LangchainEmbedding(OpenAIEmbeddings()) - elif openai.api_type == "azure": - Settings.embed_model = LangchainEmbedding( - AzureOpenAIEmbeddings( - azure_endpoint=os.getenv("OPENAI_AZURE_ENDPOINT"), - model=os.getenv("EMBEDDING_MODEL_INSTANCE") - ) - ) - - -def main(): - """ - Use an LLM to query a database in human language. - """ - - # Configure application. - load_dotenv() - configure_llm() - - # Configure database connection and query engine. - print("Connecting to CrateDB") - engine_crate = sa.create_engine(os.getenv("CRATEDB_SQLALCHEMY_URL")) - engine_crate.connect() - - print("Creating LlamaIndex QueryEngine") - sql_database = SQLDatabase(engine_crate, include_tables=[os.getenv("CRATEDB_TABLE_NAME")]) - query_engine = NLSQLTableQueryEngine( - sql_database=sql_database, - tables=[os.getenv("CRATEDB_TABLE_NAME")], - llm=Settings.llm - ) - - # Invoke an inquiry. - print("Running query") - QUERY_STR = "What is the average value for sensor 1?" - answer = query_engine.query(QUERY_STR) - print(answer.get_formatted_sources()) - print("Query was:", QUERY_STR) - print("Answer was:", answer) - print(answer.metadata) - - -if __name__ == "__main__": - main() diff --git a/topic/machine-learning/llama-index/requirements.txt b/topic/machine-learning/llama-index/requirements.txt index c4771397..7ff8171b 100644 --- a/topic/machine-learning/llama-index/requirements.txt +++ b/topic/machine-learning/llama-index/requirements.txt @@ -1,7 +1,9 @@ +cratedb-about==0.0.6 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 diff --git a/topic/machine-learning/llama-index/test.py b/topic/machine-learning/llama-index/test.py index 3c81566a..7991ae8c 100644 --- a/topic/machine-learning/llama-index/test.py +++ b/topic/machine-learning/llama-index/test.py @@ -22,9 +22,9 @@ def init_database(cratedb): cratedb.run_sql((HERE / "init.sql").read_text()) -def test_main(cratedb, capsys): +def test_nlsql(cratedb, capsys): """ - Execute `main.py` and verify outcome. + Execute `demo_nlsql.py` and verify outcome. """ # Load the standalone configuration also for software testing. @@ -32,7 +32,25 @@ def test_main(cratedb, capsys): load_dotenv("env.standalone") # Invoke the workload, in-process. - from main import main + from demo_nlsql 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 + + +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.