add ai libs check (#5)

Author: cbeagan-kx

README.md

@@ -66,6 +66,8 @@ Before installing and running the KDB-X MCP Server, ensure you have met the foll
   - `npx` comes bundled with the [nodejs](https://nodejs.org/en) installer - available on Windows/Mac/Linux/WSL
   - See [example configuration with streamable-http](#example-configuration-with-streamable-http)
 
+> Note: ⚠️ KDB-X public preview has recently been extended. If you have installed KDB-X prior to Sept 30th 2025, you will receive an email notification about this update. Please update to the latest [KDB-X](https://kdb-x.kx.com/sign-in) to ensure uninterrupted access, valid through January 4, 2026
+
 ## Quickstart
 
 To demonstrate basic usage of the KDB-X MCP Server, using an empty KDB-X database, follow the quickstart steps below.
@@ -438,6 +440,10 @@ The below tools can aid in the development, testing and debugging of new MCP too
 
 ## Troubleshooting
 
+### KDB-X license expired
+
+KDB-X public preview has recently been extended. If you have installed KDB-X prior to Sept 30th 2025, you will receive an email notification about this update. Please update to the latest [KDB-X](https://kdb-x.kx.com/sign-in) to ensure uninterrupted access, valid through January 4, 2026
+
 ### KDB-X connection error
 
 Ensure that your KDB-X database is online and accessible on the specified kdb host and port.

pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "mcp[cli]>=1.2.0",
-    "pykx>=4.0.0.b2",
+    "pykx>=4.0.0.b3",
     "pydantic-settings",
     # Optional: Manage below packages for embedding support
     # "sentence_transformers",

src/mcp_server/server.py

@@ -55,6 +55,16 @@ def __init__(self, kdbx_mcp_transport=None, kdbx_mcp_port=None, kdbx_host=None,
         self.host = self.settings.host
         self.transport = self.settings.transport
 
+# Global flag to track AI Libs availability
+_ai_libs_available = False
+
+def set_ai_libs_available(available: bool):
+    global _ai_libs_available
+    _ai_libs_available = available
+
+def is_ai_libs_available() -> bool:
+    return _ai_libs_available
+
 class McpServer:
 
     def __init__(self, config: Optional[MCPServerConfig] = None):
@@ -111,11 +121,18 @@ 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")
+
+            # check if AI libs are loaded on KDB-X service
+            ai_libs_available = conn('@[{2< count .ai};(::);{0b}]').py()
+            if not ai_libs_available:
+                self.logger.warning("KDB-X AI Libs check: NOT AVAILABLE - AI-powered tools (similarity_search, hybrid_search) will be disabled.")
+                self.logger.warning("To enable AI tools, load the KDB-X AI libraries by running: \l ai-libs/init.q in your KDB-X Session and then restart the MCP server")
+            else:
+                self.logger.info("KDB-X AI Libs check: SUCCESS - AI Libs are loaded, AI tools will be available")
+
+            set_ai_libs_available(ai_libs_available)
             conn.close()
 
         except QError as e:

src/mcp_server/tools/kdbx_sim_search.py

@@ -90,6 +90,17 @@ async def kdbx_similarity_search_impl( table_name: str,
 
 
 def register_tools(mcp_server):
+    # Check if AI Libs are available
+    try:
+        from mcp_server.server import is_ai_libs_available
+        
+        if not is_ai_libs_available():
+            logger.info("AI Libs not available - skipping...")
+            return []
+    except Exception as e:
+        logger.warning(f"Could not check AI Libs availability: {e}. Skipping AI tools.")
+        return []
+    
     @mcp_server.tool()
     async def kdbx_similarity_search(table_name: str,
                             query: str,

src/mcp_server/utils/embeddings_helpers.py

@@ -40,4 +40,4 @@ def get_embedding_config(
     row = matching_rows.iloc[0]
 
     # Choose columns based on config_type flag
-    return [row['embedding_column'], row['embedding_provider'], row['embedding_model'], row['sparse_tokenizer_provider'], row['sparse_tokenizer_model']]
+    return [row['embedding_column'], row['embedding_provider'], row['embedding_model'], row['sparse_tokenizer_provider'], row['sparse_tokenizer_model']]