69 lines
3.0 KiB
Python
69 lines
3.0 KiB
Python
"""
|
|
Tool for searching memories by entity relationships.
|
|
Allows agents to find all memories related to a specific entity (Agent, Technology, Project, etc.)
|
|
"""
|
|
from crewai.tools import BaseTool
|
|
from pydantic import BaseModel, Field
|
|
from typing import Optional
|
|
from src.memory.graph_wrapper import GraphWrapper
|
|
import logging
|
|
|
|
logger = logging.getLogger("AntigravityEntitySearch")
|
|
|
|
class EntitySearchInput(BaseModel):
|
|
entity_name: str = Field(..., description="Name of the entity to search for (e.g., 'Arthur Mendes', 'Zabbix', 'Infrastructure Project')")
|
|
entity_type: Optional[str] = Field(
|
|
default=None,
|
|
description="Type of entity: Agent, Technology, Project, Problem, Solution. If not specified, searches all types."
|
|
)
|
|
limit: Optional[int] = Field(
|
|
default=10,
|
|
description="Maximum number of memories to return"
|
|
)
|
|
|
|
class SearchByEntityTool(BaseTool):
|
|
name: str = "Search Memories by Entity"
|
|
description: str = (
|
|
"Search for all memories related to a specific entity (person, technology, project, etc.) "
|
|
"using the relationship graph. This is useful when you need to find everything related to "
|
|
"a specific agent, technology, or project, not just semantically similar content. "
|
|
"Example: 'Find all memories about Zabbix' or 'Show me everything related to Arthur Mendes'."
|
|
)
|
|
args_schema: type = EntitySearchInput
|
|
|
|
def _run(self, entity_name: str, entity_type: Optional[str] = None, limit: int = 10) -> str:
|
|
"""
|
|
Busca memórias relacionadas a uma entidade específica usando o grafo de relacionamentos.
|
|
|
|
Args:
|
|
entity_name: Nome da entidade
|
|
entity_type: Tipo da entidade (Agent, Technology, Project, etc.)
|
|
limit: Limite de resultados
|
|
"""
|
|
try:
|
|
memories = GraphWrapper.find_memories_by_entity(
|
|
entity_name=entity_name,
|
|
entity_type=entity_type,
|
|
limit=limit
|
|
)
|
|
|
|
if not memories:
|
|
return f"No memories found related to entity '{entity_name}'" + (f" of type '{entity_type}'" if entity_type else "") + "."
|
|
|
|
formatted = [f"Found {len(memories)} memories related to '{entity_name}':\n"]
|
|
|
|
for i, mem in enumerate(memories, 1):
|
|
preview = mem.get('preview', 'N/A')
|
|
mem_type = mem.get('type', 'unknown')
|
|
entity_type_display = mem.get('entity_type', 'N/A')
|
|
related_count = mem.get('related_count', 0)
|
|
|
|
formatted.append(f"{i}. [{mem_type}] {preview}")
|
|
formatted.append(f" Entity Type: {entity_type_display} | Related memories: {related_count}")
|
|
|
|
return "\n".join(formatted)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Entity search failed for '{entity_name}': {e}")
|
|
return f"Entity search temporarily unavailable: {str(e)[:100]}"
|