From ef2363cd3ab67bd06011ac8a87bb94db19beb45a Mon Sep 17 00:00:00 2001 From: drbh Date: Tue, 16 Apr 2024 15:53:27 +0000 Subject: [PATCH] fix: refactor tests to support completions snapshot --- clients/python/text_generation/types.py | 21 + integration-tests/conftest.py | 22 +- ..._llama_completion_many_prompts_stream.json | 603 +++++++++++++++++- .../models/test_completion_prompts.py | 9 +- 4 files changed, 647 insertions(+), 8 deletions(-) diff --git a/clients/python/text_generation/types.py b/clients/python/text_generation/types.py index deb987c5..5dab1f80 100644 --- a/clients/python/text_generation/types.py +++ b/clients/python/text_generation/types.py @@ -59,6 +59,17 @@ class ChatCompletionComplete(BaseModel): usage: Optional[Any] = None +class CompletionComplete(BaseModel): + # Index of the chat completion + index: int + # Message associated with the chat completion + text: str + # Log probabilities for the chat completion + logprobs: Optional[Any] + # Reason for completion + finish_reason: str + + class Function(BaseModel): name: Optional[str] arguments: str @@ -104,6 +115,16 @@ class ChatComplete(BaseModel): usage: Any +class CompletionComplete(BaseModel): + # Completion details + id: str + object: str + created: int + model: str + system_fingerprint: str + choices: List[CompletionComplete] + + class ChatRequest(BaseModel): # Model identifier model: str diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index e8ce0d84..616446ba 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -9,6 +9,7 @@ import json import math import time import random +import re from docker.errors import NotFound from typing import Optional, List, Dict @@ -26,6 +27,7 @@ from text_generation.types import ( ChatComplete, ChatCompletionChunk, ChatCompletionComplete, + CompletionComplete, ) DOCKER_IMAGE = os.getenv("DOCKER_IMAGE", None) @@ -69,12 +71,12 @@ class ResponseComparator(JSONSnapshotExtension): data = json.loads(data) if isinstance(data, Dict) and "choices" in data: choices = data["choices"] - if ( - isinstance(choices, List) - and len(choices) >= 1 - and "delta" in choices[0] - ): - return ChatCompletionChunk(**data) + print(choices) + if isinstance(choices, List) and len(choices) >= 1: + if "delta" in choices[0]: + return ChatCompletionChunk(**data) + if "text" in choices[0]: + return CompletionComplete(**data) return ChatComplete(**data) if isinstance(data, Dict): @@ -161,6 +163,9 @@ class ResponseComparator(JSONSnapshotExtension): ) ) + def eq_completion(response: ChatComplete, other: ChatComplete) -> bool: + return response.choices[0].text == other.choices[0].text + def eq_chat_complete(response: ChatComplete, other: ChatComplete) -> bool: return ( response.choices[0].message.content == other.choices[0].message.content @@ -184,6 +189,11 @@ class ResponseComparator(JSONSnapshotExtension): if not isinstance(snapshot_data, List): snapshot_data = [snapshot_data] + if isinstance(serialized_data[0], CompletionComplete): + return len(snapshot_data) == len(serialized_data) and all( + [eq_completion(r, o) for r, o in zip(serialized_data, snapshot_data)] + ) + if isinstance(serialized_data[0], ChatComplete): return len(snapshot_data) == len(serialized_data) and all( [eq_chat_complete(r, o) for r, o in zip(serialized_data, snapshot_data)] diff --git a/integration-tests/models/__snapshots__/test_completion_prompts/test_flash_llama_completion_many_prompts_stream.json b/integration-tests/models/__snapshots__/test_completion_prompts/test_flash_llama_completion_many_prompts_stream.json index b0850cc9..b3e20c30 100644 --- a/integration-tests/models/__snapshots__/test_completion_prompts/test_flash_llama_completion_many_prompts_stream.json +++ b/integration-tests/models/__snapshots__/test_completion_prompts/test_flash_llama_completion_many_prompts_stream.json @@ -1 +1,602 @@ -"\n\n" +[ + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": "\n" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": "\n" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": "\n" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": "hd" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": "2" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": "2" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": "2" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": "2" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": "0" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": "." + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": "." + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": "\n" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": "." + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": " And" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": " And" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": "def" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": " What" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": " away" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": " what" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": " mai" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": " is" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": " we" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": " is" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": "hd" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": " the" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": " go" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": " the" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": "\n" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": " temperature" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": "," + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": " capital" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": "\t" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": " in" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": " sc" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": " of" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": "p" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 0, + "logprobs": null, + "text": " New" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 1, + "logprobs": null, + "text": "rap" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 2, + "logprobs": null, + "text": " Ukraine" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + }, + { + "choices": [ + { + "finish_reason": "", + "index": 3, + "logprobs": null, + "text": " text" + } + ], + "created": 1713282708, + "id": "", + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "object": "text_completion", + "system_fingerprint": "1.4.5-native" + } +] diff --git a/integration-tests/models/test_completion_prompts.py b/integration-tests/models/test_completion_prompts.py index 496bd062..1d7af0a4 100644 --- a/integration-tests/models/test_completion_prompts.py +++ b/integration-tests/models/test_completion_prompts.py @@ -3,6 +3,10 @@ import requests import json from aiohttp import ClientSession +from text_generation.types import ( + CompletionComplete, +) + @pytest.fixture(scope="module") def flash_llama_completion_handle(launcher): @@ -82,6 +86,7 @@ async def test_flash_llama_completion_many_prompts_stream( url = f"{flash_llama_completion.base_url}/v1/completions" + chunks = [] async with ClientSession(headers=flash_llama_completion.headers) as session: async with session.post(url, json=request) as response: # iterate over the stream @@ -96,8 +101,10 @@ async def test_flash_llama_completion_many_prompts_stream( chunk = [json.loads(c) for c in chunk] for c in chunk: + print(c) + chunks.append(CompletionComplete(**c)) assert "choices" in c assert 0 <= c["choices"][0]["index"] <= 4 assert response.status == 200 - assert response == response_snapshot + assert chunks == response_snapshot