From ec13bd14c011e00c48561cc6b33ef949d5d31a97 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?= =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= Date: Thu, 18 Jun 2026 09:44:16 +0000 Subject: [PATCH] Solution published to Gitea.: update main.py --- main.py | 40 ++++++++++++++++++++-------------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/main.py b/main.py index 7afdd9c..ba4e50a 100644 --- a/main.py +++ b/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."