diff --git a/letta/server/rest_api/routers/v1/agents.py b/letta/server/rest_api/routers/v1/agents.py index e8571fa5..bd03348e 100644 --- a/letta/server/rest_api/routers/v1/agents.py +++ b/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" diff --git a/letta/server/rest_api/routers/v1/identities.py b/letta/server/rest_api/routers/v1/identities.py index 5f081600..dd48fd4e 100644 --- a/letta/server/rest_api/routers/v1/identities.py +++ b/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, diff --git a/letta/server/rest_api/routers/v1/sources.py b/letta/server/rest_api/routers/v1/sources.py index 5f08b3ea..ac91d69b 100644 --- a/letta/server/rest_api/routers/v1/sources.py +++ b/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, diff --git a/letta/server/rest_api/routers/v1/tools.py b/letta/server/rest_api/routers/v1/tools.py index b1d95386..06482175 100644 --- a/letta/server/rest_api/routers/v1/tools.py +++ b/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(...), diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index e68ac443..e71e1083 100644 --- a/letta/services/agent_manager.py +++ b/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.""" diff --git a/letta/services/identity_manager.py b/letta/services/identity_manager.py index e6bf881f..798b01a0 100644 --- a/letta/services/identity_manager.py +++ b/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, diff --git a/letta/services/source_manager.py b/letta/services/source_manager.py index 21a36ded..c872f490 100644 --- a/letta/services/source_manager.py +++ b/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]: """ diff --git a/letta/services/tool_manager.py b/letta/services/tool_manager.py index 90dbdcfa..571e67d4 100644 --- a/letta/services/tool_manager.py +++ b/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."""