mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-15 01:21:48 +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
|
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
|
||||||
|
@ -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):
|
||||||
|
@ -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)
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user