feat: integrate ChromaDB for lightweight vector search on Raspberry Pi
This commit is contained in:
rodolpho
+43
-54
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user