#!/usr/bin/env python3
"""
question_and_answer: Tool for performing Q&A on PDF documents using retrieval augmented generation.
This module provides functionality to extract text from PDF binary data, split it into
chunks, retrieve relevant segments via a vector store, and generate an answer to a
user-provided question using a language model chain.
"""
import io
import logging
from typing import Annotated, Dict, Any, List
from PyPDF2 import PdfReader
from pydantic import BaseModel, Field
import hydra
from langchain.chains.question_answering import load_qa_chain
from langchain.docstore.document import Document
from langchain.text_splitter import CharacterTextSplitter
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_core.messages import ToolMessage
from langchain_core.tools import tool
from langchain_core.tools.base import InjectedToolCallId
from langchain_core.embeddings import Embeddings
from langchain_community.vectorstores import Annoy
from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import OpenAIEmbeddings
from langgraph.types import Command
from langgraph.prebuilt import InjectedState
# Set up logging.
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# Load configuration using Hydra.
with hydra.initialize(version_base=None, config_path="../../configs"):
cfg = hydra.compose(
config_name="config", overrides=["tools/question_and_answer=default"]
)
cfg = cfg.tools.question_and_answer
logger.info("Loaded Question and Answer tool configuration.")
class QuestionAndAnswerInput(BaseModel):
"""
Input schema for the PDF Question and Answer tool.
Attributes:
question (str): The question to ask regarding the PDF content.
tool_call_id (str): Unique identifier for the tool call, injected automatically.
"""
question: str = Field(description="The question to ask regarding the PDF content.")
tool_call_id: Annotated[str, InjectedToolCallId]
state: Annotated[dict, InjectedState]
def extract_text_from_pdf_data(pdf_bytes: bytes) -> str:
"""
Extract text content from PDF binary data.
This function uses PyPDF2 to read the provided PDF bytes and concatenates the text
extracted from each page.
Args:
pdf_bytes (bytes): The binary data of the PDF document.
Returns:
str: The complete text extracted from the PDF.
"""
reader = PdfReader(io.BytesIO(pdf_bytes))
text = ""
for page in reader.pages:
page_text = page.extract_text() or ""
text += page_text
return text
def generate_answer(
question: str, pdf_bytes: bytes, llm_model: BaseChatModel
) -> Dict[str, Any]:
"""
Generate an answer for a question using retrieval augmented generation on PDF content.
This function extracts text from the PDF data, splits the text into manageable chunks,
performs a similarity search to retrieve the most relevant segments, and then uses a
question-answering chain (built using the provided llm_model) to generate an answer.
Args:
question (str): The question to be answered.
pdf_bytes (bytes): The binary content of the PDF document.
llm_model (BaseChatModel): The language model instance to use for answering.
Returns:
Dict[str, Any]: A dictionary containing the answer generated by the language model.
"""
text = extract_text_from_pdf_data(pdf_bytes)
logger.info("Extracted text from PDF.")
text_splitter = CharacterTextSplitter(
separator="\n", chunk_size=cfg.chunk_size, chunk_overlap=cfg.chunk_overlap
)
chunks = text_splitter.split_text(text)
documents: List[Document] = [Document(page_content=chunk) for chunk in chunks]
logger.info("Split PDF text into %d chunks.", len(documents))
embeddings = OpenAIEmbeddings(openai_api_key=cfg.openai_api_key)
vector_store = Annoy.from_documents(documents, embeddings)
search_results = vector_store.similarity_search(question, k=cfg.num_retrievals)
logger.info("Retrieved %d relevant document chunks.", len(search_results))
# Use the provided llm_model to build the QA chain.
qa_chain = load_qa_chain(llm_model, chain_type=cfg.qa_chain_type)
answer = qa_chain.invoke(
input={"input_documents": search_results, "question": question}
)
return answer
def generate_answer2(
question: str, pdf_url: str, text_embedding_model: Embeddings
) -> Dict[str, Any]:
"""
Generate an answer for a question using retrieval augmented generation on PDF content.
This function extracts text from the PDF data, splits the text into manageable chunks,
performs a similarity search to retrieve the most relevant segments, and then uses a
question-answering chain (built using the provided llm_model) to generate an answer.
Args:
question (str): The question to be answered.
pdf_bytes (bytes): The binary content of the PDF document.
llm_model (BaseChatModel): The language model instance to use for answering.
Returns:
Dict[str, Any]: A dictionary containing the answer generated by the language model.
"""
# text = extract_text_from_pdf_data(pdf_bytes)
# logger.info("Extracted text from PDF.")
logger.log(logging.INFO, "searching the article with the question: %s", question)
# Load the article
# loader = PyPDFLoader(state['pdf_file_name'])
# loader = PyPDFLoader("https://arxiv.org/pdf/2310.08365")
loader = PyPDFLoader(pdf_url)
# Load the pages of the article
pages = []
for page in loader.lazy_load():
pages.append(page)
# Set up text embedding model
# text_embedding_model = state['text_embedding_model']
# text_embedding_model = OpenAIEmbeddings(openai_api_key=cfg.openai_api_key)
logging.info("Loaded text embedding model %s", text_embedding_model)
# Create a vector store from the pages
vector_store = InMemoryVectorStore.from_documents(pages, text_embedding_model)
# Search the article with the question
docs = vector_store.similarity_search(question)
# Return the content of the pages
return "\n".join([doc.page_content for doc in docs])
# return answer
@tool(args_schema=QuestionAndAnswerInput)
def question_and_answer_tool(
question: str,
tool_call_id: Annotated[str, InjectedToolCallId],
state: Annotated[dict, InjectedState],
) -> Dict[str, Any]:
"""
Answer a question using PDF content stored in the state via retrieval augmented generation.
This tool retrieves the PDF binary data from the state (under the key "pdf_data"), extracts its
textual content, and generates an answer to the specified question. It also extracts the
llm_model (of type BaseChatModel) from the state to use for answering.
Args:
question (str): The question regarding the PDF content.
tool_call_id (str): Unique identifier for the current tool call.
state (dict): A dictionary representing the current state, expected to contain PDF data
under the key "pdf_data" with a sub-key "pdf_object" for the binary content,
and a key "llm_model" holding the language model instance.
Returns:
Dict[str, Any]: A dictionary containing the generated answer or an error message.
"""
logger.info("Starting PDF Question and Answer tool using PDF data from state.")
# print (state['text_embedding_model'])
text_embedding_model = state["text_embedding_model"]
pdf_state = state.get("pdf_data")
if not pdf_state:
error_msg = "No pdf_data found in state."
logger.error(error_msg)
return Command(
update={
"messages": [ToolMessage(content=error_msg, tool_call_id=tool_call_id)]
}
)
pdf_bytes = pdf_state.get("pdf_object")
if not pdf_bytes:
error_msg = "PDF binary data is missing in the pdf_data from state."
logger.error(error_msg)
return Command(
update={
"messages": [ToolMessage(content=error_msg, tool_call_id=tool_call_id)]
}
)
pdf_url = pdf_state.get("pdf_url")
# Retrieve llm_model from state; use a default if not provided.
llm_model = state.get("llm_model")
if not llm_model:
logger.error("Missing LLM model instance in state.")
return {"error": "No LLM model found in state."}
# answer = generate_answer(question, pdf_bytes, llm_model)
print(pdf_url)
answer = generate_answer2(question, pdf_url, text_embedding_model)
# logger.info("Generated answer: %s", answer)
return answer
Datasets
Models
