Publish solution for task 6a02e23da6fe2e4ac16acf65: update main.py

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."