fix: entities list only show 100 entities (#142)

* fix: entities list only show 100 entities

* fix: update Rust CLI for entities pagination API changes

Author: nicoloboschi

hindsight-api/hindsight_api/api/http.py

@@ -188,12 +188,18 @@ class EntityListResponse(BaseModel):
                         "first_seen": "2024-01-15T10:30:00Z",
                         "last_seen": "2024-02-01T14:00:00Z",
                     }
-                ]
+                ],
+                "total": 150,
+                "limit": 100,
+                "offset": 0,
             }
         }
     )
 
     items: list[EntityListItem]
+    total: int
+    limit: int
+    offset: int
 
 
 class EntityDetailResponse(BaseModel):
@@ -1516,19 +1522,27 @@ async def api_stats(
         "/v1/default/banks/{bank_id}/entities",
         response_model=EntityListResponse,
         summary="List entities",
-        description="List all entities (people, organizations, etc.) known by the bank, ordered by mention count.",
+        description="List all entities (people, organizations, etc.) known by the bank, ordered by mention count. Supports pagination.",
         operation_id="list_entities",
         tags=["Entities"],
     )
     async def api_list_entities(
         bank_id: str,
         limit: int = Query(default=100, description="Maximum number of entities to return"),
+        offset: int = Query(default=0, description="Offset for pagination"),
         request_context: RequestContext = Depends(get_request_context),
     ):
-        """List entities for a memory bank."""
+        """List entities for a memory bank with pagination."""
         try:
-            entities = await app.state.memory.list_entities(bank_id, limit=limit, request_context=request_context)
-            return EntityListResponse(items=[EntityListItem(**e) for e in entities])
+            data = await app.state.memory.list_entities(
+                bank_id, limit=limit, offset=offset, request_context=request_context
+            )
+            return EntityListResponse(
+                items=[EntityListItem(**e) for e in data["items"]],
+                total=data["total"],
+                limit=data["limit"],
+                offset=data["offset"],
+            )
         except (AuthenticationError, HTTPException):
             raise
         except Exception as e:

hindsight-api/hindsight_api/engine/interface.py

@@ -406,18 +406,20 @@ async def list_entities(
         bank_id: str,
         *,
         limit: int = 100,
+        offset: int = 0,
         request_context: "RequestContext",
-    ) -> list[dict[str, Any]]:
+    ) -> dict[str, Any]:
         """
-        List entities for a bank.
+        List entities for a bank with pagination.
 
         Args:
             bank_id: The memory bank ID.
             limit: Maximum results.
+            offset: Offset for pagination.
             request_context: Request context for authentication.
 
         Returns:
-            List of entity dicts.
+            Dict with items, total, limit, offset.
         """
         ...
 

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -3601,32 +3601,47 @@ async def list_entities(
         bank_id: str,
         *,
         limit: int = 100,
+        offset: int = 0,
         request_context: "RequestContext",
-    ) -> list[dict[str, Any]]:
+    ) -> dict[str, Any]:
         """
-        List all entities for a bank.
+        List all entities for a bank with pagination.
 
         Args:
             bank_id: bank IDentifier
             limit: Maximum number of entities to return
+            offset: Offset for pagination
             request_context: Request context for authentication.
 
         Returns:
-            List of entity dicts with id, canonical_name, mention_count, first_seen, last_seen
+            Dict with items, total, limit, offset
         """
         await self._authenticate_tenant(request_context)
         pool = await self._get_pool()
         async with acquire_with_retry(pool) as conn:
+            # Get total count
+            total_row = await conn.fetchrow(
+                f"""
+                SELECT COUNT(*) as total
+                FROM {fq_table("entities")}
+                WHERE bank_id = $1
+                """,
+                bank_id,
+            )
+            total = total_row["total"] if total_row else 0
+
+            # Get paginated entities
             rows = await conn.fetch(
                 f"""
                 SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
                 FROM {fq_table("entities")}
                 WHERE bank_id = $1
                 ORDER BY mention_count DESC, last_seen DESC
-                LIMIT $2
+                LIMIT $2 OFFSET $3
                 """,
                 bank_id,
                 limit,
+                offset,
             )
 
             entities = []
@@ -3653,7 +3668,12 @@ async def list_entities(
                         "metadata": metadata,
                     }
                 )
