84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
|
import asyncio
|
||
|
|
import os
|
||
|
|
from mcp import ClientSession, StdioServerParameters
|
||
|
|
from mcp.client.stdio import stdio_client
|
||
|
|
|
||
|
|
|
||
|
|
async def main():
|
||
|
|
server = StdioServerParameters(
|
||
|
|
command="uv",
|
||
|
|
args=[
|
||
|
|
"run",
|
||
|
|
"--directory",
|
||
|
|
"/home/df/projects/rags/lightrag-mcp",
|
||
|
|
"python",
|
||
|
|
"main.py",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
|
||
|
|
async with stdio_client(server) as (read, write):
|
||
|
|
async with ClientSession(read, write) as session:
|
||
|
|
await session.initialize()
|
||
|
|
|
||
|
|
print("--- INSERT ---")
|
||
|
|
result = await session.call_tool(
|
||
|
|
"insert_documents",
|
||
|
|
arguments={
|
||
|
|
"documents": [
|
||
|
|
"Python is a high-level programming language known for its simplicity and readability.",
|
||
|
|
"JavaScript was created in 1995 by Brendan Eich at Netscape.",
|
||
|
|
"Machine learning is a subset of artificial intelligence that enables systems to learn from data.",
|
||
|
|
]
|
||
|
|
},
|
||
|
|
)
|
||
|
|
print(f"Insert result: {result.content[0].text[:200]}")
|
||
|
|
|
||
|
|
print("\n--- QUERY (mix) ---")
|
||
|
|
result = await session.call_tool(
|
||
|
|
"query_documents",
|
||
|
|
arguments={
|
||
|
|
"query": "Tell me about programming languages",
|
||
|
|
"mode": "mix",
|
||
|
|
"top_k": 60,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
import json
|
||
|
|
|
||
|
|
data = json.loads(result.content[0].text)
|
||
|
|
d = data.get("data", {})
|
||
|
|
print(f"Entities: {len(d.get('entities', []))}")
|
||
|
|
print(f"Relationships: {len(d.get('relationships', []))}")
|
||
|
|
print(f"Chunks: {len(d.get('chunks', []))}")
|
||
|
|
for c in d.get("chunks", [])[:2]:
|
||
|
|
print(f" - {c.get('content', '')[:100]}")
|
||
|
|
|
||
|
|
print("\n--- QUERY (local) ---")
|
||
|
|
result = await session.call_tool(
|
||
|
|
"query_documents",
|
||
|
|
arguments={"query": "What is Python?", "mode": "local", "top_k": 60},
|
||
|
|
)
|
||
|
|
data = json.loads(result.content[0].text)
|
||
|
|
d = data.get("data", {})
|
||
|
|
print(f"Entities: {len(d.get('entities', []))}")
|
||
|
|
print(f"Chunks: {len(d.get('chunks', []))}")
|
||
|
|
|
||
|
|
print("\n--- QUERY (global) ---")
|
||
|
|
result = await session.call_tool(
|
||
|
|
"query_documents",
|
||
|
|
arguments={
|
||
|
|
"query": "What topics are covered?",
|
||
|
|
"mode": "global",
|
||
|
|
"top_k": 60,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
data = json.loads(result.content[0].text)
|
||
|
|
d = data.get("data", {})
|
||
|
|
print(f"Entities: {len(d.get('entities', []))}")
|
||
|
|
print(f"Relationships: {len(d.get('relationships', []))}")
|
||
|
|
|
||
|
|
print("\nDone!")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
asyncio.run(main())
|