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> | 🚀 This PR
description was created by [Ellipsis](https://www.ellipsis.dev) for
commit ee5fdf70f6. |
|--------|--------|

### 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 ❤️ by [ellipsis.dev](https://www.ellipsis.dev)

<!--
ELLIPSIS_HIDDEN
-->
This commit is contained in:
Stan Girard 2024-04-16 09:14:38 -07:00 committed by GitHub
parent c4106dbb86
commit 1931848262
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 49 additions and 1 deletions

View File

@ -82,9 +82,15 @@ class Outputs(BaseModel):
brain: Optional[OutputBrain] = None brain: Optional[OutputBrain] = None
class Pricing(BaseModel):
cost: int = 20
description: str = "Credits per use"
class AssistantOutput(BaseModel): class AssistantOutput(BaseModel):
name: str name: str
description: str description: str
pricing: Optional[Pricing] = Pricing()
tags: Optional[List[str]] = [] tags: Optional[List[str]] = []
input_description: str input_description: str
output_description: str output_description: str

View File

@ -2,11 +2,13 @@ import random
from abc import abstractmethod from abc import abstractmethod
from io import BytesIO from io import BytesIO
from tempfile import NamedTemporaryFile from tempfile import NamedTemporaryFile
from typing import List from typing import List, Optional
from fastapi import UploadFile from fastapi import UploadFile
from logger import get_logger from logger import get_logger
from models.user_usage import UserUsage
from modules.assistant.dto.inputs import InputAssistant 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.contact_support.controller.settings import ContactsSettings
from modules.upload.controller.upload_routes import upload_file from modules.upload.controller.upload_routes import upload_file
from modules.user.entity.user_identity import UserIdentity from modules.user.entity.user_identity import UserIdentity
@ -20,6 +22,40 @@ class ITO(BaseModel):
input: InputAssistant input: InputAssistant
files: List[UploadFile] files: List[UploadFile]
current_user: UserIdentity 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 @abstractmethod
async def process_assistant(self): async def process_assistant(self):

View File

@ -65,6 +65,12 @@ class SummaryAssistant(ITO):
async def process_assistant(self): 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 # 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) tmp_file = tempfile.NamedTemporaryFile(delete=False)