From feea5528db61cc577cd37ad5ec8f8e61432456df Mon Sep 17 00:00:00 2001 From: fhennerkes <162764597+fhennerkes@users.noreply.github.com> Date: Fri, 13 Dec 2024 14:23:12 -0800 Subject: [PATCH] variable pricing: Add support for variable priced server bots (#129) * variable pricing: Add support for variable priced server bots - Added authorize and capture functions - Added custom_rate_card to the bot settings - Updated version to 0.0.50 --- pyproject.toml | 2 +- src/fastapi_poe/__init__.py | 6 +- src/fastapi_poe/base.py | 118 +++++++++++++++++++++++++++++++++++- src/fastapi_poe/types.py | 17 +++++- 4 files changed, 138 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index aaf0b92..7dac22a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "fastapi_poe" -version = "0.0.49" +version = "0.0.50" authors = [ { name="Lida Li", email="lli@quora.com" }, { name="Jelle Zijlstra", email="jelle@quora.com" }, diff --git a/src/fastapi_poe/__init__.py b/src/fastapi_poe/__init__.py index 3e3b9ba..13373fd 100644 --- a/src/fastapi_poe/__init__.py +++ b/src/fastapi_poe/__init__.py @@ -24,9 +24,12 @@ "ToolResultDefinition", "MessageFeedback", "sync_bot_settings", + "CostItem", + "InsufficientFundError", + "CostRequestError", ] -from .base import PoeBot, make_app, run +from .base import CostRequestError, InsufficientFundError, PoeBot, make_app, run from .client import ( BotError, BotErrorNoRetry, @@ -37,6 +40,7 @@ ) from .types import ( Attachment, + CostItem, ErrorResponse, MessageFeedback, MetaResponse, diff --git a/src/fastapi_poe/base.py b/src/fastapi_poe/base.py index aae2fc1..1e35cea 100644 --- a/src/fastapi_poe/base.py +++ b/src/fastapi_poe/base.py @@ -14,12 +14,14 @@ BinaryIO, Callable, Dict, + List, Optional, Sequence, Union, ) import httpx +import httpx_sse from fastapi import Depends, FastAPI, HTTPException, Request, Response from fastapi.exceptions import RequestValidationError from fastapi.responses import HTMLResponse, JSONResponse @@ -38,6 +40,7 @@ from fastapi_poe.types import ( AttachmentUploadResponse, ContentType, + CostItem, ErrorResponse, Identifier, MetaResponse, @@ -63,6 +66,14 @@ class AttachmentUploadError(Exception): pass +class CostRequestError(Exception): + pass + + +class InsufficientFundError(Exception): + pass + + class LoggingMiddleware(BaseHTTPMiddleware): async def set_body(self, request: Request) -> None: receive_ = await request._receive() @@ -182,8 +193,11 @@ async def get_response_with_context( response to the Poe servers. This is what gets displayed to the user. """ - async for event in self.get_response(request): - yield event + try: + async for event in self.get_response(request): + yield event + except InsufficientFundError: + yield ErrorResponse(error_type="insufficient_fund", text="") async def get_settings(self, setting: SettingsRequest) -> SettingsResponse: """ @@ -629,6 +643,106 @@ def make_prompt_author_role_alternated( return new_messages + async def capture_cost( + self, + request: QueryRequest, + amounts: Union[List[CostItem], CostItem], + base_url: str = "https://api.poe.com/", + ) -> None: + """ + + Used to capture variable costs for monetized and eligible bot creators. + Visit https://creator.poe.com/docs/creator-monetization for more information. + + #### Parameters: + - `request` (`QueryRequest`): The currently handlded QueryRequest object. + - `amounts` (`Union[List[CostItem], CostItem]`): The to be captured amounts. + + """ + + if not self.access_key: + raise CostRequestError( + "Please provide the bot access_key when make_app is called." + ) + + if not request.bot_query_id: + raise InvalidParameterError( + "bot_query_id is required to make cost requests." + ) + + url = f"{base_url}bot/cost/{request.bot_query_id}/capture" + result = await self._cost_requests_inner( + amounts=amounts, access_key=self.access_key, url=url + ) + if not result: + raise InsufficientFundError() + + async def authorize_cost( + self, + request: QueryRequest, + amounts: Union[List[CostItem], CostItem], + base_url: str = "https://api.poe.com/", + ) -> None: + """ + + Used to authorize a cost for monetized and eligible bot creators. + Visit https://creator.poe.com/docs/creator-monetization for more information. + + #### Parameters: + - `request` (`QueryRequest`): The currently handlded QueryRequest object. + - `amounts` (`Union[List[CostItem], CostItem]`): The to be authorized amounts. + + """ + + if not self.access_key: + raise CostRequestError( + "Please provide the bot access_key when make_app is called." + ) + + if not request.bot_query_id: + raise InvalidParameterError( + "bot_query_id is required to make cost requests." + ) + + url = f"{base_url}bot/cost/{request.bot_query_id}/authorize" + result = await self._cost_requests_inner( + amounts=amounts, access_key=self.access_key, url=url + ) + if not result: + raise InsufficientFundError() + + async def _cost_requests_inner( + self, amounts: Union[List[CostItem], CostItem], access_key: str, url: str + ) -> bool: + amounts = [amounts] if isinstance(amounts, CostItem) else amounts + amounts = [amount.model_dump() for amount in amounts] + data = {"amounts": amounts, "access_key": access_key} + try: + async with httpx.AsyncClient(timeout=300) as client, httpx_sse.aconnect_sse( + client, method="POST", url=url, json=data + ) as event_source: + if event_source.response.status_code != 200: + error_pieces = [ + json.loads(event.data).get("message", "") + async for event in event_source.aiter_sse() + ] + raise CostRequestError( + f"{event_source.response.status_code} " + f"{event_source.response.reason_phrase}: {''.join(error_pieces)}" + ) + + async for event in event_source.aiter_sse(): + if event.event == "result": + event_data = json.loads(event.data) + result = event_data["status"] + return result == "success" + return False + except httpx.HTTPError: + logger.error( + "An HTTP error occurred when attempting to send a cost request." + ) + raise + @staticmethod def text_event(text: str) -> ServerSentEvent: return ServerSentEvent(data=json.dumps({"text": text}), event="text") diff --git a/src/fastapi_poe/types.py b/src/fastapi_poe/types.py index 1e45eaa..c4a8b14 100644 --- a/src/fastapi_poe/types.py +++ b/src/fastapi_poe/types.py @@ -7,7 +7,7 @@ Identifier: TypeAlias = str FeedbackType: TypeAlias = Literal["like", "dislike"] ContentType: TypeAlias = Literal["text/markdown", "text/plain"] -ErrorType: TypeAlias = Literal["user_message_too_long"] +ErrorType: TypeAlias = Literal["user_message_too_long", "insufficient_fund"] class MessageFeedback(BaseModel): @@ -24,6 +24,20 @@ class MessageFeedback(BaseModel): reason: Optional[str] +class CostItem(BaseModel): + """ + + An object representing a cost item used for authorization and charge request. + #### Fields: + - `amount_usd_milli_cents` (`int`) + - `description` (`str`) + + """ + + amount_usd_milli_cents: int + description: Optional[str] = None + + class Attachment(BaseModel): """ @@ -215,6 +229,7 @@ class SettingsResponse(BaseModel): enable_image_comprehension: Optional[bool] = None enforce_author_role_alternation: Optional[bool] = None enable_multi_bot_chat_prompting: Optional[bool] = None + custom_rate_card: Optional[str] = None class AttachmentUploadResponse(BaseModel):