create embeddings of ppt and store metadata of files in tables

migrations/versions/{c649dcba2de3_create_tables.py → 3292a4761b48_files_and_vector_records.py}

@@ -1,8 +1,8 @@
-"""create_tables
+"""files_and_vector_records
 
-Revision ID: c649dcba2de3
+Revision ID: 3292a4761b48
 Revises: 8652e8501339
-Create Date: 2024-10-25 14:28:36.463641
+Create Date: 2024-10-28 13:03:53.442309
 
 """
 from typing import Sequence, Union
@@ -12,7 +12,7 @@ import sqlalchemy as sa
 
 
 # revision identifiers, used by Alembic.
-revision: str = 'c649dcba2de3'
+revision: str = '3292a4761b48'
 down_revision: Union[str, None] = '8652e8501339'
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -20,11 +20,29 @@ depends_on: Union[str, Sequence[str], None] = None
 
 def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    pass
+    op.create_table('files',
+    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+    sa.Column('file_name', sa.String(), nullable=False),
+    sa.Column('uploaded_at', sa.DateTime(), nullable=False),
+    sa.Column('processed_at', sa.DateTime(), nullable=True),
+    sa.Column('slack_file_id', sa.String(), nullable=True),
+    sa.Column('user_id', sa.Integer(), nullable=False),
+    sa.Column('mime_type', sa.String(), nullable=False),
+    sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_table('vector_records',
+    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+    sa.Column('file_id', sa.Integer(), nullable=False),
+    sa.Column('record_id', sa.Integer(), nullable=False),
+    sa.ForeignKeyConstraint(['file_id'], ['files.id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    pass
+    op.drop_table('vector_records')
+    op.drop_table('files')
     # ### end Alembic commands ###

migrations/versions/37cddaae80e7_create_tables.py

@@ -1,30 +0,0 @@
-"""create_tables
-
-Revision ID: 37cddaae80e7
-Revises: c649dcba2de3
-Create Date: 2024-10-28 11:29:18.878826
-
-"""
-from typing import Sequence, Union
-
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision: str = '37cddaae80e7'
-down_revision: Union[str, None] = 'c649dcba2de3'
-branch_labels: Union[str, Sequence[str], None] = None
-depends_on: Union[str, Sequence[str], None] = None
-
-
-def upgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
-    pass
-    # ### end Alembic commands ###
-
-
-def downgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
-    pass
-    # ### end Alembic commands ###

src/controllers/_bot_controller.py

@@ -7,7 +7,7 @@ from src.utils import SlackClient
 from src.services import SlackBotService, SlackTeamService, UserService
 from src.models import ContentDeliveryFrequency
 from src.utils import logger
-
+from src.services import FileService
 
 class BotController:
     def __init__(self):
@@ -23,23 +23,68 @@ class BotController:
         try:
             if "challenge" in event:
                 return {"challenge": event["challenge"]}
+            
             if "event" in event:
                 actual_event = event["event"]
                 logger.info(f"Actual event: {actual_event}")
+                
                 if "file" in actual_event:
                     file = actual_event["file"]
                     file_id = file["id"]
                     logger.info(f"file id is {file_id}")
+                    
                     async with SlackClient() as slack_client:
+                        # Get file info
                         file_info = await slack_client.get_file_info(file_id)
                         logger.info(f"File info: {file_info}")
+                        
+                        # Get file URL
+                        file_url = file_info.get("url_private")
+                        file_name = file_info.get("name", "downloaded_file")
+                        
+                        if file_url:
+                            # Download the file content
+                            headers = {"Authorization": f"Bearer {slack_client.token}"}
+                            async with aiohttp.ClientSession() as session:
+                                async with session.get(file_url, headers=headers) as response:
+                                    if response.status == 200:
+                                        content = await response.read()
+                                        # Create UploadFile object from the content
+                                        upload_file = UploadFile(
+                                            filename=file_name,
+                                            file=BytesIO(content)
+                                        )
+                                        
+                                        # Process the file
+                                        async with self.user_repository as user_repo:
+                                            user = await user_repo.get_user_by_slack_id(user_slack_id)
+                                            if not user:
+                                                logger.error(f"User not found for slack_id: {user_slack_id}")
+                                                return {"message": "User not found"}
+
+                                        async with self.file_service as service:
+                                            result = await service.parse_document(
+                                                file=upload_file,
+                                                user_id=user.id,
+                                                slack_file_id=slack_file_id
+                                            )
+                                        logger.info(f"File processing result: {result}")
+                                    else:
+                                        logger.error(f"Failed to download file. Status: {response.status}")
+                                        raise HTTPException(
+                                            status_code=response.status,
+                                            detail="Failed to download file from Slack"
+                                        )
+                        
                         users = await slack_client.get_users_list()
                         logger.info(f"Users: {users}")
                 else:
                     logger.info("file not shared")
             else:
                 logger.info("event not found")
+                
             return {"message": "Event received successfully"}
+            
         except HTTPException as e:
             logger.warning(e)
             raise e

src/models/__init__.py

@@ -1,7 +1,9 @@
 from ._base import Base
 from ._user import User, ContentDeliveryFrequency
 from ._slack_team import SlackTeam
+from ._files import File
+from ._vector_records import VectorRecord
 
-__all__ = ["Base", "User", "ContentDeliveryFrequency", "SlackTeam"]
+__all__ = ["Base", "User", "ContentDeliveryFrequency", "SlackTeam", "File", "VectorRecord"]
 __version__ = "0.1.0"
 __author__ = "Kanha Upadhyay"

src/repositories/__init__.py

@@ -2,7 +2,9 @@ from ._base_repository import BaseRepository
 from ._config import DatabaseConfig
 from ._user_repository import UserRepository
 from ._slack_team_repository import SlackTeamRepository
+from ._file_repository import FileRepository
+from ._vector_record_repository import VectorRecordRepository
 
-__all__ = ["DatabaseConfig", "BaseRepository", "UserRepository", "SlackTeamRepository"]
+__all__ = ["DatabaseConfig", "BaseRepository", "UserRepository", "SlackTeamRepository", "FileRepository", "VectorRecordRepository"]
 __version__ = "0.1.0"
 __author__ = "Kanha Upadhyay"

src/repositories/_file_repository.py

@@ -0,0 +1,32 @@
+from sqlalchemy import select
+from datetime import datetime
+from src.models import File
+from ._base_repository import BaseRepository
+
+class FileRepository(BaseRepository):
+    def __init__(self):
+        super().__init__(model=File)
+    
+    async def create_file(self, file_name: str, user_id: int, mime_type: str, slack_file_id: str = None) -> File:
+        """Create a new file record"""
+        file = File(
+            file_name=file_name,
+            user_id=user_id,
+            mime_type=mime_type,
+            slack_file_id=slack_file_id,
+            uploaded_at=datetime.utcnow()
+        )
+        return await self.add(file)
+    
+    async def mark_as_processed(self, file_id: int) -> None:
+        """Mark file as processed"""
+        file = await self.get_by_id(file_id)
+        if file:
+            file.processed_at = datetime.utcnow()
+            await self.update(file)
+    
+    async def get_by_slack_file_id(self, slack_file_id: str) -> File:
+        """Get file by Slack file ID"""
+        query = select(File).where(File.slack_file_id == slack_file_id)
+        result = await self.execute(query)
+        return result.scalar_one_or_none()

src/repositories/_vector_record_repository.py

@@ -0,0 +1,21 @@
+from src.models import VectorRecord
+from ._base_repository import BaseRepository
+from typing import List
+
+class VectorRecordRepository(BaseRepository):
+    def __init__(self):
+        super().__init__(model=VectorRecord)
+    
+    async def create_records(self, file_id: int, record_ids: List[int]) -> List[VectorRecord]:
+        """Create multiple vector records for a file"""
+        records = [
+            VectorRecord(file_id=file_id, record_id=record_id)
+            for record_id in record_ids
+        ]
+        return await self.add_all(records)
+    
+    async def get_by_file_id(self, file_id: int) -> List[VectorRecord]:
+        """Get all vector records for a file"""
+        query = select(VectorRecord).where(VectorRecord.file_id == file_id)
+        result = await self.execute(query)
+        return result.scalars().all()

src/services/__init__.py

@@ -4,6 +4,6 @@ from ._slack_bot_service import SlackBotService
 from ._slack_team_service import SlackTeamService
 from ._file_service import FileService
 
-__all__ = ["ContentDeliveryService", "UserService", "SlackBotService", "SlackTeamService","FileService"]
+__all__ = ["ContentDeliveryService", "UserService", "SlackBotService", "SlackTeamService", "FileService"]
 __version__ = "0.1.0"
 __author__ = "Ashok Bhati"

src/services/_file_service.py

@@ -36,18 +36,75 @@ class FileService:
             shutil.copyfileobj(file.file, buffer)
         return file_path
 
-    async def parse_document(self, file: UploadFile) -> Dict:
+    async def filter_ppt_content(self, parsed_content: Tuple[Tuple[int, str], ...]) -> Tuple[Tuple[int, str], ...]:
+        """
+        Filter PowerPoint content based on specific target phrases.
+
+        Args:
+            parsed_content: Tuple of (page_number, content) pairs
+
+        Returns:
+            Filtered tuple of (page_number, content) pairs
+        """
+        target_phrases = [
+            "My Strategy for Managing Myself",
+            "My Strategy For Active Listening",
+            "My Strategy for Managing Others",
+            "What It Concretely Looks Like",
+            "My People",
+            "My Philosophy",
+            "My Preferences",
+            "My Priorities for Q1 2025"
+        ]
+
+        filtered_content = []
+
+        for page_num, content in parsed_content:
+            content_lower = content.lower()
+            for phrase in target_phrases:
+                if phrase.lower() in content_lower:
+                    filtered_content.append((page_num, content))
+                    break  # Move to next page once we find a matching phrase
+
+        logger.info(f"Total pages before filtering: {len(parsed_content)}")
+        logger.info(f"Total pages after filtering: {len(filtered_content)}")
+
+        return tuple(filtered_content)
+
+    async def parse_document(self, file: UploadFile, user_id: int, slack_file_id: str = None) -> Dict:
         try:
             file_path = await self.save_uploaded_file(file)
             
+            # Create file record in database
+            mime_type = file.content_type or f"application/{file.filename.split('.')[-1]}"
+            db_file = await self.file_repository.create_file(
+                file_name=file.filename,
+                user_id=user_id,
+                mime_type=mime_type,
+                slack_file_id=slack_file_id
+            )
+            
             parsed_content = await parse_file_content(str(file_path))
+            
+            file_extension = file.filename.lower().split('.')[-1]
+            if file_extension in ['ppt', 'pptx']:
+                parsed_content = await self.filter_ppt_content(parsed_content)
 
             documents = create_documents(parsed_content, file.filename)
             
-            await store_documents(documents)
+            # Store documents in vector store and get record IDs
+            record_ids = await store_documents(documents)
+            
+            # Store vector record mappings in database
+            if record_ids:
+                await self.vector_repository.create_records(db_file.id, record_ids)
+            
+            # Mark file as processed
+            await self.file_repository.mark_as_processed(db_file.id)
 
             return {
                 "message": "Document's embeddings created successfully",
+                "file_id": db_file.id
             }
         except Exception as e:
             logger.error(f"Error in parse_document: {str(e)}")