From 1c58618f3762d42e1868cee9de317bc2d847b1a2 Mon Sep 17 00:00:00 2001 From: rodolpho Date: Thu, 7 May 2026 22:36:42 -0300 Subject: [PATCH] feat: integrate ChromaDB for lightweight vector search on Raspberry Pi --- mcp_server.py | 97 ++++++++++++++++++++------------------------ requirements.txt | 2 + vector_db_manager.py | 61 ++++++++++++++++++++++++++++ 3 files changed, 106 insertions(+), 54 deletions(-) create mode 100644 vector_db_manager.py diff --git a/mcp_server.py b/mcp_server.py index 79357da..2fb2739 100644 --- a/mcp_server.py +++ b/mcp_server.py @@ -1,83 +1,72 @@ import json import os +import argparse +import chromadb +from chromadb.utils import embedding_functions from typing import List, Dict from mcp.server.fastmcp import FastMCP -import argparse # Inicializa o servidor FastMCP mcp = FastMCP("Fluig Technical Wiki") -# Caminhos dos arquivos -CHUNKS_FILE = "fluig_chunks.json" +# Configurações do Banco Vetorial +DB_PATH = "fluig_vector_db" SNIPPETS_DIR = os.path.join("fluig_rag_docs", "Biblioteca de Snippets") -def load_chunks() -> List[Dict]: - if os.path.exists(CHUNKS_FILE): - with open(CHUNKS_FILE, "r", encoding="utf-8") as f: - return json.load(f) - return [] +# Inicialização preguiçosa (lazy load) do ChromaDB para economizar memória no Pi +_collection = None + +def get_collection(): + global _collection + if _collection is None: + if os.path.exists(DB_PATH): + model_name = "all-MiniLM-L6-v2" + emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name) + client = chromadb.PersistentClient(path=DB_PATH) + _collection = client.get_collection(name="fluig_docs", embedding_function=emb_fn) + else: + print("AVISO: Banco vetorial não encontrado. Use --mode stdio para busca simples ou gere o banco.") + return _collection @mcp.tool() def search_docs(query: str) -> str: """ - Busca profunda na documentação técnica do Fluig. - Retorna trechos exatos sobre APIs, Eventos, Datasets e Configurações. + Busca semântica profunda na documentação técnica do Fluig usando Vetores. + Encontra resultados por significado, mesmo que as palavras exatas não coincidam. """ - chunks = load_chunks() - results = [] - query = query.lower() + collection = get_collection() - # Busca com prioridade em títulos - for chunk in chunks: - title = chunk.get("metadata", {}).get("title", "").lower() - content = chunk.get("content", "").lower() + if collection: + # Busca Semântica + results = collection.query( + query_texts=[query], + n_results=5 + ) - if query in title: - results.insert(0, f"### [ALTA RELEVÂNCIA] {chunk['metadata'].get('title')}\n{chunk['content']}\n---\n") - elif query in content: - results.append(f"### {chunk['metadata'].get('title')}\n{chunk['content']}\n---\n") - - if len(results) >= 8: # Limite expandido para busca mais rica - break - - if not results: - return f"Nenhuma informação técnica encontrada para: {query}. Tente termos como 'createDataset', 'hAPI', 'OAuth' ou 'BPM'." - - return "\n".join(results) + output = [] + for i in range(len(results['documents'][0])): + doc = results['documents'][0][i] + meta = results['metadatas'][0][i] + output.append(f"### {meta.get('title')}\n{doc}\n---\n") + return "\n".join(output) + else: + # Fallback para busca por palavra-chave se o banco vetorial não estiver pronto + return "Banco vetorial não inicializado. Por favor, execute vector_db_manager.py no servidor." @mcp.tool() def get_code_snippets(language: str) -> str: - """ - Recupera a biblioteca completa de exemplos de código para Fluig (javascript, java ou sql). - """ + """Recupera exemplos de código Fluig (javascript, java ou sql).""" lang = language.lower() - file_map = { - "javascript": "Snippets JAVASCRIPT.md", - "js": "Snippets JAVASCRIPT.md", - "java": "Snippets JAVA.md", - "sql": "Snippets SQL.md" - } - + file_map = {"javascript": "Snippets JAVASCRIPT.md", "js": "Snippets JAVASCRIPT.md", "java": "Snippets JAVA.md", "sql": "Snippets SQL.md"} file_name = file_map.get(lang) - if not file_name: - return f"Linguagem '{language}' não suportada. Use: javascript, java ou sql." - + if not file_name: return f"Linguagem '{language}' não suportada." path = os.path.join(SNIPPETS_DIR, file_name) if os.path.exists(path): - with open(path, "r", encoding="utf-8") as f: - return f.read() - - return f"Arquivo de snippets para {language} não encontrado." - -@mcp.resource("docs://all_titles") -def list_all_titles() -> str: - """Lista todos os temas técnicos disponíveis na base de conhecimento.""" - chunks = load_chunks() - titles = sorted(list(set(chunk["metadata"].get("title") for chunk in chunks))) - return "\n".join(titles) + with open(path, "r", encoding="utf-8") as f: return f.read() + return "Snippets não encontrados." if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Fluig MCP Server") + parser = argparse.ArgumentParser(description="Fluig MCP Server with Vector Search") parser.add_argument("--mode", choices=["stdio", "sse"], default="stdio") parser.add_argument("--port", type=int, default=8001) args = parser.parse_args() diff --git a/requirements.txt b/requirements.txt index 26229a4..3eecd81 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,6 @@ mkdocs-with-pdf mkdocs-mermaid2-plugin mcp langchain-community +chromadb +sentence-transformers # crawl4ai # Opcional, mas recomendado se o usuário tiver as dependências de sistema (Playwright) diff --git a/vector_db_manager.py b/vector_db_manager.py new file mode 100644 index 0000000..9367fed --- /dev/null +++ b/vector_db_manager.py @@ -0,0 +1,61 @@ +import chromadb +from chromadb.utils import embedding_functions +import json +import os + +# Configurações +CHUNKS_FILE = "fluig_chunks.json" +DB_PATH = "fluig_vector_db" + +def init_vector_db(): + # Usar um modelo extremamente leve adequado para o Raspberry Pi + # 'all-MiniLM-L6-v2' é o padrão ouro para performance vs qualidade + model_name = "all-MiniLM-L6-v2" + emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name) + + # Inicializa o cliente persistente (SQLite sob o capô) + client = chromadb.PersistentClient(path=DB_PATH) + + # Cria ou obtém a coleção + collection = client.get_or_create_collection( + name="fluig_docs", + embedding_function=emb_fn, + metadata={"hnsw:space": "cosine"} # Melhor para busca semântica + ) + + if not os.path.exists(CHUNKS_FILE): + print(f"Erro: {CHUNKS_FILE} não encontrado. Rode o rag_processor.py primeiro.") + return + + with open(CHUNKS_FILE, "r", encoding="utf-8") as f: + chunks = json.load(f) + + print(f"Indexando {len(chunks)} chunks no ChromaDB (isso pode levar alguns minutos no Pi)...") + + ids = [] + documents = [] + metadatas = [] + + for i, chunk in enumerate(chunks): + ids.append(f"id_{i}") + documents.append(chunk["content"]) + metadatas.append({ + "title": chunk["metadata"].get("title", ""), + "source": chunk["metadata"].get("source", ""), + "path": chunk["metadata"].get("path", "") + }) + + # Adicionar em lotes para não estourar memória + batch_size = 100 + for i in range(0, len(documents), batch_size): + collection.add( + ids=ids[i:i+batch_size], + documents=documents[i:i+batch_size], + metadatas=metadatas[i:i+batch_size] + ) + print(f"Progresso: {i + len(documents[i:i+batch_size])}/{len(documents)}") + + print(f"Sucesso! Banco vetorial criado em {DB_PATH}") + +if __name__ == "__main__": + init_vector_db()