feat(api, sdk, docs): add session token count (#34) (#44)

* feat(api): add endpoint to retrieve token counts for sessions and implement tokenizer functionality

* feat(sdk-ts): implement getTokenCounts method to retrieve total token counts for sessions

* feat(sdk-py): add get_token_counts method to retrieve total token counts for sessions

* feat(docs): add new endpoint for retrieving total token counts in a session with detailed response schema

Author: GenerQAQ

docs/api-reference/openapi.json

@@ -968,6 +968,47 @@
         } ]
       }
     },
+    "/session/{session_id}/token_counts" : {
+      "get" : {
+        "description" : "Get total token counts for all text and tool-call parts in a session",
+        "parameters" : [ {
+          "description" : "Session ID",
+          "in" : "path",
+          "name" : "session_id",
+          "required" : true,
+          "schema" : {
+            "format" : "uuid",
+            "type" : "string"
+          }
+        } ],
+        "responses" : {
+          "200" : {
+            "content" : {
+              "application/json" : {
+                "schema" : {
+                  "$ref" : "#/components/schemas/_session__session_id__token_counts_get_200_response"
+                }
+              }
+            },
+            "description" : "OK"
+          }
+        },
+        "security" : [ {
+          "BearerAuth" : [ ]
+        } ],
+        "summary" : "Get token counts for session",
+        "tags" : [ "session" ],
+        "x-code-samples" : [ {
+          "label" : "Python",
+          "lang" : "python",
+          "source" : "from acontext import AcontextClient\n\nclient = AcontextClient(api_key='sk_project_token')\n\n# Get token counts\nresult = client.sessions.get_token_counts(session_id='session-uuid')\nprint(f\"Total tokens: {result.total_tokens}\")\n"
+        }, {
+          "label" : "JavaScript",
+          "lang" : "javascript",
+          "source" : "import { AcontextClient } from '@acontext/acontext';\n\nconst client = new AcontextClient({ apiKey: 'sk_project_token' });\n\n// Get token counts\nconst result = await client.sessions.getTokenCounts('session-uuid');\nconsole.log(`Total tokens: ${result.total_tokens}`);\n"
+        } ]
+      }
+    },
     "/space" : {
       "get" : {
         "description" : "Get all spaces under a project",
@@ -2006,6 +2047,14 @@
         "required" : [ "blob" ],
         "type" : "object"
       },
+      "handler.TokenCountsResp" : {
+        "properties" : {
+          "total_tokens" : {
+            "type" : "integer"
+          }
+        },
+        "type" : "object"
+      },
       "handler.ToolRenameItem" : {
         "properties" : {
           "new_name" : {
@@ -2658,6 +2707,18 @@
           "type" : "object"
         } ]
       },
+      "_session__session_id__token_counts_get_200_response" : {
+        "allOf" : [ {
+          "$ref" : "#/components/schemas/serializer.Response"
+        }, {
+          "properties" : {
+            "data" : {
+              "$ref" : "#/components/schemas/handler.TokenCountsResp"
+            }
+          },
+          "type" : "object"
+        } ]
+      },
       "_space_get_200_response" : {
         "allOf" : [ {
           "$ref" : "#/components/schemas/serializer.Response"

src/client/acontext-py/src/acontext/resources/async_sessions.py

@@ -15,6 +15,7 @@
     ListSessionsOutput,
     Message,
     Session,
+    TokenCounts,
 )
 from ..uploads import FileUpload, normalize_file_upload
 from pydantic import BaseModel
@@ -321,3 +322,17 @@ async def get_learning_status(self, session_id: str) -> LearningStatus:
             "GET", f"/session/{session_id}/get_learning_status"
         )
         return LearningStatus.model_validate(data)
+
+    async def get_token_counts(self, session_id: str) -> TokenCounts:
+        """Get total token counts for all text and tool-call parts in a session.
+
+        Args:
+            session_id: The UUID of the session.
+
+        Returns:
+            TokenCounts object containing total_tokens.
+        """
+        data = await self._requester.request(
+            "GET", f"/session/{session_id}/token_counts"
+        )
+        return TokenCounts.model_validate(data)

src/client/acontext-py/src/acontext/resources/sessions.py

@@ -15,6 +15,7 @@
     ListSessionsOutput,
     Message,
     Session,
+    TokenCounts,
 )
 from ..uploads import FileUpload, normalize_file_upload
 from pydantic import BaseModel
@@ -321,3 +322,17 @@ def get_learning_status(self, session_id: str) -> LearningStatus:
             "GET", f"/session/{session_id}/get_learning_status"
         )
         return LearningStatus.model_validate(data)
+
+    def get_token_counts(self, session_id: str) -> TokenCounts:
+        """Get total token counts for all text and tool-call parts in a session.
+
+        Args:
+            session_id: The UUID of the session.
+
+        Returns:
+            TokenCounts object containing total_tokens.
+        """
+        data = self._requester.request(
+            "GET", f"/session/{session_id}/token_counts"
+        )
+        return TokenCounts.model_validate(data)

