Unlock real-time market intelligence with KDB-X MCP server

import logging
import pykx as kx
import json
from typing import Dict, Any
from mcp_server.utils.kdbx import get_kdb_connection

logger = logging.getLogger(__name__)
MAX_ROWS_RETURNED = 1000

async def run_query_impl(sqlSelectQuery: str) -> Dict[str, Any]:
    try:
        dangerous_keywords = ['INSERT', 'DROP', 'DELETE', 'TRUNCATE', 'ALTER', 'CREATE']
        query_upper = sqlSelectQuery.upper().strip()

        for keyword in dangerous_keywords:
            if keyword in query_upper and not query_upper.startswith('SELECT'):
                raise ValueError(f"Query contains dangerous keyword: {keyword}")

        conn = get_kdb_connection()
        # below query gets kdbx table data back as json for correct conversion of different datatypes
        result = conn('{r:.s.e x;`rowCount`data!(count r;.j.j y sublist r)}', kx.CharVector(sqlSelectQuery), MAX_ROWS_RETURNED)
        total = int(result['rowCount'])
        if 0==total:
            return {"status": "success", "data": [], "message": "No rows returned"}
        # parse json result
        rows = json.loads(result['data'].py().decode('utf-8'))
        if total > MAX_ROWS_RETURNED:
            logger.info(f"Table has {total} rows. Query returned truncated data to {MAX_ROWS_RETURNED} rows.")
            return {
                "status": "success",
                "data": rows,
                "message": f"Showing first {MAX_ROWS_RETURNED} of {total} rows",
            }

        logger.info(f"Query returned {total} rows.")
        return {"status": "success", "data": rows}

    except Exception as e:
        logger.error(f"Query failed: {e}")
        if ".s.e" in str(e):
            logger.error(f"It looks like the SQL interface is not loaded. You can load it manually by running .s.init[]:")
            return {
                "status": "error",
                "error_type": "sql_interface_not_loaded",
                "message": "It looks like the SQL interface is not loaded in the KDB-X database. Please initialize it by running `.s.init[]` in your KDB-X session, or contact your system administrator.",
                "technical_details": str(e)
            }
        return {"status": "error", "message": str(e)}


def register_tools(mcp_server):
    @mcp_server.tool()
    async def kdbx_run_sql_query(query: str) -> Dict[str, Any]:
        """
        Execute a SQL query and return structured results only to be used on kdb and not on kdbai.

        This function processes SQL SELECT statements to retrieve data from the underlying
        database. It parses the query, executes it against the data source, and returns
        the results in a structured format suitable for further analysis or display.

        Use the kdbx_sql_query_guidance resource when creating queries


        Supported query types:
            - SELECT statements with column specifications
            - WHERE clauses for filtering
            - ORDER BY for result sorting
            - LIMIT for result pagination
            - Basic aggregation functions (COUNT, SUM, AVG, etc.)

        For query syntax and examples, see: file://guidance/kdbx-sql-queries

        Args:
            query (str): SQL SELECT query string to execute. Must be a valid SQL statement
                        following standard SQL syntax conventions.

        Returns:
            Dict[str, Any]: Query execution results.
        """
        return await run_query_impl(sqlSelectQuery=query)

    return ['kdbx_run_sql_query']

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"]

from typing import List, Union
import pandas as pd
import logging
from datetime import date
from mcp_server.utils.kdbx import get_kdb_connection
from mcp_server.settings import default_kdb_config

logger = logging.getLogger(__name__)
async def compute_stock_indicator_impl(symbols: List[str],
                                       start_date: str,
                                       end_date: str,
                                       indicator: str,
                                       window: int = 14) -> Union[str, pd.Series]:
    """
    Compute specified technical indicators for given stock symbols over a date range.
    This function retrieves end-of-day (EOD) price data for the provided stock symbols
    between the specified start and end dates. It then calculates the requested technical
    indicator using the retrieved data.
    Supported indicators:
        - 'sma': Simple Moving Average
        - 'ema': Exponential Moving Average
        - 'volatility': Rolling standard deviation of returns
        - 'vwap': Volume Weighted Average Price
    Args:
        symbols (List[str]): List of stock symbols to analyze.
        start_date (str): Start date in 'YYYY-MM-DD' format.
        end_date (str): End date in 'YYYY-MM-DD' format.
        indicator (str): Technical indicator to compute ('sma', 'ema', 'volatility', 'vwap').
        window (int, optional): Window size for rolling calculations. Defaults to 14.
    Returns:
        Union[str, pd.Series]: Calculated indicator values or an error message.
    """
    conn = get_kdb_connection()
    try:
        stock_data = conn('{[x;y;z] select from stocks where (Symbol in x) and (Date within (y;z))}',
               symbols,
               date.fromisoformat(start_date),
               date.fromisoformat(end_date)).pd()
    
        if stock_data.empty:
            return "No data found for the given parameters."
    
    except Exception as e:
        logger.error(f"Error querying data from kdb: {e}")
        return "Error retrieving data."
    if indicator.lower() == "sma":
        return stock_data["Close"].rolling(window).mean().dropna()
    
    elif indicator.lower() == "ema":
        return stock_data["Close"].ewm(span=window, adjust=False).mean().dropna()
    elif indicator.lower() == "volatility":
        returns = stock_data["Close"].pct_change()
        return returns.rolling(window).std().dropna()
    
    elif indicator.lower() == "vwap":
        return (stock_data['Close'] * stock_data['Volume']).cumsum() / stock_data['Volume'].cumsum()
    else:
        raise ValueError("Metric must be 'SMA', 'EMA', 'volatility' or 'vwap'")
def register_tools(mcp_server):
    @mcp_server.tool()
    async def compute_stock_indicator(
        symbols: List[str], start_date: str, end_date: str, indicator: str, window: int = 14
    ) -> Union[str, pd.Series]:
        return await compute_stock_indicator_impl(symbols, start_date, end_date, indicator, window)
    
    return ['compute_stock_indicator']