From 9755993d7d32cab04abc6c61d6b36d40dcd8ef5a Mon Sep 17 00:00:00 2001
From: noble-varghese <noblekvarghese96@gmail.com>
Date: Fri, 1 Dec 2023 19:25:01 +0530
Subject: [PATCH] fix: updated feedbacks response, added streaming in post and
 added over ride function in utils

---
 portkey_ai/api_resources/apis/feedback.py |  3 +-
 portkey_ai/api_resources/apis/post.py     | 35 +++++++++++++++++++++--
 portkey_ai/api_resources/utils.py         | 18 ++++++------
 3 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/portkey_ai/api_resources/apis/feedback.py b/portkey_ai/api_resources/apis/feedback.py
index 16325348..effaf048 100644
--- a/portkey_ai/api_resources/apis/feedback.py
+++ b/portkey_ai/api_resources/apis/feedback.py
@@ -17,8 +17,7 @@ def create(
         weight: Optional[float] = None,
         metadata: Optional[Dict[str, Any]] = None
     ) -> None:
-        body = dict(trace_id=trace_id, value=value,
-                    weight=weight, metadata=metadata)
+        body = dict(trace_id=trace_id, value=value, weight=weight, metadata=metadata)
         return self._post(
             PortkeyApiPaths.FEEDBACK_API,
             body=body,
diff --git a/portkey_ai/api_resources/apis/post.py b/portkey_ai/api_resources/apis/post.py
index 153a68c4..743fe6e2 100644
--- a/portkey_ai/api_resources/apis/post.py
+++ b/portkey_ai/api_resources/apis/post.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, Union
+from typing import Any, Dict, Union, overload, Literal
 
 from portkey_ai.api_resources.base_client import APIClient
 
@@ -10,10 +10,41 @@ class Post(APIResource):
     def __init__(self, client: APIClient) -> None:
         super().__init__(client)
 
+    @overload
     def create(
         self,
         *,
         url: str,
+        stream: Literal[True],
+        **kwargs,
+    ) -> Stream[Dict[str, Any]]:
+        ...
+
+    @overload
+    def create(
+        self,
+        *,
+        url: str,
+        stream: Literal[False] = False,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        ...
+
+    @overload
+    def create(
+        self,
+        *,
+        url: str,
+        stream: bool = False,
+        **kwargs,
+    ) -> Union[Dict[str, Any], Stream[Dict[str, Any]]]:
+        ...
+
+    def create(
+        self,
+        *,
+        url: str,
+        stream: bool = False,
         **kwargs,
     ) -> Union[Dict[str, Any], Stream[Dict[str, Any]]]:
         return self._post(
@@ -22,6 +53,6 @@ def create(
             params=None,
             cast_to=dict,
             stream_cls=Stream[dict],
-            stream=False,
+            stream=stream,
             headers={},
         )
diff --git a/portkey_ai/api_resources/utils.py b/portkey_ai/api_resources/utils.py
index db74e1ed..732c3ded 100644
--- a/portkey_ai/api_resources/utils.py
+++ b/portkey_ai/api_resources/utils.py
@@ -258,7 +258,7 @@ def __str__(self):
 
 
 # Models for Chat Stream
-class Delta(BaseModel):
+class Delta(BaseModel, extra="allow"):
     role: Optional[str] = None
     content: Optional[str] = ""
 
@@ -272,7 +272,7 @@ def get(self, key: str, default: Optional[Any] = None):
         return getattr(self, key, None) or default
 
 
-class StreamChoice(BaseModel):
+class StreamChoice(BaseModel, extra="allow"):
     index: Optional[int] = None
     delta: Union[Delta, Dict[Any, Any]] = {}
     finish_reason: Optional[str] = None
@@ -287,7 +287,7 @@ def __getitem__(self, key):
         return getattr(self, key, None)
 
 
-class ChatCompletionChunk(BaseModel):
+class ChatCompletionChunk(BaseModel, extra="allow"):
     id: Optional[str] = None
     object: Optional[str] = None
     created: Optional[int] = None
@@ -305,7 +305,7 @@ def get(self, key: str, default: Optional[Any] = None):
 
 
 # Models for Chat Non-stream
-class ChatChoice(BaseModel):
+class ChatChoice(BaseModel, extra="allow"):
     index: Optional[int] = None
     message: Optional[Message] = None
     finish_reason: Optional[str] = None
@@ -320,7 +320,7 @@ def get(self, key: str, default: Optional[Any] = None):
         return getattr(self, key, None) or default
 
 
-class Usage(BaseModel):
+class Usage(BaseModel, extra="allow"):
     prompt_tokens: Optional[int] = None
     completion_tokens: Optional[int] = None
     total_tokens: Optional[int] = None
@@ -335,7 +335,7 @@ def get(self, key: str, default: Optional[Any] = None):
         return getattr(self, key, None) or default
 
 
-class ChatCompletions(BaseModel):
+class ChatCompletions(BaseModel, extra="allow"):
     id: Optional[str] = None
     object: Optional[str] = None
     created: Optional[int] = None
@@ -354,7 +354,7 @@ def get(self, key: str, default: Optional[Any] = None):
 
 
 # Models for text completion Non-stream
-class TextChoice(BaseModel):
+class TextChoice(BaseModel, extra="allow"):
     index: Optional[int] = None
     text: Optional[str] = None
     logprobs: Any
@@ -370,7 +370,7 @@ def get(self, key: str, default: Optional[Any] = None):
         return getattr(self, key, None) or default
 
 
-class TextCompletion(BaseModel):
+class TextCompletion(BaseModel, extra="allow"):
     id: Optional[str] = None
     object: Optional[str] = None
     created: Optional[int] = None
@@ -389,7 +389,7 @@ def get(self, key: str, default: Optional[Any] = None):
 
 
 # Models for text completion stream
-class TextCompletionChunk(BaseModel):
+class TextCompletionChunk(BaseModel, extra="allow"):
     id: Optional[str] = None
     object: Optional[str] = None
     created: Optional[int] = None