add usage message

capeprivacy · Sep 26, 2023 · ee288a5 · ee288a5
1 parent e4a36de
commit ee288a5
Showing 1 changed file with 15 additions and 4 deletions.
diff --git a/pycape/llms/llms.py b/pycape/llms/llms.py
@@ -113,11 +113,16 @@ async def completions(
 
         async for msg in self.ctx.websocket:
             msg = WSMessage.model_validate_json(msg)
-            if msg.msg_type != WSMessageType.STREAM_CHUNK:
+            if msg.msg_type not in [WSMessageType.STREAM_CHUNK, WSMessageType.USAGE]:
                 raise Exception(
-                    f"expected {WSMessageType.STREAM_CHUNK} not {msg.msg_type}"
+                    f"expected {WSMessageType.STREAM_CHUNK} or "
+                    f"{WSMessageType.USAGE} not {msg.msg_type}"
                 )
 
+            if msg.msg_type == WSMessageType.USAGE:
+                yield msg.data
+                continue
+
             dec = crypto.aes_decrypt(
                 base64.b64decode(msg.data["data"].encode()), aes_key
             )
@@ -155,11 +160,16 @@ async def chat_completions(
 
         async for msg in self.ctx.websocket:
             msg = WSMessage.model_validate_json(msg)
-            if msg.msg_type != WSMessageType.STREAM_CHUNK:
+            if msg.msg_type not in [WSMessageType.STREAM_CHUNK, WSMessageType.USAGE]:
                 raise Exception(
-                    f"expected {WSMessageType.STREAM_CHUNK} not {msg.msg_type}"
+                    f"expected {WSMessageType.STREAM_CHUNK} or "
+                    f"{WSMessageType.USAGE} not {msg.msg_type}"
                 )
 
+            if msg.msg_type == WSMessageType.USAGE:
+                yield msg.data
+                continue
+
             dec = crypto.aes_decrypt(
                 base64.b64decode(msg.data["data"].encode()), aes_key
             )
@@ -195,6 +205,7 @@ class WSMessageType(str, Enum):
     STREAM_CHUNK = "stream_chunk"
     CHAT_COMPLETIONS_REQUEST = "chat_completion_request"
     COMPLETIONS_REQUEST = "completions_request"
+    USAGE = "usage"
 
 
 class WSMessage(BaseModel):