initial

Dockerfile

@@ -0,0 +1,40 @@
+FROM python:3.10-slim
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    gcc \
+    python3-dev \
+    openssl \
+    curl \
+    ca-certificates \
+    gnupg \
+    build-essential && \
+    rm -rf /var/lib/apt/lists/*
+
+# Set up MariaDB's Python connector dependencies
+RUN curl -LsSO https://r.mariadb.com/downloads/mariadb_repo_setup && \
+    echo "c4a0f3dade02c51a6a28ca3609a13d7a0f8910cccbb90935a2f218454d3a914a  mariadb_repo_setup" | sha256sum -c - && \
+    chmod +x mariadb_repo_setup && \
+    ./mariadb_repo_setup --mariadb-server-version="mariadb-11.7" && \
+    rm mariadb_repo_setup && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+    libmariadb3 \
+    libmariadb-dev && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install uv package manager
+RUN pip install --no-cache-dir uv
+
+WORKDIR /app
+
+# Copy project files
+COPY . /app
+
+# Install project dependencies
+RUN uv sync
+
+EXPOSE 8000
+
+CMD ["uv", "run", "mcp-server-mariadb-vector", "--transport", "sse", "--host", "0.0.0.0"]

app.py

@@ -0,0 +1,7 @@
+from fastapi import FastAPI
+
+app = FastAPI()
+
+@app.get("/")
+def greet_json():
+    return {"Hello": "World!"}

requirements.txt

@@ -0,0 +1,2 @@
+fastapi
+uvicorn[standard]

src/mcp_server_mariadb_vector/app_context.py

@@ -0,0 +1,33 @@
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import AsyncIterator
+
+import mariadb
+from mcp.server.fastmcp import FastMCP
+
+from mcp_server_mariadb_vector.settings import DatabaseSettings
+
+
+@dataclass
+class AppContext:
+    conn: mariadb.Connection
+
+
+@asynccontextmanager
+async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
+    """Open a MariaDB connection for the duration of the FastMCP session."""
+
+    cfg = DatabaseSettings()
+    conn = mariadb.connect(
+        host=cfg.host,
+        port=cfg.port,
+        user=cfg.user,
+        password=cfg.password,
+        database=cfg.database,
+    )
+    conn.autocommit = True
+
+    try:
+        yield AppContext(conn=conn)
+    finally:
+        conn.close()

src/mcp_server_mariadb_vector/embeddings/base.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import List
+
+
+class EmbeddingProviderType(Enum):
+    OPENAI = "openai"
+    TEST = "test"
+    # SENTENCE_TRANSFORMERS = "sentence-transformers"
+
+
+class EmbeddingProvider(ABC):
+    """Abstract base class for embedding providers."""
+
+    @abstractmethod
+    def length_of_embedding(self) -> int:
+        """Get the length of the embedding for a given model."""
+        pass
+
+    @abstractmethod
+    def embed_documents(self, documents: List[str]) -> List[List[float]]:
+        """Embed a list of documents into vectors."""
+        pass
+
+    @abstractmethod
+    def embed_query(self, query: str) -> List[float]:
+        """Embed a query into a vector."""
+        pass

src/mcp_server_mariadb_vector/embeddings/factory.py

@@ -0,0 +1,22 @@
+from mcp_server_mariadb_vector.embeddings.base import (
+    EmbeddingProvider,
+    EmbeddingProviderType,
+)
+from mcp_server_mariadb_vector.embeddings.openai import OpenAIEmbeddingProvider
+from mcp_server_mariadb_vector.embeddings.test import TestEmbeddingProvider
+from mcp_server_mariadb_vector.settings import EmbeddingSettings
+
+
+def create_embedding_provider(settings: EmbeddingSettings) -> EmbeddingProvider:
+    """
+    Create an instance of the specified embedding provider.
+
+    Args:
+        settings: The settings for the embedding provider.
+    """
+    if settings.provider == EmbeddingProviderType.OPENAI:
+        return OpenAIEmbeddingProvider(settings.model, settings.openai_api_key)
+    elif settings.provider == EmbeddingProviderType.TEST:
+        return TestEmbeddingProvider()
+    else:
+        raise ValueError(f"Unsupported embedding provider: {settings.provider}")

src/mcp_server_mariadb_vector/embeddings/openai.py

@@ -0,0 +1,48 @@
+from typing import List
+
+from openai import OpenAI
+
+from mcp_server_mariadb_vector.embeddings.base import EmbeddingProvider
+
+
+class OpenAIEmbeddingProvider(EmbeddingProvider):
+    """
+    OpenAI implementation of the embedding provider.
+
+    Args:
+        model: The name of the OpenAI model to use.
+    """
+
+    def __init__(self, model: str, api_key: str):
+        self.model = model
+        self.client = OpenAI(api_key=api_key)
+
+    def length_of_embedding(self) -> int:
+        """Get the length of the embedding for a given model."""
+        if self.model == "text-embedding-3-small":
+            return 1536
+        elif self.model == "text-embedding-3-large":
+            return 3072
+        else:
+            raise ValueError(f"Unknown embedding model: {self.model}")
+
+    def embed_documents(self, documents: List[str]) -> List[List[float]]:
+        """Embed a list of documents into vectors."""
+        embeddings = [
+            self.client.embeddings.create(
+                model=self.model,
+                input=document,
+            )
+            .data[0]
+            .embedding
+            for document in documents
+        ]
+        return embeddings
+
+    def embed_query(self, query: str) -> List[float]:
+        """Embed a query into a vector."""
+        embedding = self.client.embeddings.create(
+            model=self.model,
+            input=query,
+        )
+        return embedding.data[0].embedding

src/mcp_server_mariadb_vector/embeddings/test.py

@@ -0,0 +1,18 @@
+from typing import List
+
+from mcp_server_mariadb_vector.embeddings.base import EmbeddingProvider
+
+
+class TestEmbeddingProvider(EmbeddingProvider):
+    """
+    Embedding provider for testing.
+    """
+
+    def length_of_embedding(self) -> int:
+        return 3
+
+    def embed_documents(self, documents: List[str]) -> List[List[float]]:
+        return [[0.1, 0.2, 0.3]] * len(documents)
+
+    def embed_query(self, query: str) -> List[float]:
+        return [0.1, 0.2, 0.3]

src/mcp_server_mariadb_vector/server.py

@@ -0,0 +1,257 @@
+import argparse
+import json
+from typing import Annotated, List, Literal
+
+import mariadb
+from fastmcp import Context, FastMCP
+from pydantic import Field
+
+from mcp_server_mariadb_vector.app_context import app_lifespan
+from mcp_server_mariadb_vector.embeddings.factory import create_embedding_provider
+from mcp_server_mariadb_vector.settings import EmbeddingSettings
+
+mcp = FastMCP(
+    "Mariadb Vector",
+    lifespan=app_lifespan,
+    dependencies=["mariadb", "openai", "pydantic", "pydantic-settings"],
+)
+
+
+embedding_provider = create_embedding_provider(EmbeddingSettings())
+
+
+@mcp.tool()
+def mariadb_create_vector_store(
+    ctx: Context,
+    vector_store_name: Annotated[
+        str,
+        Field(description="The name of the vector store to create"),
+    ],
+    distance_function: Annotated[
+        Literal["euclidean", "cosine"],
+        Field(description="The distance function to use."),
+    ] = "euclidean",
+) -> str:
+    """Create a vector store in the MariaDB database."""
+
+    embedding_length = embedding_provider.length_of_embedding()
+
+    schema_query = f"""
+    CREATE TABLE `{vector_store_name}` (
+        id BIGINT UNSIGNED PRIMARY KEY AUTO_INCREMENT,
+        document LONGTEXT NOT NULL,
+        embedding VECTOR({embedding_length}) NOT NULL,
+        metadata JSON NOT NULL,
+        VECTOR INDEX (embedding) DISTANCE={distance_function}
+    )
+    """
+
+    try:
+        conn = ctx.request_context.lifespan_context.conn
+        with conn.cursor() as cursor:
+            cursor.execute(schema_query)
+    except mariadb.Error as e:
+        return f"Error creating vector store `{vector_store_name}`: {e}"
+
+    return f"Vector store `{vector_store_name}` created successfully."
+
+
+def is_vector_store(conn, table: str, embedding_length: int) -> bool:
+    """
+    True if `table` has the right schema, with vectors of the correct length, and a VECTOR index.
+    """
+
+    with conn.cursor(dictionary=True) as cur:
+        # check columns
+        cur.execute(f"SHOW COLUMNS FROM `{table}`")
+        rows = {r["Field"]: r for r in cur}
+
+        if set(rows) != {"id", "document", "embedding", "metadata"}:
+            return False
+
+        # id
+        id_type = rows["id"]["Type"].lower()
+        if id_type != "bigint(20) unsigned":
+            return False
+        if (
+            rows["id"]["Null"] != "NO"
+            or rows["id"]["Key"] != "PRI"
+            or "auto_increment" not in rows["id"]["Extra"].lower()
+        ):
+            return False
+
+        # document
+        if (
+            rows["document"]["Type"].lower() != "longtext"
+            or rows["document"]["Null"] != "NO"
+        ):
+            return False
+
+        # embedding
+        if (
+            rows["embedding"]["Type"].lower() != f"vector({embedding_length})"
+            or rows["embedding"]["Null"] != "NO"
+        ):
+            return False
+
+        # metadata
+        if (
+            rows["metadata"]["Type"].lower() != "longtext"
+            or rows["metadata"]["Null"] != "NO"
+        ):
+            return False
+
+        # check vector index
+        cur.execute(f"""
+            SHOW INDEX FROM `{table}`
+            WHERE Index_type = 'VECTOR' AND Column_name = 'embedding'
+        """)
+        if cur.fetchone() is None:
+            return False
+
+    return True
+
+
+@mcp.tool()
+def mariadb_list_vector_stores(ctx: Context) -> str:
+    """List all vector stores in a MariaDB database."""
+    try:
+        conn = ctx.request_context.lifespan_context.conn
+        with conn.cursor() as cursor:
+            cursor.execute("SHOW TABLES")
+            tables = [table[0] for table in cursor]
+    except mariadb.Error as e:
+        return f"Error listing vector stores: {e}"
+
+    embedding_length = embedding_provider.length_of_embedding()
+    vector_stores = [
+        table for table in tables if is_vector_store(conn, table, embedding_length)
+    ]
+
+    return "Vector stores: " + ", ".join(vector_stores)
+
+
+@mcp.tool()
+def mariadb_delete_vector_store(
+    ctx: Context,
+    vector_store_name: Annotated[
+        str, Field(description="The name of the vector store to delete.")
+    ],
+) -> str:
+    """Delete a vector store in the MariaDB database."""
+
+    try:
+        conn = ctx.request_context.lifespan_context.conn
+        with conn.cursor() as cursor:
+            cursor.execute(f"DROP TABLE `{vector_store_name}`")
+    except mariadb.Error as e:
+        return f"Error deleting vector store `{vector_store_name}`: {e}"
+
+    return f"Vector store `{vector_store_name}` deleted successfully."
+
+
+@mcp.tool()
+def mariadb_insert_documents(
+    ctx: Context,
+    vector_store_name: Annotated[
+        str, Field(description="The name of the vector store to insert documents into.")
+    ],
+    documents: Annotated[
+        List[str], Field(description="The documents to insert into the vector store.")
+    ],
+    metadata: Annotated[
+        List[dict], Field(description="The metadata of the documents to insert.")
+    ],
+) -> str:
+    """Insert a document into a vector store."""
+
+    embeddings = embedding_provider.embed_documents(documents)
+
+    metadata_json = [json.dumps(metadata) for metadata in metadata]
+
+    insert_query = f"""
+    INSERT INTO `{vector_store_name}` (document, embedding, metadata) VALUES (%s, VEC_FromText(%s), %s)
+    """
+    try:
+        conn = ctx.request_context.lifespan_context.conn
+        with conn.cursor() as cursor:
+            cursor.executemany(
+                insert_query, list(zip(documents, embeddings, metadata_json))
+            )
+    except mariadb.Error as e:
+        return f"Error inserting documents`{vector_store_name}`: {e}"
+
+    return f"Documents inserted into `{vector_store_name}` successfully."
+
+
+@mcp.tool()
+def mariadb_search_vector_store(
+    ctx: Context,
+    query: Annotated[str, Field(description="The query to search for.")],
+    vector_store_name: Annotated[
+        str, Field(description="The name of the vector store to search.")
+    ],
+    k: Annotated[int, Field(gt=0, description="The number of results to return.")] = 5,
+) -> str:
+    """Search a vector store for the most similar documents to a query."""
+
+    embedding = embedding_provider.embed_query(query)
+
+    search_query = f"""
+    SELECT 
+        document,
+        metadata,
+        VEC_DISTANCE_EUCLIDEAN(embedding, VEC_FromText(%s)) AS distance
+    FROM `{vector_store_name}`
+    ORDER BY distance ASC
+    LIMIT %s
+    """
+
+    try:
+        conn = ctx.request_context.lifespan_context.conn
+        with conn.cursor(buffered=True) as cursor:
+            cursor.execute(
+                search_query,
+                (str(embedding), k),
+            )
+            rows = cursor.fetchall()
+    except mariadb.Error as e:
+        return f"Error searching vector store`{vector_store_name}`: {e}"
+
+    if not rows:
+        return "No similar context found."
+
+    return "\n\n".join(
+        f"Document: {row[0]}\nMetadata: {json.loads(row[1])}\nDistance: {row[2]}"
+        for row in rows
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--transport",
+        choices=["stdio", "sse"],
+        default="stdio",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="127.0.0.1",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+    )
+
+    args = parser.parse_args()
+
+    if args.transport == "sse":
+        mcp.run(transport=args.transport, host=args.host, port=args.port)
+    else:
+        mcp.run(transport=args.transport)
+
+
+if __name__ == "__main__":
+    main()

src/mcp_server_mariadb_vector/settings.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+from mcp_server_mariadb_vector.embeddings.base import EmbeddingProviderType
+
+
+class DatabaseSettings(BaseSettings):
+    host: str = Field(default="127.0.0.1", alias="MARIADB_HOST")
+    port: int = Field(default=3306, alias="MARIADB_PORT")
+    user: str = Field(..., alias="MARIADB_USER")
+    password: str = Field(..., alias="MARIADB_PASSWORD")
+    database: str = Field(..., alias="MARIADB_DATABASE")
+
+
+class EmbeddingSettings(BaseSettings):
+    provider: EmbeddingProviderType = Field(
+        default=EmbeddingProviderType.OPENAI, alias="EMBEDDING_PROVIDER"
+    )
+    model: str = Field(default="text-embedding-3-small", alias="EMBEDDING_MODEL")
+    openai_api_key: Optional[str] = Field(default=None, alias="OPENAI_API_KEY")