1import os
2from openai import OpenAI
3import chromadb
4from dotenv import load_dotenv
5
6load_dotenv()
7client = OpenAI()
8
9# Sample knowledge base
10knowledge_base = [
11 "MinAI Learning Platform offers courses in AI, Data Science, and Automation.",
12 "Course pricing starts at 500,000 VND for basic courses.",
13 "Premium subscription costs 2,000,000 VND per year with unlimited access.",
14 "Refunds are available within 7 days of purchase.",
15 "Contact support at support@minai.vn for assistance.",
16]
17
18# Initialize vector database
19chroma_client = chromadb.Client()
20collection = chroma_client.create_collection("minai_kb")
21
22# Index documents
23print("Indexing documents...")
24for i, doc in enumerate(knowledge_base):
25 embedding = client.embeddings.create(
26 model="text-embedding-3-small",
27 input=doc
28 ).data[0].embedding
29
30 collection.add(
31 ids=[f"doc_{i}"],
32 embeddings=[embedding],
33 documents=[doc]
34 )
35
36def ask_rag(question):
37 # Get query embedding
38 query_embedding = client.embeddings.create(
39 model="text-embedding-3-small",
40 input=question
41 ).data[0].embedding
42
43 # Retrieve relevant docs
44 results = collection.query(
45 query_embeddings=[query_embedding],
46 n_results=2
47 )
48 context = "\n".join(results["documents"][0])
49
50 # Generate answer
51 response = client.chat.completions.create(
52 model="gpt-4o-mini",
53 messages=[
54 {"role": "system", "content": "Answer based only on the provided context. If you don't know, say so."},
55 {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
56 ]
57 )
58
59 return response.choices[0].message.content
60
61# Test
62questions = [
63 "How much does a premium subscription cost?",
64 "What's the refund policy?",
65 "How do I contact support?",
66]
67
68for q in questions:
69 print(f"\nQ: {q}")
70 print(f"A: {ask_rag(q)}")
