class QuestionGeneration(Enum):"""Enum class to specify the level of question generation for document processing.
Attributes:DOCUMENT_LEVEL (int): Represents question generation at the entire document level.FRAGMENT_LEVEL (int): Represents question generation at the individual text fragment level."""DOCUMENT_LEVEL = 1FRAGMENT_LEVEL = 2

def generate_questions(text: str) -> List[str]:"""Generates a list of questions based on the provided text using OpenAI.
Args:text (str): The context data from which questions are generated.
Returns:List[str]: A list of unique, filtered questions."""llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)prompt = PromptTemplate(input_variables=["context", "num_questions"],template="Using the context data: {context}\n\nGenerate a list of at least {num_questions} " "possible questions that can be asked about this context. Ensure the questions are " "directly answerable within the context and do not include any answers or headers. " "Separate the questions with a new line character.")chain = prompt | llm.with_structured_output(QuestionList)input_data = {"context": text, "num_questions": QUESTIONS_PER_DOCUMENT}result = chain.invoke(input_data)# Extract the list of questions from the QuestionList objectquestions = result.question_listfiltered_questions = clean_and_filter_questions(questions)return list(set(filtered_questions))

def process_documents(content: str, embedding_model: OpenAIEmbeddings):"""Process the document content, split it into fragments, generate questions,create a FAISS vector store, and return a retriever.
Args:content (str): The content of the document to process.embedding_model (OpenAIEmbeddings): The embedding model to use for vectorization.
Returns:VectorStoreRetriever: A retriever for the most relevant FAISS document."""# Split the whole text content into text documentstext_documents = split_document(content, DOCUMENT_MAX_TOKENS, DOCUMENT_OVERLAP_TOKENS)print(f'Text content split into: {len(text_documents)} documents')
documents = []counter = 0for i, text_document in enumerate(text_documents):text_fragments = split_document(text_document, FRAGMENT_MAX_TOKENS, FRAGMENT_OVERLAP_TOKENS)print(f'Text document {i} - split into: {len(text_fragments)} fragments')for j, text_fragment in enumerate(text_fragments):documents.append(Document(page_content=text_fragment,metadata={"type": "ORIGINAL", "index": counter, "text": text_document}))counter += 1if QUESTION_GENERATION == QuestionGeneration.FRAGMENT_LEVEL:questions = generate_questions(text_fragment)documents.extend([Document(page_content=question, metadata={"type": "AUGMENTED", "index": counter + idx, "text": text_document})for idx, question in enumerate(questions)])counter += len(questions)print(f'Text document {i} Text fragment {j} - generated: {len(questions)} questions')if QUESTION_GENERATION == QuestionGeneration.DOCUMENT_LEVEL:questions = generate_questions(text_document)documents.extend([Document(page_content=question, metadata={"type": "AUGMENTED", "index": counter + idx, "text": text_document})for idx, question in enumerate(questions)])counter += len(questions)print(f'Text document {i} - generated: {len(questions)} questions')
for document in documents:print_document("Dataset", document)
print(f'Creating store, calculating embeddings for {len(documents)} FAISS documents')vectorstore = FAISS.from_documents(documents, embedding_model)
print("Creating retriever returning the most relevant FAISS document")return vectorstore.as_retriever(search_kwargs={"k": 1})