From 12a21d81b6d4bdabfc39a5edb0bebb76f97034a3 Mon Sep 17 00:00:00 2001
From: phernandez <phernandez@basicmachines.co>
Date: Mon, 30 Dec 2024 15:05:35 -0600
Subject: [PATCH] add select relations to list_by_type tool

---
 .../repository/entity_repository.py           | 29 +++++--
 tests/repository/test_entity_repository.py    | 77 +++++++++++++++++++
 2 files changed, 100 insertions(+), 6 deletions(-)

diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py
index e7fc8b02..d1bb7bf5 100644
--- a/src/basic_memory/repository/entity_repository.py
+++ b/src/basic_memory/repository/entity_repository.py
@@ -24,17 +24,34 @@ class EntityRepository(Repository[Entity]):
         return await self.find_one(query)
 
     async def list_entities(
-        self,
-        entity_type: Optional[str] = None,
-        doc_id: Optional[int] = None,
-        sort_by: Optional[str] = "updated_at",
+            self,
+            entity_type: Optional[str] = None,
+            doc_id: Optional[int] = None,
+            sort_by: Optional[str] = "updated_at",
+            include_related: bool = False,
     ) -> Sequence[Entity]:
         """List all entities, optionally filtered by type and sorted."""
-        query = self.select().options(*self.get_load_options())
+        query = self.select()
+
+        # Always load base relations
+        query = query.options(*self.get_load_options())
 
         # Apply filters
         if entity_type:
-            query = query.where(Entity.entity_type == entity_type)
+            # When include_related is True, get both:
+            # 1. Entities of the requested type
+            # 2. Entities that have relations with entities of the requested type
+            if include_related:
+                query = query.where(
+                    or_(
+                        Entity.entity_type == entity_type,
+                        Entity.from_relations.any(Relation.to_entity.has(entity_type=entity_type)),
+                        Entity.to_relations.any(Relation.from_entity.has(entity_type=entity_type))
+                    )
+                )
+            else:
+                query = query.where(Entity.entity_type == entity_type)
+
         if doc_id:
             query = query.where(Entity.doc_id == doc_id)
 
diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py
index 0c56cc1b..77624e08 100644
--- a/tests/repository/test_entity_repository.py
+++ b/tests/repository/test_entity_repository.py
@@ -461,3 +461,80 @@ async def test_delete_by_path_ids_with_observations(
         result = await session.execute(query)
         remaining_observations = result.scalars().all()
         assert len(remaining_observations) == 0
+
+
+@pytest.mark.asyncio
+async def test_list_entities_with_related(entity_repository: EntityRepository, session_maker):
+    """Test listing entities with related entities included."""
+    
+    # Create test entities
+    async with db.scoped_session(session_maker) as session:
+        # Core entities
+        core = Entity(
+            name="core_service",
+            entity_type="service",
+            path_id="service/core",
+            file_path="service/core.md",
+            description="Core service"
+        )
+        dbe = Entity(
+            name="db_service",
+            entity_type="service",
+            path_id="service/db",
+            file_path="service/db.md",
+            description="Database service"
+        )
+        # Related entity of different type
+        config = Entity(
+            name="service_config",
+            entity_type="configuration",
+            path_id="config/service",
+            file_path="config/service.md",
+            description="Service configuration"
+        )
+        session.add_all([core, dbe, config])
+        await session.flush()
+
+        # Create relations in both directions
+        relations = [
+            # core -> db (depends_on)
+            Relation(from_id=core.id, to_id=dbe.id, relation_type="depends_on"),
+            # config -> core (configures)
+            Relation(from_id=config.id, to_id=core.id, relation_type="configures")
+        ]
+        session.add_all(relations)
+
+    # Test 1: List services without related entities
+    services = await entity_repository.list_entities(
+        entity_type="service",
+        include_related=False
+    )
+    assert len(services) == 2
+    service_names = {s.name for s in services}
+    assert service_names == {"core_service", "db_service"}
+
+    # Test 2: List services with related entities
+    services_and_related = await entity_repository.list_entities(
+        entity_type="service",
+        include_related=True
+    )
+    assert len(services_and_related) == 3
+    # Should include both services and the config
+    entity_names = {e.name for e in services_and_related}
+    assert entity_names == {"core_service", "db_service", "service_config"}
+
+    # Test 3: Verify relations are loaded
+    core_service = next(e for e in services_and_related if e.name == "core_service")
+    assert len(core_service.from_relations) > 0  # Has incoming relation from config
+    assert len(core_service.to_relations) > 0    # Has outgoing relation to db
+
+    # Test 4: List configurations with related
+    configs = await entity_repository.list_entities(
+        entity_type="configuration",
+        include_related=True,
+        sort_by="name"
+    )
+    config_names = {c.name for c in configs}
+    # Should include both config and the services it relates to
+    assert "service_config" in config_names
+    assert "core_service" in config_names  # Related via configures relation
\ No newline at end of file