query_service.py

from typing import List
from models import Question, FinalAnswer
from document_manager import DocumentManager
from retriever import VectorStoreManager
from workflow import RAGWorkflow
from rich import print as rprint
from rich.panel import Panel
import time

class QueryService:
    def __init__(self):
        self.workflow = RAGWorkflow()
        rprint(Panel("[bold green]QueryService Initialized[/bold green]", subtitle="Ready to process requests.", border_style="green"))

    def process_queries(
        self,
        document_url: str,
        questions: List[Question]
    ) -> List[FinalAnswer]:
        
        start_time = time.time()
        try:
            rprint(Panel(f"Fetching document from URL:", title="[yellow]Step 1: Document Manager[/yellow]", border_style="yellow"))
            print(f'URL: \n{document_url}')
            print(f'Questions:')
            print('{')
            for question in questions:
                print(question.question, end=",\n")
            print('}')

            doc_manager = DocumentManager(document_url)
            file_path = doc_manager.get_filepath()

            rprint(f"[green]Document located at:[/] {file_path}")
            rprint(f"[blue]Elapsed Time: {(time.time()-start_time):.2f}")

            rprint(Panel("Initializing vector store and retriever...", title="[yellow]Step 2: Vector Store Manager[/yellow]", border_style="yellow"))
            store_manager = VectorStoreManager(file_path)
            retriever = store_manager.retriever
            rprint(f"[green]Retriever created successfully. Indexed [bold]{len(store_manager.documents)}[/bold] chunks.[/green]")
            rprint(f"[blue]Elapsed Time: {(time.time()-start_time):.2f}")

            rprint(Panel(f"Invoking RAG workflow for {len(questions)} questions...", title="[yellow]Step 3: RAG Workflow[/yellow]", border_style="yellow"))
            results = self.workflow.invoke(questions, retriever)
            rprint("[bold green]Workflow finished successfully.[/bold green]")
            rprint(f"[blue]Elapsed Time: {(time.time()-start_time):.2f}")
            
            return results

        except Exception as e:
            rprint(Panel(f"[bold red]An error occurred during query processing:[/] {e}", title="[red]FATAL ERROR[/red]"))
            raise