문서 기반 질의응답 봇

from langchain.text_splitter import RecursiveCharacterTextSplitter
from pypdf import PdfReader

text = ''
for page in PdfReader('handbook.pdf').pages:
    text += page.extract_text() + '\n'

splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
chunks = splitter.split_text(text)
print(f'{len(chunks)} chunks')

import chromadb
from openai import OpenAI

oai = OpenAI()
chroma = chromadb.PersistentClient(path='./db')
collection = chroma.get_or_create_collection('handbook')

resp = oai.embeddings.create(model='text-embedding-3-small', input=chunks)
vectors = [d.embedding for d in resp.data]
collection.add(ids=[f'c{i}' for i in range(len(chunks))], embeddings=vectors, documents=chunks)

def query(question, k=4):
    qvec = oai.embeddings.create(model='text-embedding-3-small', input=question).data[0].embedding
    res = collection.query(query_embeddings=[qvec], n_results=k)
    return res['documents'][0]

def make_prompt(question, chunks):
    context = '\n\n'.join(f'[{i+1}] {c}' for i, c in enumerate(chunks))
    return f'''
Using only the context below, answer the question.
If the answer is not present, say so.
Cite the source like [1], [2], etc.

Context:
{context}

Question: {question}
Answer:
'''

print(make_prompt("What is the return policy?", ["Returns accepted within 30 days.", "Item must be unused."]))

def answer(question):
    chunks = query(question)
    prompt = make_prompt(question, chunks)
    response = oai.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': prompt}],
        temperature=0,
    )
    return response.choices[0].message.content

print(answer('What is our refund policy?'))

# pip install fastapi uvicorn
from fastapi import FastAPI
app = FastAPI()

@app.post('/ask')
def ask(payload: dict):
    return {'answer': answer(payload['question'])}

messages = [{'role': 'system', 'content': 'You are a doc Q&A assistant.'}]

def chat(user_input):
    messages.append({'role': 'user', 'content': user_input})
    chunks = query(user_input)
    context_msg = {'role': 'system', 'content': f'Context:\n{chr(10).join(chunks)}'}
    response = oai.chat.completions.create(
        model='gpt-4o-mini',
        messages=[messages[0], context_msg] + messages[1:]
    )
    answer = response.choices[0].message.content
    messages.append({'role': 'assistant', 'content': answer})
    return answer

metadata = [{'source': 'handbook.pdf', 'page': i} for i in range(len(chunks))]
# In the prompt, include 'source: <url>' so the LLM can produce clickable references.