1
You are a knowledge assistant. Please answer the user's question based on the following reference materials.
2
If you cannot find the answer in the reference materials, please state so honestly.
3

4
Reference materials:
5
{retrieved_chunks}
6

7
User question: {user_query}

1
import os
2
from openai import OpenAI
3
import chromadb
4
from chromadb.utils import embedding_functions
5

6
# ============================================================
7
# 1. Initialization
8
# ============================================================
9
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
10

11
# Use OpenAI's Embedding function
12
embedding_fn = embedding_functions.OpenAIEmbeddingFunction(
13
    api_key=os.environ["OPENAI_API_KEY"],
14
    model_name="text-embedding-3-small"
15
)
16

17
# Create or connect to a ChromaDB collection
18
chroma_client = chromadb.PersistentClient(path="./chroma_db")
19
collection = chroma_client.get_or_create_collection(
20
    name="my_knowledge_base",
21
    embedding_function=embedding_fn
22
)
23

24

25
# ============================================================
26
# 2. Document chunking function
27
# ============================================================
28
def chunk_text(text: str, chunk_size: int = 500, overlap: int = 100) -> list[str]:
29
    """
30
    Simple fixed-size chunking with overlap to maintain context continuity.
31

32
    Args:
33
        text: Input text
34
        chunk_size: Number of characters per chunk
35
        overlap: Number of overlapping characters between adjacent chunks
36

37
    Returns:
38
        List of text chunks
39
    """
40
    chunks = []
41
    start = 0
42
    while start < len(text):
43
        end = start + chunk_size
44
        chunk = text[start:end]
45
        chunks.append(chunk)
46
        start += chunk_size - overlap  # Next chunk start = current start + chunk_size - overlap
47
    return chunks
48

49

50
# ============================================================
51
# 3. Index documents: chunk documents and store them in the vector database
52
# ============================================================
53
def index_documents(documents: list[dict[str, str]]) -> None:
54
    """
55
    Index a list of documents into the vector database.
56

57
    Args:
58
        documents: List of documents, each containing "id", "title", "content" fields
59
    """
60
    for doc in documents:
61
        chunks = chunk_text(doc["content"])
62
        for i, chunk in enumerate(chunks):
63
            chunk_id = f"{doc['id']}_chunk_{i}"
64
            collection.add(
65
                ids=[chunk_id],
66
                documents=[chunk],
67
                metadatas=[{
68
                    "title": doc["title"],
69
                    "chunk_index": i,
70
                    "source": doc["id"]
71
                }]
72
            )
73
    print(f"Indexed {len(documents)} documents, {collection.count()} text chunks in total")
74

75

76
# ============================================================
77
# 4. Retrieval: return the most relevant text chunks based on the query
78
# ============================================================
79
def retrieve(query: str, top_k: int = 3) -> list[str]:
80
    """
81
    Retrieve the text chunks most relevant to the query.
82

83
    Args:
84
        query: User query
85
        top_k: Number of text chunks to return
86

87
    Returns:
88
        List of relevant text chunks
89
    """
90
    results = collection.query(
91
        query_texts=[query],
92
        n_results=top_k
93
    )
94
    # results["documents"] is a nested list: [[chunk1, chunk2, ...]]
95
    return results["documents"][0] if results["documents"] else []
96

97

98
# ============================================================
99
# 5. Augmented generation: the complete RAG question-answering workflow
100
# ============================================================
101
SYSTEM_PROMPT = """You are a knowledge assistant. Please answer the user's question strictly based on the following reference materials.
102

103
Rules:
104
1. If the reference materials contain the answer, answer based on the materials and cite the source
105
2. If the reference materials are insufficient to answer the question, honestly say "My current knowledge base does not contain enough information to answer this question"
106
3. Do not fabricate information not present in the reference materials
107

108
Reference materials:
109
{context}"""
110

111

112
def ask(query: str, top_k: int = 3) -> str:
113
    """
114
    Execute the complete RAG question-answering pipeline.
115

116
    Args:
117
        query: User question
118
        top_k: Number of text chunks to retrieve
119

120
    Returns:
121
        LLM-generated answer
122
    """
123
    # Step 1: Retrieve
124
    retrieved_chunks = retrieve(query, top_k=top_k)
125

126
    if not retrieved_chunks:
