Class DocumentQAEngine

document_qa/document_qa_engine.py:254–633

End-to-end RAG engine for scientific PDF documents. Orchestrates the full pipeline: 1. **PDF parsing** via a GROBID server (structured text + coordinates). 2. **Chunking** — paragraphs kept as-is or merged with :class:`TextMerger`. 3. **Embedding and storage** chunks are embedded

class DocumentQAEngine:
    """End-to-end RAG engine for scientific PDF documents.

    Orchestrates the full pipeline:

    1. **PDF parsing** via a GROBID server (structured text + coordinates).
    2. **Chunking** — paragraphs kept as-is or merged with :class:`TextMerger`.
    3. **Embedding and storage**  chunks are embedded and stored.
    4. **Retrieval + LLM** — relevant chunks are retrieved and fed to an LLM
       to produce an answer.

    Args:
        llm: A LangChain chat model (e.g. ``ChatOpenAI``).
        data_storage: A `DataStorage` instance for managing embeddings.
        grobid_url: URL of the GROBID server.
        memory: Optional ``ConversationBufferMemory`` for multi-turn context.

    """

    llm = None
    qa_chain_type = None

    default_prompts = {
        "stuff": stuff_prompt,
        "refine": refine_prompts,
        "map_reduce": map_reduce_prompt,
        "map_rerank": map_rerank_prompt,
    }

    def __init__(self, llm, data_storage: DataStorage, grobid_url=None, memory=None, ping_grobid_server: bool = True):

        self.llm = llm
        self.memory = memory
        self.chain = create_stuff_documents_chain(llm, self.default_prompts["stuff"].PROMPT)
        self.text_merger = TextMerger()
        self.data_storage = data_storage

        if grobid_url:
            self.grobid_processor = GrobidProcessor(grobid_url, ping_server=ping_grobid_server)

    def query_document(
        self, query: str, doc_id, output_parser=None, context_size=4, extraction_schema=None, verbose=False
    ) -> tuple[Any, str, list]:
        """Ask a question and get an LLM-generated answer.

        Retrieves the most relevant chunks from the vector store, feeds
        them as context to the LLM, and returns the response.

        Args:
            query: The natural-language question.
            doc_id: Document identifier returned by create_memory_embeddings`.
            output_parser: Optional LangChain output parser.  If provided the
                raw LLM response is re-processed into structured output.
            context_size: Number of chunks to retrieve as context (default 4).
            extraction_schema: Optional extraction schema.
            verbose: Print debug information.

        Returns: