update to support loading ai-libs as kdbx module (#7)

Author: cbeagan-kx

README.md

@@ -56,7 +56,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://developer.kx.com/products/kdb-x/install) - 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)
@@ -67,7 +67,7 @@ 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
+> 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://developer.kx.com/products/kdb-x/install) to ensure uninterrupted access, valid through 31st Dec 2025
 
 ## Quickstart
 
@@ -88,7 +88,7 @@ To demonstrate basic usage of the KDB-X MCP Server, using an empty KDB-X databas
 2. Load the ai and sql interfaces.
 
    ```q
-   \l ai-libs/init.q
+   .ai:use`kx.ai
    .s.init[]
    ```
 
@@ -519,9 +519,19 @@ The below tools can aid in the development, testing and debugging of new MCP too
 
 ## Troubleshooting
 
+### Failed to import pykx
+
+The KDB-X MCP Server requires a valid KDB-X license to operate.
+
+If you see an error like "Failed to import pykx", verify the following:
+
+- The `QLIC` environment variable is set and points to your license directory
+- Your license directory contains a valid license file
+- Your license has not expired
+
 ### 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 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://developer.kx.com/products/kdb-x/install) to ensure uninterrupted access, valid through 31st Dec 2025
 
 ### KDB-X connection error
 

src/mcp_server/server.py

@@ -1,13 +1,18 @@
+import os
 import sys
 import logging
 import socket
+from packaging import version
 from mcp.server.fastmcp import FastMCP
 from mcp_server.utils.logging import setup_logging
 from mcp_server.settings import AppSettings
 from mcp_server.tools import register_tools
 from mcp_server.prompts import register_prompts
 from mcp_server.resources import register_resources
 
+# Ensure pykx throws error on import if license is not valid
+os.environ["PYKX_LICENSED"] = "true"
+
 # Global flag to track AI Libs availability
 _ai_libs_available = False
 
@@ -66,6 +71,13 @@ def _check_kdb_connection(self):
         try:
             import pykx as kx
             from pykx.exceptions import QError
+        except Exception as e:
+            self.logger.error(f"Failed to import pykx: {e}")
+            self.logger.error("KDB-X Python is required to connect to KDB-X. Please ensure you have a valid license.")
+            self.logger.error("Set the QLIC environment variable to point to your license directory.")
+            sys.exit(1)
+
+        try:
             conn = kx.SyncQConnection(
                 host=self.db_config.host,
                 port=self.db_config.port,
@@ -75,7 +87,15 @@ def _check_kdb_connection(self):
                 tls=self.db_config.tls,
 
             )
-            self.logger.info(f"KDB-X connectivity check with 'tls={self.db_config.tls}': SUCCESS - {self.db_config.host}:{self.db_config.port} is accessible")
+            # Try to get kdbx version first, fall back to kdb+ version
+            try:
+                kdb_version = conn('.z.v`version').py().decode('utf-8')
+                kdb_type = "KDB-X"
+            except (QError, AttributeError):
+                kdb_version = str(conn('.z.K').py())
+                kdb_type = "KDB+"
+
+            self.logger.info(f"KDB-X connectivity check with 'tls={self.db_config.tls}': SUCCESS - {self.db_config.host}:{self.db_config.port} is accessible. You are running {kdb_type} version: {kdb_version}")
 
             # check if sql interface is loaded on KDB-X service
             if not conn('@[{2< count .s};(::);{0b}]').py():
@@ -87,8 +107,16 @@ def _check_kdb_connection(self):
             # 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")
+
+                # check if KDB-X version supports loading AI libs as a module
+                if kdb_type == "KDB-X" and version.parse(kdb_version) < version.parse("0.1.2"):
+                    self.logger.warning("KDB-X AI Libs check: NOT AVAILABLE - AI-powered tools (similarity_search, hybrid_search) will be disabled.")
+                    self.logger.warning(f"To use AI tools, you need at least KDB-X version '0.1.2'. Your version is '{kdb_version}'. Please update to the latest KDB-X version.")
+                elif kdb_type == "KDB+":
+                    self.logger.warning("KDB-X AI Libs check: NOT AVAILABLE - AI-powered tools (similarity_search, hybrid_search) are only available in KDB-X.")
+                else:
+                    self.logger.warning("KDB-X AI Libs check: NOT LOADED - AI-powered tools (similarity_search, hybrid_search) will be disabled.")
+                    self.logger.warning(r"To enable AI tools, load the KDB-X AI libraries by running: .ai:use`kx.ai 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")
 
@@ -148,7 +176,7 @@ def _preload_embedding_models(self):
         try:
             import asyncio
             from mcp_server.utils.embeddings import preload_models_from_config
-            
+
             # Run the async preload function
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
@@ -157,7 +185,7 @@ def _preload_embedding_models(self):
                 self.logger.info("Embedding models preloaded successfully")
             finally:
                 loop.close()
-                
+
         except Exception as e:
             self.logger.warning(f"Failed to preload embedding models: {e}")
             self.logger.warning("Models will be loaded on first use, which may cause delays")