-            return entities
+            return {
+                "items": entities,
+                "total": total,
+                "limit": limit,
+                "offset": offset,
+            }
 
     async def get_entity_state(
         self,

hindsight-api/tests/test_http_api_integration.py

@@ -250,11 +250,34 @@ async def test_full_api_workflow(api_client, test_bank_id):
     # 8. Test Entity Endpoints
     # ================================================================
 
-    # List entities
+    # List entities with pagination
     response = await api_client.get(f"/v1/default/banks/{test_bank_id}/entities")
     assert response.status_code == 200
     entities_data = response.json()
     assert "items" in entities_data
+    assert "total" in entities_data
+    assert "limit" in entities_data
+    assert "offset" in entities_data
+    assert entities_data["offset"] == 0
+    assert entities_data["limit"] == 100  # default limit
+
+    # Test pagination with custom limit and offset
+    response = await api_client.get(f"/v1/default/banks/{test_bank_id}/entities?limit=5&offset=0")
+    assert response.status_code == 200
+    paginated_data = response.json()
+    assert paginated_data["limit"] == 5
+    assert paginated_data["offset"] == 0
+    assert len(paginated_data["items"]) <= 5
+
+    # Test offset
+    if entities_data["total"] > 1:
+        response = await api_client.get(f"/v1/default/banks/{test_bank_id}/entities?limit=1&offset=1")
+        assert response.status_code == 200
+        offset_data = response.json()
+        assert offset_data["offset"] == 1
+        # With offset=1, we should get different entity than first one (if there are multiple)
+        if len(offset_data["items"]) > 0 and len(entities_data["items"]) > 1:
+            assert offset_data["items"][0]["id"] != entities_data["items"][0]["id"]
 
     # Get specific entity if any exist
     if len(entities_data['items']) > 0:

hindsight-cli/src/api.rs

@@ -241,9 +241,9 @@ impl ApiClient {
         })
     }
 
