Compare commits

...
17 changed files with 318 additions and 40 deletions
+19 -5
View File
@@ -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 @@ def _register_routes(app: FastAPI):
"/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:
@@ -406,18 +406,20 @@ class MemoryEngineInterface(ABC):
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.
"""
...
@@ -3587,32 +3587,47 @@ Guidelines:
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 = []
@@ -3639,7 +3654,12 @@ Guidelines:
"metadata": metadata,
}
)
return entities
return {
"items": entities,
"total": total,
"limit": limit,
"offset": offset,
}
async def get_entity_state(
self,
@@ -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:
+3 -3
View File
@@ -103,7 +103,7 @@ impl ApiClient {
pub fn get_stats(&self, agent_id: &str, _verbose: bool) -> Result<AgentStats> {
self.runtime.block_on(async {
let response = self.client.get_agent_stats(agent_id).await?;
let response = self.client.get_agent_stats(agent_id, None).await?;
let value = response.into_inner();
// Convert to JSON Value first, then parse into our type
let json_value = serde_json::to_value(&value)?;
@@ -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())
})
}
+1 -1
View File
@@ -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();
+1 -1
View File
@@ -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() {
@@ -939,6 +939,7 @@ class BanksApi:
async def get_agent_stats(
self,
bank_id: StrictStr,
authorization: Optional[StrictStr] = None,
_request_timeout: Union[
None,
Annotated[StrictFloat, Field(gt=0)],
@@ -958,6 +959,8 @@ class BanksApi:
:param bank_id: (required)
:type bank_id: str
:param authorization:
:type authorization: str
:param _request_timeout: timeout setting for this request. If one
number provided, it will be total request
timeout. It can also be a pair (tuple) of
@@ -982,6 +985,7 @@ class BanksApi:
_param = self._get_agent_stats_serialize(
bank_id=bank_id,
authorization=authorization,
_request_auth=_request_auth,
_content_type=_content_type,
_headers=_headers,
@@ -1007,6 +1011,7 @@ class BanksApi:
async def get_agent_stats_with_http_info(
self,
bank_id: StrictStr,
authorization: Optional[StrictStr] = None,
_request_timeout: Union[
None,
Annotated[StrictFloat, Field(gt=0)],
@@ -1026,6 +1031,8 @@ class BanksApi:
:param bank_id: (required)
:type bank_id: str
:param authorization:
:type authorization: str
:param _request_timeout: timeout setting for this request. If one
number provided, it will be total request
timeout. It can also be a pair (tuple) of
@@ -1050,6 +1057,7 @@ class BanksApi:
_param = self._get_agent_stats_serialize(
bank_id=bank_id,
authorization=authorization,
_request_auth=_request_auth,
_content_type=_content_type,
_headers=_headers,
@@ -1075,6 +1083,7 @@ class BanksApi:
async def get_agent_stats_without_preload_content(
self,
bank_id: StrictStr,
authorization: Optional[StrictStr] = None,
_request_timeout: Union[
None,
Annotated[StrictFloat, Field(gt=0)],
@@ -1094,6 +1103,8 @@ class BanksApi:
:param bank_id: (required)
:type bank_id: str
:param authorization:
:type authorization: str
:param _request_timeout: timeout setting for this request. If one
number provided, it will be total request
timeout. It can also be a pair (tuple) of
@@ -1118,6 +1129,7 @@ class BanksApi:
_param = self._get_agent_stats_serialize(
bank_id=bank_id,
authorization=authorization,
_request_auth=_request_auth,
_content_type=_content_type,
_headers=_headers,
@@ -1138,6 +1150,7 @@ class BanksApi:
def _get_agent_stats_serialize(
self,
bank_id,
authorization,
_request_auth,
_content_type,
_headers,
@@ -1163,6 +1176,8 @@ class BanksApi:
_path_params['bank_id'] = bank_id
# process the query parameters
# process the header parameters
if authorization is not None:
_header_params['authorization'] = authorization
# process the form parameters
# process the body parameter
@@ -338,6 +338,7 @@ class EntitiesApi:
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 @@ class EntitiesApi:
) -> 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 @@ class EntitiesApi:
_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 @@ class EntitiesApi:
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 @@ class EntitiesApi:
) -> 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 @@ class EntitiesApi:
_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 @@ class EntitiesApi:
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 @@ class EntitiesApi:
) -> 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 @@ class EntitiesApi:
_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 @@ class EntitiesApi:
self,
bank_id,
limit,
offset,
authorization,
_request_auth,
_content_type,
@@ -590,6 +603,10 @@ class EntitiesApi:
_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
@@ -17,7 +17,7 @@ import pprint
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 @@ class EntityListResponse(BaseModel):
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
@@ -449,6 +449,38 @@ class TestEntities:
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."""
@@ -236,7 +236,7 @@ export const getAgentStats = <ThrowOnError extends boolean = false>(
/**
* 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.
*/
export const listEntities = <ThrowOnError extends boolean = false>(
options: Options<ListEntitiesData, ThrowOnError>,
@@ -495,6 +495,18 @@ export type EntityListResponse = {
* Items
*/
items: Array<EntityListItem>;
/**
* Total
*/
total: number;
/**
* Limit
*/
limit: number;
/**
* Offset
*/
offset: number;
};
/**
@@ -1320,6 +1332,12 @@ export type ListBanksResponse = ListBanksResponses[keyof ListBanksResponses];
export type GetAgentStatsData = {
body?: never;
headers?: {
/**
* Authorization
*/
authorization?: string | null;
};
path: {
/**
* Bank Id
@@ -1370,6 +1388,12 @@ export type ListEntitiesData = {
* Maximum number of entities to return
*/
limit?: number;
/**
* Offset
*
* Offset for pagination
*/
offset?: number;
};
url: "/v1/default/banks/{bank_id}/entities";
};
@@ -11,11 +11,12 @@ export async function GET(request: NextRequest) {
}
const limit = searchParams.get("limit") ? Number(searchParams.get("limit")) : undefined;
const offset = searchParams.get("offset") ? Number(searchParams.get("offset")) : undefined;
const response = await sdk.listEntities({
client: lowLevelClient,
path: { bank_id: bankId },
query: { limit },
query: { limit, offset },
});
if (response.error) {
@@ -4,6 +4,7 @@ import { useState, useEffect } from "react";
import { client } from "@/lib/api";
import { useBank } from "@/lib/bank-context";
import { Button } from "@/components/ui/button";
import { ChevronLeft, ChevronRight, ChevronsLeft, ChevronsRight } from "lucide-react";
import {
Table,
TableBody,
@@ -29,6 +30,8 @@ interface EntityDetail extends Entity {
}>;
}
const ITEMS_PER_PAGE = 50;
export function EntitiesView() {
const { currentBank } = useBank();
const [entities, setEntities] = useState<Entity[]>([]);
@@ -37,16 +40,26 @@ export function EntitiesView() {
const [loadingDetail, setLoadingDetail] = useState(false);
const [regenerating, setRegenerating] = useState(false);
const loadEntities = async () => {
// Pagination state
const [currentPage, setCurrentPage] = useState(1);
const [total, setTotal] = useState(0);
const totalPages = Math.ceil(total / ITEMS_PER_PAGE);
const offset = (currentPage - 1) * ITEMS_PER_PAGE;
const loadEntities = async (page: number = 1) => {
if (!currentBank) return;
setLoading(true);
try {
const result: any = await client.listEntities({
const pageOffset = (page - 1) * ITEMS_PER_PAGE;
const result = await client.listEntities({
bank_id: currentBank,
limit: 100,
limit: ITEMS_PER_PAGE,
offset: pageOffset,
});
setEntities(result.items || []);
setTotal(result.total || 0);
} catch (error) {
console.error("Error loading entities:", error);
alert("Error loading entities: " + (error as Error).message);
@@ -86,9 +99,16 @@ export function EntitiesView() {
}
};
// Handle page change
const handlePageChange = (newPage: number) => {
setCurrentPage(newPage);
loadEntities(newPage);
};
useEffect(() => {
if (currentBank) {
loadEntities();
setCurrentPage(1);
loadEntities(1);
setSelectedEntity(null);
}
}, [currentBank]);
@@ -105,13 +125,15 @@ export function EntitiesView() {
{loading ? (
<div className="flex items-center justify-center py-20">
<div className="text-center">
<div className="text-4xl mb-2"></div>
<div className="text-4xl mb-2">...</div>
<div className="text-sm text-muted-foreground">Loading entities...</div>
</div>
</div>
) : entities.length > 0 ? (
<>
<div className="mb-4 text-sm text-muted-foreground">{entities.length} entities</div>
<div className="mb-4 text-sm text-muted-foreground">
{total} {total === 1 ? "entity" : "entities"}
</div>
<div className="overflow-x-auto">
<Table>
<TableHeader>
@@ -146,11 +168,61 @@ export function EntitiesView() {
</TableBody>
</Table>
</div>
{/* Pagination Controls */}
{totalPages > 1 && (
<div className="flex items-center justify-between mt-3 pt-3 border-t">
<div className="text-xs text-muted-foreground">
{offset + 1}-{Math.min(offset + ITEMS_PER_PAGE, total)} of {total}
</div>
<div className="flex items-center gap-1">
<Button
variant="outline"
size="sm"
onClick={() => handlePageChange(1)}
disabled={currentPage === 1 || loading}
className="h-7 w-7 p-0"
>
<ChevronsLeft className="h-3 w-3" />
</Button>
<Button
variant="outline"
size="sm"
onClick={() => handlePageChange(currentPage - 1)}
disabled={currentPage === 1 || loading}
className="h-7 w-7 p-0"
>
<ChevronLeft className="h-3 w-3" />
</Button>
<span className="text-xs px-2">
{currentPage} / {totalPages}
</span>
<Button
variant="outline"
size="sm"
onClick={() => handlePageChange(currentPage + 1)}
disabled={currentPage === totalPages || loading}
className="h-7 w-7 p-0"
>
<ChevronRight className="h-3 w-3" />
</Button>
<Button
variant="outline"
size="sm"
onClick={() => handlePageChange(totalPages)}
disabled={currentPage === totalPages || loading}
className="h-7 w-7 p-0"
>
<ChevronsRight className="h-3 w-3" />
</Button>
</div>
</div>
)}
</>
) : (
<div className="flex items-center justify-center py-20">
<div className="text-center">
<div className="text-4xl mb-2">👥</div>
<div className="text-4xl mb-2">...</div>
<div className="text-sm text-muted-foreground">No entities found</div>
<div className="text-xs text-muted-foreground mt-1">
Entities are extracted from facts when memories are added.
@@ -178,7 +250,7 @@ export function EntitiesView() {
onClick={() => setSelectedEntity(null)}
className="h-8 w-8 p-0"
>
<span className="text-lg">×</span>
<span className="text-lg">x</span>
</Button>
</div>
+8 -2
View File
@@ -127,11 +127,17 @@ export class ControlPlaneClient {
/**
* List entities
*/
async listEntities(params: { bank_id: string; limit?: number }) {
async listEntities(params: { bank_id: string; limit?: number; offset?: number }) {
const queryParams = new URLSearchParams();
queryParams.append("bank_id", params.bank_id);
if (params.limit) queryParams.append("limit", params.limit.toString());
return this.fetchApi(`/api/entities?${queryParams}`);
if (params.offset) queryParams.append("offset", params.offset.toString());
return this.fetchApi<{
items: any[];
total: number;
limit: number;
offset: number;
}>(`/api/entities?${queryParams}`);
}
/**
+49 -3
View File
@@ -454,6 +454,22 @@
"type": "string",
"title": "Bank Id"
}
},
{
"name": "authorization",
"in": "header",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Authorization"
}
}
],
"responses": {
@@ -486,7 +502,7 @@
"Entities"
],
"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.",
"operationId": "list_entities",
"parameters": [
{
@@ -510,6 +526,18 @@
},
"description": "Maximum number of entities to return"
},
{
"name": "offset",
"in": "query",
"required": false,
"schema": {
"type": "integer",
"description": "Offset for pagination",
"default": 0,
"title": "Offset"
},
"description": "Offset for pagination"
},
{
"name": "authorization",
"in": "header",
@@ -2407,11 +2435,26 @@
},
"type": "array",
"title": "Items"
},
"total": {
"type": "integer",
"title": "Total"
},
"limit": {
"type": "integer",
"title": "Limit"
},
"offset": {
"type": "integer",
"title": "Offset"
}
},
"type": "object",
"required": [
"items"
"items",
"total",
"limit",
"offset"
],
"title": "EntityListResponse",
"description": "Response model for entity list endpoint.",
@@ -2424,7 +2467,10 @@
"last_seen": "2024-02-01T14:00:00Z",
"mention_count": 15
}
]
],
"limit": 100,
"offset": 0,
"total": 150
}
},
"EntityObservationResponse": {