Update mcp_server.py

mcp_server.py

@@ -1,78 +1,68 @@
-"""
-mcp_server.py – MCP tool server exposing:
-  • query_database    – SQLite sample
-  • query_hubspot     – HubSpot via hubspot-api-client
-
-Secrets expected in env:
-  HUBSPOT_TOKEN
-"""
-
 import os
 import json
 import sqlite3
-
 from mcp.server.fastmcp import FastMCP
 
-# HubSpot
+# HubSpot SDK
 from hubspot import HubSpot
 from hubspot.crm.contacts import ApiException as HSContactsError
 
-# ──────────────────────────────────────────────────────────────
+# ─────────────────────────── MCP Server Init ─────────────────────────────
 mcp = FastMCP("EnterpriseData")
 
-# ─── 1) In-memory SQLite sample ────────────────────────────────────────
+# ───────────────────── In-Memory SQLite Sample ───────────────────────────
 conn = sqlite3.connect(":memory:", check_same_thread=False)
-cur  = conn.cursor()
+cur = conn.cursor()
 cur.execute("""
 CREATE TABLE Customers (
-    CustomerID      INTEGER PRIMARY KEY AUTOINCREMENT,
-    Name            TEXT,
-    Region          TEXT,
-    LastOrderDate   TEXT
+    CustomerID INTEGER PRIMARY KEY AUTOINCREMENT,
+    Name TEXT,
+    Region TEXT,
+    LastOrderDate TEXT
 )
 """)
 cur.executemany(
-    "INSERT INTO Customers (Name, Region, LastOrderDate) VALUES (?,?,?)",
+    "INSERT INTO Customers (Name, Region, LastOrderDate) VALUES (?, ?, ?)",
     [
-        ("Acme Corp",  "Northeast", "2024-12-01"),
-        ("Beta Inc",   "West",      "2025-06-01"),
-        ("Gamma Co",   "Northeast", "2023-09-15"),
-        ("Delta LLC",  "South",     "2025-03-20"),
-        ("Epsilon Ltd","Northeast", "2025-07-10"),
+        ("Acme Corp", "Northeast", "2024-12-01"),
+        ("Beta Inc", "West", "2025-06-01"),
+        ("Gamma Co", "Northeast", "2023-09-15"),
+        ("Delta LLC", "South", "2025-03-20"),
+        ("Epsilon Ltd", "Northeast", "2025-07-10"),
     ],
 )
 conn.commit()
 
 @mcp.tool()
 def query_database(sql: str) -> str:
-    """Run SQL against the in-memory Customers table and return JSON rows."""
+    """Run SQL query on Customers table and return results."""
     try:
         cur.execute(sql)
-        cols = [d[0] for d in cur.description or []]
-        rows = [dict(zip(cols, r)) for r in cur.fetchall()]
-        return json.dumps(rows)
+        cols = [desc[0] for desc in cur.description or []]
+        rows = [dict(zip(cols, row)) for row in cur.fetchall()]
+        return json.dumps(rows, indent=2)
     except Exception as e:
         return json.dumps({"error": str(e)})
 
-# ─── 2) HubSpot tool ───────────────────────────────────────────────
-hs_client = HubSpot(access_token=os.getenv("HUBSPOT_TOKEN"))
+# ───────────────────────── HubSpot Integration ───────────────────────────
+hs_token = os.getenv("HUBSPOT_TOKEN")
+hs_client = HubSpot(access_token=hs_token)
 
 @mcp.tool()
-def query_hubspot(object_type: str, limit: int = 100) -> str:
+def query_hubspot_contacts(limit: int = 100) -> str:
     """
-    Fetch up to `limit` objects of type `contacts`, `companies`, or `deals`.
-    Example: object_type="contacts"
+    Fetch up to `limit` HubSpot contacts.
+    Example: query_hubspot_contacts(limit=50)
     """
     try:
-        api = getattr(hs_client.crm, object_type)
-        page = api.basic_api.get_page(limit=limit)
-        items = [r.to_dict() for r in page.results]
-        return json.dumps(items)
-    except HSContactsError as he:
-        return json.dumps({"error": he.body})
+        page = hs_client.crm.contacts.basic_api.get_page(limit=limit)
+        contacts = [contact.to_dict() for contact in page.results]
+        return json.dumps(contacts, indent=2)
+    except HSContactsError as e:
+        return json.dumps({"error": e.body})
     except Exception as e:
         return json.dumps({"error": str(e)})
 
-# ─── 3) Start the MCP server ─────────────────────────────────────
+# ──────────────────────── Start the MCP Server ───────────────────────────
 if __name__ == "__main__":
     mcp.run(transport="stdio")