import os
import httpx
from pydantic import BaseModel, Field
from langchain_core.tools import StructuredTool
class MCPSession:
def __init__(self, base_url: str, bot_token: str):
self.endpoint = base_url.rstrip("/") + "/mcp"
self.headers = {
"Authorization": f"Bearer {bot_token}",
"Content-Type": "application/json",
"Accept": "application/json, text/event-stream",
}
self.client = httpx.Client(timeout=30)
self._next_id = 1
self._initialize()
def _request(self, method: str, params: dict | None = None, notify: bool = False) -> dict | None:
payload = {"jsonrpc": "2.0", "method": method}
if not notify:
payload["id"] = self._next_id
self._next_id += 1
if params is not None:
payload["params"] = params
response = self.client.post(self.endpoint, headers=self.headers, json=payload)
response.raise_for_status()
return response.json() if response.content else None
def _initialize(self) -> None:
self._request(
"initialize",
{
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "langchain-agent", "version": "1.0.0"},
},
)
self._request("notifications/initialized", notify=True)
def call_tool(self, name: str, arguments: dict) -> str:
response = self._request("tools/call", {"name": name, "arguments": arguments})
assert response is not None, "tools/call must not be sent as a notification"
result = response["result"]
text = result["content"][0]["text"]
if result.get("isError"):
raise RuntimeError(f"{name} failed: {text}")
return text
session = MCPSession(
base_url=os.environ["OM_SERVER_URL"],
bot_token=os.environ["OM_BOT_TOKEN"],
)
class SearchMetadataInput(BaseModel):
query: str = Field(description="Keywords to search for")
entity_type: str | None = Field(default=None, description="Optional entity type filter, e.g. 'table'")
def search_metadata(query: str, entity_type: str | None = None) -> str:
arguments = {"query": query}
if entity_type:
arguments["entityType"] = entity_type
return session.call_tool("search_metadata", arguments)
search_metadata_tool = StructuredTool.from_function(
func=search_metadata,
name="search_metadata",
description="Search the OpenMetadata catalog by keyword.",
args_schema=SearchMetadataInput,
)