struc outputs support

parea-ai · Aug 7, 2024 · 33a5374 · 33a5374
1 parent dfea695
commit 33a5374
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/parea/utils/universal_encoder.py b/parea/utils/universal_encoder.py
@@ -87,8 +87,6 @@ def handle_openai_not_given(self, obj) -> Any:
         except ImportError:
             return None
 
-        from openai import NotGiven
-
         if isinstance(obj, NotGiven):
             return {"not_given": None}
         return None

diff --git a/parea/wrapper/openai/openai.py b/parea/wrapper/openai/openai.py
@@ -17,8 +17,6 @@
     from openai.util import convert_to_openai_object
 else:
     from openai.types.chat import ChatCompletion as OpenAIObject
-    from openai.types.chat import ParsedChatCompletion as OpenAIObjectParsed
-    from openai.types.chat import ParsedChatCompletionMessage
 
     def convert_to_openai_object(kwargs) -> OpenAIObject:
         if "id" not in kwargs:
@@ -57,9 +55,16 @@ def get_original_methods(self, module_client=openai):
             original_methods = {"ChatCompletion.create": module_client.ChatCompletion.create, "ChatCompletion.acreate": module_client.ChatCompletion.acreate}
         else:
             try:
-                original_methods = {"chat.completions.create": module_client.chat.completions.create, "beta.chat.completions.parse": module_client.beta.chat.completions.parse}
+                original_methods = {"chat.completions.create": module_client.chat.completions.create}
             except openai.OpenAIError:
                 original_methods = {}
+
+            try:
+                latest_methods = {"beta.chat.completions.parse": module_client.beta.chat.completions.parse}
+                original_methods.update(latest_methods)
+            except Exception:
+                pass
+
         return list(original_methods.keys())
 
     def init(self, log: Callable, cache: Cache = None, module_client=openai):
@@ -271,7 +276,13 @@ def _kwargs_to_llm_configuration(kwargs):
 
     @staticmethod
     def _get_output(result: Any, model: Optional[str] = None) -> str:
-        if not isinstance(result, (OpenAIObject, OpenAIObjectParsed)) and isinstance(result, dict):
+        try:
+            from openai.types.chat import ParsedChatCompletion, ParsedChatCompletionMessage
+        except ImportError:
+            ParsedChatCompletion = None
+            ParsedChatCompletionMessage = None
+
+        if not isinstance(result, (OpenAIObject, ParsedChatCompletion)) and isinstance(result, dict):
             result = convert_to_openai_object(
                 {
                     "choices": [