Implementa pipeline de chunking, geração de embeddings e indexação no ChromaDB - #3
Merged
Merged
Conversation
There was a problem hiding this comment.
Pull request overview
Esta PR implementa a pipeline de preparação de documentos para RAG no projeto (chunking → embeddings Gemini → indexação local no ChromaDB), adicionando componentes de app/vectorstore/, um script de carga em lote e testes/documentação de suporte.
Changes:
- Adiciona o
DocumentChunker(RecursiveCharacterTextSplitter) para gerar chunks com metadados. - Implementa geração de embeddings com fallback/rotação de chaves e instanciação do ChromaDB.
- Inclui script
populate.py, endpoint/healthaprimorado, testes e documentação da arquitetura.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 14 comments.
Show a summary per file
| File | Description |
|---|---|
app/vectorstore/chunker.py |
Implementa chunking recursivo e propagação de metadados para Document. |
app/vectorstore/database.py |
Adiciona gerenciador de chaves Gemini e wrapper de embeddings com retry/rotação; fornece helpers para Chroma. |
app/vectorstore/populate.py |
Script CLI para fatiar documentos processados e indexar em lote no ChromaDB. |
app/vectorstore/__init__.py |
Expõe APIs públicas do pacote vectorstore. |
app/main.py |
Atualiza /health para reportar conectividade/contagem de chunks no ChromaDB. |
tests/test_chunker.py |
Testes unitários do chunker (texto normal e vazio). |
tests/test_database.py |
Testes do gerenciador de chaves e da rotação/fallback de embeddings (com mocks). |
tests/test_vectorstore_integration.py |
Testes de integração de inserção e busca por similaridade com Chroma usando embeddings mockados. |
docs/vectorstore.md |
Documenta arquitetura/decisões e instruções de uso/testes. |
docs/backlog.md |
Marca Epic 3 como concluída no backlog. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+130
to
+134
| except Exception as e: | ||
| err_str = str(e) | ||
| # Se for limite de cota/autorização e temos outras chaves, rotacionamos sem gastar retentativas locais | ||
| if ("RESOURCE_EXHAUSTED" in err_str or "429" in err_str or "403" in err_str) and self.key_manager.count > 1: | ||
| break |
Comment on lines
+174
to
+177
| except Exception as e: | ||
| err_str = str(e) | ||
| if ("RESOURCE_EXHAUSTED" in err_str or "429" in err_str or "403" in err_str) and self.key_manager.count > 1: | ||
| break |
Comment on lines
+40
to
+44
| from app.vectorstore.database import get_vectorstore | ||
| db = get_vectorstore() | ||
| count = db._collection.count() | ||
| db_status = f"connected ({count} chunks)" | ||
| except Exception as e: |
Comment on lines
+25
to
+42
| # 2. Limpa coleção anterior para evitar duplicações | ||
| print("🧹 Limpando dados antigos da coleção...") | ||
| try: | ||
| # Chroma do LangChain não tem um método direto 'clear' de alto nível, | ||
| # mas podemos deletar por ids ou usar a API do client interno para deletar | ||
| # a coleção e criá-la novamente. | ||
| # Uma forma limpa usando a API do LangChain Chroma é: | ||
| # obter todos os documentos e deletar por IDs. | ||
| all_docs = vectorstore.get() | ||
| if all_docs and "ids" in all_docs and all_docs["ids"]: | ||
| print(f"🗑️ Removendo {len(all_docs['ids'])} registros existentes...") | ||
| vectorstore.delete(ids=all_docs["ids"]) | ||
| print("✅ Coleção limpa com sucesso.") | ||
| else: | ||
| print("ℹ️ Coleção já estava vazia.") | ||
| except Exception as e: | ||
| print(f"⚠️ Erro ao limpar a coleção (pode ser a primeira execução): {str(e)}") | ||
|
|
| print(f"⏳ Limite de requisições excedido. Aguardando 35 segundos para recuperar cota... (Tentativas restantes: {retries})") | ||
| time.sleep(35) | ||
| else: | ||
| raise e |
| assert manager.current_key == "good_key" | ||
| assert mock_google_embeddings.call_count == 3 | ||
|
|
||
| @patch("app.vectorstore.database.GoogleGenerativeAIEmbeddings") |
Comment on lines
+52
to
+53
| # Cria uma cópia profunda/estendida dos metadados originais | ||
| metadata = doc_metadata.copy() |
Comment on lines
+209
to
+227
| def get_embedding_model() -> FallbackGeminiEmbeddings: | ||
| """ | ||
| Cria e retorna a instância de embeddings configurada com suporte a fallback. | ||
| """ | ||
| key_manager = get_key_manager() | ||
| return FallbackGeminiEmbeddings(key_manager=key_manager) | ||
|
|
||
| def get_vectorstore() -> Chroma: | ||
| """ | ||
| Retorna a instância do ChromaDB configurada com os embeddings suportados por fallback. | ||
| """ | ||
| chroma_db_path = os.getenv("CHROMA_DB_PATH", "data/db") | ||
| embeddings = get_embedding_model() | ||
|
|
||
| return Chroma( | ||
| collection_name="brasil_copa_2026", | ||
| embedding_function=embeddings, | ||
| persist_directory=chroma_db_path | ||
| ) |
| * **Modelo**: `models/gemini-embedding-001`. Este modelo foi selecionado por possuir excelente custo-benefício e suporte nativo robusto a idiomas multilíngues, em especial o Português do Brasil. | ||
| * **Gerenciador de Chaves (`GeminiAPIKeyManager`)**: Permite o carregamento de múltiplas chaves de API a partir da variável de ambiente `GEMINI_API_KEYS` (lista separada por vírgula) ou fallback automático para a chave tradicional `GEMINI_API_KEY`. | ||
| * **Resiliência e Rotação**: Criamos a classe wrapper `FallbackGeminiEmbeddings` que encapsula o cliente de embeddings e intercepta erros de cota ou limites de taxa (`429 RESOURCE_EXHAUSTED`). Caso uma chave falhe, ela é rotacionada automaticamente para a próxima chave configurada na lista. | ||
| * **Retentativas locais com Backoff**: Para suportar falhas temporárias de rede (ex: `WinError 10060`), implementamos uma lógica de retentativas locais (até 3 tentativas) com tempo de espera exponencial crescente (2s, 4s, 8s) antes de desistir ou rotacionar a chave. |
| Criamos um script utilitário CLI para processar a base de dados de `data/processed/` e indexar tudo no banco vetorial. | ||
| * **Carga Fracionada**: Os 502 chunks são enviados ao banco em lotes de 10 chunks com 4 segundos de intervalo entre envios. Isso garante que a API Free Tier do Gemini nunca exceda o limite de requisições por minuto (RPM) e previne timeouts por tamanho excessivo do payload. | ||
| * **Uso**: `python -m app.vectorstore.populate` | ||
| * **Limpeza Automática**: O script limpa registros duplicados da coleção antes de iniciar uma nova carga, evitando duplicidade de dados caso seja executado mais de uma vez. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Descrição da PR
Esta Pull Request conclui a Epic 3 — Chunking e Embeddings do backlog, realizando o merge da branch correspondente para
develop.Ela implementa toda a pipeline responsável pela preparação dos documentos para o sistema RAG, incluindo a divisão dos textos em chunks, geração de embeddings utilizando o Google Gemini e indexação em uma base vetorial local com ChromaDB. A implementação adota uma arquitetura modular, com tratamento de erros, rotação automática de chaves da API e utilitários para facilitar a operação e manutenção da base vetorial.
Implementação da pipeline
DocumentChunker(chunker.py) para realizar a divisão recursiva dos documentos em chunks utilizando Recursive Character Text Splitting, preservando os metadados de cada documento.FallbackGeminiEmbeddingseGeminiAPIKeyManager(database.py) para geração de embeddings com o Google Gemini, incluindo rotação automática de chaves da API, tentativas de recuperação (retry/backoff) e acesso singleton às instâncias de embeddings e do banco vetorial.populate.py, responsável pela indexação em lote dos documentos processados no ChromaDB, incluindo limpeza opcional da coleção, processamento em batches, tratamento de erros e gerenciamento de limites de requisição da API.Integração e melhorias na API
__init__.pypara expor os componentes de chunking, embeddings e banco vetorial, facilitando sua reutilização em outras partes do projeto./health, que agora valida a conectividade com o ChromaDB e informa, em tempo real, a quantidade de chunks atualmente indexados.Documentação
vectorstore.md, detalhando a arquitetura, decisões de implementação e instruções de utilização da pipeline de chunking, embeddings e banco vetorial.backlog.md, marcando todas as funcionalidades da Epic 3 como concluídas.Epic atendida
Epic 3 — Chunking e Embeddings
Responsável pela preparação e armazenamento dos documentos na base vetorial.
Feature 3.1 — Chunking
Feature 3.2 — Embeddings
Feature 3.3 — Banco Vetorial
✅ Definition of Done