mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-14 17:03:29 +03:00
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:
parent
c4106dbb86
commit
1931848262
@ -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
|
||||
|
@ -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):
|
||||
|
@ -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)
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user