Andrey Lutsenko
86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
|
|
import requests
|
|
from fastmcp import FastMCP
|
|
from qdrant_client import QdrantClient
|
|
|
|
MCP_NAME = os.getenv("MCP_NAME", "local-rag")
|
|
QDRANT_URL = os.getenv("QDRANT_URL", "http://127.0.0.1:6333")
|
|
COLLECTION = os.getenv("QDRANT_COLLECTION", "docs")
|
|
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
|
|
GEN_MODEL = os.getenv("GEN_MODEL", "qwen2.5-coder:7b")
|
|
|
|
mcp = FastMCP(MCP_NAME)
|
|
qdrant = QdrantClient(url=QDRANT_URL)
|
|
|
|
|
|
def embed(text: str) -> list[float]:
|
|
r = requests.post(
|
|
"http://127.0.0.1:11434/api/embeddings",
|
|
json={"model": EMBED_MODEL, "prompt": text},
|
|
timeout=120,
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()["embedding"]
|
|
|
|
|
|
def retrieve(query: str, top_k: int = 5):
|
|
query_vec = embed(query)
|
|
return qdrant.search(collection_name=COLLECTION, query_vector=query_vec, limit=top_k)
|
|
|
|
|
|
def sanitize(text: str, max_len: int = 1800) -> str:
|
|
return (text or "").strip()[:max_len]
|
|
|
|
|
|
@mcp.tool
|
|
def search_docs(query: str, top_k: int = 5) -> str:
|
|
hits = retrieve(query, top_k=top_k)
|
|
if not hits:
|
|
return "No relevant docs found."
|
|
lines = []
|
|
for i, hit in enumerate(hits, start=1):
|
|
src = hit.payload.get("source", "unknown")
|
|
txt = sanitize(hit.payload.get("text", ""))
|
|
lines.append(f"[{i}] source={src}\n{txt}")
|
|
return "\n\n---\n\n".join(lines)
|
|
|
|
|
|
@mcp.tool
|
|
def answer_with_citations(query: str, top_k: int = 5) -> str:
|
|
hits = retrieve(query, top_k=top_k)
|
|
if not hits:
|
|
return "Не нашел релевантных документов."
|
|
|
|
context = []
|
|
citations = []
|
|
for i, hit in enumerate(hits, start=1):
|
|
src = hit.payload.get("source", "unknown")
|
|
txt = sanitize(hit.payload.get("text", ""))
|
|
context.append(f"[{i}] SOURCE: {src}\n{txt}")
|
|
citations.append(f"[{i}] {src}")
|
|
|
|
prompt = (
|
|
"Ты отвечаешь только на основе контекста.\n"
|
|
"Если данных недостаточно, явно скажи об этом.\n"
|
|
"Добавь ссылки на источники в формате [номер].\n\n"
|
|
f"Вопрос: {query}\n\n"
|
|
"Контекст:\n"
|
|
+ "\n\n".join(context)
|
|
)
|
|
|
|
r = requests.post(
|
|
"http://127.0.0.1:11434/api/generate",
|
|
json={"model": GEN_MODEL, "prompt": prompt, "stream": False},
|
|
timeout=240,
|
|
)
|
|
r.raise_for_status()
|
|
answer = r.json().get("response", "").strip()
|
|
|
|
return f"{answer}\n\nИсточники:\n" + "\n".join(citations)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
mcp.run()
|