advanced filtering on dataset collection

Author: jalexanderII

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,7 +55,7 @@ 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 = {}
         return list(original_methods.keys())
@@ -105,7 +103,7 @@ def resolver(self, trace_id: str, _args: Sequence[Any], kwargs: Dict[str, Any],
         trace_data.get()[trace_id].output_tokens = output_tokens
         trace_data.get()[trace_id].total_tokens = total_tokens
         trace_data.get()[trace_id].cost = _compute_cost(input_tokens, output_tokens, model)
-        trace_data.get()[trace_id].output = json_dumps(output) if not isinstance(output, str) else output
+        trace_data.get()[trace_id].output = output
         return response
 
     def gen_resolver(self, trace_id: str, _args: Sequence[Any], kwargs: Dict[str, Any], response, final_log):
@@ -271,7 +269,7 @@ 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):
+        if not isinstance(result, OpenAIObject) and isinstance(result, dict):
             result = convert_to_openai_object(
                 {
                     "choices": [
@@ -284,9 +282,7 @@ def _get_output(result: Any, model: Optional[str] = None) -> str:
                 }
             )
         response_message = result.choices[0].message
-        if isinstance(response_message, ParsedChatCompletionMessage):
-            completion = response_message.parsed.model_dump_json() if response_message.parsed else ""
-        elif not response_message.get("content", None) if is_old_openai else not response_message.content:
+        if not response_message.get("content", None) if is_old_openai else not response_message.content:
             completion = OpenAIWrapper._format_function_call(response_message)
         else:
             completion = response_message.content