tests/gerar_relatorio_modelo_usuario/test_contextual_retriever.py

import pytest
from _utils.gerar_documento_utils.contextual_retriever import (
    ContextualRetriever,
)
from _utils.models.gerar_documento import RetrievalConfig
from _utils.models.gerar_documento import (
    ContextualizedChunk,
    DocumentChunk,
    RetrievalConfig,
)
from tests.fixtures.processo_civil_01 import response_auxiliary_summary
from tests.gerar_relatorio_modelo_usuario.fixture_chunks import (
    ChunkResponseFromLLM,
    cenario_01_lista_com_20_chunks,
    cenario_01_resumo_auxiliar,
    resposta_errada_ids,
)
from tests.gerar_relatorio_modelo_usuario.fixture_lista_respostas_contextuals import (
    lista_respostas_contextuals,
    lista_respostas,
    lista_respostas_contextuals_cenario_02,
    lista_20_chunks,
    lista_20_chunks_cenario_02,
    Listas20Chunks
)

from unittest.mock import Mock, AsyncMock, patch


config = RetrievalConfig(
    num_chunks=20,
    embedding_weight=0.5,
    bm25_weight=0.5,
    context_window=3,
    chunk_overlap=800,
    chunk_size=1000,
)

lista_chunks_errados = [
    DocumentChunk(
        content="AO JUÍZO DA VARA DO SISTEMA DE JUIZADO ESPECIAL CÍVEL DA \nCOMARCA DE MARAÚ/ITACARÉ-BA. \n \n \n \n \nProcesso nº: 8000630-48.2023.8.05.0114 \n \n \nGABRIELA SALVATORE TEBET PIEDADE PETROCELLO , já \nqualificada nos autos do processo em epígrafe, vem mui respeitosamente, através de seus \nprocuradores in fine assinado, instrumento procuratório anexo aos autos, apresentar  \nIMPUGNAÇÃO À CONTESTAÇÃO , pelos motivos de fato e de direito expostos a \nseguir: \n \nI. DAS PRELIMINARES EM SEDE DE CONTESTAÇÃO \nA. DA IMPUGNAÇÃO DA PRELIMINAR INCOMPETÊNCIA \nABSOLUTA DO JUÍZO \n \nA empresa requerida, aduz que o Juízo da comarca de Itacaré, não é competente \npara julgar as demandas de residentes na com arca de Maraú (Barra Grande), sob o \nargumento de que a demandante não anexou comprovante de residência, no entanto, por \num infortúnio , a autora n ão colacionou o comprovante de residência, porém, como é \ndocumento obrigatório e o sistema não se atentou a inf ormar a ausência, junta-se agora, \npara sanar quaisquer dúvidas. \n \n \nNum. 391025923 - Pág. 2\nAssinado eletronicamente por: ROGERIO DANTAS OLIVEIRA - 30/05/2023 10:58:20\nhttps://pje.tjba.jus.br/pje/Processo/ConsultaDocumento/listView.seam?x=23053010581937600000380220036\nNúmero do documento: 23053010581937600000380220036\nEste documento foi gerado pelo usuário 048.***.***-01 em 01/02/2025 18:32:20\n \n \nNão obstante, a autora reside em Barra Grande, na pen ínsula de Maraú, o nde a \ncomarca de Itacar é atua como seção judiciária, ou seja, as ações d essa região , são \nNum. 391025923 - Pág. 3\nAssinado eletronicamente por: ROGERIO DANTAS OLIVEIRA - 30/05/2023 10:58:20\nhttps://pje.tjba.jus.br/pje/Processo/ConsultaDocumento/listView.seam?x=23053010581937600000380220036\nNúmero do documento: 23053010581937600000380220036\nEste documento foi gerado pelo usuário 048.***.***-01 em 01/02/2025 18:32:20\n \nenglobadas à comarca de Itacaré, tendo em vista qu e o Município de Maraú não possui \ncomarca. \nPosto isso, não há o que se falar em incompetência do juízo, devendo a preliminar \nser afastada. \nB. IMPUGNAÇÃO DA PRELIMINAR DE CONEXÃO COM OUTRO \nPROCESSO \n \nInsta dizer, que a empresa ré, informa que há conexão entre o processo da autora \ncom o de seu esposo, com mesma causa de pedir e pedido, o que não é verdade. \nO caso da autora, foi que ela teve prejuízo moral,  onde teve por obr igação da \nempresa ré, a modificação de seu voo, para chegar depois do almoço de páscoa o que a \nfez perder a páscoa com sua família, com uma criança, diante da comprovada falha no \nserviço prestado, causando um transtorno absurdo. Já no caso do marido da autora, os \nfatos são diferentes, embora ocasionados pela mesma falha de serviço da ré. \nO marido da autora teve que pedir o cancelamento, uma vez que não podia viajar \nna data, e teve prejuízo, tanto moral como material, e foi obrigado a pagar um valor e não \nteve o reembolso, portanto, os fatos não são os mesmos, as partes não são as mesmas e os \nprejuízos causados são diversos. Devendo, portanto, a presente preliminar ser afastada. \n \nII. IMPUGNAÇÃO DAS TELAS E DOCUMENTOS \nCOLACIONADOS AOS AUTOS \n \nA tela colacionada no bojo da peça em que mostra os fatos narrados na exordial, \nsequer foram co nfrontados pelo requerido, uma vez que apenas trazem alegações sem \nprovas do que realmente aconteceu, usando do argu mento de que a demandante foi \ninformada com antecedência.  Num. 391025923 - Pág. 4\nAssinado eletronicamente por: ROGERIO DANTAS OLIVEIRA - 30/05/2023 10:58:20",
        page_number=89,
        chunk_id="fee5636c-92fb-4614-889f-ba981b62cbd0",
        start_char=138901,
        end_char=142372,
        id_do_processo=0,
        contextual_summary="",
    ),
    DocumentChunk(
        content="O marido da autora teve que pedir o cancelamento, uma vez que não podia viajar \nna data, e teve prejuízo, tanto moral como material, e foi obrigado a pagar um valor e não \nteve o reembolso, portanto, os fatos não são os mesmos, as partes não são as mesmas e os \nprejuízos causados são diversos. Devendo, portanto, a presente preliminar ser afastada. \n \nII. IMPUGNAÇÃO DAS TELAS E DOCUMENTOS \nCOLACIONADOS AOS AUTOS \n \nA tela colacionada no bojo da peça em que mostra os fatos narrados na exordial, \nsequer foram co nfrontados pelo requerido, uma vez que apenas trazem alegações sem \nprovas do que realmente aconteceu, usando do argu mento de que a demandante foi \ninformada com antecedência.  Num. 391025923 - Pág. 4\nAssinado eletronicamente por: ROGERIO DANTAS OLIVEIRA - 30/05/2023 10:58:20\nhttps://pje.tjba.jus.br/pje/Processo/ConsultaDocumento/listView.seam?x=23053010581937600000380220036\nNúmero do documento: 23053010581937600000380220036\nEste documento foi gerado pelo usuário 048.***.***-01 em 01/02/2025 18:32:20\n \nImportante mencionar que, os documentos colacionados aos autos não possuem o \ncondão de confrontar as alegações autorais, uma vez que não impugnam especificamente \nos direitos da Autora.  \nOs documentos acostados, são apenas documentos de representação, não tendo \nqualquer embasamento para trazer prova do contrário. \nDesse modo, resta evidente que o requerido não se desincumbiu do ônus de provar \nfato impeditivo, modificativo e extintivo. Sendo assim, a única conclusão lógica é a total \nprocedência dos pedidos autorais. Por fim, impugna todos os documentos apresentados \npelo réu. \n \nIII. DOS PEDIDOS \n \nAnte exposto, diante dos argumentos trazidos na peça contestatória se revelam \ninsuficientes e ineficazes para rechaçar os pedidos formulados no petitório inaugural, \nportanto, impugnado a defesa requer:  \na) A manutenção de todos os argumentos explanados na inicial, sendo julgada \ntotalmente procedente a ação;  \n \nNestes termos,  \nPede deferimento. \n \nIlhéus-Ba, 30 de maio de 2023. \n \n \nROGÉRIO DANTAS OLIVEIRA          MARIA CLARA NASCIMENTO BARROS \n               OAB/BA 64.702                                                 OAB/BA 63.039 \n Num. 391055932 - Pág. 1\nAssinado eletronicamente por: GABRIELLE CAROLINA LOPES PEREIRA - 30/05/2023 12:50:30\nhttps://pje.tjba.jus.br/pje/Processo/ConsultaDocumento/listView.seam?x=23053012503024400000380248733\nNúmero do documento: 23053012503024400000380248733\nEste documento foi gerado pelo usuário 048.***.***-01 em 01/02/2025 18:32:21\n\xa0\n\xa0\nPODER JUDICIÁRIO DO ESTADO DA BAHIA  \nComarca de Itacaré \n\xa0\nEnd: Rua Joaquim Vieira, sn, Itacaré-Bahia – Fone: (73) 3251-2158\n\xa0\n\xa0\nTERMO DE AUDIÊNCIA\xa0\nProcesso nº :\xa0 \xa08000630-48.2023.8.05.0114\nClasse - Assunto :\xa0[Transporte Aéreo, Cancelamento de vôo]\xa0\nAUTOR: GABRIELA SALVATORE TEBET PIEDADE  \n\xa0\nREU: GOL LINHAS AEREAS INTELIGENTES S.A.  \n\xa0\n \nNesta cidade Itacaré, Estado da Bahia, às 12:40, neste fórum, na sala de audiências,\xa0onde\xa0se encontrava Gabrielle Carolina Lopes \nPereira, conciliadora desta comarca, foi aberta a\xa0audiência de conciliação TELEPRESENCIAL no aplicativo LIFESIZE, sala \nAudiências\xa0Preliminares\xa0(9835079).\xa0Estavam presentes a autora da ação\xa0GABRIELA SALVATORE TEBET PIEDADE, \nacompanhada do seu advogado\xa0\xa0ROGERIO DANTAS OLIVEIRA - OAB BA64702 \xa0, bem como o\xa0réu GOL LINHAS AEREAS \nINTELIGENTES S.A. , representado pela preposta\xa0Juliana Scofano de Araujo Gonçalves CPF 096.147.127-16\xa0, desacompanhada \nde advogado.",
        page_number=89,
        chunk_id="bfb957f3-f98f-4680-95c3-e0232592c04a",
        start_char=142373,
        end_char=145831,
        id_do_processo=0,
        contextual_summary="",
    ),
    DocumentChunk(
        content="Processo nº :\xa0 \xa08000630-48.2023.8.05.0114\nClasse - Assunto :\xa0[Transporte Aéreo, Cancelamento de vôo]\xa0\nAUTOR: GABRIELA SALVATORE TEBET PIEDADE  \n\xa0\nREU: GOL LINHAS AEREAS INTELIGENTES S.A.  \n\xa0\n \nNesta cidade Itacaré, Estado da Bahia, às 12:40, neste fórum, na sala de audiências,\xa0onde\xa0se encontrava Gabrielle Carolina Lopes \nPereira, conciliadora desta comarca, foi aberta a\xa0audiência de conciliação TELEPRESENCIAL no aplicativo LIFESIZE, sala \nAudiências\xa0Preliminares\xa0(9835079).\xa0Estavam presentes a autora da ação\xa0GABRIELA SALVATORE TEBET PIEDADE, \nacompanhada do seu advogado\xa0\xa0ROGERIO DANTAS OLIVEIRA - OAB BA64702 \xa0, bem como o\xa0réu GOL LINHAS AEREAS \nINTELIGENTES S.A. , representado pela preposta\xa0Juliana Scofano de Araujo Gonçalves CPF 096.147.127-16\xa0, desacompanhada \nde advogado.\xa0\n\xa0\nAberta a audiência com as formalidades de estilo e iniciada a gravação, foi infrutífera a tentativa de acordo.\n \nManifestou-se a parte requerida nos seguintes termos:\nReitera a defesa sob ID\xa0390345882. Requer habilitação exclusiva do Dr. GUSTAVO ANTONIO FERES PAIXAO \nOAB BA 55666. Sem mais. Pede deferimento.\n \nManifestou-se a parte requerente nos seguintes termos:\nReitera a Impugnação à Contestação já juntada aos autos sob ID 391025923 e requer a total procedência da ação, nos \ntermos da Inicial. Pede deferimento.\xa0\n\xa0\nAo final voltam-se os\xa0autos conclusos para MM Juíza. Nada mais havendo encerra-se este termo.\nLink para gravação da audiência:\xa0 https://playback.lifesize.com/#/publicvideo/5555e592-32f9-4f1a-bea5-",
        page_number=89,
        chunk_id="f4c1c853-8d2f-4aae-9d04-d725066b6bf8",
        start_char=145832,
        end_char=147339,
        id_do_processo=0,
        contextual_summary="",
    ),
]


