feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

vectorstore.py

@@ -0,0 +1,125 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from langchain_ollama import OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from qdrant_client import QdrantClient
+
+
+def create_vectorstore(
+    persist_directory: str = "./qdrant_db",
+    collection_name: str = "documents",
+    embedding_model: str = "nomic-embed-text",
+) -> QdrantVectorStore:
+    """
+    Create a Qdrant vector store backed by Ollama embeddings.
+
+    Parameters
+    ----------
+    persist_directory : str
+        Path to the directory where Qdrant will store its data.
+    collection_name : str
+        Name of the collection to use.
+    embedding_model : str
+        Ollama model name for embeddings.
+
+    Returns
+    -------
+    QdrantVectorStore
+        Initialized vector store.
+    """
+    # Ensure the persistence directory exists
+    Path(persist_directory).mkdir(parents=True, exist_ok=True)
+
+    # Create a local Qdrant client that stores data on disk
+    client = QdrantClient(path=persist_directory)
+
+    # Initialize embeddings
+    embeddings = OllamaEmbeddings(model=embedding_model)
+
+    # Create or load the collection
+    vectorstore = QdrantVectorStore(
+        client=client,
+        embeddings=embeddings,
+        collection_name=collection_name,
+    )
+    return vectorstore
+
+
+def load_documents(
+    directory: str,
+    vectorstore: QdrantVectorStore,
+    chunk_size: int = 1000,
+    chunk_overlap: int = 200,
+) -> None:
+    """
+    Load all .txt and .md files from a directory, split them into chunks,
+    embed, and add to the vector store.
+
+    Parameters
+    ----------
+    directory : str
+        Path to the directory containing documents.
+    vectorstore : QdrantVectorStore
+        The vector store to populate.
+    chunk_size : int
+        Maximum size of each chunk in characters.
+    chunk_overlap : int
+        Number of overlapping characters between chunks.
+    """
+    from langchain.text_splitter import RecursiveCharacterTextSplitter
+    from langchain.docstore.document import Document
+
+    # Prepare text splitter
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+    )
+
+    # Collect all documents
+    docs = []
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.lower().endswith((".txt", ".md")):
+                file_path = os.path.join(root, file)
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+                # Split content into chunks
+                chunks = splitter.split_text(content)
+                # Create Document objects with metadata
+                for i, chunk in enumerate(chunks):
+                    doc = Document(
+                        page_content=chunk,
+                        metadata={
+                            "source": file_path,
+                            "chunk_index": i,
+                        },
+                    )
+                    docs.append(doc)
+
+    if docs:
+        vectorstore.add_documents(docs)
+    else:
+        print("No documents found to load.")
+
+
+def collection_exists(vectorstore: QdrantVectorStore) -> bool:
+    """
+    Check if the collection already exists in Qdrant.
+
+    Parameters
+    ----------
+    vectorstore : QdrantVectorStore
+        The vector store to check.
+
+    Returns
+    -------
+    bool
+        True if collection exists, False otherwise.
+    """
+    try:
+        vectorstore.client.get_collection(vectorstore.collection_name)
+        return True
+    except Exception:
+        return False