Solution published to Gitea.: update main.py

main.py

@@ -8,7 +8,8 @@ from langchain_core.documents import Document
 from langchain_core.tools import tool
 from langchain_ollama import ChatOllama, OllamaEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from chromadb import Client
+from langchain_qdrant import QdrantVectorStore
+from qdrant_client import QdrantClient
 
 OLLAMA_CHAT_MODEL = os.getenv("OLLAMA_CHAT_MODEL", "llama3")
 OLLAMA_EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
@@ -21,11 +22,14 @@ def get_embeddings() -> OllamaEmbeddings:
 
 
 @lru_cache(maxsize=1)
-def get_vector_store() -> Client:
-    client = Client()
+def get_vector_store() -> QdrantVectorStore:
+    client = QdrantClient()
     # Ensure collection exists
-    client.get_or_create_collection(name="rag_memory")
-    return client
+    client.recreate_collection(
+        collection_name="rag_memory",
+        vectors_config={"size": EMBEDDING_SIZE, "distance": "cosine"},
+    )
+    return QdrantVectorStore(client=client, collection_name="rag_memory")
 
 
 def chunk_document(content: str, title: str) -> list[Document]:
@@ -35,32 +39,28 @@ def chunk_document(content: str, title: str) -> list[Document]:
 
 @tool
 def search_knowledge_base(query: str, max_results: int = 3) -> str:
-    """Search relevant chunks in the local ChromaDB knowledge base."""
-    client = get_vector_store()
-    collection = client.get_collection(name="rag_memory")
-    results = collection.query(query_texts=[query], n_results=max_results)
-    if not results.get("documents"):
+    """Search relevant chunks in the local Qdrant knowledge base."""
+    vector_store = get_vector_store()
+    results = vector_store.similarity_search_with_score(query, k=max_results)
+    if not results:
         return "No relevant documents found."
 
     lines: list[str] = []
-    for idx, (doc, dist) in enumerate(zip(results["documents"], results["distances"]), start=1):
-        ids = results["ids"]
-        metadata = collection.get(ids=[ids[idx-1]])["metadatas"][0]
-        title = metadata.get("title", "untitled") if metadata else "untitled"
-        snippet = doc.replace("\n", " ")[:300]
-        lines.append(f"{idx}. {title} (score={dist:.4f}): {snippet}")
+    for idx, (doc, score) in enumerate(results, start=1):
+        title = doc.metadata.get("title", "untitled") if doc.metadata else "untitled"
+        snippet = doc.page_content.replace("\n", " ")[:300]
+        lines.append(f"{idx}. {title} (score={score:.4f}): {snippet}")
     return "\n".join(lines)
 
 
 @tool
 def add_to_knowledge_base(content: str, title: str) -> str:
-    """Split content into chunks and store it in the local ChromaDB knowledge base."""
+    """Split content into chunks and store it in the local Qdrant knowledge base."""
     documents = chunk_document(content, title)
     ids = [str(uuid4()) for _ in documents]
     embeddings = get_embeddings().embed_documents([doc.page_content for doc in documents])
-    client = get_vector_store()
-    collection = client.get_collection(name="rag_memory")
-    collection.add(ids=ids, documents=[doc.page_content for doc in documents], embeddings=embeddings, metadatas=[doc.metadata for doc in documents])
+    vector_store = get_vector_store()
+    vector_store.add_documents(documents, ids=ids, embeddings=embeddings)
     return f"Added {len(documents)} chunk(s) from '{title}' to the knowledge base."