from mcp.server.fastmcp import FastMCP
import chromadb
import os

# On définit l'hôte sur 0.0.0.0 et le port sur 8080 pour le container Docker
mcp = FastMCP("Chroma-Server", host="0.0.0.0", port=8080)

# Initialisation du client ChromaDB
CHROMA_URL = os.getenv("CHROMA_URL", "http://chroma-db:8000")
client = chromadb.HttpClient(host="chroma-db", port=8000)

@mcp.tool()
async def query_chroma(collection_name: str, query_text: str, n_results: int = 5):
    """Interroge une collection ChromaDB pour trouver des documents similaires."""
    collection = client.get_collection(name=collection_name)
    results = collection.query(query_texts=[query_text], n_results=n_results)
    return str(results)

if __name__ == "__main__":
    # On lance le serveur en mode SSE (Server-Sent Events)
    # C'est ce que LM Studio et la plupart des clients attendent
    mcp.run(transport="sse")