From 732388d2361321d06ccbdc6d0d4d81b5dc0779ef Mon Sep 17 00:00:00 2001 From: joshua Date: Wed, 26 Nov 2025 21:09:40 -0500 Subject: [PATCH] WIP testing --- .env.example | 13 +++ PLAN.md | 160 +++++++++++++++++++++++++++++++ README.md | 14 ++- pyproject.toml | 73 ++++++++++++++ requirements-dev.txt | 10 ++ requirements.txt | 8 ++ typedb_mcp/__init__.py | 4 + typedb_mcp/resources/__init__.py | 2 + typedb_mcp/resources/schema.py | 53 ++++++++++ typedb_mcp/server.py | 118 +++++++++++++++++++++++ typedb_mcp/utils/__init__.py | 2 + typedb_mcp/utils/config.py | 48 ++++++++++ typedb_mcp/utils/connection.py | 76 +++++++++++++++ 13 files changed, 576 insertions(+), 5 deletions(-) create mode 100644 .env.example create mode 100644 PLAN.md create mode 100644 pyproject.toml create mode 100644 requirements-dev.txt create mode 100644 requirements.txt create mode 100644 typedb_mcp/__init__.py create mode 100644 typedb_mcp/resources/__init__.py create mode 100644 typedb_mcp/resources/schema.py create mode 100644 typedb_mcp/server.py create mode 100644 typedb_mcp/utils/__init__.py create mode 100644 typedb_mcp/utils/config.py create mode 100644 typedb_mcp/utils/connection.py diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..9738eff --- /dev/null +++ b/.env.example @@ -0,0 +1,13 @@ +# TypeDB Connection Configuration +# Copy this file to .env and fill in your actual values + +# TypeDB Server Configuration +TYPEDB_HOST=localhost +TYPEDB_PORT=1729 + +# Database Configuration +TYPEDB_DATABASE=your_database_name + +# Authentication (optional, only if TypeDB server has authentication enabled) +TYPEDB_USERNAME= +TYPEDB_PASSWORD= diff --git a/PLAN.md b/PLAN.md new file mode 100644 index 0000000..d43e1c7 --- /dev/null +++ b/PLAN.md @@ -0,0 +1,160 @@ +# TypeDB MCP Server - Implementation Plan + +## Overview + +This document outlines the plan for building a Model Context Protocol (MCP) server for TypeDB, enabling AI assistants and other MCP clients to interact with TypeDB databases through a standardized interface. + +## Architecture + +### Dependencies +- **TypeDB Python Driver**: Official TypeDB Python SDK for database connectivity +- **MCP Python SDK**: The `mcp` Python package for MCP server implementation +- **Deployment**: Docker container (implementation details TBD - see Future Considerations) + +## MCP Resources + +Resources provide read-only access to schema information: + +1. **get schema** + - Returns the complete database schema + - Format: TypeQL schema definition + - Supports multi-tenancy: database name can be passed as a query parameter in the URI + - URI format: `typedb://schema?database=` + - If database parameter is omitted, uses default from configuration + +## MCP Tools + +### Core Tools + +These tools provide direct query execution and transaction management: + +1. **run_query** + - **Args**: `query` (string) - TypeQL query to execute + - **Returns**: Query response/result + - **Description**: Execute a TypeQL query in a read transaction + +2. **analyze_query** + - **Args**: `query` (string) - TypeQL query to analyze + - **Returns**: Analysis response (query plan, optimization info, etc.) + - **Description**: Analyze a query without executing it + +3. **open_transaction** + - **Args**: `type` (string) - Transaction type ("read" or "write") + - **Returns**: Transaction ID (string) + - **Description**: Open a new transaction and return its identifier + +4. **commit_transaction** + - **Args**: `transaction_id` (string) - Transaction ID to commit + - **Returns**: Commit result (success/failure) + - **Description**: Commit a write transaction + +5. **close_transaction** + - **Args**: `transaction_id` (string) - Transaction ID to close + - **Returns**: Success confirmation + - **Description**: Close a transaction (read or write) + +6. **transaction_run_query** + - **Args**: `transaction_id` (string), `query` (string) - Transaction ID and TypeQL query + - **Returns**: Query response/result + - **Description**: Execute a query within a specific transaction context + +7. **transaction_run_analyze** + - **Args**: `transaction_id` (string), `query` (string) - Transaction ID and TypeQL query + - **Returns**: Analysis response + - **Description**: Analyze a query within a specific transaction context + +### Exploration Tools + +These tools provide convenient ways to explore the graph structure: + +1. **get_concept_with_attribute** + - **Args**: + - `attribute_type` (string) - Type of the attribute + - `attribute_value` (string) - Value of the attribute + - **Returns**: List of IIDs (Instance IDs) matching the criteria + - **Description**: Find concepts (entities/relations) that have a specific attribute value + +2. **get_attributes_by_concept** + - **Args**: `iid` (string) - Instance ID of the concept + - **Returns**: JSON object with all attributes of the concept + - **Description**: Retrieve all attributes owned by a specific concept + +3. **get_relations_of_concept** + - **Args**: `iid` (string) - Instance ID of the concept + - **Returns**: List of relation IIDs along with the role played by the concept + - **Description**: Find all relations where a concept participates and the role it plays + +4. **get_players_of_relation** + - **Args**: `iid` (string) - Instance ID of the relation + - **Returns**: List of role+player pairs + - **Description**: Get all players (concepts) in a relation and their respective roles + +## Implementation Phases + +### Phase 1: Core Infrastructure +- [ ] Set up project structure +- [ ] Initialize MCP server with Python SDK +- [ ] Configure TypeDB connection +- [ ] Basic error handling and logging + +### Phase 2: MCP Resources +- [ ] Implement `get_schema` resource +- [ ] Implement `get_schema_functions` resource +- [ ] Implement `get_schema_types` resource +- [ ] Test resource endpoints + +### Phase 3: Core Tools +- [ ] Implement `run_query` tool +- [ ] Implement `analyze_query` tool +- [ ] Implement transaction management tools: + - [ ] `open_transaction` + - [ ] `commit_transaction` + - [ ] `close_transaction` +- [ ] Implement transaction-scoped query tools: + - [ ] `transaction_run_query` + - [ ] `transaction_run_analyze` +- [ ] Add transaction state management + +### Phase 4: Exploration Tools +- [ ] Implement `get_concept_with_attribute` +- [ ] Implement `get_attributes_by_concept` +- [ ] Implement `get_relations_of_concept` +- [ ] Implement `get_players_of_relation` + +### Phase 5: Testing & Documentation +- [ ] Unit tests for all tools and resources +- [ ] Integration tests with TypeDB +- [ ] API documentation +- [ ] Usage examples + +### Phase 6: Docker Deployment +- [ ] Create Dockerfile +- [ ] Configure container orchestration (TBD - discuss with team) +- [ ] Environment variable configuration +- [ ] Health checks +- [ ] Documentation for deployment + +## Future Considerations + +### Docker Deployment +- **Status**: Pending discussion +- **Questions to resolve**: + - Container orchestration approach (Docker Compose, Kubernetes, etc.) + - Configuration management (environment variables, config files) + - Connection pooling and resource management + - Scaling strategy + - Health check endpoints + +### Additional Features (Potential) +- Query result caching +- Batch query execution +- Schema validation tools +- Query optimization suggestions +- Performance monitoring/metrics + +## Notes + +- The exploration tools provide a GraphQL-like interface for traversing the TypeDB graph, which may be useful for AI assistants that need to explore the database structure. +- Transaction management allows for multi-step operations and write transactions. +- All tools should include proper error handling and validation of inputs. + diff --git a/README.md b/README.md index 2e4d47c..5053b6a 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ This MCP server provides tools and resources for: ## Features ### MCP Resources -- **Schema Access**: Retrieve complete database schemas as TypeQL +- **Schema Access**: Retrieve complete database schemas as TypeQL (supports multi-tenancy via database parameter) ### MCP Tools - **Query Execution**: Run and analyze TypeQL queries @@ -54,7 +54,9 @@ Set the following environment variables. You can either: Required variables: - `TYPEDB_HOST`: TypeDB server host (default: `localhost`) - `TYPEDB_PORT`: TypeDB server port (default: `1729`) -- `TYPEDB_DATABASE`: Database name to connect to + +Optional variables: +- `TYPEDB_DATABASE`: Default database name (can also be specified per resource request) Optional variables (only if TypeDB server has authentication enabled): - `TYPEDB_USERNAME`: Username for authentication @@ -139,9 +141,11 @@ pytest tests/ ## MCP Resources Reference -- `schema`: Complete database schema -- `schema_functions`: All defined functions -- `schema_types`: All type definitions +- `schema`: Complete database schema in TypeQL format + - URI: `typedb://schema` + - Query parameter: `database` (optional) - database name to query + - Example: `typedb://schema?database=my_database` + - If database parameter is not provided, uses the default from `TYPEDB_DATABASE` environment variable ## Contributing diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b788d38 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,73 @@ +[project] +name = "typedb-mcp" +version = "0.1.0" +description = "Model Context Protocol server for TypeDB" +readme = "README.md" +requires-python = ">=3.9" +license = {text = "Apache-2.0"} +authors = [ + {name = "Vaticle", email = "hello@vaticle.com"} +] +keywords = ["typedb", "mcp", "model-context-protocol", "database", "graph"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", +] + +dependencies = [ + "typedb-driver>=2.28.0", + "mcp>=0.1.0", + "python-dotenv>=1.0.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "pytest-asyncio>=0.21.0", + "pytest-cov>=4.0.0", + "black>=23.0.0", + "ruff>=0.1.0", + "mypy>=1.0.0", +] + +[project.urls] +Homepage = "https://github.com/vaticle/typedb-mcp" +Documentation = "https://github.com/vaticle/typedb-mcp#readme" +Repository = "https://github.com/vaticle/typedb-mcp" +Issues = "https://github.com/vaticle/typedb-mcp/issues" + +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.black] +line-length = 100 +target-version = ['py39', 'py310', 'py311', 'py312'] + +[tool.ruff] +line-length = 100 +target-version = "py39" + +[tool.ruff.lint] +select = ["E", "F", "I", "N", "W", "UP"] +ignore = ["E501"] # Line too long (handled by black) + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +asyncio_mode = "auto" + +[tool.mypy] +python_version = "3.9" +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = false + diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..298aca3 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,10 @@ +# Development dependencies +-r requirements.txt + +pytest>=7.0.0 +pytest-asyncio>=0.21.0 +pytest-cov>=4.0.0 +black>=23.0.0 +ruff>=0.1.0 +mypy>=1.0.0 + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..36fbc21 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +# Core dependencies +typedb-driver>=2.28.0 +mcp>=0.1.0 +python-dotenv>=1.0.0 + +# Development dependencies (install with: pip install -r requirements.txt -r requirements-dev.txt) +# See pyproject.toml for dev dependencies + diff --git a/typedb_mcp/__init__.py b/typedb_mcp/__init__.py new file mode 100644 index 0000000..cd872ef --- /dev/null +++ b/typedb_mcp/__init__.py @@ -0,0 +1,4 @@ +"""TypeDB MCP Server - Model Context Protocol server for TypeDB.""" + +__version__ = "0.1.0" + diff --git a/typedb_mcp/resources/__init__.py b/typedb_mcp/resources/__init__.py new file mode 100644 index 0000000..b9a03f6 --- /dev/null +++ b/typedb_mcp/resources/__init__.py @@ -0,0 +1,2 @@ +"""MCP resources for TypeDB schema access.""" + diff --git a/typedb_mcp/resources/schema.py b/typedb_mcp/resources/schema.py new file mode 100644 index 0000000..f110ea3 --- /dev/null +++ b/typedb_mcp/resources/schema.py @@ -0,0 +1,53 @@ +"""Schema resources for MCP server.""" + +from typing import Optional + +from typedb.common.exception import TypeDBDriverException + +from typedb_mcp.utils.connection import TypeDBConnection + + +class SchemaResource: + """MCP resource for accessing TypeDB schema information.""" + + def __init__(self, connection: TypeDBConnection): + """Initialize schema resource. + + Args: + connection: TypeDB connection manager + """ + self.connection = connection + + def get_schema(self, database_name: Optional[str] = None) -> str: + """Get the complete database schema as TypeQL. + + Args: + database_name: Name of the database (if None, uses config default) + + Returns: + Complete schema definition in TypeQL format + + Raises: + ValueError: If database is not found or not specified + Exception: If schema retrieval fails + """ + # Use provided database name or fall back to config default + db_name = database_name or self.connection.config.database + + if not db_name: + raise ValueError("Database name must be provided either as parameter or in config") + + driver = self.connection.get_driver() + database = driver.databases.get(db_name) + + if not database: + raise ValueError(f"Database '{db_name}' not found") + + try: + # Get schema from database + schema = database.schema() + return schema + except TypeDBDriverException as e: + raise Exception(f"Failed to retrieve schema from database '{db_name}': {e}") from e + + diff --git a/typedb_mcp/server.py b/typedb_mcp/server.py new file mode 100644 index 0000000..1216bc0 --- /dev/null +++ b/typedb_mcp/server.py @@ -0,0 +1,118 @@ +"""Main MCP server for TypeDB.""" + +import asyncio +import sys +from urllib.parse import parse_qs, urlparse + +from mcp.server import Server +from mcp.server.stdio import stdio_server +from mcp.types import Resource, Uri + +from typedb_mcp.resources.schema import SchemaResource +from typedb_mcp.utils.config import TypeDBConfig +from typedb_mcp.utils.connection import TypeDBConnection + + +class TypeDBMCPServer: + """MCP server for TypeDB.""" + + def __init__(self, config: TypeDBConfig | None = None): + """Initialize the TypeDB MCP server. + + Args: + config: TypeDB configuration (if None, loads from environment) + """ + if config is None: + config = TypeDBConfig() + + self.config = config + self.connection = TypeDBConnection(config) + self.schema_resource = SchemaResource(self.connection) + self.server = Server("typedb-mcp") + + # Register resources + self._register_resources() + + def _register_resources(self) -> None: + """Register MCP resources.""" + # Register schema resource + @self.server.list_resources() + async def list_resources() -> list[Resource]: + """List available resources.""" + return [ + Resource( + uri=Uri("typedb://schema"), + name="Schema", + description="Complete database schema in TypeQL format. Use query parameter 'database' to specify database name (e.g., typedb://schema?database=my_db)", + mimeType="text/plain", + ), + ] + + @self.server.read_resource() + async def read_resource(uri: Uri) -> str: + """Read a resource by URI. + + Args: + uri: Resource URI (may include query parameter 'database') + + Returns: + Resource content as string + + Raises: + ValueError: If URI is not recognized + """ + uri_str = str(uri) + parsed = urlparse(uri_str) + + # Extract database name from query parameters + database_name = None + if parsed.query: + query_params = parse_qs(parsed.query) + if "database" in query_params: + database_name = query_params["database"][0] + + try: + if parsed.scheme == "typedb" and parsed.netloc == "schema": + return self.schema_resource.get_schema(database_name=database_name) + else: + raise ValueError(f"Unknown resource URI: {uri_str}") + except Exception as e: + raise ValueError(f"Failed to read resource {uri_str}: {e}") from e + + async def run(self) -> None: + """Run the MCP server.""" + # Connect to TypeDB + try: + self.connection.connect() + print(f"Connected to TypeDB at {self.config.address}", file=sys.stderr) + except Exception as e: + print(f"Failed to connect to TypeDB: {e}", file=sys.stderr) + sys.exit(1) + + # Run the server + async with stdio_server() as (read_stream, write_stream): + await self.server.run( + read_stream, + write_stream, + self.server.create_initialization_options(), + ) + + def cleanup(self) -> None: + """Clean up resources.""" + self.connection.close() + + +async def main() -> None: + """Main entry point for the MCP server.""" + server = TypeDBMCPServer() + try: + await server.run() + except KeyboardInterrupt: + pass + finally: + server.cleanup() + + +if __name__ == "__main__": + asyncio.run(main()) + diff --git a/typedb_mcp/utils/__init__.py b/typedb_mcp/utils/__init__.py new file mode 100644 index 0000000..ae0b66a --- /dev/null +++ b/typedb_mcp/utils/__init__.py @@ -0,0 +1,2 @@ +"""Utility modules for TypeDB MCP server.""" + diff --git a/typedb_mcp/utils/config.py b/typedb_mcp/utils/config.py new file mode 100644 index 0000000..1cd3053 --- /dev/null +++ b/typedb_mcp/utils/config.py @@ -0,0 +1,48 @@ +"""Configuration management for TypeDB MCP server.""" + +import os +from typing import Optional + +from dotenv import load_dotenv + +# Load environment variables from .env file if it exists +load_dotenv() + + +class TypeDBConfig: + """Configuration for TypeDB connection.""" + + def __init__( + self, + host: Optional[str] = None, + port: Optional[int] = None, + database: Optional[str] = None, + username: Optional[str] = None, + password: Optional[str] = None, + ): + """Initialize TypeDB configuration from parameters or environment variables. + + Args: + host: TypeDB server host (defaults to TYPEDB_HOST env var or 'localhost') + port: TypeDB server port (defaults to TYPEDB_PORT env var or 1729) + database: Database name (defaults to TYPEDB_DATABASE env var) + username: Username for authentication (defaults to TYPEDB_USERNAME env var) + password: Password for authentication (defaults to TYPEDB_PASSWORD env var) + """ + self.host = host or os.getenv("TYPEDB_HOST", "localhost") + self.port = port or int(os.getenv("TYPEDB_PORT", "1729")) + self.database = database or os.getenv("TYPEDB_DATABASE") + self.username = username or os.getenv("TYPEDB_USERNAME") + self.password = password or os.getenv("TYPEDB_PASSWORD") + + # Database is optional since it can be passed per request + + @property + def address(self) -> str: + """Get the TypeDB server address.""" + return f"{self.host}:{self.port}" + + def has_credentials(self) -> bool: + """Check if credentials are provided.""" + return bool(self.username and self.password) + diff --git a/typedb_mcp/utils/connection.py b/typedb_mcp/utils/connection.py new file mode 100644 index 0000000..820903f --- /dev/null +++ b/typedb_mcp/utils/connection.py @@ -0,0 +1,76 @@ +"""TypeDB connection management.""" + +from typing import Optional + +import typedb.driver as typedb +from typedb.api.connection.driver import Driver +from typedb.api.connection.options import TypeDBOptions +from typedb.common.exception import TypeDBDriverException + +from typedb_mcp.utils.config import TypeDBConfig + + +class TypeDBConnection: + """Manages TypeDB driver connection.""" + + def __init__(self, config: TypeDBConfig): + """Initialize TypeDB connection manager. + + Args: + config: TypeDB configuration + """ + self.config = config + self._driver: Optional[Driver] = None + + def connect(self) -> Driver: + """Create and return a TypeDB driver connection. + + Returns: + TypeDB Driver instance + + Raises: + TypeDBDriverException: If connection fails + """ + if self._driver is not None and self._driver.is_open(): + return self._driver + + try: + # Create credentials if username/password are provided + credentials = None + if self.config.has_credentials(): + credentials = typedb.TypeDBCredential( + self.config.username, + self.config.password, + tls_enabled=False, # Set to True for TypeDB Cloud + ) + + # Create driver options + driver_options = TypeDBOptions() + + # Create driver + self._driver = typedb.driver( + self.config.address, + credentials=credentials, + driver_options=driver_options, + ) + + return self._driver + except TypeDBDriverException as e: + raise ConnectionError(f"Failed to connect to TypeDB at {self.config.address}: {e}") from e + + def close(self) -> None: + """Close the TypeDB driver connection.""" + if self._driver is not None and self._driver.is_open(): + self._driver.close() + self._driver = None + + def is_connected(self) -> bool: + """Check if the driver is connected.""" + return self._driver is not None and self._driver.is_open() + + def get_driver(self) -> Driver: + """Get the current driver, connecting if necessary.""" + if not self.is_connected(): + self.connect() + return self._driver +