From 5c432727862443b2bad02609ef6ebb3071b96532 Mon Sep 17 00:00:00 2001 From: TenzDelek Date: Wed, 25 Feb 2026 13:00:48 +0530 Subject: [PATCH 1/3] base --- api/llm/token_counter.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 api/llm/token_counter.py diff --git a/api/llm/token_counter.py b/api/llm/token_counter.py new file mode 100644 index 0000000..d7f3f28 --- /dev/null +++ b/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) \ No newline at end of file From f51cab2263de73295c32d46e21749c17aa0e4360 Mon Sep 17 00:00:00 2001 From: TenzDelek Date: Wed, 25 Feb 2026 14:29:42 +0530 Subject: [PATCH 2/3] added new token field --- api/Assistant/assistant_model.py | 3 +- ...added_tokencount_field_in_context_table.py | 32 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) create mode 100644 migrations/versions/635acdc7f036_added_tokencount_field_in_context_table.py diff --git a/api/Assistant/assistant_model.py b/api/Assistant/assistant_model.py index 2d6d306..f2e59f2 100644 --- a/api/Assistant/assistant_model.py +++ b/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)) diff --git a/migrations/versions/635acdc7f036_added_tokencount_field_in_context_table.py b/migrations/versions/635acdc7f036_added_tokencount_field_in_context_table.py new file mode 100644 index 0000000..e5b54f0 --- /dev/null +++ b/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 ### From 9365a25c4a4bc6e9a047728fdc62198af87edebd Mon Sep 17 00:00:00 2001 From: TenzDelek Date: Wed, 25 Feb 2026 14:39:47 +0530 Subject: [PATCH 3/3] token counter --- api/Assistant/assistant_service.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/api/Assistant/assistant_service.py b/api/Assistant/assistant_service.py index 41b2649..fa16f12 100644 --- a/api/Assistant/assistant_service.py +++ b/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