-    pub fn list_entities(&self, bank_id: &str, limit: Option<i64>, _verbose: bool) -> Result<types::EntityListResponse> {
+    pub fn list_entities(&self, bank_id: &str, limit: Option<i64>, offset: Option<i64>, _verbose: bool) -> Result<types::EntityListResponse> {
         self.runtime.block_on(async {
-            let response = self.client.list_entities(bank_id, limit, None).await?;
+            let response = self.client.list_entities(bank_id, limit, offset, None).await?;
             Ok(response.into_inner())
         })
     }

hindsight-cli/src/commands/entity.rs

@@ -16,7 +16,7 @@ pub fn list(
         None
     };
 
-    let response = client.list_entities(bank_id, Some(limit), verbose)?;
+    let response = client.list_entities(bank_id, Some(limit), None, verbose)?;
 
     if let Some(mut sp) = spinner {
         sp.finish();

hindsight-cli/src/commands/explore.rs

@@ -283,7 +283,7 @@ impl App {
     }
 
     fn load_entities(&mut self, bank_id: &str) -> Result<()> {
-        let response = self.client.list_entities(bank_id, Some(100), false)?;
+        let response = self.client.list_entities(bank_id, Some(100), None, false)?;
         self.entities = response.items;
 
         if !self.entities.is_empty() && self.entities_state.selected().is_none() {

hindsight-clients/python/hindsight_client_api/api/entities_api.py

@@ -338,6 +338,7 @@ async def list_entities(
         self,
         bank_id: StrictStr,
         limit: Annotated[Optional[StrictInt], Field(description="Maximum number of entities to return")] = None,
+        offset: Annotated[Optional[StrictInt], Field(description="Offset for pagination")] = None,
         authorization: Optional[StrictStr] = None,
         _request_timeout: Union[
             None,
@@ -354,12 +355,14 @@ async def list_entities(
     ) -> EntityListResponse:
         """List entities
 
-        List all entities (people, organizations, etc.) known by the bank, ordered by mention count.
+        List all entities (people, organizations, etc.) known by the bank, ordered by mention count. Supports pagination.
 
         :param bank_id: (required)
         :type bank_id: str
         :param limit: Maximum number of entities to return
         :type limit: int
+        :param offset: Offset for pagination
+        :type offset: int
         :param authorization:
         :type authorization: str
         :param _request_timeout: timeout setting for this request. If one
@@ -387,6 +390,7 @@ async def list_entities(
         _param = self._list_entities_serialize(
             bank_id=bank_id,
             limit=limit,
+            offset=offset,
             authorization=authorization,
             _request_auth=_request_auth,
             _content_type=_content_type,
@@ -414,6 +418,7 @@ async def list_entities_with_http_info(
         self,
         bank_id: StrictStr,
         limit: Annotated[Optional[StrictInt], Field(description="Maximum number of entities to return")] = None,
+        offset: Annotated[Optional[StrictInt], Field(description="Offset for pagination")] = None,
         authorization: Optional[StrictStr] = None,
         _request_timeout: Union[
             None,
@@ -430,12 +435,14 @@ async def list_entities_with_http_info(
     ) -> ApiResponse[EntityListResponse]:
         """List entities
 
-        List all entities (people, organizations, etc.) known by the bank, ordered by mention count.
+        List all entities (people, organizations, etc.) known by the bank, ordered by mention count. Supports pagination.
 
         :param bank_id: (required)
         :type bank_id: str
         :param limit: Maximum number of entities to return
         :type limit: int
+        :param offset: Offset for pagination
+        :type offset: int
         :param authorization:
         :type authorization: str
         :param _request_timeout: timeout setting for this request. If one
@@ -463,6 +470,7 @@ async def list_entities_with_http_info(
         _param = self._list_entities_serialize(
             bank_id=bank_id,
             limit=limit,
+            offset=offset,
             authorization=authorization,
             _request_auth=_request_auth,
             _content_type=_content_type,
@@ -490,6 +498,7 @@ async def list_entities_without_preload_content(
         self,
         bank_id: StrictStr,
         limit: Annotated[Optional[StrictInt], Field(description="Maximum number of entities to return")] = None,
+        offset: Annotated[Optional[StrictInt], Field(description="Offset for pagination")] = None,
         authorization: Optional[StrictStr] = None,
         _request_timeout: Union[
             None,
@@ -506,12 +515,14 @@ async def list_entities_without_preload_content(
     ) -> RESTResponseType:
         """List entities
 
-        List all entities (people, organizations, etc.) known by the bank, ordered by mention count.
+        List all entities (people, organizations, etc.) known by the bank, ordered by mention count. Supports pagination.
 
         :param bank_id: (required)
         :type bank_id: str
         :param limit: Maximum number of entities to return
         :type limit: int
+        :param offset: Offset for pagination
+        :type offset: int
         :param authorization:
         :type authorization: str
         :param _request_timeout: timeout setting for this request. If one
@@ -539,6 +550,7 @@ async def list_entities_without_preload_content(
         _param = self._list_entities_serialize(
             bank_id=bank_id,
             limit=limit,
+            offset=offset,
             authorization=authorization,
             _request_auth=_request_auth,
             _content_type=_content_type,
@@ -561,6 +573,7 @@ def _list_entities_serialize(
         self,
         bank_id,
         limit,
+        offset,
         authorization,
         _request_auth,
         _content_type,
@@ -590,6 +603,10 @@ def _list_entities_serialize(
 
             _query_params.append(('limit', limit))
 
+        if offset is not None:
+            
+            _query_params.append(('offset', offset))
+            
         # process the header parameters
         if authorization is not None:
             _header_params['authorization'] = authorization

hindsight-clients/python/hindsight_client_api/models/entity_list_response.py

@@ -17,7 +17,7 @@
 import re  # noqa: F401
 import json
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, StrictInt
 from typing import Any, ClassVar, Dict, List
 from hindsight_client_api.models.entity_list_item import EntityListItem
 from typing import Optional, Set
@@ -28,7 +28,10 @@ class EntityListResponse(BaseModel):
     Response model for entity list endpoint.
     """ # noqa: E501
     items: List[EntityListItem]
-    __properties: ClassVar[List[str]] = ["items"]
+    total: StrictInt
+    limit: StrictInt
+    offset: StrictInt
+    __properties: ClassVar[List[str]] = ["items", "total", "limit", "offset"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -88,7 +91,10 @@ def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
             return cls.model_validate(obj)
 
         _obj = cls.model_validate({
-            "items": [EntityListItem.from_dict(_item) for _item in obj["items"]] if obj.get("items") is not None else None
+            "items": [EntityListItem.from_dict(_item) for _item in obj["items"]] if obj.get("items") is not None else None,
+            "total": obj.get("total"),
+            "limit": obj.get("limit"),
+            "offset": obj.get("offset")
         })
         return _obj
 

hindsight-clients/python/tests/test_main_operations.py

@@ -449,6 +449,38 @@ async def do_list():
         assert response is not None
         assert response.items is not None
         assert isinstance(response.items, list)
+        # Verify pagination fields
+        assert response.total is not None
+        assert response.limit is not None
+        assert response.offset is not None
+        assert response.offset == 0
+        assert response.limit == 100  # default limit
+
+    def test_list_entities_with_pagination(self, client, bank_id):
+        """Test listing entities with pagination parameters."""
+        import asyncio
+        from hindsight_client_api import ApiClient, Configuration
+        from hindsight_client_api.api import EntitiesApi
+
+        async def do_list_paginated():
+            config = Configuration(host=HINDSIGHT_API_URL)
+            api_client = ApiClient(config)
+            api = EntitiesApi(api_client)
+
+            # Test with custom limit
+            response = await api.list_entities(bank_id=bank_id, limit=5, offset=0)
+            assert response.limit == 5
+            assert response.offset == 0
+            assert len(response.items) <= 5
+
+            # Test with offset
+            response_offset = await api.list_entities(bank_id=bank_id, limit=1, offset=1)
+            assert response_offset.offset == 1
+            assert response_offset.limit == 1
+
+            return response
+
+        asyncio.get_event_loop().run_until_complete(do_list_paginated())
 
     def test_get_entity(self, client, bank_id):
         """Test getting a specific entity."""