Merge branch 'COG-970-refactor-tokenizing' into feat/cog-1071-input-t…

…oken-counting

cognee/modules/data/models/Data.py

@@ -1,6 +1,6 @@
 from datetime import datetime, timezone
 from uuid import uuid4
-from sqlalchemy import UUID, Column, DateTime, String, JSON
+from sqlalchemy import UUID, Column, DateTime, String, JSON, Integer
 from sqlalchemy.orm import relationship
 
 from cognee.infrastructure.databases.relational import Base
@@ -20,6 +20,7 @@ class Data(Base):
     owner_id = Column(UUID, index=True)
     content_hash = Column(String)
     external_metadata = Column(JSON)
+    token_count = Column(Integer)
     created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
     updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
 

cognee/tasks/documents/extract_chunks_from_documents.py

@@ -17,9 +17,7 @@ async def extract_chunks_from_documents(
     """
     for document in documents:
         document_token_count = 0
-        for document_chunk in document.read(
-            chunk_size=chunk_size, chunker=chunker, max_tokens=max_tokens
-        ):
+        for document_chunk in document.read(chunk_size=chunk_size, chunker=chunker):
             document_token_count += document_chunk.token_count
             yield document_chunk
         document.token_count = document_token_count

cognee/tasks/ingestion/ingest_data.py

@@ -107,6 +107,7 @@ async def data_storing(data: Any, dataset_name: str, user: User):
                             owner_id=user.id,
                             content_hash=file_metadata["content_hash"],
                             external_metadata=get_external_metadata_dict(data_item),
+                            token_count=-1,
                         )
 
                     # Check if data is already in dataset