zentreopy_mcp_server.py

"""
ZeroEntropy MCP Server
"""

import os
from mcp.server.fastmcp import FastMCP
from pydantic import Field
import mcp.types as types
from zeroentropy import ZeroEntropy

# Initialize ZeroEntropy client
ZEROENTROPY_API_KEY = os.getenv("ZEROENTROPY_API_KEY")
client = ZeroEntropy(api_key=ZEROENTROPY_API_KEY)

mcp = FastMCP("ZeroEntropy Server", port=3000, stateless_http=True, debug=True)

@mcp.tool(
    title="Search African History",
    description="Search the African history book collection using ZeroEntropy",
)
def search_african_history(
    query: str = Field(description="The search query"),
    k: int = Field(description="Number of results to return", default=21),
    reranker: str = Field(description="Reranker to use", default="zerank-1")
) -> str:
    """
    Performs a search on the african history collection.
    """
    try:
        response = client.queries.top_snippets(
            collection_name="african_history_book",
            query=query,
            k=k,
            precise_responses=True,
            reranker=reranker
        )
        return str(response.results)
    except Exception as e:
        return f"Error performing search: {str(e)}"

@mcp.resource(
    uri="search://{query}",
    description="Get search results for a query",
    name="Search Results",
)
def get_search_results(query: str) -> str:
    try:
        response = client.queries.top_snippets(
            collection_name="african_history_book",
            query=query,
            k=5,
            precise_responses=True,
            reranker="zerank-1"
        )
        return str(response.results)
    except Exception as e:
        return f"Error: {str(e)}"

@mcp.prompt("")
def search_prompt(
    topic: str = Field(description="The topic to search for"),
    focus: str = Field(description="Specific focus area", default="general"),
) -> str:
    """Generate a search prompt for African history"""
    return f"Search for information about {topic} in African history, focusing on {focus} aspects."

if __name__ == "__main__":
    mcp.run(transport="streamable-http")