Spaces:
Runtime error
Runtime error
| """ | |
| AnswerGenerator: orchestrates retrieval, re-ranking, and answer generation. | |
| This module contains: | |
| - Retriever: Hybrid BM25 + dense retrieval over parsed chunks | |
| - Reranker: Cross-encoder based re-ranking of candidate chunks | |
| - AnswerGenerator: ties together retrieval, re-ranking, and LLM generation | |
| Each component is modular and can be swapped or extended (e.g., add HyDE retriever). | |
| """ | |
| import os | |
| import random | |
| from typing import List, Dict, Any, Tuple | |
| from src import logger, RetrieverConfig | |
| from src.utils import LLMClient | |
| from src.retriever import Retriever | |
| class AnswerGenerator: | |
| """ | |
| Generates answers by retrieving documents from a vector store | |
| and using them to build a context for an LLM. | |
| This version is optimized for low latency by skipping the reranking step. | |
| """ | |
| def __init__(self, collection_name: str): | |
| self.retriever = Retriever(collection_name, RetrieverConfig) | |
| self.context_chunks_count = 5 # Use top 5 chunks for the final prompt | |
| self.greetings = [ | |
| "Hello! I'm ready to answer your questions about the document. What would you like to know?", | |
| "Hi there! How can I help you with your document today?", | |
| "Hey! I've got the document open and I'm ready for your questions.", | |
| "Greetings! Ask me anything about the document, and I'll do my best to find the answer for you." | |
| ] | |
| def _truncate_to_last_sentence(self, text: str) -> str: | |
| """Finds the last period or newline and truncates the text to that point.""" | |
| # Find the last period | |
| last_period = text.rfind('.') | |
| # Find the last newline | |
| last_newline = text.rfind('\n') | |
| # Find the last of the two | |
| last_marker = max(last_period, last_newline) | |
| if last_marker != -1: | |
| return text[:last_marker + 1].strip() | |
| # If no sentence-ending punctuation, return the text as is (or a portion) | |
| return text | |
| def answer(self, question: str) -> Tuple[str, List[Dict[str, Any]]]: | |
| """ | |
| Retrieves documents, builds a context, and generates an answer. | |
| Handles simple greetings separately to improve user experience. | |
| """ | |
| # Handle simple greetings to avoid a failed retrieval | |
| normalized_question = question.lower().strip().rstrip('.,!') | |
| greeting_triggers = ["hi", "hello", "hey", "hallo", "hola"] | |
| if normalized_question in greeting_triggers: | |
| return random.choice(self.greetings), [] | |
| # Retrieve candidate documents from the vector store | |
| candidates = self.retriever.retrieve(question) | |
| if not candidates: | |
| logger.warning("No candidates retrieved from vector store.") | |
| return "The document does not contain information on this topic.", [] | |
| # Use the top N chunks for context, without reranking | |
| top_chunks = candidates[:self.context_chunks_count] | |
| context = "\n\n".join(f"- {c['narration']}" for c in top_chunks) | |
| # A more robust prompt that encourages a natural, conversational tone | |
| prompt = ( | |
| "You are a helpful and friendly AI assistant for document analysis. " | |
| "Your user is asking a question about a document. " | |
| "Based *only* on the context provided below, formulate a clear and conversational answer. " | |
| "Adopt a helpful and slightly informal tone, as if you were a knowledgeable colleague.\n\n" | |
| "CONTEXT:\n" | |
| "---------------------\n" | |
| f"{context}\n" | |
| "---------------------\n\n" | |
| "USER'S QUESTION: " | |
| f'"{question}"\n\n' | |
| "YOUR TASK:\n" | |
| "1. Carefully read the provided context.\n" | |
| "2. If the context contains the answer, explain it to the user in a natural, conversational way. Do not just repeat the text verbatim.\n" | |
| "3. If the context does not contain the necessary information, respond with: " | |
| "'I've checked the document, but I couldn't find any information on that topic.'\n" | |
| "4. **Crucially, do not use any information outside of the provided context.**\n\n" | |
| "Answer:" | |
| ) | |
| answer, finish_reason = LLMClient.generate(prompt, max_tokens=256) | |
| # Handle cases where the response might be cut off | |
| if finish_reason == 'length': | |
| logger.warning("LLM response was truncated due to token limit.") | |
| truncated_answer = self._truncate_to_last_sentence(answer) | |
| answer = truncated_answer + " ... (response shortened)" | |
| return answer, top_chunks | |