From 69c168d1e0b8e8e99d8154cbeee9d70879fe30a7 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Fri, 6 Sep 2024 18:50:52 +0200 Subject: [PATCH] Fix parsing --- integration-tests/conftest.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index 22a61a7c..94143058 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -102,15 +102,16 @@ class ResponseComparator(JSONSnapshotExtension): def _convert_data(data): if isinstance(data, Dict): - choices = data["choices"] - if isinstance(choices, List) and len(choices) >= 1: - if "delta" in choices[0]: - return ChatCompletionChunk(**data) - if "text" in choices[0]: - return Completion(**data) - return ChatComplete(**data) - if isinstance(data, Dict): - return Response(**data) + if "choices" in data: + choices = data["choices"] + if isinstance(choices, List) and len(choices) >= 1: + if "delta" in choices[0]: + return ChatCompletionChunk(**data) + if "text" in choices[0]: + return Completion(**data) + return ChatComplete(**data) + else: + return Response(**data) if isinstance(data, List): return [_convert_data(d) for d in data] raise NotImplementedError