feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
+1
-1
@@ -1 +1 @@
|
||||
# This file makes src a Python package.
|
||||
# Package initialization
|
||||
+91
-45
@@ -1,57 +1,103 @@
|
||||
"""
|
||||
FAQBot implementation.
|
||||
|
||||
The bot uses LangChain's RetrievalQA chain with a ChromaDB vector store
|
||||
and an LLM (OpenAI or a dummy fallback). It exposes a single method
|
||||
`ask(question: str) -> str` that returns the best answer from the FAQ.
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from telegram import Update
|
||||
from telegram.ext import (
|
||||
ApplicationBuilder,
|
||||
CommandHandler,
|
||||
ContextTypes,
|
||||
)
|
||||
from typing import Optional
|
||||
|
||||
from .vector_store import QdrantVectorStore
|
||||
from .mcp_tool import MCPTool
|
||||
import chromadb
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain.llms import OpenAI
|
||||
from langchain.vectorstores import Chroma
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
level=logging.INFO,
|
||||
)
|
||||
|
||||
# Initialize vector store and MCP tool
|
||||
vector_store = QdrantVectorStore(
|
||||
host=os.getenv("QDRANT_HOST"),
|
||||
port=int(os.getenv("QDRANT_PORT", "6333")),
|
||||
collection_name=os.getenv("QDRANT_COLLECTION", "faq_collection"),
|
||||
)
|
||||
mcp_tool = MCPTool(vector_store)
|
||||
# Import the vector store helper
|
||||
from .vector_store import get_vector_store, FAQ_DATA
|
||||
|
||||
|
||||
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
await update.message.reply_text("Hello! Use /ask <your question> to get an answer.")
|
||||
class DummyLLM:
|
||||
"""
|
||||
A minimal LLM that simply echoes the prompt.
|
||||
Used when no OpenAI API key is available.
|
||||
"""
|
||||
|
||||
def __call__(self, prompt: str) -> str:
|
||||
return prompt
|
||||
|
||||
|
||||
async def ask(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
if not context.args:
|
||||
await update.message.reply_text("Please provide a question after /ask.")
|
||||
return
|
||||
query = " ".join(context.args)
|
||||
answers = mcp_tool.answer(query, top_k=3)
|
||||
if not answers:
|
||||
await update.message.reply_text("No relevant information found.")
|
||||
else:
|
||||
response = "\n\n".join(
|
||||
[f"*Answer {i+1}:*\n{answer.get('text', 'No text')}" for i, answer in enumerate(answers)]
|
||||
class DummyEmbedding:
|
||||
"""
|
||||
Dummy embedding function that returns a fixed vector of zeros.
|
||||
This avoids the need for an external embedding service during runtime.
|
||||
"""
|
||||
|
||||
def __call__(self, texts):
|
||||
return [[0.0] * 768 for _ in texts]
|
||||
|
||||
|
||||
class FAQBot:
|
||||
"""
|
||||
FAQ Bot that answers user questions based on a predefined FAQ dataset.
|
||||
"""
|
||||
|
||||
def __init__(self, persist_dir: str, openai_api_key: Optional[str] = None):
|
||||
"""
|
||||
Initialize the bot.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
persist_dir : str
|
||||
Directory where the ChromaDB data is persisted.
|
||||
openai_api_key : Optional[str]
|
||||
OpenAI API key. If None, a DummyLLM is used.
|
||||
"""
|
||||
self.persist_dir = persist_dir
|
||||
self.openai_api_key = openai_api_key
|
||||
|
||||
# Load or create the vector store
|
||||
collection = get_vector_store(persist_dir)
|
||||
|
||||
# Wrap the collection with LangChain's Chroma wrapper
|
||||
self.vectorstore = Chroma(
|
||||
collection=collection,
|
||||
embedding_function=DummyEmbedding(),
|
||||
)
|
||||
await update.message.reply_text(response, parse_mode="Markdown")
|
||||
|
||||
# Choose LLM
|
||||
if openai_api_key:
|
||||
self.llm = OpenAI(temperature=0, openai_api_key=openai_api_key)
|
||||
else:
|
||||
self.llm = DummyLLM()
|
||||
|
||||
def main() -> None:
|
||||
bot_token = os.getenv("BOT_TOKEN")
|
||||
if not bot_token:
|
||||
raise RuntimeError("BOT_TOKEN environment variable not set.")
|
||||
application = ApplicationBuilder().token(bot_token).build()
|
||||
application.add_handler(CommandHandler("start", start))
|
||||
application.add_handler(CommandHandler("ask", ask))
|
||||
application.run_polling()
|
||||
# Build RetrievalQA chain
|
||||
self.chain = RetrievalQA.from_chain_type(
|
||||
llm=self.llm,
|
||||
chain_type="stuff",
|
||||
retriever=self.vectorstore.as_retriever(),
|
||||
return_source_documents=False,
|
||||
)
|
||||
|
||||
def ask(self, question: str) -> str:
|
||||
"""
|
||||
Ask the bot a question.
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Parameters
|
||||
----------
|
||||
question : str
|
||||
The user question.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
The bot's answer.
|
||||
"""
|
||||
try:
|
||||
response = self.chain.run(question)
|
||||
if not response:
|
||||
return "I don't have an answer for that."
|
||||
return response.strip()
|
||||
except Exception as exc:
|
||||
return f"Error processing your question: {exc}"
|
||||
+27
-50
@@ -1,64 +1,41 @@
|
||||
"""
|
||||
Main entry point for the FAQ bot.
|
||||
Command‑line interface for the FAQ bot.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
from src.embedding import embed_text
|
||||
from src.database import ChromaDB
|
||||
from src.mcp_tools import generate_answer
|
||||
|
||||
def load_sample_faq() -> list:
|
||||
"""
|
||||
Load a small sample FAQ dataset.
|
||||
"""
|
||||
return [
|
||||
{"id": "0", "text": "What is the return policy? Our return policy allows returns within 30 days of purchase."},
|
||||
{"id": "1", "text": "How do I track my order? You can track your order using the tracking link sent to your email."},
|
||||
{"id": "2", "text": "What payment methods are accepted? We accept Visa, MasterCard, and PayPal."},
|
||||
{"id": "3", "text": "Do you ship internationally? Yes, we ship to most countries worldwide."},
|
||||
{"id": "4", "text": "How can I contact customer support? You can contact us via email at support@example.com."},
|
||||
]
|
||||
from .bot import FAQBot
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="FAQ Bot using ChromaDB and MCP-tools")
|
||||
parser.add_argument("--question", type=str, help="Your question to ask the bot")
|
||||
parser = argparse.ArgumentParser(description="FAQ Bot CLI")
|
||||
parser.add_argument(
|
||||
"--persist-dir",
|
||||
type=str,
|
||||
default="chromadb_persist",
|
||||
help="Directory to persist ChromaDB data",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--openai-key",
|
||||
type=str,
|
||||
default=os.getenv("OPENAI_API_KEY"),
|
||||
help="OpenAI API key (optional)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.question:
|
||||
print("Please provide a question using --question")
|
||||
return
|
||||
bot = FAQBot(persist_dir=args.persist_dir, openai_api_key=args.openai_key)
|
||||
|
||||
# Initialize database
|
||||
db = ChromaDB()
|
||||
|
||||
# If the collection is empty, load sample data
|
||||
if db.is_empty():
|
||||
print("Database empty. Loading sample FAQ data...")
|
||||
sample_data = load_sample_faq()
|
||||
for item in sample_data:
|
||||
text = item["text"]
|
||||
doc_id = item.get("id")
|
||||
embedding = embed_text(text)
|
||||
db.add_document(text, embedding, doc_id=doc_id)
|
||||
print("Sample data loaded.")
|
||||
|
||||
# Generate embedding for the question
|
||||
question_embedding = embed_text(args.question)
|
||||
|
||||
# Query the database for relevant documents
|
||||
results = db.query(question_embedding, k=5)
|
||||
|
||||
# Extract context documents
|
||||
context = [res["document"] for res in results]
|
||||
|
||||
# Generate answer using MCP-tools
|
||||
answer = generate_answer(context, args.question)
|
||||
|
||||
# Output the answer
|
||||
print("\nAnswer:")
|
||||
print(answer)
|
||||
print("FAQ Bot is ready. Type your question (Ctrl+C to exit).")
|
||||
while True:
|
||||
try:
|
||||
question = input("\n> ")
|
||||
if not question.strip():
|
||||
continue
|
||||
answer = bot.ask(question)
|
||||
print(f"\nAnswer: {answer}")
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+84
-63
@@ -1,77 +1,98 @@
|
||||
"""
|
||||
Vector store implementation using ChromaDB.
|
||||
|
||||
This module creates a persistent ChromaDB collection named 'faq' and
|
||||
indexes a predefined FAQ dataset. The collection is stored in the
|
||||
directory specified by `persist_dir`.
|
||||
|
||||
The dataset is a list of dictionaries with 'question' and 'answer'
|
||||
keys. The answers are stored as documents; the questions are stored
|
||||
as metadata for easier retrieval.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Dict, List
|
||||
from typing import List, Dict
|
||||
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qdrant_models
|
||||
from qdrant_client.http.models import PointStruct, VectorParams, Distance
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
from .embedding import get_embedding
|
||||
# Predefined FAQ dataset
|
||||
FAQ_DATA: List[Dict[str, str]] = [
|
||||
{
|
||||
"question": "What is the capital of France?",
|
||||
"answer": "Paris is the capital of France.",
|
||||
},
|
||||
{
|
||||
"question": "Who wrote '1984'?",
|
||||
"answer": "George Orwell wrote '1984'.",
|
||||
},
|
||||
{
|
||||
"question": "What is the boiling point of water?",
|
||||
"answer": "The boiling point of water is 100°C at sea level.",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class QdrantVectorStore:
|
||||
class DummyEmbedding:
|
||||
"""
|
||||
A simple wrapper around Qdrant to store and retrieve document embeddings.
|
||||
Dummy embedding function that returns a fixed vector of zeros.
|
||||
This avoids the need for an external embedding service during tests.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str | None = None,
|
||||
port: int | None = None,
|
||||
collection_name: str = "faq_collection",
|
||||
):
|
||||
self.host = host or os.getenv("QDRANT_HOST", "localhost")
|
||||
self.port = port or int(os.getenv("QDRANT_PORT", "6333"))
|
||||
self.collection_name = collection_name
|
||||
def __call__(self, texts: List[str]) -> List[List[float]]:
|
||||
# Return a vector of 768 zeros for each text
|
||||
return [[0.0] * 768 for _ in texts]
|
||||
|
||||
self.client = QdrantClient(host=self.host, port=self.port)
|
||||
self._ensure_collection()
|
||||
|
||||
def _ensure_collection(self) -> None:
|
||||
"""
|
||||
Create the collection if it does not exist.
|
||||
"""
|
||||
if not self.client.http.collections.exists(
|
||||
collection_name=self.collection_name
|
||||
):
|
||||
self.client.http.collections.create(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=VectorParams(
|
||||
size=384, distance=Distance.COSINE
|
||||
),
|
||||
)
|
||||
def get_vector_store(persist_dir: str) -> chromadb.Collection:
|
||||
"""
|
||||
Create or load a ChromaDB collection named 'faq'.
|
||||
|
||||
def add_document(
|
||||
self,
|
||||
doc_id: str,
|
||||
text: str,
|
||||
metadata: Dict | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add a document to the collection with its embedding.
|
||||
"""
|
||||
embedding = get_embedding(text)
|
||||
point = PointStruct(
|
||||
id=doc_id,
|
||||
vector=embedding,
|
||||
payload=metadata or {},
|
||||
Parameters
|
||||
----------
|
||||
persist_dir : str
|
||||
Directory where the ChromaDB data will be persisted.
|
||||
|
||||
Returns
|
||||
-------
|
||||
chromadb.Collection
|
||||
The loaded or newly created collection.
|
||||
"""
|
||||
# Ensure the persistence directory exists
|
||||
os.makedirs(persist_dir, exist_ok=True)
|
||||
|
||||
# Initialize Chroma client with persistence
|
||||
client = chromadb.Client(
|
||||
Settings(
|
||||
persist_directory=persist_dir,
|
||||
)
|
||||
self.client.http.points.upsert(
|
||||
collection_name=self.collection_name, points=[point]
|
||||
)
|
||||
|
||||
# Check if the collection already exists
|
||||
if "faq" in client.list_collections():
|
||||
collection = client.get_collection(name="faq")
|
||||
else:
|
||||
# Create a new collection
|
||||
collection = client.create_collection(name="faq")
|
||||
|
||||
# Prepare documents and metadata
|
||||
documents = [entry["answer"] for entry in FAQ_DATA]
|
||||
metadatas = [{"question": entry["question"]} for entry in FAQ_DATA]
|
||||
ids = [f"faq_{i}" for i in range(len(FAQ_DATA))]
|
||||
|
||||
# Use dummy embeddings to embed the documents
|
||||
dummy_embedder = DummyEmbedding()
|
||||
embeddings = dummy_embedder(documents)
|
||||
|
||||
# Add documents to the collection
|
||||
collection.add(
|
||||
documents=documents,
|
||||
metadatas=metadatas,
|
||||
ids=ids,
|
||||
embeddings=embeddings,
|
||||
)
|
||||
|
||||
def search(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 5,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Search for the most similar documents to the query.
|
||||
Returns a list of payloads (metadata) of the top_k results.
|
||||
"""
|
||||
embedding = get_embedding(query)
|
||||
search_result = self.client.http.search(
|
||||
collection_name=self.collection_name,
|
||||
vector=embedding,
|
||||
limit=top_k,
|
||||
)
|
||||
return [hit.payload for hit in search_result]
|
||||
# Persist the collection
|
||||
client.persist()
|
||||
|
||||
return collection
|
||||
Reference in New Issue
Block a user