class TestContextualRetriever:
    context_retrieval = ContextualRetriever(config, "claude-3-haiku-20240307")

    @pytest.mark.asyncio
    async def test_contextualize_uma_lista_de_20_chunks(self):
        cenarios = Listas20Chunks()
        resultado = await self.context_retrieval.contextualize_uma_lista_de_20_chunks(
            cenarios.cenario_03_lista, cenarios.cenario_03_resumo_contextual
        )

        print('\n\nresultado[0]', resultado[0])

    @pytest.mark.asyncio
    async def test_llm_call_uma_lista_de_20_chunks(self):
        resultado = await self.context_retrieval.llm_call_uma_lista_de_20_chunks(
            cenario_01_lista_com_20_chunks, cenario_01_resumo_auxiliar
        )
        print("resultado: ", resultado)
        assert len(resultado) == len(cenario_01_lista_com_20_chunks)
        assert all(len(item) == 3 for item in resultado)

    @pytest.mark.asyncio
    async def test_validate_many_chunks_in_one_request(self):
        fixture = ChunkResponseFromLLM()
        lista_fixtures = [
            [fixture.resposta_cenario_01, fixture.lista_ids_cenario_01],
            [fixture.resposta_cenario_02, fixture.lista_ids_cenario_02],
        ]

        for cenario in lista_fixtures:
            resultado = validate_many_chunks_in_one_request(
                cenario[0],
                cenario[1],
            )

            assert resultado and len(resultado) == len(cenario[1])
            assert all( isinstance(item[0], int) and item[0] != 0 for item in resultado)
            assert all( isinstance(item[1], str) and item[1] != "" for item in resultado)
            assert all( isinstance(item[2], str) and item[2] != "" for item in resultado)

            print("\n\nresultado", resultado)

    @pytest.mark.asyncio
    async def test_get_info_from_validated_chunks(self):
        contador = 1
        print(len(lista_respostas_contextuals_cenario_02))
        # lista_respostas + lista_respostas_contextuals_cenario_02
        for response in (lista_respostas_contextuals_cenario_02 + lista_respostas):
            matches = validate_many_chunks_in_one_request(response, resposta_errada_ids)
            try:
                assert matches and len(matches) > 2
                assert all( isinstance(item[0], int) and item[0] != 0 for item in matches)
                assert all( isinstance(item[1], str) and item[1] != "" for item in matches)
                assert all( isinstance(item[2], str) and item[2] != "" for item in matches)
            except:
                print('\n\nresponse', response)
                raise BaseException("deu ruim")

            resultado = self.context_retrieval.get_info_from_validated_chunks(matches)
            print("\n\nresultado", resultado)
            print("\n\ncontador", contador)
            contador += 1
            assert all(len(item) == 3 for item in resultado)
            assert all(isinstance(item[0], int) for item in resultado)
            assert all(isinstance(item[1], str) for item in resultado)
            assert all(isinstance(item[2], str) for item in resultado)

        # COMENTADO ABAIXO PARA RODAR APENAS UMA RESPOSTA
        # print("lista_respostas[143]: ", lista_respostas[143])
        # matches = validate_many_chunks_in_one_request(
        #     lista_respostas[143], resposta_errada_ids
        # )

        # resultado = self.context_retrieval.get_info_from_validated_chunks(matches)
        # # print("\n\nresultado", resultado)
        # print("\n\ncontador", contador)
        # contador += 1
        # assert all(len(item) == 3 for item in resultado)

    @pytest.mark.asyncio
    async def test_valida_retries_automaticos_do_contextual(self, monkeypatch):
        call_count = 0  # Track how many times the function is called

        # Define the custom side effect
        def custom_side_effect(response, document_ids):
            nonlocal call_count
            print(f"MOCK CALL {call_count+1} RECEIVED")
            if call_count < 2:
                call_count += 1
                print("MOCK RETORNADO FALSE")
                return False  # Mocked return value for first two calls
            return validate_many_chunks_in_one_request(response, document_ids) # Call the real function after two mocks

        monkeypatch.setattr(
            "_utils.gerar_documento_utils.contextual_retriever.validate_many_chunks_in_one_request",
            custom_side_effect,
        )

        mock_agpt_answer = AsyncMock(return_value=lista_respostas_contextuals_cenario_02[0])
        monkeypatch.setattr(
            "_utils.gerar_documento_utils.contextual_retriever.agpt_answer",
            mock_agpt_answer,
        )

        resultado = await self.context_retrieval.llm_call_uma_lista_de_20_chunks(
            cenario_01_lista_com_20_chunks, cenario_01_resumo_auxiliar
        )
        assert all(len(item) == 3 for item in resultado)
        assert call_count == 2  # Ensure only the first two calls were mocked