decodingtrust-agent-sdk 0.1.0__py3-none-any.whl

dt_arena/mcp_server/snowflake/main.py

@@ -0,0 +1,612 @@
+#!/usr/bin/env python3
+"""
+Snowflake MCP Server (local, no OAuth/account)
+- Implements the Snowflake-documented tool names exactly:
+  - product-search (CORTEX_SEARCH_SERVICE_QUERY)
+  - revenue-semantic-view (CORTEX_ANALYST_MESSAGE)
+  - sql_exec_tool (SYSTEM_EXECUTE_SQL)
+  - agent_1 (CORTEX_AGENT_RUN)
+
+Data storage: local PostgreSQL (via sync psycopg to localhost; can be docker-compose).
+product-search supports two backends, selectable via env SEARCH_MODE:
+  - "faiss": vector search using FAISS with simple hashing embeddings
+  - "simple": n-gram/fuzzy style text similarity without external model
+"""
+import os
+import json
+import asyncio
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+try:
+    import faiss  # type: ignore
+except Exception:
+    faiss = None  # Will be checked when SEARCH_MODE=faiss
+
+import psycopg
+try:
+    from openai import OpenAI  # Optional OpenAI integration
+except Exception:
+    OpenAI = None
+
+from fastmcp import FastMCP
+
+mcp = FastMCP("Snowflake MCP (Local)")
+
+# -----------------------------------------------------------------------------
+# Environment configuration
+# -----------------------------------------------------------------------------
+PORT = int(os.getenv("PORT", "8842"))
+HOST = os.getenv("HOST", "0.0.0.0")
+
+# Local PostgreSQL DSN (e.g. postgresql://snow:snow@127.0.0.1:5452/snowdb)
+POSTGRES_DSN = os.getenv(
+    "POSTGRES_DSN",
+    "postgresql://snow:snow@127.0.0.1:5452/snowdb",
+)
+
+# Search behavior
+SEARCH_MODE = os.getenv("SEARCH_MODE", "simple").lower()  # "faiss" or "simple"
+SEARCH_TABLE = os.getenv("SEARCH_TABLE", "products")
+# Comma-separated list of text columns to search over
+SEARCH_COLUMNS = [c.strip() for c in os.getenv("SEARCH_COLUMNS", "name,description").split(",") if c.strip()]
+
+# Optional row limit for indexing
+INDEX_ROW_LIMIT = int(os.getenv("INDEX_ROW_LIMIT", "10000"))
+
+# -----------------------------------------------------------------------------
+# Globals for connections and search index
+# -----------------------------------------------------------------------------
+_pg_conn: Optional["psycopg.Connection"] = None
+
+_faiss_index: Optional["faiss.Index"] = None
+_faiss_id_to_row: List[Dict[str, Any]] = []
+_faiss_dim: int = int(os.getenv("FAISS_DIM", "384"))
+_faiss_norm: bool = True
+
+# Optional: OpenAI configuration
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "").strip()
+OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "").strip()
+OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini").strip()
+_openai_client: Optional["OpenAI"] = None
+
+def get_openai() -> Optional["OpenAI"]:
+    global _openai_client
+    if not OPENAI_API_KEY or OpenAI is None:
+        return None
+    if _openai_client is None:
+        if OPENAI_BASE_URL:
+            _openai_client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)
+        else:
+            _openai_client = OpenAI(api_key=OPENAI_API_KEY)
+    return _openai_client
+
+def _connect_pg_with_retry(
+    dsn: str,
+    attempts: int = 3,
+    delay: float = 0.75,
+) -> "psycopg.Connection":
+    """
+    Connect to Postgres with simple retry.
+
+    In per-task eval runs the snowflake Postgres container may still be
+    starting up when the MCP server begins handling requests, so the first
+    few connection attempts can see \"connection refused\". Other envs have
+    explicit HTTP health waits in env_up.py; for Postgres we handle this at
+    the MCP layer with a short retry loop.
+    """
+    last_exc: Optional[Exception] = None
+    for _ in range(max(1, attempts)):
+        try:
+            return psycopg.connect(dsn, autocommit=True)
+        except Exception as e:  # connection refused, DNS issues, etc.
+            last_exc = e
+            time.sleep(delay)
+    if last_exc is not None:
+        raise last_exc
+    return psycopg.connect(dsn, autocommit=True)
+
+
+def get_pg() -> "psycopg.Connection":
+    global _pg_conn
+    if _pg_conn is None or _pg_conn.closed:
+        _pg_conn = _connect_pg_with_retry(POSTGRES_DSN)
+        _ensure_postgres_initialized(_pg_conn)
+    return _pg_conn
+
+
+# -----------------------------------------------------------------------------
+# Utility: simple hashing-based embedding for text
+# -----------------------------------------------------------------------------
+def _hashing_embed(text: str, dim: int) -> np.ndarray:
+    # Simple n-gram hashing into fixed-size vector
+    text = (text or "").lower()
+    vec = np.zeros(dim, dtype=np.float32)
+    if not text:
+        return vec
+    n = 3
+    for i in range(len(text) - n + 1):
+        gram = text[i : i + n]
+        h = (hash(gram) % dim + dim) % dim
+        vec[h] += 1.0
+    # Normalize
+    norm = np.linalg.norm(vec)
+    if norm > 0:
+        vec /= norm
+    return vec
+
+
+def _concat_columns(row: Dict[str, Any], columns: List[str]) -> str:
+    parts: List[str] = []
+    for c in columns:
+        v = row.get(c)
+        if v is None:
+            continue
+        parts.append(str(v))
+    return " ".join(parts)
+
+
+def _simple_similarity(a: str, b: str) -> float:
+    # Lightweight token-based similarity (Jaccard of tokens)
+    ta = set(a.lower().split())
+    tb = set(b.lower().split())
+    if not ta or not tb:
+        return 0.0
+    inter = len(ta & tb)
+    union = len(ta | tb)
+    return inter / union if union else 0.0
+
+
+# -----------------------------------------------------------------------------
+# PostgreSQL initialization and helpers
+# -----------------------------------------------------------------------------
+def _ensure_postgres_initialized(conn: "psycopg.Connection") -> None:
+    with conn.cursor() as cur:
+        cur.execute(
+            """
+            CREATE TABLE IF NOT EXISTS products (
+              id SERIAL PRIMARY KEY,
+              name TEXT,
+              description TEXT,
+              category TEXT
+            );
+            """
+        )
+        cur.execute(
+            """
+            CREATE TABLE IF NOT EXISTS revenue (
+              id SERIAL PRIMARY KEY,
+              date DATE,
+              product_id INT,
+              product_name TEXT,
+              revenue NUMERIC
+            );
+            """
+        )
+        # Seed minimal data if empty
+        cur.execute("SELECT COUNT(*) FROM products;")
+        cnt = cur.fetchone()[0]
+        if cnt == 0:
+            cur.execute(
+                """
+                INSERT INTO products (name, description, category) VALUES
+                ('Laptop Pro 14', 'High-performance laptop with retina display', 'Electronics'),
+                ('Noise Cancelling Headphones', 'Over-ear ANC headphones with long battery life', 'Audio'),
+                ('Smartwatch X', 'Fitness tracking and notifications', 'Wearables');
+                """
+            )
+        cur.execute("SELECT COUNT(*) FROM revenue;")
+        cnt2 = cur.fetchone()[0]
+        if cnt2 == 0:
+            cur.execute(
+                """
+                INSERT INTO revenue (date, product_id, product_name, revenue) VALUES
+                ('2025-10-01', 1, 'Laptop Pro 14', 250000),
+                ('2025-10-01', 2, 'Noise Cancelling Headphones', 85000),
+                ('2025-10-01', 3, 'Smartwatch X', 120000),
+                ('2025-10-02', 1, 'Laptop Pro 14', 265000),
+                ('2025-10-02', 2, 'Noise Cancelling Headphones', 82000),
+                ('2025-10-02', 3, 'Smartwatch X', 110000);
+                """
+            )
+
+def _get_table_columns(conn: "psycopg.Connection", table: str) -> List[str]:
+    with conn.cursor() as cur:
+        cur.execute(
+            """
+            SELECT column_name
+            FROM information_schema.columns
+            WHERE table_schema = 'public' AND table_name = %s
+            ORDER BY ordinal_position;
+            """,
+            (table,),
+        )
+        rows = cur.fetchall()
+    return [r[0] for r in rows] if rows else []
+
+
+# -----------------------------------------------------------------------------
+# FAISS Index building/query
+# -----------------------------------------------------------------------------
+async def _ensure_faiss_index() -> None:
+    global _faiss_index, _faiss_id_to_row
+    if _faiss_index is not None:
+        return
+    if faiss is None:
+        raise RuntimeError("FAISS not installed but SEARCH_MODE=faiss was requested.")
+
+    conn = get_pg()
+    cols = ", ".join([f"{c}" for c in SEARCH_COLUMNS if c])
+    query = f"SELECT id, {cols} FROM {SEARCH_TABLE} LIMIT %s;"
+    with conn.cursor() as cur:
+        cur.execute(query, (INDEX_ROW_LIMIT,))
+        rows = cur.fetchall()
+        colnames = [desc[0] for desc in cur.description]
+
+    _faiss_id_to_row = []
+    vectors: List[np.ndarray] = []
+    for r in rows:
+        row = {colnames[i]: r[i] for i in range(len(colnames))}
+        text = _concat_columns(row, SEARCH_COLUMNS)
+        emb = _hashing_embed(text, _faiss_dim)
+        vectors.append(emb)
+        _faiss_id_to_row.append(row)
+
+    if not vectors:
+        # Empty index
+        _faiss_index = faiss.IndexFlatIP(_faiss_dim)
+        return
+
+    xb = np.stack(vectors, axis=0)
+    if _faiss_norm:
+        faiss.normalize_L2(xb)
+    _faiss_index = faiss.IndexFlatIP(_faiss_dim)
+    _faiss_index.add(xb)
+
+
+async def _faiss_search(query: str, k: int) -> List[Tuple[Dict[str, Any], float]]:
+    await _ensure_faiss_index()
+    assert _faiss_index is not None
+    if not _faiss_id_to_row:
+        return []
+    qv = _hashing_embed(query, _faiss_dim).reshape(1, -1)
+    if _faiss_norm:
+        faiss.normalize_L2(qv)
+    scores, idxs = _faiss_index.search(qv, min(k, len(_faiss_id_to_row)))
+    out: List[Tuple[Dict[str, Any], float]] = []
+    for i, s in zip(idxs[0].tolist(), scores[0].tolist()):
+        if i < 0 or i >= len(_faiss_id_to_row):
+            continue
+        out.append((_faiss_id_to_row[i], float(s)))
+    return out
+
+
+# -----------------------------------------------------------------------------
+# Mock product data (when database is not available)
+# -----------------------------------------------------------------------------
+MOCK_PRODUCTS = [
+    {"id": 1, "name": "Laptop Pro 14", "description": "High-performance laptop with 14-inch display", "price": 1299.99},
+    {"id": 2, "name": "Noise Cancelling Headphones", "description": "Premium wireless headphones with ANC", "price": 349.99},
+    {"id": 3, "name": "Smartwatch X", "description": "Advanced smartwatch with health tracking", "price": 299.99},
+    {"id": 4, "name": "Laptop Air 13", "description": "Ultralight laptop for everyday use", "price": 999.99},
+    {"id": 5, "name": "Wireless Mouse Pro", "description": "Ergonomic wireless mouse for productivity", "price": 79.99},
+    {"id": 6, "name": "Mechanical Keyboard", "description": "RGB mechanical keyboard for gaming", "price": 149.99},
+    {"id": 7, "name": "4K Monitor 27", "description": "Ultra HD monitor with HDR support", "price": 449.99},
+    {"id": 8, "name": "USB-C Hub", "description": "Multi-port USB-C hub with HDMI", "price": 59.99},
+]
+
+
+def _mock_search(query: str, limit: int) -> List[Tuple[Dict[str, Any], float]]:
+    """Fallback search using mock data when database is not available."""
+    scored = []
+    for product in MOCK_PRODUCTS:
+        text = f"{product.get('name', '')} {product.get('description', '')}"
+        score = _simple_similarity(query, text)
+        scored.append((product, score))
+    scored.sort(key=lambda x: x[1], reverse=True)
+    return scored[:limit]
+
+
+# -----------------------------------------------------------------------------
+# Simple fuzzy search over DB (no FAISS)
+# -----------------------------------------------------------------------------
+async def _simple_search(
+    query: str,
+    limit: int,
+    filter_obj: Optional[Dict[str, Any]],
+) -> List[Tuple[Dict[str, Any], float]]:
+    try:
+        conn = get_pg()
+        table_cols = _get_table_columns(conn, SEARCH_TABLE)
+        where = ""
+        params: List[Any] = []
+        if filter_obj:
+            clauses: List[str] = []
+            for k, v in filter_obj.items():
+                if k in table_cols:
+                    clauses.append(f"{k} = %s")
+                    params.append(v)
+            if clauses:
+                where = " WHERE " + " AND ".join(clauses)
+        cols = ", ".join([f"{c}" for c in SEARCH_COLUMNS if c])
+        sql = f"SELECT id, {cols} FROM {SEARCH_TABLE}{where} LIMIT %s;"
+        params.append(max(limit * 5, 50))
+        with conn.cursor() as cur:
+            cur.execute(sql, params)
+            rows = cur.fetchall()
+            colnames = [desc[0] for desc in cur.description]
+
+        scored: List[Tuple[Dict[str, Any], float]] = []
+        for r in rows:
+            row = {colnames[i]: r[i] for i in range(len(colnames))}
+            text = _concat_columns(row, SEARCH_COLUMNS)
+            scored.append((row, _simple_similarity(query, text)))
+
+        scored.sort(key=lambda x: x[1], reverse=True)
+        return scored[:limit]
+    except Exception:
+        # Fallback to mock data when database is not available
+        return _mock_search(query, limit)
+
+
+# -----------------------------------------------------------------------------
+# Tools
+# -----------------------------------------------------------------------------
+@mcp.tool(name="product-search")
+async def product_search(
+    query: str,
+    columns: Optional[List[str]] = None,
+    filter: Optional[Dict[str, Any]] = None,
+    limit: int = 10,
+    access_token: Optional[str] = None,
+) -> str:
+    """
+    Search tool (CORTEX_SEARCH_SERVICE_QUERY equivalent).
+    Args match Snowflake docs: query, columns, filter, limit.
+    """
+    try:
+        columns = columns or []
+        limit = max(1, min(200, int(limit)))
+
+        if SEARCH_MODE == "faiss":
+            results = await _faiss_search(query, limit)
+        else:
+            results = await _simple_search(query, limit, filter)
+
+        formatted = []
+        for row, score in results:
+            record = {}
+            if columns:
+                for c in columns:
+                    if c in row:
+                        record[c] = row[c]
+            else:
+                # Return all known fields
+                record = {k: v for k, v in row.items()}
+            record["_score"] = score
+            formatted.append(record)
+
+        return json.dumps(
+            {
+                "results": formatted,
+                "request_id": f"req-{abs(hash(query)) % (10**9)}",
+                "mode": SEARCH_MODE,
+            },
+            ensure_ascii=False,
+            indent=2,
+            default=str,
+        )
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+@mcp.tool(name="revenue-semantic-view")
+async def revenue_semantic_view(message: str, access_token: Optional[str] = None) -> str:
+    """
+    Analyst tool (CORTEX_ANALYST_MESSAGE equivalent).
+    We return a text response that includes a suggested SQL for the local Postgres.
+    If OpenAI is configured via OPENAI_API_KEY, we generate AI suggestions; otherwise fallback.
+    """
+    try:
+        client = get_openai()
+        if client is not None:
+            # Minimal schema context to help the model
+            conn = get_pg()
+            revenue_cols = _get_table_columns(conn, "revenue")
+            search_cols = _get_table_columns(conn, SEARCH_TABLE)
+            sys_prompt = (
+                "You are an analyst assistant for a local Postgres database.\n"
+                "Your job is to propose a helpful SQL for the user's request.\n"
+                f"Tables available include revenue({', '.join(revenue_cols) or 'unknown'}) "
+                f"and {SEARCH_TABLE}({', '.join(search_cols) or 'unknown'}).\n"
+                "Constraints:\n"
+                "- Only produce a single SELECT/WITH query that runs on Postgres.\n"
+                "- Limit output rows to ~50.\n"
+                "- If user asks for top revenue, aggregate and order by total revenue desc.\n"
+            )
+            user_prompt = (
+                f"User message:\n{message or ''}\n\n"
+                "Return a concise explanation (2-3 lines) and a SQL fenced in triple backticks."
+            )
+            try:
+                resp = client.chat.completions.create(
+                    model=OPENAI_MODEL,
+                    messages=[
+                        {"role": "system", "content": sys_prompt},
+                        {"role": "user", "content": user_prompt},
+                    ],
+                    temperature=0.2,
+                    max_tokens=700,
+                )
+                ai_text = resp.choices[0].message.content or ""
+                return json.dumps({"content": [{"type": "text", "text": ai_text}]}, indent=2)
+            except Exception:
+                # Fall through to deterministic suggestion
+                pass
+
+        # Fallback deterministic suggestion
+        msg = (message or "").lower()
+        if "top" in msg and "revenue" in msg:
+            sql = (
+                'SELECT product_id, product_name, SUM(revenue) AS total_revenue\n'
+                'FROM "revenue"\n'
+                "GROUP BY product_id, product_name\n"
+                "ORDER BY total_revenue DESC\n"
+                "LIMIT 10;"
+            )
+        else:
+            sql = (
+                'SELECT date, product_id, product_name, revenue\n'
+                'FROM "revenue"\n'
+                "ORDER BY date DESC\n"
+                "LIMIT 20;"
+            )
+        text = (
+            "CORTEX Analyst (local simulation)\n"
+            "Proposed SQL based on your message:\n\n"
+            f"{sql}\n\n"
+            "Note: This is a textual response and does not execute the SQL."
+        )
+        return json.dumps({"content": [{"type": "text", "text": text}]}, indent=2)
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+@mcp.tool(name="sql_exec_tool")
+async def sql_exec_tool(sql: str, access_token: Optional[str] = None) -> str:
+    """
+    SQL execution tool (SYSTEM_EXECUTE_SQL equivalent).
+    Read-only guard: only allow SELECTs.
+    """
+    try:
+        if not sql or sql.strip() == "":
+            return json.dumps({"error": "sql is required"})
+        first = sql.strip().split(None, 1)[0].upper()
+        if first not in {"SELECT", "WITH"}:
+            return json.dumps({"error": "Only read-only SELECT queries are allowed."})
+
+        conn = get_pg()
+        with conn.cursor() as cur:
+            cur.execute(sql)
+            rows = cur.fetchall()
+            colnames = [desc[0] for desc in cur.description]
+        out = [{colnames[i]: r[i] for i in range(len(colnames))} for r in rows] if rows else []
+        return json.dumps({"rows": out, "row_count": len(out)}, indent=2, default=str)
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+@mcp.tool(name="agent_1")
+async def agent_1(message: str, access_token: Optional[str] = None) -> str:
+    """
+    Agent tool (CORTEX_AGENT_RUN equivalent).
+    If OpenAI is configured, use it to suggest the next tool with brief reasoning.
+    Otherwise return a simple heuristic suggestion.
+    """
+    try:
+        client = get_openai()
+        if client is not None:
+            sys_prompt = (
+                "You are a tool-routing assistant in an MCP server.\n"
+                "Available tools:\n"
+                "- product-search: text search over products (FAISS/simple)\n"
+                "- sql_exec_tool: run read-only SELECT/WITH SQL\n"
+                "- revenue-semantic-view: produce helpful SQL suggestion text\n"
+                "Choose a single best next tool and give 1-2 lines of reasoning. "
+                "Respond with a short JSON first, then a brief explanation."
+            )
+            user_prompt = (
+                f"User message:\n{message or ''}\n\n"
+                "Return JSON like {\"suggested_tool\":\"product-search\"} followed by a brief explanation."
+            )
+            try:
+                resp = client.chat.completions.create(
+                    model=OPENAI_MODEL,
+                    messages=[
+                        {"role": "system", "content": sys_prompt},
+                        {"role": "user", "content": user_prompt},
+                    ],
+                    temperature=0.2,
+                    max_tokens=300,
+                )
+                ai_text = resp.choices[0].message.content or ""
+                return json.dumps({"content": [{"type": "text", "text": ai_text}]}, indent=2)
+            except Exception:
+                pass
+
+        suggestion = "sql_exec_tool" if "sql" in (message or "").lower() else "product-search"
+        text = (
+            "Agent V2 (local simulation)\n"
+            f"Received: {message}\n"
+            f"Suggested next tool: {suggestion}"
+        )
+        return json.dumps({"content": [{"type": "text", "text": text}]}, indent=2)
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+# -----------------------------------------------------------------------------
+# Internal reset helper (used by env reset flows, not exposed as an MCP tool)
+# -----------------------------------------------------------------------------
+async def admin_reset(access_token: Optional[str] = None) -> str:
+    """
+    Reset Snowflake environment to clean state.
+    Clears all data from products and revenue tables, then re-seeds with defaults.
+    Called during container reset to prepare for a new task.
+    """
+    global _faiss_index, _faiss_id_to_row
+    
+    try:
+        conn = get_pg()
+        with conn.cursor() as cur:
+            # Clear existing data
+            cur.execute("DELETE FROM revenue;")
+            cur.execute("DELETE FROM products;")
+            
+            # Re-seed with default data
+            cur.execute(
+                """
+                INSERT INTO products (name, description, category) VALUES
+                ('Laptop Pro 14', 'High-performance laptop with retina display', 'Electronics'),
+                ('Noise Cancelling Headphones', 'Over-ear ANC headphones with long battery life', 'Audio'),
+                ('Smartwatch X', 'Fitness tracking and notifications', 'Wearables');
+                """
+            )
+            cur.execute(
+                """
+                INSERT INTO revenue (date, product_id, product_name, revenue) VALUES
+                ('2025-10-01', 1, 'Laptop Pro 14', 250000),
+                ('2025-10-01', 2, 'Noise Cancelling Headphones', 85000),
+                ('2025-10-01', 3, 'Smartwatch X', 120000),
+                ('2025-10-02', 1, 'Laptop Pro 14', 265000),
+                ('2025-10-02', 2, 'Noise Cancelling Headphones', 82000),
+                ('2025-10-02', 3, 'Smartwatch X', 110000);
+                """
+            )
+        
+        # Reset FAISS index
+        _faiss_index = None
+        _faiss_id_to_row = []
+        
+        return json.dumps({
+            "ok": True,
+            "message": "Snowflake environment reset"
+        })
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+def main() -> None:
+    print(f"Starting Snowflake MCP (Local) on http://{HOST}:{PORT}/mcp")
+    print(f"- POSTGRES_DSN: {POSTGRES_DSN}")
+    print(f"- SEARCH_MODE : {SEARCH_MODE}  (faiss available: {faiss is not None})")
+    print(f"- SEARCH_TABLE: {SEARCH_TABLE}, columns={SEARCH_COLUMNS}")
+    mcp.run(transport="http", host=HOST, port=PORT)
+
+
+if __name__ == "__main__":
+    main()
+
+