feat: Integrate Neo4j for enhanced relationship mapping in memory system

- Added Neo4j service to docker-compose for relationship graph management.
- Updated requirements.txt to include Neo4j dependency.
- Enhanced MemoryWrapper to support hybrid search using vector and graph data.
- Introduced optional parameters for relationship extraction in memory save and search tools.
- Updated README.md to reflect new hybrid memory capabilities and Neo4j configuration.
This commit is contained in:
João Pedro Toledo Goncalves 2026-01-08 20:44:28 -03:00
parent b75f73f46a
commit 3a967f4fe6
5 changed files with 230 additions and 22 deletions

View File

@ -10,7 +10,8 @@ Uma aplicação modular **CrewAI** com memória compartilhada, 26 agentes de IA
## 🚀 Funcionalidades ## 🚀 Funcionalidades
- **26 Agentes de IA Especializados** - De infraestrutura (Arthur Mendes, Gus Fring) a vendas (Ari Gold, Don Draper) até gestão de crises (Olivia Pope, Saul Goodman) - **26 Agentes de IA Especializados** - De infraestrutura (Arthur Mendes, Gus Fring) a vendas (Ari Gold, Don Draper) até gestão de crises (Olivia Pope, Saul Goodman)
- **Memória Compartilhada** - Agentes compartilham conhecimento via Mem0 + banco de dados vetorial Qdrant - **Memória Híbrida Inteligente** - Combinação de banco vetorial (Qdrant) + grafo de relacionamentos (Neo4j) para busca semântica e contextualizada
- **Agentes Escolhem Estratégia de Busca** - Podem usar busca rápida (vetorial) ou busca contextualizada (vetor + relacionamentos)
- **Roteamento Inteligente** - Classificação automática de solicitações para a equipe apropriada - **Roteamento Inteligente** - Classificação automática de solicitações para a equipe apropriada
- **Suporte Multi-Provedor LLM** - Funciona com Gemini, OpenAI, Anthropic, ou Ollama local - **Suporte Multi-Provedor LLM** - Funciona com Gemini, OpenAI, Anthropic, ou Ollama local
- **Interface Web** - Interface de chat moderna powered by Chainlit - **Interface Web** - Interface de chat moderna powered by Chainlit
@ -32,7 +33,9 @@ minions-da-itguys/
│ ├── knowledge/ │ ├── knowledge/
│ │ └── standards/ # Base de conhecimento corporativo │ │ └── standards/ # Base de conhecimento corporativo
│ ├── memory/ │ ├── memory/
│ │ └── wrapper.py # Integração Mem0 com rate limiting │ │ ├── wrapper.py # Integração Mem0 + busca híbrida
│ │ ├── graph_wrapper.py # Gerenciamento de relacionamentos Neo4j
│ │ └── entity_search_tool.py # Busca por entidades específicas
│ └── tools/ # Ferramentas customizadas (Zabbix, Evolution, etc.) │ └── tools/ # Ferramentas customizadas (Zabbix, Evolution, etc.)
├── docker-compose.yml # Orquestração de containers ├── docker-compose.yml # Orquestração de containers
├── Dockerfile # Container da aplicação ├── Dockerfile # Container da aplicação
@ -89,6 +92,11 @@ GEMINI_API_KEY=sua-chave-api
# Memória: qdrant (local) ou mem0 (nuvem) # Memória: qdrant (local) ou mem0 (nuvem)
MEMORY_PROVIDER=qdrant MEMORY_PROVIDER=qdrant
MEMORY_EMBEDDING_PROVIDER=local MEMORY_EMBEDDING_PROVIDER=local
# Neo4j Graph Database (para relacionamentos)
NEO4J_URI=bolt://localhost:7687
NEO4J_USER=neo4j
NEO4J_PASSWORD=antigravity2024
``` ```
## 🤖 Crews Disponíveis ## 🤖 Crews Disponíveis
@ -123,6 +131,7 @@ Usuário: "Analise a segurança da nossa página de login"
- `chainlit` - Interface Web - `chainlit` - Interface Web
- `mem0ai` - Memória compartilhada - `mem0ai` - Memória compartilhada
- `qdrant-client` - Banco de dados vetorial - `qdrant-client` - Banco de dados vetorial
- `neo4j` - Banco de dados de grafos para relacionamentos
- `litellm` - Suporte multi-provedor LLM - `litellm` - Suporte multi-provedor LLM
- `sentence-transformers` - Embeddings locais - `sentence-transformers` - Embeddings locais