127
        return "No relevant reference materials found; unable to answer this question."
128

129
    # Step 2: Concatenate context
130
    context = "\n\n---\n\n".join(
131
        f"[Source {i+1}]\n{chunk}"
132
        for i, chunk in enumerate(retrieved_chunks)
133
    )
134

135
    # Step 3: Call the LLM to generate an answer
136
    response = client.chat.completions.create(
137
        model="gpt-4o-mini",
138
        messages=[
139
            {"role": "system", "content": SYSTEM_PROMPT.format(context=context)},
140
            {"role": "user", "content": query}
141
        ],
142
        temperature=0.3,  # Low temperature to reduce hallucinations
143
        max_tokens=1024
144
    )
145

146
    return response.choices[0].message.content
147

148

149
# ============================================================
150
# 6. Demo: index some documents and test question answering
151
# ============================================================
152
if __name__ == "__main__":
153
    # Prepare sample documents
154
    documents = [
155
        {
156
            "id": "python_history",
157
            "title": "A Brief History of the Python Language",
158
            "content": """
159
Python was started by Guido van Rossum during Christmas 1989, and the first public version was released in 1991.
160
Python's design philosophy emphasizes code readability, using significant whitespace indentation to delimit code blocks.
161
Python 2.0 was released in 2000, introducing garbage collection and Unicode support.
162
Python 3.0 was released in 2008 as a major backward-incompatible release.
163
Major changes in Python 3.0 include print changing from a statement to a function, and all strings using Unicode by default.
164
As of 2026, Python 3.14 is the latest stable version.
165
Python is widely used in web development (Django, FastAPI), data science (NumPy, Pandas), and AI/machine learning (PyTorch, TensorFlow), among other fields.
166
            """.strip()
167
        },
168
        {
169
            "id": "rag_overview",
170
            "title": "RAG Technology Overview",
171
            "content": """
172
RAG (Retrieval-Augmented Generation) was first proposed by Meta AI (then Facebook AI) in 2020.
173
RAG combines information retrieval and text generation technologies, aiming to solve two core problems of large language models: knowledge timeliness and hallucination.
174
A typical RAG architecture includes two stages: the indexing stage and the query stage.
175
Indexing stage: chunk knowledge base documents, embed them, and store them in a vector database.
176
Query stage: after the user's question is embedded, retrieve from the vector database, and send the retrieval results together with the question to the LLM to generate an answer.
177
The advantages of RAG are: knowledge can be updated in real time (no need to retrain the model), answers are traceable (with explicit citation sources), and costs are relatively low.
178
Common RAG optimization techniques include: HyDE (Hypothetical Document Embeddings), Multi-Query Retrieval, Re-ranking, Self-RAG, etc.
179
            """.strip()
180
        },
181
        {
182
            "id": "vector_db",
183
            "title": "Introduction to Vector Databases",
184
            "content": """
185
A vector database is a database system specifically designed for storing and retrieving vectors (high-dimensional floating-point arrays).
186
Unlike traditional databases, the core operation of a vector database is approximate nearest neighbor search (ANN), i.e., quickly finding the K vectors most similar to a query vector.
187
Commonly used vector databases include: ChromaDB (lightweight, suitable for prototyping), Pinecone (fully managed cloud service), Weaviate (open source, supports hybrid search), Milvus (high performance, suitable for large-scale production environments), Qdrant (written in Rust, excellent performance).
188
Factors to consider when choosing a vector database: data scale, query latency requirements, whether filtered search is needed, operational costs, etc.
189
For prototypes and small-scale applications, ChromaDB is usually the most suitable choice; for large-scale production applications, Milvus and Pinecone are more mature choices.
190
            """.strip()
191
        }
192
    ]
193

194
    # Index documents
195
    index_documents(documents)
196

197
    # Test question answering
198
    test_questions = [
199
        "When was Python created? By whom?",
200
        "What problems can RAG technology solve?",
201
        "What are some commonly used vector databases?",
202
        "What's the weather like today?"  # Content not in the knowledge base
203
    ]
204

205
    for q in test_questions:
206
        print(f"\n{'='*60}")
207
        print(f"Q: {q}")
208
        print(f"{'='*60}")
209
        answer = ask(q)
210
        print(f"A: {answer}")