VectorDB (ChromaDB) partially added, README adapted

Author: SailingDirk

README.md

@@ -52,10 +52,12 @@ open source AI Desktops like Jan.ai.
 ## Features
 
 - Checks the relevance of a natural language question for a requested database schema
+- Checks the VectorDB for similar questions and give hints if feasible
 - Transforms the question into an SQL statement
 - Checks if the SQL statement is allowed; currently, we only allow read-only statements.
 - Executes the SQL statement
 - Checks if SQL statement is valid
+  - If a result set is returned, the question, the SQL statement, and some metadata is stored in a VectorDB
 - If required, rewrites the question
 - Generates a result
 
@@ -212,6 +214,7 @@ MCP_OPENAI_SERVER_URL=http://localhost:1234/v1
 MCP_OPENAI_SERVER_API_KEY=<API-Key of your LLM Server>
 MCP_OPENAI_SERVER_MODEL_NAME=<your selexted model>
 MCP_VECTORDB_FILE=/Users/dirk/Temp/mcp_exasol_t2s.vectordb
+MCP_VECTORDB_SIMILARITY_SEARCH_DISTANCE=0.3
 ```
 
 They secret key and the encrypted password shall be created with the 
@@ -230,7 +233,7 @@ with good quality. We have made some good experience with the
 - qwen/qwn-coder-30b 
 
 Large Language Model. In case you decide to use an AI Desktop where you can configure the  
-LLM to be utilized you need to check, if the LLM is trained for tool usage. The more parameters  
+LLM to be utilized, you need to check, if the LLM is trained for tool usage. The more parameters  
 the LLM features the higher is the performance requirement for a timely answer. A 70B parameter  
 LLM on an Apple Macbook Pro with M4MAX and 40 GPU cores and *LM-Studio* or *ollama* is already  
 consuming quite some time. Having a dedicated LLM server (on premise) is definitely a plus.

exasol/ai/mcp/server/text_to_sql.py

@@ -4,33 +4,28 @@
 ## Version 0.1 DirkB : Initial version             ##
 #####################################################
 
-import os
-import pyexasol
-#from mypy.state import state
-from pyexasol import ExaError
-
-import sys
+########################
+## Required Libraries ##
+########################
 
+import chromadb
+from chromadb.config import Settings
 from cryptography.fernet import Fernet
+from datetime import datetime
 from dotenv import load_dotenv
+from exasol.ai.mcp.server.load_prompt import load_prompt
+from exasol.ai.mcp.server.server_settings import ExaDbResult
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, START, END
+import os
 from pydantic import BaseModel, Field
-from typing_extensions import TypedDict
-
-from langchain_openai import ChatOpenAI
-from langchain_core.prompts import ChatPromptTemplate
-
-from sqlglot import (
-    exp,
-    parse_one,
-)
+import pyexasol
+from pyexasol import ExaError
+from sqlglot import exp, parse_one
 from sqlglot.errors import ParseError
-
-from exasol.ai.mcp.server.load_prompt import load_prompt
-from exasol.ai.mcp.server.server_settings import ExaDbResult
-
-
-
+import sys
+from typing_extensions import TypedDict
 
 
 ############################################################################
@@ -63,7 +58,7 @@ def get_environment() -> dict:
     load_dotenv()
 
     secret_key = os.getenv("MCP_SERVER_EXASOL_SECRET_KEY")
-    assert secret_key is not None, "Please set SECRET_KEY environment variable"
+    assert secret_key is not None, "Please set 'MCP_SERVER_EXASOL_SECRET_KEY' environment variable"
     fernet = Fernet(secret_key)
     stored_password = os.getenv("MCP_EXASOL_DATABASE_PASSWORD")
     db_password =  fernet.decrypt(stored_password).decode()
@@ -76,30 +71,30 @@ def get_environment() -> dict:
             "llm_server_api_token": os.getenv("MCP_OPENAI_SERVER_API_KEY"),
             "llm_server_model_check": os.getenv("MCP_OPENAI_SERVER_MODEL_NAME"),
             "llm_server_sql_transform": os.getenv("MCP_OPENAI_SERVER_MODEL_NAME"),
+            "vectordb_persistent_storage": os.getenv("MCP_VECTORDB_FILE"),
+            "vectordb_similarity_distance": os.getenv("MCP_VECTORDB_SIMILARITY_DISTANCE"),
         }
 
     return env
 
 
-
-
 #######################################################
 ## Working status of Text2SQL transformation process ##
 #######################################################
 
 class GraphState(TypedDict):
-    question: str
-    db_schema: str
-    sql_statement: str
-    query_num_rows: int
-    query_result: str
-    display_result: str
-    num_of_attempts: int
-    is_allowed: str
-    is_relevant: str
-    sql_is_valid: str
-    sql_error: str
-    info: str
+    question: str                 # The natural language question
+    db_schema: str                # The database schema to be used
+    sql_statement: str            # The generated SQL statement
+    query_num_rows: int           # The number of rows returned
+    query_result: str             # The result of the generated SQL statement
+    display_result: str           # The transformed result into a visual version
+    num_of_attempts: int          # The number of attempts to generate a valid SQL statement
+    is_allowed: str               # Is the generated SQL statement allowed (READ-ONLY, currently)
+    is_relevant: str              # Does the natural language fit to the underlying database schema
+    sql_is_valid: str             # SQL statements accepted by the Exasol database
+    sql_error: str                # The SQL error returned by the Exasol database, if any
+    info: str                     # Additional INFO field
 
 
 ############################################################
@@ -226,6 +221,28 @@ def t2s_human_language_to_sql(state: GraphState):
 
     system_prompt = load_prompt(db_schema=db_schema, schema=schema)
 
+    ##
+    ## Check VectorDB for a similar question and SQL Statement
+    ##
+
+    try:
+        vectordb_client = chromadb.PersistentClient(path=env['vectordb_persistent_storage'])
+        sql_collection = vectordb_client.get_or_create_collection(name="Questions_SQL_History")
+        tmp = sql_collection.query(query_texts=state['question'], n_results=1, include=["distances", "documents", "metadatas"])
+
+        print(f"VectorDB Result in T2S ::: {tmp}", file=sys.stderr)
+        print(f"Vector-DB-Result-Distance in T2S ::: {tmp["distances"][0][0]}", file=sys.stderr)
+
+        if float(tmp["distances"][0][0]) <= float(env['vectordb_similarity_distance']):
+            system_prompt += f"""
+                                For a similar natural language question you have created the following SQL statement:
+                                
+                                {tmp['metadatas'][0][0]['sql']}
+                                
+                            """
+    except Exception as e:
+        print(f"ChromaDB - Error: {e}", file=sys.stderr)
+
     print(f"Prompt: {system_prompt}", file=sys.stderr)
 
 
@@ -306,7 +323,7 @@ def t2s_check_sql_router(state: GraphState):
 
 def t2s_execute_query(state: GraphState):
 
-    print(f"#### Beginning of SQL Execution ::: state['sql_statement']", file=sys.stderr)
+    print(f"#### Beginning of SQL Execution -01::: {state['sql_statement']}", file=sys.stderr)
 
     env = get_environment()
     try:
@@ -315,6 +332,7 @@ def t2s_execute_query(state: GraphState):
             #rows = c.export_to_pandas(state['sql_statement'])
 
             state['query_result'] = str(ExaDbResult(rows))
+            state['query_num_rows'] = c.last_statement().rowcount()
 
     except ExaError as e:
         state['sql_is_valid'] = "NO"
@@ -323,7 +341,28 @@ def t2s_execute_query(state: GraphState):
         state['sql_is_valid'] = "YES"
         state['sql_error'] = "None"
 
-    print("#### End of SQL Execution", file=sys.stderr)
+        ## Store the generated SQL statement and the natural language question into a VectorDB
+        ## We will use it for similarity search and may add this query to the prompt for future
+        ## natural language questions
+
+        if state['query_num_rows'] > 0:
+            print(f"Storing into VectorDB", file=sys.stderr)
+            load_dotenv()
+            vectordb_client = chromadb.PersistentClient(path=env['vectordb_persistent_storage'])
+            sql_collection = vectordb_client.get_or_create_collection(name="Questions_SQL_History")
+
+            new_idx = sql_collection.count() + 1
+            sql_collection.add(
+                documents=[state['question']],
+                metadatas=[{"sql": state['sql_statement'],
+                            "execution_date": str(datetime.now()),
+                            "db_schema": state['db_schema'],
+                            "user": env['db_user'].lower(),
+                            "origin": "text-to-sql"}],
+                ids=[f"{new_idx}"]
+            )
+
+    print("#### End of SQL Execution-02", file=sys.stderr)
 
     return state