View File

@ -1,5 +1,3 @@
version: '3.8'
services: services:
app: app:
build: . build: .
@ -13,7 +11,10 @@ services:
env_file: env_file:
- .env - .env
depends_on: depends_on:
- qdrant # Only strictly needed if using local Qdrant qdrant:
condition: service_started
neo4j:
condition: service_healthy
networks: networks:
- antigravity_net - antigravity_net
@ -28,6 +29,30 @@ services:
networks: networks:
- antigravity_net - antigravity_net
# Graph Database for Relationship Mapping
neo4j:
image: neo4j:5.15
container_name: antigravity_neo4j
ports:
- "7474:7474" # Browser
- "7687:7687" # Bolt
environment:
- NEO4J_AUTH=neo4j/${NEO4J_PASSWORD:-antigravity2024}
- NEO4J_PLUGINS=["apoc"]
- NEO4J_dbms_security_procedures_unrestricted=apoc.*
- NEO4J_dbms_security_procedures_allowlist=apoc.*
volumes:
- neo4j_data:/data
- neo4j_logs:/logs
networks:
- antigravity_net
healthcheck:
test: ["CMD-SHELL", "wget --no-verbose --tries=1 --spider http://localhost:7474 || exit 1"]
interval: 10s
timeout: 5s
retries: 5
start_period: 30s
# Telegram Listener Service (Runs the bot in background) # Telegram Listener Service (Runs the bot in background)
telegram_listener: telegram_listener:
build: . build: .
@ -44,6 +69,8 @@ services:
volumes: volumes:
qdrant_data: qdrant_data:
neo4j_data:
neo4j_logs:
networks: networks:
antigravity_net: antigravity_net:

View File

@ -14,3 +14,5 @@ fastapi-sso
google-generativeai google-generativeai
# For local embeddings if needed # For local embeddings if needed
sentence-transformers sentence-transformers
# Graph database for relationship mapping
neo4j

View File

@ -162,3 +162,16 @@ class Config:
return [int(c.strip()) for c in chats if c.strip()] return [int(c.strip()) for c in chats if c.strip()]
except ValueError: except ValueError:
return [] return []
@staticmethod
def get_neo4j_config():
"""
Returns Neo4j connection configuration.
Used for relationship graph management (complements Qdrant vector store).
"""
return {
"uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"),
"user": os.getenv("NEO4J_USER", "neo4j"),
"password": os.getenv("NEO4J_PASSWORD", "antigravity2024"),
"database": os.getenv("NEO4J_DATABASE", "neo4j")
}

View File

