Teradata MCP Server

import json import logging from datetime import date, datetime from decimal import Decimal from typing import Any from pathlib import Path import yaml from teradatasql import TeradataConnection logger = logging.getLogger("teradata_mcp_server") # Load RAG configuration def load_rag_config(): """Load RAG configuration from rag_config.yml""" try: # Get the directory path current_dir = Path(__file__).parent # Go to config/ config_path = current_dir.parent.parent / 'config' / 'rag_config.yml' with open(config_path, 'r') as file: logger.info(f"Loading RAG config from: {config_path}") return yaml.safe_load(file) except FileNotFoundError: logger.warning(f"RAG config file not found: {config_path}, using defaults") return get_default_rag_config() except Exception as e: logger.error(f"Error loading RAG config: {e}") return get_default_rag_config() def get_default_rag_config(): """Default RAG configuration as fallback""" return { 'version': 'ivsm', 'databases': { 'query_db': 'demo_db', 'model_db': 'demo_db', 'vector_db': 'demo_db' }, 'tables': { 'query_table': 'user_query', 'query_embedding_store': 'user_query_embeddings', 'vector_table': 'icici_fr_embeddings_store', 'model_table': 'embeddings_models', 'tokenizer_table': 'embeddings_tokenizers' }, 'model': { 'model_id': 'bge-small-en-v1.5' }, 'retrieval': { 'default_k': 10, 'max_k': 50 }, 'vector_store_schema': { 'required_fields': ['txt'], 'metadata_fields_in_vector_store': ['chunk_num', 'section_title', 'doc_name'] }, 'embedding': { 'vector_length': 384, 'vector_column_prefix': 'emb_', 'distance_measure': 'cosine', 'feature_columns': '[emb_0:emb_383]' } } # Load config at module level RAG_CONFIG = load_rag_config() def build_search_query(vector_db, dst_table, chunk_embed_table, k, config): """Build dynamic search query based on available metadata fields in vector store""" # Get metadata fields from config metadata_fields = config['vector_store_schema']['metadata_fields_in_vector_store'] or [] feature_columns = config['embedding']['feature_columns'] # Build SELECT clause dynamically - txt is always required select_fields = ["e_ref.txt AS reference_txt"] # Add all metadata fields from vector store for field in metadata_fields: if field != 'txt': select_fields.append(f"e_ref.{field} AS {field}") # Add similarity select_fields.append("(1.0 - dt.distance) AS similarity") select_clause = ",\n ".join(select_fields) return f""" SELECT {select_clause} FROM TD_VECTORDISTANCE ( ON {vector_db}.{dst_table} AS TargetTable ON {vector_db}.{chunk_embed_table} AS ReferenceTable DIMENSION USING TargetIDColumn('id') TargetFeatureColumns('{feature_columns}') RefIDColumn('id') RefFeatureColumns('{feature_columns}') DistanceMeasure('cosine') TopK({k}) ) AS dt JOIN {vector_db}.{chunk_embed_table} e_ref ON e_ref.id = dt.reference_id ORDER BY similarity DESC; """ def serialize_teradata_types(obj: Any) -> Any: """Convert Teradata-specific types to JSON serializable formats""" if isinstance(obj, date | datetime): return obj.isoformat() if isinstance(obj, Decimal): return float(obj) return str(obj) def rows_to_json(cursor_description: Any, rows: list[Any]) -> list[dict[str, Any]]: """Convert database rows to JSON objects using column names as keys""" if not cursor_description or not rows: return [] columns = [col[0] for col in cursor_description] return [ { col: serialize_teradata_types(value) for col, value in zip(columns, row) } for row in rows ] def create_response(data: Any, metadata: dict[str, Any] | None = None) -> str: """Create a standardized JSON response structure""" if metadata: response = { "status": "success", "metadata": metadata, "results": data } else: response = { "status": "success", "results": data } return json.dumps(response, default=serialize_teradata_types) def handle_rag_Execute_Workflow( conn: TeradataConnection, question: str, k: int | None = None, *args, **kwargs, ): """ Execute complete RAG workflow to answer user questions based on document context. This tool handles the entire RAG pipeline in a single step when a user query is tagged with /rag. WORKFLOW STEPS (executed automatically): 1. Configuration setup using configurable values from rag_config.yml 2. Store user query with '/rag ' prefix stripping 3. Generate query embeddings using either BYOM (ONNXEmbeddings) or IVSM functions based on config 4. Perform semantic search against precomputed chunk embeddings 5. Return context chunks for answer generation CONFIGURATION VALUES (from rag_config.yml): - version: 'ivsm' or 'byom' to select embedding approach - All database names, table names, and model settings are configurable - Vector store metadata fields are dynamically detected - Embedding parameters are configurable - Default chunk retrieval count is configurable - Default values are provided as fallback TECHNICAL DETAILS: - Strips the '/rag ' prefix if present from user questions - Creates query table if it does not exist (columns: id, txt, created_ts) - BYOM approach: Uses mldb.ONNXEmbeddings UDF for tokenization and embedding - IVSM approach: Uses ivsm.tokenizer_encode and ivsm.IVSM_score functions - Both approaches store embeddings in configured output table - Uses cosine similarity via TD_VECTORDISTANCE for semantic search - Returns the top-k matching chunks from the configured vector store - Each result includes chunk text, similarity score, and metadata fields CRITICAL ANSWERING RULES: - Answer ONLY using retrieved chunks - no external knowledge, speculation, or inference - Quote source content directly without paraphrasing, summarizing, or rewriting - If no relevant context found: "Not enough information found in the provided context" - If partial context: "The available context does not fully answer the question" - Include document/page references when available (e.g., "On page 2 of 'demo_policy.pdf'...") - Execute entire workflow silently without showing function calls to user - only show final answer LANGUAGE RESTRICTIONS: - Do not say "According to the context" or "The context says" - just quote directly - Do not say "It can be inferred that..." - no inference allowed - Use exact or near-verbatim quotes only - No paraphrasing, summarizing, or adding transitions between quotes - Copy content precisely - no expansion or interpretation BEHAVIORAL REQUIREMENTS: - Execute all RAG workflow steps silently without showing tool calls to users - Only display final answers based strictly on retrieved context chunks - Do not use external knowledge or speculation - only quote retrieved content directly - RAG mode only triggers with explicit /rag prefix - do not enter RAG mode implicitly - Hide all tool execution steps from user - Provide clean, direct response with only the answer based on retrieved chunks - This is a complete end-to-end RAG workflow - no additional tools needed EXECUTION: Run completely silently - user only sees their question and the final context-grounded answer. """ # Use configuration from loaded config config = RAG_CONFIG # Check version and delegate to appropriate implementation version = config.get('version', 'ivsm').lower() if version == 'ivsm': return _execute_rag_workflow_ivsm(conn, question, k, config, *args, **kwargs) elif version == 'byom': return _execute_rag_workflow_byom(conn, question, k, config, *args, **kwargs) else: raise ValueError(f"Unsupported RAG version: {version}. Supported versions: 'byom', 'ivsm'") def _execute_rag_workflow_byom(conn: TeradataConnection, question: str, k: int | None, config: dict, *args, **kwargs): """Execute RAG workflow using BYOM (ONNXEmbeddings)""" # Use config default if k not provided if k is None: k = config['retrieval']['default_k'] # Optional: Enforce max limit max_k = config['retrieval'].get('max_k', 50) if k > max_k: logger.warning(f"Requested k={k} exceeds max_k={max_k}, using max_k") k = max_k logger.debug(f"handle_rag_executeWorkflow (BYOM): question={question[:60]}..., k={k}") # Extract config values database_name = config['databases']['query_db'] table_name = config['tables']['query_table'] dst_table = config['tables']['query_embedding_store'] model_id = config['model']['model_id'] model_db = config['databases']['model_db'] model_table = config['tables']['model_table'] tokenizer_table = config['tables']['tokenizer_table'] vector_db = config['databases']['vector_db'] chunk_embed_table = config['tables']['vector_table'] with conn.cursor() as cur: # Store user query logger.debug(f"Step 2: Storing user query in {database_name}.{table_name}") # Create table if it doesn't exist ddl = f""" CREATE TABLE {database_name}.{table_name} ( id INTEGER GENERATED ALWAYS AS IDENTITY (START WITH 1 INCREMENT BY 1) NOT NULL, txt VARCHAR(5000), created_ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY (id) ) """ try: cur.execute(ddl) logger.debug(f"Table {database_name}.{table_name} created") except Exception as e: error_msg = str(e).lower() if "already exists" in error_msg or "3803" in error_msg: logger.debug(f"Table {database_name}.{table_name} already exists, skipping creation") else: logger.error(f"Error creating table: {e}") raise # Insert cleaned question insert_sql = f""" INSERT INTO {database_name}.{table_name} (txt) SELECT CASE WHEN TRIM(?) LIKE '/rag %' THEN SUBSTRING(TRIM(?) FROM 6) ELSE TRIM(?) END """ cur.execute(insert_sql, [question, question, question]) # Get inserted ID and cleaned text cur.execute(f"SELECT MAX(id) AS id FROM {database_name}.{table_name}") new_id = cur.fetchone()[0] cur.execute(f"SELECT txt FROM {database_name}.{table_name} WHERE id = ?", [new_id]) cleaned_txt = cur.fetchone()[0] logger.debug(f"Stored query with ID {new_id}: {cleaned_txt[:60]}...") # Generate query embeddings logger.debug(f"Step 3: Generating embeddings in {database_name}.{dst_table}") # Drop existing embeddings table drop_sql = f"DROP TABLE {database_name}.{dst_table}" try: cur.execute(drop_sql) logger.debug(f"Dropped existing table {database_name}.{dst_table}") except Exception as e: logger.debug(f"DROP failed or table not found: {e}") # Create embeddings table create_sql = f""" CREATE TABLE {database_name}.{dst_table} AS ( SELECT * FROM mldb.ONNXEmbeddings( ON (SELECT id, txt FROM {database_name}.{table_name}) ON (SELECT model_id, model FROM {model_db}.{model_table} WHERE model_id = '{model_id}') DIMENSION ON (SELECT model AS tokenizer FROM {model_db}.{tokenizer_table} WHERE model_id = '{model_id}') DIMENSION USING Accumulate('id', 'txt') ModelOutputTensor('sentence_embedding') OutputFormat('FLOAT32({config["embedding"]["vector_length"]})') ) AS a ) WITH DATA """ cur.execute(create_sql) logger.debug(f"Created embeddings table {database_name}.{dst_table}") # Perform semantic search logger.debug(f"Step 4: Performing semantic search with k={k}") search_sql = build_search_query(vector_db, dst_table, chunk_embed_table, k, config) rows = cur.execute(search_sql) data = rows_to_json(cur.description, rows.fetchall()) logger.debug(f"Retrieved {len(data)} chunks for semantic search") # Return metadata metadata = { "tool_name": "rag_executeWorkflow", "workflow_type": "BYOM", "workflow_steps": ["config_set", "query_stored", "embeddings_generated", "semantic_search_completed"], "query_id": new_id, "cleaned_question": cleaned_txt, "database": database_name, "query_table": table_name, "embedding_table": dst_table, "vector_table": chunk_embed_table, "model_id": model_id, "chunks_retrieved": len(data), "topk_requested": k, "topk_configured_default": config['retrieval']['default_k'], "metadata_fields": config['vector_store_schema']['metadata_fields_in_vector_store'], "description": "Complete RAG workflow executed using BYOM: config - store query - generate embeddings - semantic search" } logger.debug(f"Tool: handle_rag_executeWorkflow (BYOM): metadata: {metadata}") return create_response(data, metadata) def _execute_rag_workflow_ivsm(conn: TeradataConnection, question: str, k: int | None, config: dict, *args, **kwargs): """Execute RAG workflow using IVSM functions""" # Use config default if k not provided if k is None: k = config['retrieval']['default_k'] # Optional: Enforce max limit max_k = config['retrieval'].get('max_k', 50) if k > max_k: logger.warning(f"Requested k={k} exceeds max_k={max_k}, using max_k") k = max_k logger.debug(f"handle_rag_executeWorkflow (IVSM): question={question[:60]}..., k={k}") # Extract config values database_name = config['databases']['query_db'] table_name = config['tables']['query_table'] dst_table = config['tables']['query_embedding_store'] model_id = config['model']['model_id'] model_db = config['databases']['model_db'] model_table = config['tables']['model_table'] tokenizer_table = config['tables']['tokenizer_table'] vector_db = config['databases']['vector_db'] chunk_embed_table = config['tables']['vector_table'] with conn.cursor() as cur: # Store user query logger.debug(f"Step 2: Storing user query in {database_name}.{table_name}") # Create table if it doesn't exist ddl = f""" CREATE TABLE {database_name}.{table_name} ( id INTEGER GENERATED ALWAYS AS IDENTITY (START WITH 1 INCREMENT BY 1) NOT NULL, txt VARCHAR(5000), created_ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY (id) ) """ try: cur.execute(ddl) logger.debug(f"Table {database_name}.{table_name} created") except Exception as e: error_msg = str(e).lower() if "already exists" in error_msg or "3803" in error_msg: logger.debug(f"Table {database_name}.{table_name} already exists, skipping creation") else: logger.error(f"Error creating table: {e}") raise # Insert cleaned question insert_sql = f""" INSERT INTO {database_name}.{table_name} (txt) SELECT CASE WHEN TRIM(?) LIKE '/rag %' THEN SUBSTRING(TRIM(?) FROM 6) ELSE TRIM(?) END """ cur.execute(insert_sql, [question, question, question]) # Get inserted ID and cleaned text cur.execute(f"SELECT MAX(id) AS id FROM {database_name}.{table_name}") new_id = cur.fetchone()[0] cur.execute(f"SELECT txt FROM {database_name}.{table_name} WHERE id = ?", [new_id]) cleaned_txt = cur.fetchone()[0] logger.debug(f"Stored query with ID {new_id}: {cleaned_txt[:60]}...") # Tokenize query logger.debug("Step 3: Tokenizing query using ivsm.tokenizer_encode") cur.execute(f""" REPLACE VIEW v_topics_tokenized AS ( SELECT id, txt, IDS AS input_ids, attention_mask FROM ivsm.tokenizer_encode( ON ( SELECT * FROM {database_name}.{table_name} QUALIFY ROW_NUMBER() OVER (ORDER BY created_ts DESC) = 1 ) ON ( SELECT model AS tokenizer FROM {model_db}.{tokenizer_table} WHERE model_id = '{model_id}' ) DIMENSION USING ColumnsToPreserve('id','txt') OutputFields('IDS','ATTENTION_MASK') MaxLength(1024) PadToMaxLength('True') TokenDataType('INT64') ) AS t ); """) logger.debug("Tokenized view v_topics_tokenized created") # Create embedding view logger.debug("Step 4: Creating embedding view using ivsm.IVSM_score") cur.execute(f""" REPLACE VIEW v_topics_embeddings AS ( SELECT * FROM ivsm.IVSM_score( ON v_topics_tokenized ON ( SELECT * FROM {model_db}.{model_table} WHERE model_id = '{model_id}' ) DIMENSION USING ColumnsToPreserve('id','txt') ModelType('ONNX') BinaryInputFields('input_ids','attention_mask') BinaryOutputFields('sentence_embedding') Caching('inquery') ) AS s ); """) logger.debug("Embedding view v_topics_embeddings created") # Create query embedding table logger.debug("Step 5: Creating query embedding table using ivsm.vector_to_columns") # Drop existing embeddings table drop_sql = f"DROP TABLE {database_name}.{dst_table}" try: cur.execute(drop_sql) logger.debug(f"Dropped existing table {database_name}.{dst_table}") except Exception as e: logger.debug(f"DROP failed or table not found: {e}") # Create embeddings table create_sql = f""" CREATE TABLE {database_name}.{dst_table} AS ( SELECT * FROM ivsm.vector_to_columns( ON v_topics_embeddings USING ColumnsToPreserve('id', 'txt') VectorDataType('FLOAT32') VectorLength({config['embedding']['vector_length']}) OutputColumnPrefix('{config['embedding']['vector_column_prefix']}') InputColumnName('sentence_embedding') ) a ) WITH DATA """ cur.execute(create_sql) logger.debug(f"Created embeddings table {database_name}.{dst_table}") # Perform semantic search logger.debug(f"Step 6: Performing semantic search with k={k}") search_sql = build_search_query(vector_db, dst_table, chunk_embed_table, k, config) rows = cur.execute(search_sql) data = rows_to_json(cur.description, rows.fetchall()) logger.debug(f"Retrieved {len(data)} chunks for semantic search") # Return metadata metadata = { "tool_name": "rag_executeWorkflow", "workflow_type": "IVSM", "workflow_steps": ["config_set", "query_stored", "query_tokenized", "embedding_view_created", "embedding_table_created", "semantic_search_completed"], "query_id": new_id, "cleaned_question": cleaned_txt, "database": database_name, "query_table": table_name, "embedding_table": dst_table, "vector_table": chunk_embed_table, "model_id": model_id, "chunks_retrieved": len(data), "topk_requested": k, "topk_configured_default": config['retrieval']['default_k'], "views_created": ["v_topics_tokenized", "v_topics_embeddings"], "metadata_fields": config['vector_store_schema']['metadata_fields_in_vector_store'], "description": "Complete RAG workflow executed using IVSM functions: config - store query - tokenize - create embedding view - create embedding table - semantic search" } logger.debug(f"Tool: handle_rag_executeWorkflow (IVSM): metadata: {metadata}") return create_response(data, metadata)