From 193184826203f195a27ed31738a69017b3b7e340 Mon Sep 17 00:00:00 2001
From: Stan Girard <girard.stanislas@gmail.com>
Date: Tue, 16 Apr 2024 09:14:38 -0700
Subject: [PATCH] feat(assistants): Add user usage update and pricing
 calculation to ITO assistant (#2433)

This pull request adds functionality to update user usage and calculate
pricing in the ITO assistant. It includes a new method
`increase_usage_user()` that raises an error if the user has consumed
all of their credits, and a new method `calculate_pricing()` that
returns a fixed pricing value of 20.

<!--
ELLIPSIS_HIDDEN
-->
----

| <a href="https://ellipsis.dev" target="_blank"><img
src="https://avatars.githubusercontent.com/u/80834858?s=400&u=31e596315b0d8f7465b3ee670f25cea677299c96&v=4"
alt="Ellipsis" width="30px" height="30px"/></a> | :rocket: This PR
description was created by [Ellipsis](https://www.ellipsis.dev) for
commit ee5fdf70f6670b30edbf9c782f4fa97dcd36eed3. |
|--------|--------|

### Summary:
This PR adds functionality to update user usage and calculate pricing in
the ITO assistant, with new methods in the `ITO` class and an update to
where the `increase_usage_user()` method is called.

**Key points**:
- Added `increase_usage_user()` and `calculate_pricing()` methods to
`ITO` class in `/backend/modules/assistant/ito/ito.py`
- `increase_usage_user()` updates user usage and raises an error if all
credits are consumed
- `calculate_pricing()` returns a fixed pricing value of 20
- `increase_usage_user()` is now called in the `ITO` class constructor


----
Generated with :heart: by [ellipsis.dev](https://www.ellipsis.dev)

<!--
ELLIPSIS_HIDDEN
-->
---
 backend/modules/assistant/dto/outputs.py |  6 ++++
 backend/modules/assistant/ito/ito.py     | 38 +++++++++++++++++++++++-
 backend/modules/assistant/ito/summary.py |  6 ++++
 3 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/backend/modules/assistant/dto/outputs.py b/backend/modules/assistant/dto/outputs.py
index 3254e030b..cf6398ad8 100644
--- a/backend/modules/assistant/dto/outputs.py
+++ b/backend/modules/assistant/dto/outputs.py
@@ -82,9 +82,15 @@ class Outputs(BaseModel):
     brain: Optional[OutputBrain] = None
 
 
+class Pricing(BaseModel):
+    cost: int = 20
+    description: str = "Credits per use"
+
+
 class AssistantOutput(BaseModel):
     name: str
     description: str
+    pricing: Optional[Pricing] = Pricing()
     tags: Optional[List[str]] = []
     input_description: str
     output_description: str
diff --git a/backend/modules/assistant/ito/ito.py b/backend/modules/assistant/ito/ito.py
index 44611fb80..e2658718a 100644
--- a/backend/modules/assistant/ito/ito.py
+++ b/backend/modules/assistant/ito/ito.py
@@ -2,11 +2,13 @@ import random
 from abc import abstractmethod
 from io import BytesIO
 from tempfile import NamedTemporaryFile
-from typing import List
+from typing import List, Optional
 
 from fastapi import UploadFile
 from logger import get_logger
+from models.user_usage import UserUsage
 from modules.assistant.dto.inputs import InputAssistant
+from modules.chat.controller.chat.utils import update_user_usage
 from modules.contact_support.controller.settings import ContactsSettings
 from modules.upload.controller.upload_routes import upload_file
 from modules.user.entity.user_identity import UserIdentity
@@ -20,6 +22,40 @@ class ITO(BaseModel):
     input: InputAssistant
     files: List[UploadFile]
     current_user: UserIdentity
+    user_usage: Optional[UserUsage] = None
+    user_settings: Optional[dict] = None
+
+    def __init__(
+        self,
+        input: InputAssistant,
+        files: List[UploadFile] = None,
+        current_user: UserIdentity = None,
+        **kwargs,
+    ):
+        super().__init__(
+            input=input,
+            files=files,
+            current_user=current_user,
+            **kwargs,
+        )
+        self.user_usage = UserUsage(
+            id=current_user.id,
+            email=current_user.email,
+        )
+        self.user_settings = self.user_usage.get_user_settings()
+        self.increase_usage_user()
+
+    def increase_usage_user(self):
+        # Raises an error if the user has consumed all of of his credits
+
+        update_user_usage(
+            usage=self.user_usage,
+            user_settings=self.user_settings,
+            cost=self.calculate_pricing(),
+        )
+
+    def calculate_pricing(self):
+        return 20
 
     @abstractmethod
     async def process_assistant(self):
diff --git a/backend/modules/assistant/ito/summary.py b/backend/modules/assistant/ito/summary.py
index 6bfb49706..a1da5ff6b 100644
--- a/backend/modules/assistant/ito/summary.py
+++ b/backend/modules/assistant/ito/summary.py
@@ -65,6 +65,12 @@ class SummaryAssistant(ITO):
 
     async def process_assistant(self):
 
+        try:
+            self.increase_usage_user()
+        except Exception as e:
+            logger.error(f"Error increasing usage: {e}")
+            return {"error": str(e)}
+
         # Create a temporary file with the uploaded file as a temporary file and then pass it to the loader
         tmp_file = tempfile.NamedTemporaryFile(delete=False)