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

@@ -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" },

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,

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")

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):