feat: adicionando mais testes

_utils/gerar_documento_utils/contextual_retriever.py

@@ -53,7 +53,7 @@ class ContextualRetriever:
         self.claude_api_key = os.environ.get("CLAUDE_API_KEY", "")
         self.claude_client = AsyncAnthropic(api_key=self.claude_api_key)
 
-    async def llm_call_uma_lista_de_chunks(
+    async def llm_call_uma_lista_de_20_chunks(
         self,
         lista_com_20_chunks: List[DocumentChunk],
         resumo_auxiliar,
@@ -120,13 +120,13 @@ class ContextualRetriever:
         self.lista_contador.append(0)
         print("contador: ", len(self.lista_contador))
 
-        result = await self.llm_call_uma_lista_de_chunks(
+        result = await self.llm_call_uma_lista_de_20_chunks(
             lista_com_20_chunks, response_auxiliar_summary, axiom_instance
         )
 
         lista_chunks: List[ContextualizedChunk] = []
-        for index, chunk in enumerate(lista_com_20_chunks):
-            try:
+        try:
+            for index, chunk in enumerate(lista_com_20_chunks):
                 lista_chunks.append(
                     ContextualizedChunk(
                         contextual_summary=result[index][2],
@@ -139,12 +139,11 @@ class ContextualRetriever:
                         context=result[index][1],
                     )
                 )
-            except BaseException as e:
-                print(
-                    f"ERRO EM UMA LISTA COM 20 CHUNKS CONTEXTUALS {index + 1}: {result} ------- {e}"
-                )
+        except BaseException as e:
+            axiom_instance.send_axiom(
+                f"ERRO EM UMA LISTA COM 20 CHUNKS CONTEXTUALS --------- lista: {lista_com_20_chunks} ------------ resultado da llm_call_uma_lista_de_20_chunks: {result} ------- erro: {e}"
+            )
 
-        axiom_instance.send_axiom(f"UMA LISTA COM 20 CHUNKS: {lista_chunks}")
         return lista_chunks
 
     async def contextualize_all_chunks(
@@ -161,20 +160,23 @@ class ContextualRetriever:
             )
         )
 
-        def processa_uma_lista_de_20_chunks(lista_com_20_chunks: List[DocumentChunk]):
-            coroutine = self.contextualize_uma_lista_de_20_chunks(
-                lista_com_20_chunks, response_auxiliar_summary, axiom_instance
-            )
-            return tg.create_task(coroutine)
-
         async with asyncio.TaskGroup() as tg:
+
+            def processa_uma_lista_de_20_chunks(
+                lista_com_20_chunks: List[DocumentChunk],
+            ):
+                coroutine = self.contextualize_uma_lista_de_20_chunks(
+                    lista_com_20_chunks, response_auxiliar_summary, axiom_instance
+                )
+                return tg.create_task(coroutine)
+
             tasks = [
                 processa_uma_lista_de_20_chunks(lista_com_20_chunks)
                 for lista_com_20_chunks in lista_de_listas_cada_com_20_chunks
             ]
 
         # contextualized_chunks = [task.result() for task in tasks]
-        contextualized_chunks = []
+        contextualized_chunks: List[ContextualizedChunk] = []
         for task in tasks:
             contextualized_chunks = contextualized_chunks + task.result()
 

_utils/gerar_documento_utils/tests/fixtures/mock_llm_call_uma_lista_de_20_chunks.py

@@ -0,0 +1,255 @@
+mock_llm_call_uma_lista_de_20_chunks_sucesso = [
+    [
+        0,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk apresenta os dados iniciais do processo de Apelação Criminal, incluindo o número, classe, órgão julgador, partes envolvidas (Estado da Bahia, Defensor Dativo, e Adenilton da Conceição Samuel), e os assuntos tratados (Honorários Advocatícios, Crimes de Trânsito). Ele se encaixa na descrição do resumo do processo, que trata da apelação do Estado da Bahia sobre os honorários do defensor dativo em um caso de crimes de trânsito.",
+    ],
+    [
+        76915042,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk lista uma série de intimações, sentenças, certidões e outros documentos relacionados ao processo, abrangendo um período de tempo considerável. Ele demonstra o andamento processual, desde a intimação da sentença até as alegações finais e a interposição da apelação, conforme descrito no resumo do processo.",
+    ],
+    [
+        76827506,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk apresenta o pronunciamento do Desembargador Relator Nartir Dantas Weber, que trata do recurso de apelação interposto pelo Estado da Bahia. Ele destaca que a controvérsia se refere a honorários advocatícios e que a intervenção do Ministério Público pode não ser necessária, alinhando-se com a questão central do processo: a responsabilidade do Estado pelo pagamento dos honorários.",
+    ],
+    [
+        76747604,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de publicação no Diário de Justiça Eletrônico (DJe) referente ao processo, informando a disponibilização de um acórdão, decisão ou despacho. Ele se relaciona com o resumo do processo, que menciona as movimentações processuais e a publicação de decisões.",
+    ],
+    [
+        76648181,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de publicação no DJe, informando sobre a disponibilização de um despacho. Ele menciona o artigo 224, § 2º, do CPC, que trata da contagem de prazos processuais, e se encaixa na descrição do resumo do processo, que acompanha as movimentações processuais.",
+    ],
+    [
+        76647593,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de triagem inicial realizada pela Diretoria de Distribuição do 2º Grau, que verifica a regularidade dos dados cadastrais do processo. Ele se relaciona com o resumo do processo, que descreve o andamento processual e as etapas de análise dos documentos.",
+    ],
+    [
+        76647591,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de triagem inicial, confirmando a conferência dos dados cadastrais do processo e identificando alterações necessárias. Ele se encaixa na descrição do resumo do processo, que acompanha as movimentações processuais e as etapas de análise dos documentos.",
+    ],
+    [
+        76647590,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é um despacho do Juiz de Direito Leandro Ferreira de Moraes, intimando o advogado Ailton Silva Dantas para apresentar as contrarrazões ao recurso de apelação. Ele se encaixa no resumo do processo, que menciona a intimação do apelado para apresentar contrarrazões.",
+    ],
+    [
+        76647588,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão que atesta que o advogado Ailton Silva Dantas, apesar de devidamente intimado, permaneceu inerte, não apresentando as contrarrazões recursais. Ele se relaciona com o resumo do processo, que menciona a inércia do apelado após a intimação.",
+    ],
+    [
+        76647583,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk contém um comprovante de comunicação à Justiça Eleitoral sobre a condenação criminal de Adenilton da Conceição Samuel, mencionando os artigos 303 e 306 da Lei 9.503/97. Ele se encaixa no resumo do processo, que descreve os crimes pelos quais o réu foi condenado.",
+    ],
+    [
+        76647583,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é a continuação da certidão, que informa o envio do ofício ao CEDEP (Centro de Documentação e Estatística Policial) sobre a condenação criminal de Adenilton da Conceição Samuel. Ele se relaciona com o resumo do processo, que descreve os crimes pelos quais o réu foi condenado e as comunicações oficiais.",
+    ],
+    [
+        76647582,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é a continuação do ofício de comunicação da condenação criminal de Adenilton da Conceição Samuel ao CEDEP, detalhando a pena imposta. Ele se relaciona com o resumo do processo, que descreve os crimes pelos quais o réu foi condenado e as comunicações oficiais.",
+    ],
+    [
+        76647579,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de trânsito em julgado da sentença condenatória para o Ministério Público em 21/08/2023 e para a defesa em 19/12/2023. Ele se encaixa no resumo do processo, que menciona o trânsito em julgado da sentença.",
+    ],
+    [
+        76647578,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de intimação de Adenilton da Conceição Samuel sobre a sentença proferida. Ele se relaciona com o resumo do processo, que descreve as etapas processuais, incluindo as intimações.",
+    ],
+    [
+        76647575,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma intimação judicial para o advogado Ailton Silva Dantas apresentar contrarrazões recursais. Ele se encaixa no resumo do processo, que menciona a intimação do apelado para apresentar contrarrazões.",
+    ],
+    [
+        76647575,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma decisão do juiz Leandro Ferreira de Moraes que recebe o recurso de apelação interposto pela Procuradoria Geral do Estado da Bahia, com efeito suspensivo, e determina a intimação do apelado para apresentar contrarrazões. Ele se encaixa no resumo do processo, que descreve as etapas processuais, incluindo a decisão sobre o recurso.",
+    ],
+    [
+        76647572,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão que atesta a tempestividade da interposição do recurso de apelação. Ele se relaciona com o resumo do processo, que descreve a interposição do recurso pelo Estado da Bahia.",
+    ],
+    [
+        76647571,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk apresenta as razões do recurso de apelação interposto pelo Estado da Bahia, argumentando sobre o cabimento do recurso e a ilegitimidade do Estado para recorrer. Ele se encaixa no resumo do processo, que descreve os argumentos do Estado na apelação.",
+    ],
+    [
+        76647571,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk continua as razões do recurso, detalhando a preliminar de inobservância do Tema Repetitivo 984 do STJ, que trata da ausência de obrigatoriedade de observância da tabela da OAB para fixação de honorários de defensor dativo. Ele se relaciona com o resumo do processo, que menciona a alegação do Estado sobre o Tema Repetitivo 984.",
+    ],
+    [
+        76647571,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk continua a discussão sobre o Tema Repetitivo 984 do STJ, apresentando a ementa do julgado e os argumentos do Estado sobre a fixação dos honorários de defensor dativo. Ele se relaciona com o resumo do processo, que menciona a alegação do Estado sobre o Tema Repetitivo 984.",
+    ],
+]
+
+mock_llm_call_uma_lista_de_20_chunks_sucesso_com_menos_itens = [
+    [
+        0,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk apresenta os dados iniciais do processo de Apelação Criminal, incluindo o número, classe, órgão julgador, partes envolvidas (Estado da Bahia, Defensor Dativo, e Adenilton da Conceição Samuel), e os assuntos tratados (Honorários Advocatícios, Crimes de Trânsito). Ele se encaixa na descrição do resumo do processo, que trata da apelação do Estado da Bahia sobre os honorários do defensor dativo em um caso de crimes de trânsito.",
+    ],
+    [
+        76915042,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk lista uma série de intimações, sentenças, certidões e outros documentos relacionados ao processo, abrangendo um período de tempo considerável. Ele demonstra o andamento processual, desde a intimação da sentença até as alegações finais e a interposição da apelação, conforme descrito no resumo do processo.",
+    ],
+    [
+        76827506,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk apresenta o pronunciamento do Desembargador Relator Nartir Dantas Weber, que trata do recurso de apelação interposto pelo Estado da Bahia. Ele destaca que a controvérsia se refere a honorários advocatícios e que a intervenção do Ministério Público pode não ser necessária, alinhando-se com a questão central do processo: a responsabilidade do Estado pelo pagamento dos honorários.",
+    ],
+    [
+        76747604,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de publicação no Diário de Justiça Eletrônico (DJe) referente ao processo, informando a disponibilização de um acórdão, decisão ou despacho. Ele se relaciona com o resumo do processo, que menciona as movimentações processuais e a publicação de decisões.",
+    ],
+    [
+        76648181,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de publicação no DJe, informando sobre a disponibilização de um despacho. Ele menciona o artigo 224, § 2º, do CPC, que trata da contagem de prazos processuais, e se encaixa na descrição do resumo do processo, que acompanha as movimentações processuais.",
+    ],
+    [
+        76647593,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de triagem inicial realizada pela Diretoria de Distribuição do 2º Grau, que verifica a regularidade dos dados cadastrais do processo. Ele se relaciona com o resumo do processo, que descreve o andamento processual e as etapas de análise dos documentos.",
+    ],
+    [
+        76647591,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de triagem inicial, confirmando a conferência dos dados cadastrais do processo e identificando alterações necessárias. Ele se encaixa na descrição do resumo do processo, que acompanha as movimentações processuais e as etapas de análise dos documentos.",
+    ],
+]
+
+mock_llm_call_uma_lista_de_20_chunks_falha = [[""]]
+
+mock_llm_call_uma_lista_de_20_chunks_falha_mais_de_20_chunks = [
+    [
+        0,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk apresenta os dados iniciais do processo de Apelação Criminal, incluindo o número, classe, órgão julgador, partes envolvidas (Estado da Bahia, Defensor Dativo, e Adenilton da Conceição Samuel), e os assuntos tratados (Honorários Advocatícios, Crimes de Trânsito). Ele se encaixa na descrição do resumo do processo, que trata da apelação do Estado da Bahia sobre os honorários do defensor dativo em um caso de crimes de trânsito.",
+    ],
+    [
+        76915042,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk lista uma série de intimações, sentenças, certidões e outros documentos relacionados ao processo, abrangendo um período de tempo considerável. Ele demonstra o andamento processual, desde a intimação da sentença até as alegações finais e a interposição da apelação, conforme descrito no resumo do processo.",
+    ],
+    [
+        76827506,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk apresenta o pronunciamento do Desembargador Relator Nartir Dantas Weber, que trata do recurso de apelação interposto pelo Estado da Bahia. Ele destaca que a controvérsia se refere a honorários advocatícios e que a intervenção do Ministério Público pode não ser necessária, alinhando-se com a questão central do processo: a responsabilidade do Estado pelo pagamento dos honorários.",
+    ],
+    [
+        76747604,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de publicação no Diário de Justiça Eletrônico (DJe) referente ao processo, informando a disponibilização de um acórdão, decisão ou despacho. Ele se relaciona com o resumo do processo, que menciona as movimentações processuais e a publicação de decisões.",
+    ],
+    [
+        76648181,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de publicação no DJe, informando sobre a disponibilização de um despacho. Ele menciona o artigo 224, § 2º, do CPC, que trata da contagem de prazos processuais, e se encaixa na descrição do resumo do processo, que acompanha as movimentações processuais.",
+    ],
+    [
+        76647593,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de triagem inicial realizada pela Diretoria de Distribuição do 2º Grau, que verifica a regularidade dos dados cadastrais do processo. Ele se relaciona com o resumo do processo, que descreve o andamento processual e as etapas de análise dos documentos.",
+    ],
+    [
+        76647591,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de triagem inicial, confirmando a conferência dos dados cadastrais do processo e identificando alterações necessárias. Ele se encaixa na descrição do resumo do processo, que acompanha as movimentações processuais e as etapas de análise dos documentos.",
+    ],
+    [
+        76647590,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é um despacho do Juiz de Direito Leandro Ferreira de Moraes, intimando o advogado Ailton Silva Dantas para apresentar as contrarrazões ao recurso de apelação. Ele se encaixa no resumo do processo, que menciona a intimação do apelado para apresentar contrarrazões.",
+    ],
+    [
+        76647588,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão que atesta que o advogado Ailton Silva Dantas, apesar de devidamente intimado, permaneceu inerte, não apresentando as contrarrazões recursais. Ele se relaciona com o resumo do processo, que menciona a inércia do apelado após a intimação.",
+    ],
+    [
+        76647583,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk contém um comprovante de comunicação à Justiça Eleitoral sobre a condenação criminal de Adenilton da Conceição Samuel, mencionando os artigos 303 e 306 da Lei 9.503/97. Ele se encaixa no resumo do processo, que descreve os crimes pelos quais o réu foi condenado.",
+    ],
+    [
+        76647583,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é a continuação da certidão, que informa o envio do ofício ao CEDEP (Centro de Documentação e Estatística Policial) sobre a condenação criminal de Adenilton da Conceição Samuel. Ele se relaciona com o resumo do processo, que descreve os crimes pelos quais o réu foi condenado e as comunicações oficiais.",
+    ],
+    [
+        76647582,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é a continuação do ofício de comunicação da condenação criminal de Adenilton da Conceição Samuel ao CEDEP, detalhando a pena imposta. Ele se relaciona com o resumo do processo, que descreve os crimes pelos quais o réu foi condenado e as comunicações oficiais.",
+    ],
+    [
+        76647579,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de trânsito em julgado da sentença condenatória para o Ministério Público em 21/08/2023 e para a defesa em 19/12/2023. Ele se encaixa no resumo do processo, que menciona o trânsito em julgado da sentença.",
+    ],
+    [
+        76647578,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão de intimação de Adenilton da Conceição Samuel sobre a sentença proferida. Ele se relaciona com o resumo do processo, que descreve as etapas processuais, incluindo as intimações.",
+    ],
+    [
+        76647575,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma intimação judicial para o advogado Ailton Silva Dantas apresentar contrarrazões recursais. Ele se encaixa no resumo do processo, que menciona a intimação do apelado para apresentar contrarrazões.",
+    ],
+    [
+        76647575,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma decisão do juiz Leandro Ferreira de Moraes que recebe o recurso de apelação interposto pela Procuradoria Geral do Estado da Bahia, com efeito suspensivo, e determina a intimação do apelado para apresentar contrarrazões. Ele se encaixa no resumo do processo, que descreve as etapas processuais, incluindo a decisão sobre o recurso.",
+    ],
+    [
+        76647572,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk é uma certidão que atesta a tempestividade da interposição do recurso de apelação. Ele se relaciona com o resumo do processo, que descreve a interposição do recurso pelo Estado da Bahia.",
+    ],
+    [
+        76647571,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk apresenta as razões do recurso de apelação interposto pelo Estado da Bahia, argumentando sobre o cabimento do recurso e a ilegitimidade do Estado para recorrer. Ele se encaixa no resumo do processo, que descreve os argumentos do Estado na apelação.",
+    ],
+    [
+        76647571,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk continua as razões do recurso, detalhando a preliminar de inobservância do Tema Repetitivo 984 do STJ, que trata da ausência de obrigatoriedade de observância da tabela da OAB para fixação de honorários de defensor dativo. Ele se relaciona com o resumo do processo, que menciona a alegação do Estado sobre o Tema Repetitivo 984.",
+    ],
+    [
+        76647571,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk continua a discussão sobre o Tema Repetitivo 984 do STJ, apresentando a ementa do julgado e os argumentos do Estado sobre a fixação dos honorários de defensor dativo. Ele se relaciona com o resumo do processo, que menciona a alegação do Estado sobre o Tema Repetitivo 984.",
+    ],
+    [
+        76647571,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk continua a discussão sobre o Tema Repetitivo 984 do STJ, apresentando a ementa do julgado e os argumentos do Estado sobre a fixação dos honorários de defensor dativo. Ele se relaciona com o resumo do processo, que menciona a alegação do Estado sobre o Tema Repetitivo 984.",
+    ],
+    [
+        76647571,
+        "APELAÇÃO CRIMINAL",
+        "Este chunk continua a discussão sobre o Tema Repetitivo 984 do STJ, apresentando a ementa do julgado e os argumentos do Estado sobre a fixação dos honorários de defensor dativo. Ele se relaciona com o resumo do processo, que menciona a alegação do Estado sobre o Tema Repetitivo 984.",
+    ],
+]

_utils/gerar_documento_utils/tests/test_contextual_retriever_inside.py

@@ -3,6 +3,7 @@ from _utils.gerar_documento_utils.contextual_retriever import (
     ContextualRetriever,
     ContextualRetrieverUtils,
 )
+from _utils.gerar_documento_utils.llm_calls import agemini_answer
 from _utils.models.gerar_documento import RetrievalConfig
 from _utils.models.gerar_documento import (
     ContextualizedChunk,
@@ -20,6 +21,7 @@ from .fixtures.resposta_llm_com_20_chunks import (
 )
 from .fixtures.lista_com_20_chunks import lista_com_20_chunks_fixture
 from .fixtures.resumo_auxiliar import resumo_auxiliar_fixture
+from .fixtures import mock_llm_call_uma_lista_de_20_chunks
 
 
 config = RetrievalConfig(
@@ -94,7 +96,7 @@ class TestContextualRetrieverUtils:
     #     mock_agpt_answer,
     # )
 
-    # resultado = await self.context_retrieval.llm_call_uma_lista_de_chunks(
+    # 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)
@@ -186,18 +188,25 @@ class TestContextuRetriever:
     axiom.generate_new_uuid()
 
     @pytest.mark.asyncio
-    async def test_llm_call_uma_lista_de_chunks_success(
+    async def test_llm_call_uma_lista_de_20_chunks_success(
         self,
     ):
-        resultado = await self.context_retrieval.llm_call_uma_lista_de_chunks(
+        resultado = await self.context_retrieval.llm_call_uma_lista_de_20_chunks(
             lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
         )
 
         print("resultado: ", resultado)
         assert resultado
+        assert len(resultado) == 20
+        for r in resultado:
+            assert isinstance(r[0], int)
+            assert isinstance(r[1], str)
+            assert isinstance(r[2], str)
+            assert len(r[1]) > 0
+            assert len(r[2]) > 0
 
     @pytest.mark.asyncio
-    async def test_llm_call_uma_lista_de_chunks_falha_01(self, monkeypatch):
+    async def test_llm_call_uma_lista_de_20_chunks_falha_01(self, monkeypatch):
         monkeypatch.setattr(
             self.context_retrieval.contextual_retriever_utils,
             "get_all_document_ids_and_contents",
@@ -212,9 +221,181 @@ class TestContextuRetriever:
             mock_agemini_answer,
         )
 
-        resultado = await self.context_retrieval.llm_call_uma_lista_de_chunks(
+        resultado = await self.context_retrieval.llm_call_uma_lista_de_20_chunks(
             lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
         )
 
         print("resultado: ", resultado)
         assert resultado == [[""]]
+
+    @pytest.mark.asyncio
+    async def test_llm_call_uma_lista_de_20_chunks_falha_02(self, monkeypatch):
+        count = 0
+
+        async def mock_agemini_answer(prompt, modelo):
+            nonlocal count
+            if count < 3:
+                count += 1
+                return "texto incorreto"
+            count += 1
+            return await agemini_answer(prompt, modelo)
+
+        monkeypatch.setattr(
+            "_utils.gerar_documento_utils.contextual_retriever.agemini_answer",
+            mock_agemini_answer,
+        )
+
+        resultado = await self.context_retrieval.llm_call_uma_lista_de_20_chunks(
+            lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
+        )
+
+        print("resultado: ", resultado)
+        assert resultado
+        assert len(resultado) == 20
+        for r in resultado:
+            assert isinstance(r[0], int)
+            assert isinstance(r[1], str)
+            assert isinstance(r[2], str)
+            assert len(r[1]) > 0
+            assert len(r[2]) > 0
+
+    @pytest.mark.asyncio
+    async def test_llm_call_uma_lista_de_20_chunks_falha_03(self, monkeypatch):
+        async def mock_agemini_answer(prompt, modelo):
+            return "texto incorreto"
+
+        monkeypatch.setattr(
+            "_utils.gerar_documento_utils.contextual_retriever.agemini_answer",
+            mock_agemini_answer,
+        )
+
+        resultado = await self.context_retrieval.llm_call_uma_lista_de_20_chunks(
+            lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
+        )
+
+        print("resultado: ", resultado)
+        assert resultado == [[""]]
+
+    @pytest.mark.asyncio
+    async def test_llm_call_uma_lista_de_20_chunks_falha_04(self, monkeypatch):
+        async def mock_agemini_answer(prompt, modelo):
+            raise Exception("Erro simulado")
+
+        monkeypatch.setattr(
+            "_utils.gerar_documento_utils.contextual_retriever.agemini_answer",
+            mock_agemini_answer,
+        )
+
+        resultado = await self.context_retrieval.llm_call_uma_lista_de_20_chunks(
+            lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
+        )
+
+        print("resultado: ", resultado)
+        assert resultado == [[""]]
+
+    @pytest.mark.asyncio
+    async def test_contextualize_uma_lista_de_20_chunks_success(self, monkeypatch):
+        async def mock_llm_call_uma_lista_de_20_chunks_function(p1, p2, p3):
+            return (
+                mock_llm_call_uma_lista_de_20_chunks.mock_llm_call_uma_lista_de_20_chunks_sucesso
+            )
+
+        monkeypatch.setattr(
+            self.context_retrieval,
+            "llm_call_uma_lista_de_20_chunks",
+            mock_llm_call_uma_lista_de_20_chunks_function,
+        )
+        resultado = await self.context_retrieval.contextualize_uma_lista_de_20_chunks(
+            lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
+        )
+
+        print("resultado: ", resultado)
+        assert resultado
+        assert len(resultado) == 20
+        for r in resultado:
+            assert isinstance(r, ContextualizedChunk)
+            assert len(r.contextual_summary) > 0
+            assert len(r.chunk_id) > 0
+            assert len(r.context) > 0
+            assert len(r.content) > 0
+
+    @pytest.mark.asyncio
+    async def test_contextualize_uma_lista_de_20_chunks_success_com_menos_chunks(
+        self, monkeypatch
+    ):
+        mock = (
+            mock_llm_call_uma_lista_de_20_chunks.mock_llm_call_uma_lista_de_20_chunks_sucesso_com_menos_itens
+        )
+
+        async def mock_llm_call_uma_lista_de_20_chunks_function(p1, p2, p3):
+            return mock
+
+        monkeypatch.setattr(
+            self.context_retrieval,
+            "llm_call_uma_lista_de_20_chunks",
+            mock_llm_call_uma_lista_de_20_chunks_function,
+        )
+        resultado = await self.context_retrieval.contextualize_uma_lista_de_20_chunks(
+            lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
+        )
+
+        print("resultado: ", resultado)
+        assert resultado
+        assert len(resultado) == len(mock)
+        for r in resultado:
+            assert isinstance(r, ContextualizedChunk)
+            assert len(r.contextual_summary) > 0
+            assert len(r.chunk_id) > 0
+            assert len(r.context) > 0
+            assert len(r.content) > 0
+
+    @pytest.mark.asyncio
+    async def test_contextualize_uma_lista_de_20_chunks_success_com_mais_chunks(
+        self, monkeypatch
+    ):
+        mock = (
+            mock_llm_call_uma_lista_de_20_chunks.mock_llm_call_uma_lista_de_20_chunks_falha_mais_de_20_chunks
+        )
+
+        async def mock_llm_call_uma_lista_de_20_chunks_function(p1, p2, p3):
+            return mock
+
+        monkeypatch.setattr(
+            self.context_retrieval,
+            "llm_call_uma_lista_de_20_chunks",
+            mock_llm_call_uma_lista_de_20_chunks_function,
+        )
+        resultado = await self.context_retrieval.contextualize_uma_lista_de_20_chunks(
+            lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
+        )
+
+        print("resultado: ", resultado)
+        assert resultado
+        assert len(resultado) == 20
+        for r in resultado:
+            assert isinstance(r, ContextualizedChunk)
+            assert len(r.contextual_summary) > 0
+            assert len(r.chunk_id) > 0
+            assert len(r.context) > 0
+            assert len(r.content) > 0
+
+    @pytest.mark.asyncio
+    async def test_contextualize_uma_lista_de_20_chunks_falha(self, monkeypatch):
+        mock = (
+            mock_llm_call_uma_lista_de_20_chunks.mock_llm_call_uma_lista_de_20_chunks_falha
+        )
+
+        async def mock_llm_call_uma_lista_de_20_chunks_function(p1, p2, p3):
+            return mock
+
+        monkeypatch.setattr(
+            self.context_retrieval,
+            "llm_call_uma_lista_de_20_chunks",
+            mock_llm_call_uma_lista_de_20_chunks_function,
+        )
+        resultado = await self.context_retrieval.contextualize_uma_lista_de_20_chunks(
+            lista_com_20_chunks_fixture, resumo_auxiliar_fixture, self.axiom
+        )
+
+        print("resultado: ", resultado)
+        assert resultado == []

tests/gerar_relatorio_modelo_usuario/test_contextual_retriever.py

@@ -80,8 +80,8 @@ class TestContextualRetriever:
         print('\n\nresultado[0]', resultado[0])
 
     @pytest.mark.asyncio
-    async def test_llm_call_uma_lista_de_chunks(self):
-        resultado = await self.context_retrieval.llm_call_uma_lista_de_chunks(
+    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)
@@ -171,7 +171,7 @@ class TestContextualRetriever:
             mock_agpt_answer,
         )
 
-        resultado = await self.context_retrieval.llm_call_uma_lista_de_chunks(
+        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)