Merge pull request #22 from OpenPecha/tokencounter

Feat: Token Counter using google genai

Author: TenzDelek

api/Assistant/assistant_model.py

@@ -1,5 +1,5 @@
 from sqlalchemy.orm import relationship
-from sqlalchemy import Column,DateTime,String,ForeignKey,Boolean,UUID,Text
+from sqlalchemy import Column,DateTime,String,ForeignKey,Boolean,UUID,Text,Integer
 from datetime import datetime,timezone
 from uuid import uuid4
 from api.db.pg_database import Base
@@ -30,6 +30,7 @@ class Context(Base):
     content=Column(Text,nullable=True)
     pecha_title = Column(String(255), nullable=True)
     pecha_text_id = Column(String(255), nullable=True)
+    token_count = Column(Integer, nullable=True)
     assistant_id = Column(UUID(as_uuid=True),ForeignKey("assistant.id", ondelete="CASCADE"),nullable=False)
     created_at = Column(DateTime(timezone=True), default=datetime.now(timezone.utc),nullable=False)
     updated_at = Column(DateTime(timezone=True), default=datetime.now(timezone.utc))

api/Assistant/assistant_service.py

@@ -9,6 +9,8 @@
 from uuid import UUID
 from datetime import datetime, timezone
 from fastapi import HTTPException, status, UploadFile
+from api.config import get
+from api.llm.token_counter import TokenCounter
 from api.error_constant import ErrorConstants
 from api.cache.cache_enums import CacheType
 from api.Assistant.assistant_cache_service import (
@@ -18,6 +20,8 @@
 )
 from api.utils import Utils
 
+token_counter = TokenCounter(api_key=get("GEMINI_API_KEY"))
+
 
 def _build_context_responses(contexts) -> List[ContextResponse]:
     return [
@@ -81,8 +85,9 @@ async def create_assistant_service(token: str, assistant_request: AssistantReque
             content = "\n\n".join([detail.content for detail in search_details]) if search_details else ""
         else:
             content = ctx.content
+        token_count = token_counter.count_tokens(content)
         contexts_list.append(
-            Context(content=content, pecha_title=ctx.pecha_title, pecha_text_id=ctx.pecha_text_id)
+            Context(content=content, pecha_title=ctx.pecha_title, pecha_text_id=ctx.pecha_text_id, token_count=token_count)
         )
     if files:
         for file in files:
@@ -91,7 +96,8 @@ async def create_assistant_service(token: str, assistant_request: AssistantReque
                 try:
                     Utils.validate_file(file.filename, len(file_bytes))
                     extracted_content = Utils.extract_content_from_file(file_bytes, file.filename)
-                    contexts_list.append(Context(content=extracted_content))
+                    file_token_count = token_counter.count_tokens(extracted_content)
+                    contexts_list.append(Context(content=extracted_content, token_count=file_token_count))
                 except ValueError as e:
                     raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
 
@@ -171,8 +177,9 @@ async def update_assistant_service(assistant_id: UUID, update_request: UpdateAss
                     content = "\n\n".join([detail.content for detail in search_details]) if search_details else ""
                 else:
                     content = ctx.content
+                token_count = token_counter.count_tokens(content)
                 new_contexts.append(
-                    Context(content=content, pecha_title=ctx.pecha_title, pecha_text_id=ctx.pecha_text_id)
+                    Context(content=content, pecha_title=ctx.pecha_title, pecha_text_id=ctx.pecha_text_id, token_count=token_count)
                 )
             assistant.contexts = new_contexts
 

api/llm/token_counter.py

@@ -0,0 +1,14 @@
+from langchain_google_genai import ChatGoogleGenerativeAI
+from typing import Optional
+
+class TokenCounter:
+    def __init__(self, api_key: str, model_name: str = "gemini-2.5-flash"):
+        self.model = ChatGoogleGenerativeAI(
+            model=model_name,
+            google_api_key=api_key
+        )
+    
+    def count_tokens(self, text: Optional[str]) -> int:
+        if not text:
+            return 0
+        return self.model.get_num_tokens(text)

migrations/versions/635acdc7f036_added_tokencount_field_in_context_table.py

@@ -0,0 +1,32 @@
+"""added tokencount field in context table
+
+Revision ID: 635acdc7f036
+Revises: 18cf46778420
+Create Date: 2026-02-25 14:29:23.227314
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '635acdc7f036'
+down_revision: Union[str, Sequence[str], None] = '18cf46778420'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('context', sa.Column('token_count', sa.Integer(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('context', 'token_count')
+    # ### end Alembic commands ###