src/client/acontext-py/src/acontext/types/__init__.py

@@ -20,6 +20,7 @@
     PublicURL,
     Session,
     Task,
+    TokenCounts,
 )
 from .block import Block
 from .space import (
@@ -55,6 +56,7 @@
     "PublicURL",
     "Session",
     "Task",
+    "TokenCounts",
     # Space types
     "ListSpacesOutput",
     "SearchResultBlockItem",

src/client/acontext-py/src/acontext/types/session.py

@@ -133,3 +133,9 @@ class LearningStatus(BaseModel):
     not_space_digested_count: int = Field(
         ..., description="Number of tasks that are not space digested"
     )
+
+
+class TokenCounts(BaseModel):
+    """Response model for token counts."""
+
+    total_tokens: int = Field(..., description="Total token count for all text and tool-call parts in a session")

src/client/acontext-py/tests/test_async_client.py

@@ -462,6 +462,27 @@ async def test_async_sessions_get_learning_status(
     assert result.not_space_digested_count == 3
 
 
+@patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_async_sessions_get_token_counts(
+    mock_request, async_client: AcontextAsyncClient
+) -> None:
+    mock_request.return_value = {
+        "total_tokens": 1234,
+    }
+
+    result = await async_client.sessions.get_token_counts("session-id")
+
+    mock_request.assert_called_once()
+    args, kwargs = mock_request.call_args
+    method, path = args
+    assert method == "GET"
+    assert path == "/session/session-id/token_counts"
+    # Verify it returns a Pydantic model
+    assert hasattr(result, "total_tokens")
+    assert result.total_tokens == 1234
+
+
 @patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock)
 @pytest.mark.asyncio
 async def test_async_blocks_list_without_filters(

src/client/acontext-py/tests/test_client.py

@@ -469,6 +469,24 @@ def test_sessions_get_learning_status(mock_request, client: AcontextClient) -> N
     assert result.not_space_digested_count == 3
 
 
+@patch("acontext.client.AcontextClient.request")
+def test_sessions_get_token_counts(mock_request, client: AcontextClient) -> None:
+    mock_request.return_value = {
+        "total_tokens": 1234,
+    }
+
+    result = client.sessions.get_token_counts("session-id")
+
+    mock_request.assert_called_once()
+    args, kwargs = mock_request.call_args
+    method, path = args
+    assert method == "GET"
+    assert path == "/session/session-id/token_counts"
+    # Verify it returns a Pydantic model
+    assert hasattr(result, "total_tokens")
+    assert result.total_tokens == 1234
+
+
 @patch("acontext.client.AcontextClient.request")
 def test_blocks_list_without_filters(mock_request, client: AcontextClient) -> None:
     mock_request.return_value = []

src/client/acontext-ts/src/resources/sessions.ts

@@ -19,6 +19,8 @@ import {
   MessageSchema,
   Session,
   SessionSchema,
+  TokenCounts,
+  TokenCountsSchema,
 } from '../types';
 
 export type MessageBlob = AcontextMessage | Record<string, unknown>;
@@ -226,5 +228,16 @@ export class SessionsAPI {
     );
     return LearningStatusSchema.parse(data);
   }
+
+  /**
+   * Get total token counts for all text and tool-call parts in a session.
+   *
+   * @param sessionId - The UUID of the session.
+   * @returns TokenCounts object containing total_tokens.
+   */
+  async getTokenCounts(sessionId: string): Promise<TokenCounts> {
+    const data = await this.requester.request('GET', `/session/${sessionId}/token_counts`);
+    return TokenCountsSchema.parse(data);
+  }
 }
 

src/client/acontext-ts/src/types/session.ts

@@ -105,3 +105,9 @@ export const LearningStatusSchema = z.object({
 
 export type LearningStatus = z.infer<typeof LearningStatusSchema>;
 
+export const TokenCountsSchema = z.object({
+  total_tokens: z.number(),
+});
+
+export type TokenCounts = z.infer<typeof TokenCountsSchema>;
+

src/client/acontext-ts/tests/client.test.ts

@@ -262,6 +262,17 @@ describe('AcontextClient Integration Tests', () => {
       expect(learningStatus.not_space_digested_count).toBeGreaterThanOrEqual(0);
     });
 
+    test('should get token counts', async () => {
+      if (!createdSessionId) {
+        throw new Error('Session not created');
+      }
+      const tokenCounts = await client.sessions.getTokenCounts(createdSessionId);
+      expect(tokenCounts).toBeDefined();
+      expect(tokenCounts.total_tokens).toBeDefined();
+      expect(typeof tokenCounts.total_tokens).toBe('number');
+      expect(tokenCounts.total_tokens).toBeGreaterThanOrEqual(0);
+    });
+
     test('should update session configs', async () => {
       if (!createdSessionId) {
         throw new Error('Session not created');