132 lines
4.0 KiB
Python
132 lines
4.0 KiB
Python
import os
|
|
import httpx
|
|
from mcp.server.fastmcp import FastMCP
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, VectorParams, PointStruct
|
|
import uuid
|
|
import logging
|
|
import io
|
|
from pypdf import PdfReader
|
|
|
|
# Configuration
|
|
QDRANT_URL = os.getenv("QDRANT_URL", "http://qdrant.openshift-gitops.svc:6333")
|
|
TEI_URL = os.getenv("TEI_URL", "http://text-embeddings.tei.svc.cluster.local:8080")
|
|
EMBEDDING_DIM = 768 # BAAI/bge-base-en-v1.5
|
|
|
|
# Initialize
|
|
mcp = FastMCP("knowledge-mcp")
|
|
qdrant = QdrantClient(url=QDRANT_URL)
|
|
|
|
def get_embedding(text: str) -> list[float]:
|
|
"""Get embedding from TEI."""
|
|
url = f"{TEI_URL}/embed"
|
|
try:
|
|
response = httpx.post(url, json={"inputs": text}, timeout=10.0)
|
|
response.raise_for_status()
|
|
return response.json()[0]
|
|
except Exception as e:
|
|
logging.error(f"Embedding failed: {e}")
|
|
raise
|
|
|
|
def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> list[str]:
|
|
"""Sliding window chunking."""
|
|
if len(text) <= chunk_size:
|
|
return [text]
|
|
|
|
chunks = []
|
|
start = 0
|
|
while start < len(text):
|
|
end = start + chunk_size
|
|
chunks.append(text[start:end])
|
|
start += (chunk_size - overlap)
|
|
return chunks
|
|
|
|
@mcp.tool()
|
|
def create_notebook(name: str) -> str:
|
|
"""Create a new RAG notebook (Qdrant collection)."""
|
|
clean_name = name.lower().replace(" ", "-")
|
|
|
|
# Check if exists
|
|
if qdrant.collection_exists(clean_name):
|
|
return f"Notebook '{clean_name}' already exists."
|
|
|
|
qdrant.create_collection(
|
|
collection_name=clean_name,
|
|
vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
|
|
)
|
|
return f"Notebook '{clean_name}' created successfully."
|
|
|
|
@mcp.tool()
|
|
def add_source(notebook: str, content: str, source_name: str, format: str = "text") -> str:
|
|
"""
|
|
Add content to a notebook.
|
|
format: 'text' or 'pdf_path' (local path inside container)
|
|
"""
|
|
if not qdrant.collection_exists(notebook):
|
|
return f"Error: Notebook '{notebook}' does not exist."
|
|
|
|
text_to_process = ""
|
|
|
|
if format == "pdf_path":
|
|
try:
|
|
reader = PdfReader(content)
|
|
for page in reader.pages:
|
|
text_to_process += page.extract_text() + "\n"
|
|
except Exception as e:
|
|
return f"Error reading PDF: {e}"
|
|
else:
|
|
text_to_process = content
|
|
|
|
chunks = chunk_text(text_to_process)
|
|
points = []
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
try:
|
|
vector = get_embedding(chunk)
|
|
points.append(PointStruct(
|
|
id=str(uuid.uuid4()),
|
|
vector=vector,
|
|
payload={
|
|
"source": source_name,
|
|
"text": chunk,
|
|
"chunk_index": i,
|
|
"total_chunks": len(chunks)
|
|
}
|
|
))
|
|
except Exception as e:
|
|
logging.error(f"Failed to embed chunk {i}: {e}")
|
|
continue
|
|
|
|
if points:
|
|
qdrant.upsert(collection_name=notebook, points=points)
|
|
return f"Added {len(points)} chunks from '{source_name}' to '{notebook}'."
|
|
return "No content added (empty or failed)."
|
|
|
|
@mcp.tool()
|
|
def query_notebook(notebook: str, query: str, limit: int = 5) -> str:
|
|
"""Query the notebook for relevant context."""
|
|
if not qdrant.collection_exists(notebook):
|
|
return f"Error: Notebook '{notebook}' does not exist."
|
|
|
|
try:
|
|
vector = get_embedding(query)
|
|
hits = qdrant.search(
|
|
collection_name=notebook,
|
|
query_vector=vector,
|
|
limit=limit
|
|
)
|
|
|
|
results = []
|
|
for hit in hits:
|
|
score = hit.score
|
|
text = hit.payload.get('text', '').replace('\n', ' ')
|
|
source = hit.payload.get('source', 'unknown')
|
|
results.append(f"[{score:.2f}] {source}: {text}...")
|
|
|
|
return "\n".join(results)
|
|
except Exception as e:
|
|
return f"Query failed: {e}"
|
|
|
|
if __name__ == "__main__":
|
|
mcp.run()
|