Spaces:
Sleeping
Sleeping
File size: 3,046 Bytes
b534a53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | from langchain_openai import ChatOpenAI
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
# from langchain.chains.history_aware_retriever import create_history_aware_retriever
# from langchain.chains.retrieval import create_retrieval_chain
# from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from src.rag.vector_store import get_vector_store
def build_doc_assistant(api_key: str):
"""
Constructs the conversational RAG pipeline.
"""
# 1. Initialize our LLM (temperature=0 because we want factual answers, not creative ones)
llm = ChatOpenAI(api_key=api_key, model="gpt-3.5-turbo", temperature=0)
# 2. Connect to our Vector DB (k=2 means return the top 2 most relevant chunks)
retriever = get_vector_store(api_key).as_retriever(search_kwargs={"k": 2})
# ==========================================
# STEP 1: The "Question Reformulation" Prompt
# ==========================================
contextualize_q_system_prompt = (
"Given a chat history and the latest user question "
"which might reference context in the chat history, "
"formulate a standalone question which can be understood "
"without the chat history. Do NOT answer the question, "
"just reformulate it if needed and otherwise return it as is."
)
contextualize_q_prompt = ChatPromptTemplate.from_messages([
("system", contextualize_q_system_prompt),
MessagesPlaceholder("chat_history"), # Injects our memory here
("human", "{input}"),
])
# This chain automatically handles rewriting the query before searching
history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)
# ==========================================
# STEP 2: The "Final Answer" Prompt
# ==========================================
system_prompt = (
"You are an elite AI Engineering Assistant. "
"Use the following pieces of retrieved context to answer the question. "
"If the answer is not contained in the context, say 'I don't know based on the documentation.' "
"Do not make up an answer. Keep it concise.\n\n"
"Context: {context}"
)
qa_prompt = ChatPromptTemplate.from_messages([
("system", system_prompt),
MessagesPlaceholder("chat_history"),
("human", "{input}"),
])
# This chain handles injecting the retrieved chunks into the {context} variable
question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
# ==========================================
# STEP 3: Tie it all together
# ==========================================
rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
return rag_chain |