feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
FAQ Bot entry point.
|
||||
|
||||
The bot loads FAQ documents from the `data/` directory, stores them in
|
||||
ChromaDB, and then enters an interactive loop where the user can ask
|
||||
questions. The bot returns the top 3 most relevant answers.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from src.vector_store import VectorStore
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Helper functions
|
||||
# --------------------------------------------------------------------------- #
|
||||
def load_documents(folder: Path) -> list:
|
||||
"""
|
||||
Load all .txt files from the given folder as documents.
|
||||
|
||||
Each file becomes a single document with its content as text.
|
||||
"""
|
||||
docs = []
|
||||
for file in folder.glob("*.txt"):
|
||||
text = file.read_text(encoding="utf-8")
|
||||
docs.append({"text": text, "metadata": {"source": file.name}})
|
||||
return docs
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Main logic
|
||||
# --------------------------------------------------------------------------- #
|
||||
def main() -> None:
|
||||
# Load environment variables (e.g. OPENAI_API_KEY)
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Resolve data directory relative to the project root
|
||||
project_root = Path(__file__).resolve().parent.parent
|
||||
data_dir = project_root / "data"
|
||||
|
||||
# Initialize vector store
|
||||
store = VectorStore()
|
||||
|
||||
# If the collection is empty, load documents
|
||||
if store.collection.count() == 0:
|
||||
print("Loading documents into ChromaDB...")
|
||||
docs = load_documents(data_dir)
|
||||
if not docs:
|
||||
print(f"No .txt files found in {data_dir}. Exiting.")
|
||||
sys.exit(1)
|
||||
store.add_documents(docs)
|
||||
print(f"Added {len(docs)} documents.")
|
||||
|
||||
print("\nFAQ Bot is ready. Type your question (or 'exit' to quit).")
|
||||
|
||||
while True:
|
||||
try:
|
||||
query = input("\nQ: ")
|
||||
except EOFError:
|
||||
break
|
||||
|
||||
if query.lower() in ("exit", "quit"):
|
||||
break
|
||||
|
||||
results = store.query(query, top_k=3)
|
||||
if not results:
|
||||
print("No answer found.")
|
||||
continue
|
||||
|
||||
print("\nTop answers:")
|
||||
for i, res in enumerate(results, 1):
|
||||
snippet = res["text"][:200].replace("\n", " ")
|
||||
print(f"{i}. {snippet}... (distance: {res['distance']:.4f})")
|
||||
|
||||
print("\nGoodbye!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+62
-11
@@ -1,18 +1,69 @@
|
||||
from typing import List, Dict
|
||||
"""
|
||||
MCP-tool: Simple embedding generator.
|
||||
|
||||
from .vector_store import QdrantVectorStore
|
||||
This module provides a single function `get_embedding` that returns a vector
|
||||
representation of a given text. The implementation first tries to use the
|
||||
OpenAI embeddings API. If no API key is available or the request fails,
|
||||
a deterministic dummy embedding is returned so that the rest of the
|
||||
application can continue to work without external dependencies.
|
||||
"""
|
||||
|
||||
import os
|
||||
import hashlib
|
||||
from typing import List
|
||||
|
||||
try:
|
||||
import openai
|
||||
except ImportError:
|
||||
openai = None
|
||||
|
||||
|
||||
class MCPTool:
|
||||
def _hash_embedding(text: str, dim: int = 1536) -> List[float]:
|
||||
"""
|
||||
A simple tool that uses the vector store to answer queries.
|
||||
Create a deterministic dummy embedding from a hash of the text.
|
||||
The values are in the range [0, 1).
|
||||
"""
|
||||
h = hashlib.sha256(text.encode("utf-8")).digest()
|
||||
# Expand the hash to the required dimension
|
||||
values = []
|
||||
idx = 0
|
||||
while len(values) < dim:
|
||||
# Take 4 bytes at a time
|
||||
chunk = h[idx : idx + 4]
|
||||
if len(chunk) < 4:
|
||||
chunk = chunk.ljust(4, b"\0")
|
||||
val = int.from_bytes(chunk, "big") / 2**32
|
||||
values.append(val)
|
||||
idx += 4
|
||||
return values
|
||||
|
||||
def __init__(self, vector_store: QdrantVectorStore):
|
||||
self.vector_store = vector_store
|
||||
|
||||
def answer(self, query: str, top_k: int = 3) -> List[Dict]:
|
||||
"""
|
||||
Return the top_k most relevant documents for the query.
|
||||
"""
|
||||
return self.vector_store.search(query, top_k=top_k)
|
||||
def get_embedding(text: str) -> List[float]:
|
||||
"""
|
||||
Return an embedding vector for the given text.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
The input text to embed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[float]
|
||||
The embedding vector.
|
||||
"""
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if api_key and openai:
|
||||
openai.api_key = api_key
|
||||
try:
|
||||
response = openai.Embedding.create(
|
||||
input=text,
|
||||
model="text-embedding-ada-002",
|
||||
)
|
||||
return response["data"][0]["embedding"]
|
||||
except Exception:
|
||||
# Fall back to dummy embedding on any error
|
||||
pass
|
||||
|
||||
# Dummy deterministic embedding
|
||||
return _hash_embedding(text)
|
||||
+88
-23
@@ -1,31 +1,96 @@
|
||||
from langchain_community.vectorstores import Chroma
|
||||
from langchain.schema import Document
|
||||
from src.config import settings
|
||||
from src.embeddings import ollama_embeddings
|
||||
"""
|
||||
Vector store abstraction over ChromaDB.
|
||||
|
||||
class FAQVectorStore:
|
||||
The `VectorStore` class encapsulates all interactions with the ChromaDB
|
||||
collection. It uses the MCP-tool to generate embeddings for documents
|
||||
and queries.
|
||||
"""
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
from typing import List, Dict
|
||||
|
||||
from .mcp_tool import get_embedding
|
||||
|
||||
|
||||
class VectorStore:
|
||||
"""
|
||||
Wrapper around Chroma vector store for FAQ documents.
|
||||
Wrapper around a ChromaDB collection.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
collection_name : str, optional
|
||||
Name of the collection to use. Defaults to "faq".
|
||||
"""
|
||||
def __init__(self):
|
||||
self.db = Chroma(
|
||||
collection_name=settings.chroma_collection_name,
|
||||
persist_directory=settings.chroma_db_path,
|
||||
embedding_function=ollama_embeddings
|
||||
|
||||
def __init__(self, collection_name: str = "faq"):
|
||||
self.client = chromadb.Client(Settings())
|
||||
self.collection = self.client.get_or_create_collection(name=collection_name)
|
||||
|
||||
def add_documents(self, documents: List[Dict[str, str]]) -> None:
|
||||
"""
|
||||
Add a list of documents to the collection.
|
||||
|
||||
Each document must contain a 'text' key and may optionally contain
|
||||
a 'metadata' dictionary.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
documents : List[Dict[str, str]]
|
||||
List of documents to add.
|
||||
"""
|
||||
ids = []
|
||||
texts = []
|
||||
embeddings = []
|
||||
metadatas = []
|
||||
|
||||
for i, doc in enumerate(documents):
|
||||
ids.append(str(i))
|
||||
texts.append(doc["text"])
|
||||
embeddings.append(get_embedding(doc["text"]))
|
||||
metadatas.append(doc.get("metadata", {}))
|
||||
|
||||
self.collection.add(
|
||||
ids=ids,
|
||||
documents=texts,
|
||||
embeddings=embeddings,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
|
||||
def add_documents(self, documents: list[Document]):
|
||||
def query(self, query_text: str, top_k: int = 5) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Add a list of Documents to the vector store and persist.
|
||||
"""
|
||||
self.db.add_documents(documents)
|
||||
self.db.persist()
|
||||
Retrieve the most relevant documents for a query.
|
||||
|
||||
def similarity_search(self, query: str, k: int = 4):
|
||||
"""
|
||||
Retrieve the top-k most similar documents to the query.
|
||||
"""
|
||||
return self.db.similarity_search(query, k=k)
|
||||
Parameters
|
||||
----------
|
||||
query_text : str
|
||||
The query string.
|
||||
top_k : int, optional
|
||||
Number of results to return. Defaults to 5.
|
||||
|
||||
# Singleton instance for use in the application
|
||||
vector_store = FAQVectorStore()
|
||||
Returns
|
||||
-------
|
||||
List[Dict[str, str]]
|
||||
List of result dictionaries containing 'text', 'distance',
|
||||
and 'metadata'.
|
||||
"""
|
||||
embedding = get_embedding(query_text)
|
||||
results = self.collection.query(
|
||||
query_embeddings=[embedding],
|
||||
n_results=top_k,
|
||||
)
|
||||
|
||||
output = []
|
||||
for doc, dist, meta in zip(
|
||||
results["documents"][0],
|
||||
results["distances"][0],
|
||||
results["metadatas"][0],
|
||||
):
|
||||
output.append(
|
||||
{
|
||||
"text": doc,
|
||||
"distance": dist,
|
||||
"metadata": meta,
|
||||
}
|
||||
)
|
||||
return output
|
||||
Reference in New Issue
Block a user