Создание цепочки RAG от начала до конца

from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

loader = PyPDFLoader('handbook.pdf')
docs = loader.load()
splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
chunks = splitter.split_documents(docs)
print(f'{len(chunks)} chunks loaded')

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

embeddings = OpenAIEmbeddings(model='text-embedding-3-small')
store = Chroma.from_documents(
    chunks,
    embeddings,
    persist_directory='./handbook_db'
)

retriever = store.as_retriever(
    search_type='similarity',
    search_kwargs={'k': 4}
)

from langchain.prompts import ChatPromptTemplate

rag_prompt = ChatPromptTemplate.from_template('''
You are a helpful assistant. Use the context below to answer.
If the answer is not in the context, say you do not know.
Cite sources like [1], [2].

Context:
{context}

Question: {question}

Answer:
''')

def format_docs(docs):
    return '\n\n'.join(
        f'[{i+1}] (source: {d.metadata.get("source", "?")}, page {d.metadata.get("page", "?")})\n{d.page_content}'
        for i, d in enumerate(docs)
    )

from types import SimpleNamespace
docs = [
    SimpleNamespace(metadata={'source': 'handbook.pdf', 'page': 3}, page_content='Vacation policy details.'),
    SimpleNamespace(metadata={'source': 'handbook.pdf', 'page': 5}, page_content='Sick leave details.'),
]
print(format_docs(docs))

from langchain_core.runnables import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from langchain_openai import ChatOpenAI

model = ChatOpenAI(model='gpt-4o-mini', temperature=0)

rag_chain = (
    {'context': retriever | format_docs, 'question': RunnablePassthrough()}
    | rag_prompt
    | model
    | StrOutputParser()
)

answer = rag_chain.invoke('What is our refund policy?')
print(answer)

for chunk in rag_chain.stream('What is our refund policy?'):
    print(chunk, end='', flush=True)

from langchain_core.runnables import RunnableParallel

rag_with_sources = RunnableParallel(
    context=retriever,
    answer=rag_chain
)

result = rag_with_sources.invoke('What is our refund policy?')
print(result['answer'])
for doc in result['context']:
    print(doc.metadata)

from langchain.chains import create_history_aware_retriever
# 'And what about XL sizes?' -> 'What is the refund policy for XL sizes?'