feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
+223
@@ -0,0 +1,223 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
FAQ Bot using QDrant as the vector store.
|
||||
|
||||
This script provides:
|
||||
- Data ingestion from a text file into QDrant.
|
||||
- Querying the vector store to retrieve relevant FAQ answers.
|
||||
- A simple CLI interface for ingestion and querying.
|
||||
|
||||
Author: Artur Kuzakhmetov
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple
|
||||
|
||||
import openai
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qdrant_models
|
||||
from qdrant_client.http.models import PointStruct
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Configuration
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
# Environment variables
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
|
||||
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY") # Optional, if QDrant requires auth
|
||||
QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "faq_collection")
|
||||
|
||||
# OpenAI embedding model
|
||||
EMBEDDING_MODEL = "text-embedding-ada-002"
|
||||
EMBEDDING_DIM = 1536 # Dimension of Ada-002 embeddings
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Helper functions
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def split_text_into_chunks(text: str, max_chunk_size: int = 500) -> List[str]:
|
||||
"""
|
||||
Split a large text into smaller chunks suitable for embedding.
|
||||
Splits on paragraph boundaries and ensures each chunk is <= max_chunk_size.
|
||||
"""
|
||||
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
|
||||
chunks = []
|
||||
current_chunk = ""
|
||||
for para in paragraphs:
|
||||
if len(current_chunk) + len(para) + 1 <= max_chunk_size:
|
||||
current_chunk += (" " if current_chunk else "") + para
|
||||
else:
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk.strip())
|
||||
current_chunk = para
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk.strip())
|
||||
return chunks
|
||||
|
||||
def embed_texts(texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Generate embeddings for a list of texts using OpenAI's embedding API.
|
||||
"""
|
||||
if not OPENAI_API_KEY:
|
||||
raise RuntimeError("OPENAI_API_KEY environment variable is not set.")
|
||||
openai.api_key = OPENAI_API_KEY
|
||||
embeddings = []
|
||||
for text in texts:
|
||||
response = openai.Embedding.create(
|
||||
input=text,
|
||||
model=EMBEDDING_MODEL
|
||||
)
|
||||
embeddings.append(response["data"][0]["embedding"])
|
||||
return embeddings
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# QDrant Vector Store Wrapper
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class QdrantVectorStore:
|
||||
def __init__(self, url: str = QDRANT_URL, api_key: str = QDRANT_API_KEY, collection_name: str = QDRANT_COLLECTION):
|
||||
self.client = QdrantClient(url=url, api_key=api_key)
|
||||
self.collection_name = collection_name
|
||||
self._ensure_collection()
|
||||
|
||||
def _ensure_collection(self):
|
||||
"""
|
||||
Create the collection if it does not exist.
|
||||
"""
|
||||
collections = self.client.get_collections()
|
||||
if self.collection_name not in [c.name for c in collections.collections]:
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=qdrant_models.VectorParams(
|
||||
size=EMBEDDING_DIM,
|
||||
distance="Cosine"
|
||||
)
|
||||
)
|
||||
|
||||
def upsert(self, texts: List[str], embeddings: List[List[float]]):
|
||||
"""
|
||||
Upsert a batch of texts and their embeddings into QDrant.
|
||||
"""
|
||||
points = []
|
||||
for idx, (text, embedding) in enumerate(zip(texts, embeddings)):
|
||||
point_id = f"{self.collection_name}_{idx}_{hash(text) % 1000000}"
|
||||
points.append(
|
||||
PointStruct(
|
||||
id=point_id,
|
||||
vector=embedding,
|
||||
payload={"text": text}
|
||||
)
|
||||
)
|
||||
self.client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=points
|
||||
)
|
||||
|
||||
def search(self, query_embedding: List[float], top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""
|
||||
Search the collection for the most similar vectors to the query embedding.
|
||||
Returns a list of (text, score) tuples.
|
||||
"""
|
||||
search_result = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_embedding,
|
||||
limit=top_k,
|
||||
with_payload=True,
|
||||
score=True
|
||||
)
|
||||
results = []
|
||||
for hit in search_result:
|
||||
text = hit.payload.get("text", "")
|
||||
score = hit.score
|
||||
results.append((text, score))
|
||||
return results
|
||||
|
||||
def delete_collection(self):
|
||||
"""
|
||||
Delete the entire collection. Use with caution.
|
||||
"""
|
||||
self.client.delete_collection(self.collection_name)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Bot Logic
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def ingest_data(file_path: str, vector_store: QdrantVectorStore):
|
||||
"""
|
||||
Read a text file, split into chunks, embed, and store in QDrant.
|
||||
"""
|
||||
if not Path(file_path).is_file():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
raw_text = f.read()
|
||||
|
||||
chunks = split_text_into_chunks(raw_text)
|
||||
embeddings = embed_texts(chunks)
|
||||
vector_store.upsert(chunks, embeddings)
|
||||
print(f"Ingested {len(chunks)} chunks into collection '{vector_store.collection_name}'.")
|
||||
|
||||
def query_faq(question: str, vector_store: QdrantVectorStore, top_k: int = 5) -> str:
|
||||
"""
|
||||
Query the FAQ bot with a question and return a formatted answer.
|
||||
"""
|
||||
query_embedding = embed_texts([question])[0]
|
||||
results = vector_store.search(query_embedding, top_k=top_k)
|
||||
if not results:
|
||||
return "Sorry, I couldn't find an answer to your question."
|
||||
|
||||
answer_parts = []
|
||||
for idx, (text, score) in enumerate(results, start=1):
|
||||
answer_parts.append(f"{idx}. (Score: {score:.4f})\n{text}\n")
|
||||
return "\n".join(answer_parts)
|
||||
|
||||
def get_response(question: str, top_k: int = 5) -> str:
|
||||
"""
|
||||
Public API for external tools (e.g., MCP-tool) to get a bot response.
|
||||
"""
|
||||
vector_store = QdrantVectorStore()
|
||||
return query_faq(question, vector_store, top_k=top_k)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# CLI Interface
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="FAQ Bot CLI")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
ingest_parser = subparsers.add_parser("ingest", help="Ingest a text file into QDrant")
|
||||
ingest_parser.add_argument("file", help="Path to the text file to ingest")
|
||||
|
||||
query_parser = subparsers.add_parser("query", help="Query the FAQ bot")
|
||||
query_parser.add_argument("question", help="Your question")
|
||||
query_parser.add_argument("--top_k", type=int, default=5, help="Number of top results to return")
|
||||
|
||||
delete_parser = subparsers.add_parser("delete", help="Delete the QDrant collection (use with caution)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
vector_store = QdrantVectorStore()
|
||||
|
||||
if args.command == "ingest":
|
||||
ingest_data(args.file, vector_store)
|
||||
elif args.command == "query":
|
||||
answer = query_faq(args.question, vector_store, top_k=args.top_k)
|
||||
print(answer)
|
||||
elif args.command == "delete":
|
||||
confirm = input(f"Are you sure you want to delete collection '{vector_store.collection_name}'? (yes/no): ")
|
||||
if confirm.lower() == "yes":
|
||||
vector_store.delete_collection()
|
||||
print("Collection deleted.")
|
||||
else:
|
||||
print("Deletion aborted.")
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user