feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'

vector_store_utils.py

@@ -0,0 +1,85 @@
+import os
+from pathlib import Path
+from typing import List, Dict, Any
+
+from chromadb import Client
+from chromadb import Collection
+from langchain_ollama import OllamaEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+# Constants
+COLLECTION_NAME = "faq_collection"
+EMBEDDING_MODEL = "nomic-embed-text"
+EMBEDDING_DIM = 1024  # Adjust if the model changes
+CHROMA_PATH = "./chroma_faq"
+
+def load_faq_to_chroma() -> Collection:
+    """
+    Load all markdown files from the data/ directory, split them into chunks,
+    embed them, and store them in a Chroma collection.
+    Returns the Chroma Collection instance.
+    """
+    # Initialize Chroma client with persistence
+    client = Client(path=CHROMA_PATH)
+
+    # Create or get collection
+    collection = client.get_or_create_collection(
+        name=COLLECTION_NAME,
+        metadata={"hnsw:space": "cosine"},
+    )
+
+    # Prepare text splitter
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+
+    # Prepare embeddings
+    embedder = OllamaEmbeddings(model=EMBEDDING_MODEL)
+
+    # Load markdown files
+    data_dir = Path("data")
+    docs = []
+    ids = []
+    metadatas = []
+
+    for md_file in data_dir.glob("*.md"):
+        text = md_file.read_text(encoding="utf-8")
+        chunks = splitter.split_text(text)
+        for idx, chunk in enumerate(chunks):
+            docs.append(chunk)
+            ids.append(f"{md_file.stem}_{idx}")
+            metadatas.append({"source": md_file.name})
+
+    # Embed documents
+    embeddings = embedder.embed_documents(docs)
+
+    # Add to collection
+    collection.add(
+        documents=docs,
+        ids=ids,
+        metadatas=metadatas,
+        embeddings=embeddings,
+    )
+
+    return collection
+
+def search_course_docs(collection: Collection, query: str, k: int = 3) -> List[Dict[str, Any]]:
+    """
+    Query the Chroma collection for the top k documents matching the query.
+    Returns a list of dicts with page_content and score.
+    """
+    results = collection.query(
+        query_texts=[query],
+        n_results=k,
+        include=["documents", "distances", "metadatas"],
+    )
+    docs = []
+    for doc, distance, metadata in zip(
+        results["documents"][0], results["distances"][0], results["metadatas"][0]
+    ):
+        docs.append(
+            {
+                "page_content": doc,
+                "score": 1 - distance,  # Convert distance to similarity
+                "metadata": metadata,
+            }
+        )
+    return docs