@ -5,10 +5,20 @@ from src.config import Config
import logging import logging
import time import time
from functools import wraps from functools import wraps
from typing import Optional
# Setup Logging # Setup Logging
logger = logging.getLogger("AntigravityMemory") logger = logging.getLogger("AntigravityMemory")
# Import GraphWrapper opcionalmente para não quebrar se Neo4j não estiver disponível
try:
from src.memory.graph_wrapper import GraphWrapper
GRAPH_AVAILABLE = True
except (ImportError, ModuleNotFoundError) as e:
logger.warning(f"GraphWrapper not available (Neo4j may not be installed): {e}")
GraphWrapper = None
GRAPH_AVAILABLE = False
# Rate Limiting Configuration # Rate Limiting Configuration
MAX_RETRIES = 3 MAX_RETRIES = 3
RETRY_DELAY_SECONDS = 2.0 # Start with 2 seconds RETRY_DELAY_SECONDS = 2.0 # Start with 2 seconds
@ -51,9 +61,17 @@ class RateLimiter:
class MemorySearchInput(BaseModel): class MemorySearchInput(BaseModel):
query: str = Field(..., description="The specific question or topic to search for in the shared memory.") query: str = Field(..., description="The specific question or topic to search for in the shared memory.")
use_relationships: Optional[bool] = Field(
default=True,
description="If True, expands search using relationship graph to find related memories. Use False for fast direct search only."
)
class MemorySaveInput(BaseModel): class MemorySaveInput(BaseModel):
fact: str = Field(..., description="The important fact, decision, or insight to be stored for future reference.") fact: str = Field(..., description="The important fact, decision, or insight to be stored for future reference.")
extract_relationships: Optional[bool] = Field(
default=True,
description="If True, automatically extracts and links relationships to other memories/entities. Set False to skip relationship extraction for speed."
)
class MemoryWrapper: class MemoryWrapper:
""" """
@ -85,12 +103,21 @@ class MemoryWrapper:
class SearchMemoryTool(BaseTool): class SearchMemoryTool(BaseTool):
name: str = "Search Shared Memory" name: str = "Search Shared Memory"
description: str = ( description: str = (
"Use this tool to search for past decisions, project context, or learned procedures " "Search for past decisions, project context, or learned procedures in the shared memory. "
"stored in the company's shared brain. Always check this before asking for redundant info." "Use 'use_relationships=True' (default) to find related memories via relationship graph for richer context. "
"Use 'use_relationships=False' for fast direct semantic search only. "
"Choose based on whether you need context about relationships between topics or just direct information."
) )
args_schema: type = MemorySearchInput args_schema: type = MemorySearchInput
def _run(self, query: str) -> str: def _run(self, query: str, use_relationships: bool = True) -> str:
"""
Busca híbrida: combina busca vetorial (Qdrant) com navegação de grafo (Neo4j).
Args:
query: Query de busca semântica
use_relationships: Se True, expande busca usando relacionamentos. Se False, apenas busca vetorial rápida.
"""
for attempt in range(MAX_RETRIES): for attempt in range(MAX_RETRIES):
try: try:
RateLimiter.wait_if_needed() RateLimiter.wait_if_needed()
@ -101,19 +128,66 @@ class SearchMemoryTool(BaseTool):
if client is None: if client is None:
return "Memory system is temporarily unavailable. Proceeding without memory context." return "Memory system is temporarily unavailable. Proceeding without memory context."
# Retrieve memories for the shared project user_id # 1. Busca vetorial inicial (sempre executada)
results = client.search( vector_results = client.search(
query, query,
user_id=config['user_id'], user_id=config['user_id'],
limit=5 limit=5 if use_relationships else 10 # Menos se vamos expandir
) )
if not results: if not vector_results:
return "No relevant information found in shared memory." return "No relevant information found in shared memory."
# Format results # 2. Busca por relacionamentos (se solicitado)
formatted = "\n".join([f"- {res.get('memory', 'Unknown')} (Score: {res.get('score', 0):.2f})" for res in results]) enriched_results = []
return f"Found in Memory:\n{formatted}" related_memories_map = {}
if use_relationships and GRAPH_AVAILABLE:
try:
# Verificar se Neo4j está disponível
driver = GraphWrapper.get_driver()
if driver:
# Para cada resultado vetorial, buscar relacionados
for result in vector_results:
memory_id = result.get('id') or result.get('memory_id') or result.get('memory', {}).get('id')
if memory_id:
related = GraphWrapper.find_related_memories(
memory_id=str(memory_id),
max_depth=2,
limit=5
)
if related:
related_memories_map[str(memory_id)] = related
except Exception as e:
logger.warning(f"Graph search failed, using vector results only: {e}")
# 3. Formatar resultados
formatted_parts = []
for i, res in enumerate(vector_results, 1):
memory_text = res.get('memory', res.get('content', 'Unknown'))
if isinstance(memory_text, dict):
memory_text = memory_text.get('content', str(memory_text))
score = res.get('score', 0)
memory_id = res.get('id') or res.get('memory_id') or res.get('memory', {}).get('id', 'unknown')
formatted_parts.append(f"{i}. {memory_text} (Relevance: {score:.2f})")
# Adicionar memórias relacionadas se existirem
if use_relationships and str(memory_id) in related_memories_map:
related = related_memories_map[str(memory_id)]
if related:
formatted_parts.append(f" └─ Related memories ({len(related)} found):")
for rel in related[:3]: # Limitar a 3 relacionadas
formatted_parts.append(f" • [{rel.get('type', 'unknown')}] {rel.get('preview', 'N/A')[:100]}...")
result_text = "Found in Memory (with relationship context):\n" if use_relationships else "Found in Memory:\n"
result_text += "\n".join(formatted_parts)
if use_relationships and related_memories_map:
total_related = sum(len(rel) for rel in related_memories_map.values())
result_text += f"\n\n[Expanded search found {total_related} additional related memories via relationship graph]"
return result_text
except Exception as e: except Exception as e:
error_str = str(e) error_str = str(e)
@ -132,12 +206,23 @@ class SearchMemoryTool(BaseTool):
class SaveMemoryTool(BaseTool): class SaveMemoryTool(BaseTool):
name: str = "Save to Shared Memory" name: str = "Save to Shared Memory"
description: str = ( description: str = (
"Use this tool to persist CRITICAL information, decisions, or new rules " "Persist CRITICAL information, decisions, or new rules to shared memory. "
"so that other agents (or you in the future) can access it. Do not save trivial chat." "The fact is saved to vector database (Qdrant) for semantic search. "
"If 'extract_relationships=True' (default), automatically extracts and links relationships "
"to other memories/entities in the graph for better context discovery. "
"Set 'extract_relationships=False' for faster saving without relationship extraction. "
"Do not save trivial chat or temporary information."
) )
args_schema: type = MemorySaveInput args_schema: type = MemorySaveInput
def _run(self, fact: str) -> str: def _run(self, fact: str, extract_relationships: bool = True) -> str:
"""
Salva memória no Qdrant (vetorial) e opcionalmente extrai relacionamentos para o grafo.
Args:
fact: Informação crítica a ser salva
extract_relationships: Se True, extrai e salva relacionamentos no grafo
"""
for attempt in range(MAX_RETRIES): for attempt in range(MAX_RETRIES):
try: try:
RateLimiter.wait_if_needed() RateLimiter.wait_if_needed()
@ -148,12 +233,84 @@ class SaveMemoryTool(BaseTool):
if client is None: if client is None:
return "Memory system temporarily unavailable. Fact not saved, but task can continue." return "Memory system temporarily unavailable. Fact not saved, but task can continue."
client.add( # 1. Salvar no Qdrant (vetorial) - sempre executado
metadata = {
"source": "agent_execution",
"type": "insight"
}
result = client.add(
fact, fact,
user_id=config['user_id'], user_id=config['user_id'],
metadata={"source": "agent_execution", "type": "insight"} metadata=metadata
) )
return "Successfully saved to shared memory."
# Obter ID da memória criada
# Mem0 pode retornar dict, objeto, ou lista dependendo da versão
memory_id = None
if isinstance(result, dict):
memory_id = result.get('id') or result.get('memory_id') or result.get('memory', {}).get('id')
elif isinstance(result, list) and result:
memory_id = result[0].get('id') if isinstance(result[0], dict) else getattr(result[0], 'id', None)
elif hasattr(result, 'id'):
memory_id = result.id
elif hasattr(result, 'get'):
memory_id = result.get('id')
# Fallback: buscar memórias recentes para obter ID (menos confiável)
if not memory_id:
try:
recent = client.search(fact[:50], user_id=config['user_id'], limit=1)
if recent and isinstance(recent, list) and recent:
memory_id = recent[0].get('id') or recent[0].get('memory_id')
except Exception as search_err:
logger.debug(f"Could not fetch memory ID via search: {search_err}")
response_parts = ["Successfully saved to shared memory (vector database)."]
# 2. Extrair relacionamentos e salvar no grafo (se solicitado)
if extract_relationships and memory_id and GRAPH_AVAILABLE:
try:
# Criar nó no grafo
GraphWrapper.create_memory_node(
memory_id=str(memory_id),
content=fact,
metadata=metadata
)
# Buscar memórias existentes relevantes para contexto
existing_context_ids = None
try:
# Buscar memórias similares para contexto de relacionamento
similar = client.search(fact[:100], user_id=config['user_id'], limit=5)
existing_context_ids = [
str(r.get('id') or r.get('memory_id', ''))
for r in similar
if r.get('id') or r.get('memory_id')
]
except:
pass
# Extrair e criar relacionamentos
relationship_stats = GraphWrapper.extract_and_link_relationships(
memory_id=str(memory_id),
content=fact,
existing_memory_ids=existing_context_ids
)
if relationship_stats.get("relationships_created", 0) > 0:
response_parts.append(
f"Relationship graph updated: {relationship_stats['relationships_created']} relationships, "
f"{relationship_stats['entities_extracted']} entities extracted."
)
else:
response_parts.append("Saved to relationship graph (no relationships extracted yet).")
except Exception as e:
logger.warning(f"Failed to extract relationships (memory still saved): {e}")
response_parts.append("Note: Relationship extraction failed, but memory was saved successfully.")
return " ".join(response_parts)
except Exception as e: except Exception as e:
error_str = str(e) error_str = str(e)