58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
|
import os
|
||
|
|
import httpx
|
||
|
|
from fastmcp import FastMCP
|
||
|
|
|
||
|
|
LIGHTRAG_URL = os.getenv("LIGHTRAG_URL", "http://localhost:9621")
|
||
|
|
|
||
|
|
mcp = FastMCP("LightRAG")
|
||
|
|
|
||
|
|
|
||
|
|
@mcp.tool
|
||
|
|
async def insert_documents(documents: list[str]) -> str:
|
||
|
|
"""Insert text documents into LightRAG for indexing.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
documents: List of document strings to index. Each string is treated as a separate document.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Tracking ID for the insertion operation.
|
||
|
|
"""
|
||
|
|
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
|
|
r = await client.post(
|
||
|
|
f"{LIGHTRAG_URL}/documents/texts",
|
||
|
|
json={"texts": documents},
|
||
|
|
)
|
||
|
|
r.raise_for_status()
|
||
|
|
data = r.json()
|
||
|
|
return data.get("track_id", data.get("message", "unknown"))
|
||
|
|
|
||
|
|
|
||
|
|
@mcp.tool
|
||
|
|
async def query_documents(query: str, mode: str = "mix", top_k: int = 60) -> dict:
|
||
|
|
"""Query LightRAG and retrieve relevant context without LLM generation.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
query: The search query string.
|
||
|
|
mode: Retrieval mode - "local", "global", "hybrid", "naive", "mix" (default: "mix").
|
||
|
|
top_k: Number of top results to retrieve (default: 60).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Structured retrieval data including entities, relationships, and text chunks.
|
||
|
|
"""
|
||
|
|
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
|
|
r = await client.post(
|
||
|
|
f"{LIGHTRAG_URL}/query/data",
|
||
|
|
json={
|
||
|
|
"query": query,
|
||
|
|
"mode": mode,
|
||
|
|
"only_need_context": True,
|
||
|
|
"top_k": top_k,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
r.raise_for_status()
|
||
|
|
return r.json()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
mcp.run()
|