feat: Adding usage apis (#206)

* Adding usage apis

* change date to datetime

* optimise query

---------

Co-authored-by: dhirenmathur <dhiru17398@gmail.com>

app/alembic/versions/20241127095409_625f792419e7_support_for_repo_path.py

@@ -8,7 +8,6 @@
 
 from typing import Sequence, Union
 
-import sqlalchemy as sa
 from alembic import op
 
 # revision identifiers, used by Alembic.
@@ -19,7 +18,8 @@
 
 
 def upgrade() -> None:
-    op.execute('ALTER TABLE projects ADD COLUMN repo_path TEXT DEFAULT NULL')
+    op.execute("ALTER TABLE projects ADD COLUMN repo_path TEXT DEFAULT NULL")
+
 
 def downgrade() -> None:
     op.drop_column("projects", "repo_path")

app/main.py

@@ -26,6 +26,7 @@
 )
 from app.modules.projects.projects_router import router as projects_router
 from app.modules.search.search_router import router as search_router
+from app.modules.usage.usage_router import router as usage_router
 from app.modules.users.user_router import router as user_router
 from app.modules.users.user_service import UserService
 from app.modules.utils.firebase_setup import FirebaseSetup
@@ -112,6 +113,7 @@ def include_routers(self):
         self.app.include_router(agent_router, prefix="/api/v1", tags=["Agents"])
         self.app.include_router(provider_router, prefix="/api/v1", tags=["Providers"])
         self.app.include_router(tool_router, prefix="/api/v1", tags=["Tools"])
+        self.app.include_router(usage_router, prefix="/api/v1/usage", tags=["Usage"])
         if os.getenv("isDevelopmentMode") != "enabled":
             self.app.include_router(
                 secret_manager_router, prefix="/api/v1", tags=["Secret Manager"]

app/modules/intelligence/tools/change_detection/change_detection_tool.py

@@ -22,6 +22,7 @@
 from app.modules.projects.projects_service import ProjectService
 from app.modules.search.search_service import SearchService
 
+
 class ChangeDetectionInput(BaseModel):
     project_id: str = Field(
         ..., description="The ID of the project being evaluated, this is a UUID."

app/modules/usage/usage_controller.py

@@ -0,0 +1,9 @@
+from datetime import datetime
+from app.modules.usage.usage_service import UsageService
+from app.modules.usage.usage_schema import UsageResponse
+
+class UsageController:
+    @staticmethod
+    async def get_user_usage(start_date: datetime, end_date: datetime, user_id: str) -> UsageResponse:        
+        usage_data = await UsageService.get_usage_data(start_date, end_date, user_id)
+        return usage_data 

app/modules/usage/usage_router.py

@@ -0,0 +1,18 @@
+from fastapi import APIRouter, Depends
+from datetime import datetime
+from app.modules.auth.auth_service import AuthService
+from app.modules.usage.usage_controller import UsageController
+from app.modules.usage.usage_schema import UsageResponse
+
+router = APIRouter()
+
+class UsageAPI:
+    @staticmethod
+    @router.get("/usage", response_model=UsageResponse)
+    async def get_usage(
+        start_date: datetime, 
+        end_date: datetime, 
+        user=Depends(AuthService.check_auth),
+    ):
+        user_id = user["user_id"]
+        return await UsageController.get_user_usage(start_date, end_date, user_id)

app/modules/usage/usage_schema.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+from typing import Dict
+
+class UsageResponse(BaseModel):
+    total_human_messages: int
+    agent_message_counts: Dict[str, int] 

app/modules/usage/usage_service.py

@@ -0,0 +1,45 @@
+from datetime import datetime
+from fastapi import logger
+from sqlalchemy import func
+from sqlalchemy.orm import Session
+from sqlalchemy.exc import SQLAlchemyError
+from app.core.database import SessionLocal
+from app.modules.conversations.message.message_model import Message, MessageType
+from app.modules.conversations.conversation.conversation_model import Conversation
+
+
+class UsageService:
+    @staticmethod
+    async def get_usage_data(start_date: datetime, end_date: datetime, user_id: str):
+        try:
+            with SessionLocal() as session:
+                agent_query = (
+                    session.query(
+                        func.unnest(Conversation.agent_ids).label('agent_id'),
+                        func.count(Message.id).label('message_count')
+                    )
+                    .join(Message, Message.conversation_id == Conversation.id)
+                    .filter(
+                        Conversation.user_id == user_id,
+                        Message.created_at.between(start_date, end_date),
+                        Message.type == MessageType.HUMAN,
+                    )
+                    .group_by(func.unnest(Conversation.agent_ids))
+                    .all()
+                )
+
+                agent_message_counts = {
+                    agent_id: count 
+                    for agent_id, count in agent_query
+                }
+
+                total_human_messages = sum(agent_message_counts.values())
+
+                return {
+                    "total_human_messages": total_human_messages,
+                    "agent_message_counts": agent_message_counts
+                }
+
+        except SQLAlchemyError as e:
+            logger.error(f"Failed to fetch usage data: {e}")
+            raise Exception("Failed to fetch usage data") from e