From 33a3642fa9642104029217c084016d6295580aab 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:15:08 +0000 Subject: [PATCH] Publish solution for task 6a02e23da6fe2e4ac16acf65: update main.py --- main.py | 51 ++++++++++++++++++++++----------------------------- 1 file changed, 22 insertions(+), 29 deletions(-) diff --git a/main.py b/main.py index 91987a6..7afdd9c 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,3 @@ -"""Interactive LangChain agent with local RAG memory on Qdrant and Ollama.""" - import os from functools import lru_cache from typing import Any @@ -9,13 +7,9 @@ from langchain.agents import create_agent from langchain_core.documents import Document from langchain_core.tools import tool from langchain_ollama import ChatOllama, OllamaEmbeddings -from langchain_qdrant import QdrantVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter -from qdrant_client import QdrantClient -from qdrant_client.models import Distance, VectorParams +from chromadb import Client -QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") -QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "rag_memory") OLLAMA_CHAT_MODEL = os.getenv("OLLAMA_CHAT_MODEL", "llama3") OLLAMA_EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text") EMBEDDING_SIZE = int(os.getenv("OLLAMA_EMBEDDING_SIZE", "768")) @@ -27,19 +21,11 @@ def get_embeddings() -> OllamaEmbeddings: @lru_cache(maxsize=1) -def get_vector_store() -> QdrantVectorStore: - client = QdrantClient(url=QDRANT_URL) - collections = {item.name for item in client.get_collections().collections} - if QDRANT_COLLECTION not in collections: - client.create_collection( - collection_name=QDRANT_COLLECTION, - vectors_config=VectorParams(size=EMBEDDING_SIZE, distance=Distance.COSINE), - ) - return QdrantVectorStore( - client=client, - collection_name=QDRANT_COLLECTION, - embedding=get_embeddings(), - ) +def get_vector_store() -> Client: + client = Client() + # Ensure collection exists + client.get_or_create_collection(name="rag_memory") + return client def chunk_document(content: str, title: str) -> list[Document]: @@ -49,25 +35,32 @@ 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 Qdrant knowledge base.""" - results = get_vector_store().similarity_search_with_score(query, k=max_results) - if not results: + """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"): return "No relevant documents found." lines: list[str] = [] - for index, (document, score) in enumerate(results, start=1): - title = document.metadata.get("title", "untitled") - snippet = document.page_content.replace("\n", " ")[:300] - lines.append(f"{index}. {title} (score={score:.4f}): {snippet}") + 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}") 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 Qdrant knowledge base.""" + """Split content into chunks and store it in the local ChromaDB knowledge base.""" documents = chunk_document(content, title) ids = [str(uuid4()) for _ in documents] - get_vector_store().add_documents(documents, ids=ids) + 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]) return f"Added {len(documents)} chunk(s) from '{title}' to the knowledge base."