import os
import logging
import httpx
from mcp.server.fastmcp import FastMCP
from dotenv import load_dotenv
# Configure logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("lightrag-mcp")
# 1. Initialization and Configuration
project_dir = os.path.dirname(os.path.abspath(__file__))
dotenv_path = os.path.join(project_dir, ".env")
if os.path.exists(dotenv_path):
load_dotenv(dotenv_path)
logger.info(f"Loaded .env from {dotenv_path}")
mcp = FastMCP("LightRAG")
def load_prompt(file_name: str) -> str | None:
path = os.path.join(project_dir, "prompts", file_name)
if os.path.exists(path):
try:
with open(path, "r", encoding="utf-8") as f:
return f.read().strip()
except Exception as e:
logger.error(f"Failed to read prompt file {path}: {e}")
return None
# Pre-load configurations
LIGHTRAG_URL = os.environ.get("LIGHTRAG_URL", "http://localhost:9621").rstrip("/")
DEFAULT_LANG = os.environ.get(
"LIGHTRAG_LANG", "zh"
).lower() # Default to Traditional Chinese
# 2. Define Tools
@mcp.tool()
async def query_knowledge_base(
query: str, mode: str = "hybrid", lang: str = None
) -> str:
"""
Query the internal LightRAG technical documentation.
Use this tool when the user asks about internal systems, architecture, or specific technical details.
Args:
query: The question to ask.
mode: Retrieval mode ('hybrid', 'naive', 'local', 'global'). Defaults to "hybrid".
lang: Preferred language ('zh' for Traditional Chinese, 'en' for English).
Defaults to environment variable LIGHTRAG_LANG or 'zh'.
"""
url = f"{LIGHTRAG_URL}/query"
# Determine language and system prompt
target_lang = (lang or DEFAULT_LANG).lower()
prompt_file = f"light_rag_system_{target_lang}.txt"
system_prompt = load_prompt(prompt_file) or load_prompt("light_rag_system.txt")
payload = {"query": query, "mode": mode}
if system_prompt:
payload["system_prompt"] = system_prompt
logger.info(f"Applying system prompt for language: {target_lang}")
logger.info(f"Querying LightRAG at {url} with mode={mode}, lang={target_lang}")
async with httpx.AsyncClient(timeout=60.0) as client:
try:
resp = await client.post(url, json=payload)
resp.raise_for_status()
data = resp.json()
return (
data.get("response")
or data.get("result")
or "No response content found."
)
except httpx.ConnectError:
error_msg = f"Failed to connect to LightRAG server at {LIGHTRAG_URL}. Ensure the server is running."
logger.error(error_msg)
return error_msg
except httpx.HTTPStatusError as e:
error_msg = f"LightRAG server returned error {e.response.status_code}: {e.response.text}"
logger.error(error_msg)
return error_msg
except Exception as e:
logger.exception("Unexpected error querying LightRAG")
return f"Error querying LightRAG: {str(e)}"
# 3. Execution
if __name__ == "__main__":
mcp.run()
