feat: add endpoints for counts of agents, identities, sources, & tools (#1865)

letta/server/rest_api/routers/v1/agents.py

@@ -104,6 +104,17 @@ def list_agents(
     )
 
 
+@router.get("/count", response_model=int, operation_id="count_agents")
+def count_agents(
+    server: SyncServer = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),
+):
+    """
+    Get the count of all agents associated with a given user.
+    """
+    return server.agent_manager.size(actor=server.user_manager.get_user_or_default(user_id=actor_id))
+
+
 class IndentedORJSONResponse(Response):
     media_type = "application/json"
 

letta/server/rest_api/routers/v1/identities.py

@@ -49,6 +49,24 @@ def list_identities(
     return identities
 
 
+@router.get("/count", tags=["identities"], response_model=int, operation_id="count_identities")
+def count_identities(
+    server: "SyncServer" = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),
+):
+    """
+    Get count of all identities for a user
+    """
+    try:
+        return server.identity_manager.size(actor=server.user_manager.get_user_or_default(user_id=actor_id))
+    except NoResultFound:
+        return 0
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"{e}")
+
+
 @router.get("/{identity_id}", tags=["identities"], response_model=Identity, operation_id="retrieve_identity")
 def retrieve_identity(
     identity_id: str,

letta/server/rest_api/routers/v1/sources.py

@@ -67,6 +67,17 @@ def list_sources(
     return server.list_all_sources(actor=actor)
 
 
+@router.get("/count", response_model=int, operation_id="count_sources")
+def count_sources(
+    server: "SyncServer" = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+):
+    """
+    Count all data sources created by a user.
+    """
+    return server.source_manager.size(actor=server.user_manager.get_user_or_default(user_id=actor_id))
+
+
 @router.post("/", response_model=Source, operation_id="create_source")
 def create_source(
     source_create: SourceCreate,

letta/server/rest_api/routers/v1/tools.py

@@ -80,6 +80,21 @@ def list_tools(
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@router.get("/count", response_model=int, operation_id="count_tools")
+def count_tools(
+    server: SyncServer = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),
+):
+    """
+    Get a count of all tools available to agents belonging to the org of the user
+    """
+    try:
+        return server.tool_manager.size(actor=server.user_manager.get_user_or_default(user_id=actor_id))
+    except Exception as e:
+        print(f"Error occurred: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @router.post("/", response_model=Tool, operation_id="create_tool")
 def create_tool(
     request: ToolCreate = Body(...),

letta/services/agent_manager.py

@@ -556,6 +556,16 @@ class AgentManager:
 
             return list(session.execute(query).scalars())
 
+    def size(
+        self,
+        actor: PydanticUser,
+    ) -> int:
+        """
+        Get the total count of agents for the given user.
+        """
+        with self.session_maker() as session:
+            return AgentModel.size(db_session=session, actor=actor)
+
     @enforce_types
     def get_agent_by_id(self, agent_id: str, actor: PydanticUser) -> PydanticAgentState:
         """Fetch an agent by its ID."""

letta/services/identity_manager.py

@@ -190,6 +190,17 @@ class IdentityManager:
             session.delete(identity)
             session.commit()
 
+    @enforce_types
+    def size(
+        self,
+        actor: PydanticUser,
+    ) -> int:
+        """
+        Get the total count of identities for the given user.
+        """
+        with self.session_maker() as session:
+            return IdentityModel.size(db_session=session, actor=actor)
+
     def _process_relationship(
         self,
         session: Session,

letta/services/source_manager.py

@@ -77,6 +77,17 @@ class SourceManager:
             )
             return [source.to_pydantic() for source in sources]
 
+    @enforce_types
+    def size(
+        self,
+        actor: PydanticUser,
+    ) -> int:
+        """
+        Get the total count of sources for the given user.
+        """
+        with self.session_maker() as session:
+            return SourceModel.size(db_session=session, actor=actor)
+
     @enforce_types
     def list_attached_agents(self, source_id: str, actor: Optional[PydanticUser] = None) -> List[PydanticAgentState]:
         """

letta/services/tool_manager.py

@@ -145,6 +145,17 @@ class ToolManager:
 
         return results
 
+    @enforce_types
+    def size(
+        self,
+        actor: PydanticUser,
+    ) -> int:
+        """
+        Get the total count of tools for the given user.
+        """
+        with self.session_maker() as session:
+            return ToolModel.size(db_session=session, actor=actor)
+
     @enforce_types
     def update_tool_by_id(self, tool_id: str, tool_update: ToolUpdate, actor: PydanticUser) -> PydanticTool:
         """Update a tool by its ID with the given ToolUpdate object."""