from src.internal.agents.base_agents import Agent, AgentRequest
from src.internal.rag.inference.inferencer import Inferencer
from typing import Dict, List
class RAGExecutorAgent(Agent):
    def __init__(self, inferencer : Inferencer):
        super().__init__(inferencer)
        self.inferencer = inferencer
        self.file_paths = [
            "data/documents/dokumen.pdf",
            "data/documents/iuran-bpjs.pdf",
            "data/documents/perhitungan-pph21.pdf",
            "data/documents/upah-lembur.pdf",
            "data/documents/qna.pdf",
        ]
    async def load_documents(self):
        for file_path in self.file_paths:
            await self.add_doc(file_path)
        
    async def add_doc(self, file_path):
        result = await self.inferencer.retriever.add_document_from_file(file_path)

        if result.success:
                print(f"Successfully processed: {result.document_metadata.file_name}")
                print(f"Chunks created: {result.document_metadata.chunk_count}")
        else:
                print(f"Failed to process: {result.error_message}")

    async def get_result(self, req : AgentRequest):
        if(req.question == ""):
             yield "Mohon berikan pertanyaan anda!"
        else:
            async for item in self.inferencer.infer_stream(enable_search = req.enable_search, 
                                                        enable_retrieval = req.enable_retrieval,
                                                        chat_memory = req.chat_memory, 
                                                        prompt_template = req.prompt_template, 
                                                        query = req.question, 
                                                        enable_reranking=False, 
                                                        k=3):
                    yield item