Merge pull request #4 from KxSystems/add-similarity-search-tool

add similarity-search tool

Author: cbeagan-kx

README.md

@@ -16,6 +16,7 @@ The server leverages a combination of curated resources, intelligent prompts, an
 - [MCP Server Installation](#mcp-server-installation)
 - [Transport Options](#transport-options)
 - [Command line Parameters](#command-line-parameters)
+- [Configure Embeddings](#configure-embeddings)
 - [Usage with Claude Desktop](#usage-with-claude-desktop)
 - [Prompts/Resources/Tools](#promptsresourcestools)
 - [Development](#development)
@@ -54,7 +55,7 @@ Before installing and running the KDB-X MCP Server, ensure you have met the foll
 - [Cloned this repo](#clone-the-repository)
 - A `KDB-X/KDB+` Service listening on a host and port that will be accessible to the MCP Server
   - See examples - [KDB-X Setup](#kdb-x-setup) / [KDB+ Setup](#kdb-setup)
-  - KDB-X can be installed by signing up to the [kdb-x public preview](https://kdb-x.kx.com/sign-in) - see [KDB-X documentation](https://docs.kx.com/public-preview/kdb-x/home.htm) for supporting information
+  - KDB-X can be installed by signing up to the [KDB-X public preview](https://kdb-x.kx.com/sign-in) - see [KDB-X documentation](https://docs.kx.com/public-preview/kdb-x/home.htm) for supporting information
   - Windows users can run the KDB-X MCP Server on Windows and connect to a local KDB-X database via WSL or remote KDB-X database running on Linux
   - Windows users can run a local KDB-X database by installing KDB-X on [WSL](https://learn.microsoft.com/en-us/windows/wsl/install), and use the default [streamable-http transport](#transport-options) when running the [KDB-X MCP Server](#run-the-server) - both share the same localhost network.
   - For details on KDB-X usage restrictions see [documentation](https://docs.kx.com/product/licensing/usage-restrictions.htm#kdb-x-personal-trial-download)
@@ -246,6 +247,22 @@ options:
 2. **Environment variables** (middle precedence) - `KDBX_MCP_TRANSPORT=streamable-http`, `KDBX_HOST=myhost`
 3. **Default values** (lowest precedence)
 
+## Configure Embeddings
+
+Before starting the KDB-X MCP Server, you must configure embedding models for your tables if you wish to use Similarity Search.
+The repository includes two ready-to-use embedding providers: OpenAI and SentenceTransformers.
+You can customize these implementations as needed, or add your own provider by following the steps outlined below.
+
+1. Update Dependencies - Add your required embedding providers to `pyproject.toml` dependencies section.
+
+2. Set Environment Variables - Configure required API keys for your chosen embedding providers if necessary (for example, set the environment variable `OPENAI_API_KEY` to use OpenAI's API)
+
+3. Add New Provider - The file `src/mcp_server/utils/embeddings.py` defines the base class `EmbeddingProvider` for all embedding providers.
+   To add a new provider, create a class in the same file that extends this base class and implements all required abstract methods.
+   You can use the existing implementations of OpenAI and SentenceTransformers in the same file as templates — simply copy and modify them to suit your needs. To register your provider, use the `@register_provider` decorator above your class definition. It is not compulsory for the registered provider name to follow the provider's Python package name.
+
+4. Configure Table Embeddings - Update the embeddings configuration file at `src/mcp_server/utils/embeddings.csv` with your actual database and table names, embedding providers and models. The name you provide at `embeddings.csv` should match the registered provider name specified in file `embeddings.py`.
+
 ## Usage with Claude Desktop
 
 ### Configure Claude Desktop
@@ -302,7 +319,8 @@ If you have pre-existing MCP servers see [example config with multiple mcp-serve
         "--directory",
         "/path/to/this/repo/",
         "run",
-        "mcp_server"
+        "mcp-server",
+        "--stdio"
       ]
     }
   }
@@ -311,7 +329,7 @@ If you have pre-existing MCP servers see [example config with multiple mcp-serve
 
 **Note**
 
-- Update your `<user>` to point to the absolute path of the uv executable - only required if `uv` is on your path
+- Update your `<user>` to point to the absolute path of the uv executable - only required if `uv` is not on your path
 - Update the `--directory` path to the absolute path of this repo
 - Currently `KDB-X` does not support Windows, meaning `stdio` is not an option for Windows users
 - Claude Desktop is responsible for starting/stopping the MCP server when using `stdio`
@@ -372,7 +390,7 @@ To enable Developer mode:
 
 | Name | Purpose | Params | Return |
 |-------------------|------------------------------------------|-------------------------------------------------|------------------------------------------------|
-| kdbx_table_analysis | Generate a detailed analysis prompt for a specific table. | table_name: Name of the table to analyze<br> analysis_type (optional): Type of analysis options statistical, data_quality<br> sample_size (optional): Suggested sample size for data exploration | The generated table analysis prompt |
+| kdbx_table_analysis | Generate a detailed analysis prompt for a specific table. | `table_name`: Name of the table to analyze<br> `analysis_type` (optional): Type of analysis options statistical, data_quality<br> `sample_size` (optional): Suggested sample size for data exploration | The generated table analysis prompt |
 
 ### Resources
 
@@ -385,7 +403,8 @@ To enable Developer mode:
 
 | Name | Purpose | Params | Return |
 |-------------------|------------------------------------------|-------------------------------------------------|------------------------------------------------|
-| kdbx_run_sql_query | Execute SQL SELECT against KDB-X database | query (str): SQL SELECT query string to execute | JSON object with query results (max 1000 rows) |
+| kdbx_run_sql_query | Execute SQL SELECT against KDB-X database | `query`: SQL SELECT query string to execute | JSON object with query results (max 1000 rows) |
+| kdbx_similarity_search | Perform vector similarity search on a KDB-X table | `table_name`: Name of the table to search <br> `query`: Text query to convert to vector and search <br> `n` (optional): Number of results to return | Dictionary containing search result |
 
 ## Development
 
@@ -441,7 +460,7 @@ If the MCP Server port is being used by another process you will need to specify
 
 ### Invalid transport
 
-You can only specify `streamable-http`, `stdio.`
+You can only specify `streamable-http` or `stdio.`
 
 ### Missing tools/resources
 

pyproject.toml

@@ -8,6 +8,10 @@ dependencies = [
     "mcp[cli]>=1.2.0",
     "pykx>=4.0.0.b2",
     "pydantic-settings",
+    # Optional: Manage below packages for embedding support
+    # "sentence_transformers",
+    # "openai",
+    # "tiktoken",
 ]
 
 [project.optional-dependencies]

screenshots/claude_resources.png

@@ -3,6 +3,7 @@
 from typing import List
 from mcp.types import TextContent
 from mcp_server.utils.kdbx import get_kdb_connection
+from mcp_server.utils.embeddings_helpers import get_embedding_config
 from mcp_server.server import config
 
 logger = logging.getLogger(__name__)
@@ -45,7 +46,7 @@ async def kdbx_describe_table_impl(table: str) -> List[TextContent]:
 
             output_lines.extend([
                 f"\n Data Preview ({preview_size} records):",
-                _format_data(preview_data)
+                _format_data(preview_data, table)
             ])
         else:
             output_lines.append("\n Table is empty - no data to preview")
@@ -94,7 +95,12 @@ async def kdbx_describe_tables_impl() -> List[TextContent]:
         )]
 
 
-def _format_data(data) -> str:
+def _format_data(data, table=None) -> str:
+    if table:
+        embeddings_column, _, _, _, _ = get_embedding_config(table)
+        if embeddings_column and hasattr(data, "pop"):
+            data.pop(embeddings_column, None)
+            # data = data.drop(embeddings_column, axis=1, errors="ignore")
     if hasattr(data, 'to_string'):
         return data.to_string()
     return str(data)

screenshots/claude_tools.png

@@ -111,6 +111,9 @@ def _check_kdb_connection(self):
             if not conn('@[{2< count .s};(::);{0b}]').py():
                 self.logger.error("KDB-X SQL interface check: FAILED - KDB-X service does not have the SQL interface loaded. Load it by running .s.init[] in your KDB-X Session")
                 sys.exit(1)
+            if not conn('@[{2< count .ai};(::);{0b}]').py():
+                self.logger.error("KDB-X AI Libs check: FAILED - KDB-X service does not have the AI Libs loaded. Load it by running \l ai-libs/init.q in your KDB-X Session")
+                sys.exit(1)
             else:
                 self.logger.info("KDB-X SQL interface check: SUCCESS - SQL interface is loaded")
             conn.close()

src/mcp_server/resources/kdbx_database_tables.py

@@ -10,6 +10,11 @@ class KDBConfig(BaseSettings):
     password: Optional[SecretStr] = ""
     timeout: Optional[int] = 1
     retry: Optional[int] = 2
+    
+    # Similarity Search tool
+    embedding_csv_path: str = "src/mcp_server/utils/embeddings.csv"
+    metric: str = "CS"
+    k: int = 5
 
     class Config:
         env_prefix = 'KDBX_'

src/mcp_server/server.py

@@ -0,0 +1,115 @@
+import logging
+from typing import Optional, Dict, Any, List
+from mcp_server.settings import KDBConfig
+from mcp_server.utils.embeddings import get_provider
+from mcp_server.utils.embeddings_helpers import get_embedding_config
+import numpy as np
+import pandas as pd
+import logging
+import pykx as kx
+import json
+from typing import Dict, Any
+from mcp_server.utils.kdbx import get_kdb_connection
+
+config = KDBConfig()
+logger = logging.getLogger(__name__)
+
+
+# Normalizes the result from the search operation
+def normalize_result(df: Dict)-> Any:
+    # serialize numpy ndarray type
+    df = df.map(lambda x: x.tolist() if isinstance(x, np.ndarray) else x)
+    # convert timespan type (KDB time type)
+    for col_name, col_type in df.dtypes.items():
+        timespan_type = str(col_type).lower().startswith("timedelta")
+        duration_type = str(col_type).lower().startswith("duration")
+        if timespan_type or duration_type:
+            df[col_name] = (pd.Timestamp("1970-01-01") + df[col_name]).dt.time
+        # convert to dict
+    return df.to_dict('records') if hasattr(df, 'to_dict') else df
+
+
+async def kdbx_similarity_search_impl( table_name: str,
+                                        query: str,
+                                        n: Optional[int] = None) -> Dict[str, Any]:
+    
+    try:
+        if n is None:
+            n = config.k
+
+        embeddings_column, embeddings_provider, embeddings_model, _, _ = get_embedding_config(table_name)
+        
+        dense_provider = get_provider(embeddings_provider)
+        query_vector = await dense_provider.dense_embed(query, embeddings_model)
+
+        # Build search parameters
+        search_params = {
+            "table" : table_name,
+            "vcol"  : embeddings_column,
+            "qvec"  : query_vector,
+            "metric": config.metric,
+            "n"     : int(n),
+        }
+
+        conn = get_kdb_connection()
+
+        result = conn('''{[args]
+                            c:args`vcol;
+                            $[(args`table) in .Q.pt;
+                                [
+                                res:raze{[d;args;tbl;c]
+                                    vecs:?[tbl;enlist (=;.Q.pf;d);0b;(enlist c)!enlist c]c;
+                                    if[not count vecs; :()];
+                                    res:.ai.flat.search[vecs;args`qvec;args`n;args`metric];
+                                    res:res@\:iasc res[1];
+                                    `dist xcols update dist:res[0] from ?[tbl;((=;.Q.pf;d);(in;`i;res[1]));0b;()]
+                                }[;args;get args`table;c] each .Q.pv;
+                                ![(args`n)#`dist xdesc res;();0b;enlist c]
+                                ];
+                                [
+                                res:.ai.flat.search[?[args`table;();();c];args`qvec;args`n;args`metric];
+                                ![(args`table) res[1];();0b;enlist c]
+                                ]
+                            ]}''', search_params)
+
+        result = normalize_result(result.pd())
+
+        return {
+            "status": "success",
+            "table": table_name,
+            "recordsCount": len(result),
+            "records": result
+        }
+    except Exception as e:
+        logger.error(f"Error performing search on table {table_name}: {e}")
+        return {
+            "status": "error",
+            "message": str(e),
+            "table": table_name,
+        } 
+
+
+def register_tools(mcp_server):
+    @mcp_server.tool()
+    async def kdbx_similarity_search(table_name: str,
+                            query: str,
+                            n: Optional[int] = None) -> Dict[str, Any]:
+        """
+        Perform vector similarity search on a KDB-X table.
+
+        Args:
+            table_name: Name of the table to search
+            query: Text query to convert to vector and search
+            n (Optional[int], optional): Number of results to return
+
+        Returns:
+            Dictionary containing search result.
+        """
+        results = await kdbx_similarity_search_impl(
+            table_name,
+            query, 
+            n,
+        )
+        return results
+
+    return ["kdbx_similarity_search"]

src/mcp_server/settings.py

@@ -0,0 +1,2 @@
+table,embedding_column,embedding_provider,embedding_model,sparse_tokenizer_provider,sparse_tokenizer_model
+_example_table1,vecs,sentence_transformers,all-MiniLM-L12-v2,sentence_transformers,all-MiniLM-L12-v2