Files
povtornyy-ekzamen-faq-bot-c…/src/main.py
T

139 lines
4.0 KiB
Python

import os
from typing import List, Dict
import openai
from qdrant_client import QdrantClient
from qdrant_client.http import models as qdrant_models
from dotenv import load_dotenv
load_dotenv()
# Configuration
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", None)
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
raise RuntimeError("OPENAI_API_KEY environment variable not set")
openai.api_key = OPENAI_API_KEY
# Collection name
COLLECTION_NAME = "faq_collection"
# Embedding dimension for text-embedding-ada-002
EMBEDDING_DIM = 1536
# Sample FAQ data
FAQ_DATA = [
{
"question": "What is the return policy?",
"answer": "You can return any item within 30 days of purchase."
},
{
"question": "How do I track my order?",
"answer": "Use the tracking link sent to your email after shipping."
},
{
"question": "Do you offer international shipping?",
"answer": "Yes, we ship to most countries worldwide."
},
{
"question": "What payment methods are accepted?",
"answer": "We accept credit cards, PayPal, and bank transfers."
},
{
"question": "How can I contact customer support?",
"answer": "Email us at support@example.com or call 1-800-123-4567."
}
]
def get_embedding(text: str) -> List[float]:
"""
Generate an embedding for the given text using OpenAI's embedding model.
"""
response = openai.Embedding.create(
input=text,
model="text-embedding-ada-002"
)
return response["data"][0]["embedding"]
def create_or_recreate_collection(client: QdrantClient) -> None:
"""
Create a new collection or recreate it if it already exists.
"""
client.recreate_collection(
collection_name=COLLECTION_NAME,
vectors_config=qdrant_models.VectorParams(
size=EMBEDDING_DIM,
distance=qdrant_models.Distance.COSINE
)
)
def ingest_faqs(client: QdrantClient, faqs: List[Dict[str, str]]) -> None:
"""
Ingest FAQ data into Qdrant.
"""
points = []
for idx, faq in enumerate(faqs):
vector = get_embedding(faq["question"])
point = qdrant_models.PointStruct(
id=idx,
vector=vector,
payload={
"question": faq["question"],
"answer": faq["answer"]
}
)
points.append(point)
# Upsert points in batches
batch_size = 100
for i in range(0, len(points), batch_size):
batch = points[i:i+batch_size]
client.upsert(
collection_name=COLLECTION_NAME,
points=batch
)
def query_faq(client: QdrantClient, question: str, top_k: int = 1) -> str:
"""
Query the FAQ collection for the most relevant answer.
"""
query_vector = get_embedding(question)
search_result = client.search(
collection_name=COLLECTION_NAME,
query_vector=query_vector,
limit=top_k,
with_payload=True
)
if not search_result:
return "Sorry, I couldn't find an answer to your question."
# Return the answer from the top result
return search_result[0].payload.get("answer", "Answer not found.")
def main() -> None:
client = QdrantClient(
host=QDRANT_HOST,
port=QDRANT_PORT,
api_key=QDRANT_API_KEY
)
# Ingest FAQs (only if collection is empty or you want to refresh)
print("Ingesting FAQ data into Qdrant...")
create_or_recreate_collection(client)
ingest_faqs(client, FAQ_DATA)
print("Ingestion complete.")
print("\nFAQ Bot is ready. Type your question (or 'exit' to quit).")
while True:
user_input = input("\nYour question: ").strip()
if user_input.lower() in {"exit", "quit"}:
print("Goodbye!")
break
answer = query_faq(client, user_input)
print(f"Answer: {answer}")
if __name__ == "__main__":
main()