gabi.py•2.21 kB
from typing import Any
import os
import httpx
from mcp.server.fastmcp import FastMCP
# Initialize FastMCP server
mcp = FastMCP("gabi")
# Constants
USER_AGENT = "gabi-mcp"
def build_gabi_url(path: str) -> str:
"""Build the endpoint for GABI"""
return os.environ.get("GABI_ENDPOINT") + path
def get_access_token() -> str:
"""Get a valid token from the environment"""
return os.environ["ACCESS_TOKEN"]
def build_headers() -> dict:
"""Build headers for HTTP requests"""
return {
"User-Agent": USER_AGENT,
"Authorization": f"Bearer {get_access_token()}",
"Content-Type": "application/json"
}
async def make_gabi_get_request(url: str) -> dict[str, Any] | None:
"""Make a GET request to GABI."""
headers = {}
try:
headers = build_headers()
except KeyError:
return "ERROR: token not defined!"
async with httpx.AsyncClient() as client:
try:
response = await client.get(url, headers=headers, timeout=30.0)
response.raise_for_status()
return response.json()
except Exception as e:
return e
async def make_gabi_post_request(url: str, query: str) -> dict[str, Any] | None:
"""Make a POST request to GABI."""
headers = {}
try:
headers = build_headers()
except KeyError:
return "ERROR: token not defined!"
data = {
"query": query
}
async with httpx.AsyncClient() as client:
try:
response = await client.post(url, headers=headers, timeout=30.0, json=data)
response.raise_for_status()
return response.json()
except Exception as e:
return e
@mcp.tool()
async def get_db_name() -> str:
"""Get the name of the current database.
Args:
"""
return await make_gabi_get_request(build_gabi_url("/dbname"))
@mcp.tool()
async def get_query_result(query: str) -> str:
"""Executes an SQL query on a single line against the current database.
Args: query
"""
return await make_gabi_post_request(build_gabi_url("/query"), query)
if __name__ == "__main__":
# Initialize and run the server
mcp.run(transport='stdio')