From fe77a5d11ca927f6aa11d54f6f4619d89b026d0e Mon Sep 17 00:00:00 2001 From: Dttbd Date: Tue, 14 May 2024 11:05:22 +0800 Subject: [PATCH 1/2] feat: add models object property --- taskingai/client/models/entities/action.py | 1 + taskingai/client/models/entities/assistant.py | 1 + taskingai/client/models/entities/chat.py | 1 + taskingai/client/models/entities/chat_completion.py | 1 + taskingai/client/models/entities/chat_completion_chunk.py | 1 + taskingai/client/models/entities/chunk.py | 1 + taskingai/client/models/entities/collection.py | 1 + taskingai/client/models/entities/message.py | 1 + taskingai/client/models/entities/message_chunk.py | 1 + taskingai/client/models/entities/message_generation_log.py | 1 + taskingai/client/models/entities/record.py | 1 + taskingai/client/models/schemas/chat_completion_request.py | 2 +- 12 files changed, 12 insertions(+), 1 deletion(-) diff --git a/taskingai/client/models/entities/action.py b/taskingai/client/models/entities/action.py index d2b83e5..9ec79b9 100644 --- a/taskingai/client/models/entities/action.py +++ b/taskingai/client/models/entities/action.py @@ -21,6 +21,7 @@ class Action(BaseModel): + object: str = Field("Action") action_id: str = Field(..., min_length=20, max_length=30) name: str = Field(..., min_length=1, max_length=128) operation_id: str = Field(...) diff --git a/taskingai/client/models/entities/assistant.py b/taskingai/client/models/entities/assistant.py index 3d39e7a..3e5889b 100644 --- a/taskingai/client/models/entities/assistant.py +++ b/taskingai/client/models/entities/assistant.py @@ -22,6 +22,7 @@ class Assistant(BaseModel): + object: str = Field("Assistant") assistant_id: str = Field(..., min_length=20, max_length=30) model_id: str = Field(..., min_length=8, max_length=8) name: str = Field("", min_length=0, max_length=256) diff --git a/taskingai/client/models/entities/chat.py b/taskingai/client/models/entities/chat.py index ecfe120..5ec1bed 100644 --- a/taskingai/client/models/entities/chat.py +++ b/taskingai/client/models/entities/chat.py @@ -19,6 +19,7 @@ class Chat(BaseModel): + object: str = Field("Chat") chat_id: str = Field(..., min_length=20, max_length=30) assistant_id: str = Field(..., min_length=20, max_length=30) name: str = Field("", min_length=0, max_length=127) diff --git a/taskingai/client/models/entities/chat_completion.py b/taskingai/client/models/entities/chat_completion.py index 27483d1..497d532 100644 --- a/taskingai/client/models/entities/chat_completion.py +++ b/taskingai/client/models/entities/chat_completion.py @@ -20,6 +20,7 @@ class ChatCompletion(BaseModel): + object: str = Field("ChatCompletion") finish_reason: ChatCompletionFinishReason = Field(...) message: ChatCompletionAssistantMessage = Field(...) created_timestamp: int = Field(...) diff --git a/taskingai/client/models/entities/chat_completion_chunk.py b/taskingai/client/models/entities/chat_completion_chunk.py index add19ed..e1b4765 100644 --- a/taskingai/client/models/entities/chat_completion_chunk.py +++ b/taskingai/client/models/entities/chat_completion_chunk.py @@ -18,6 +18,7 @@ class ChatCompletionChunk(BaseModel): + object: str = Field("ChatCompletionChunk") role: ChatCompletionRole = Field("assistant") index: int = Field(...) delta: str = Field(...) diff --git a/taskingai/client/models/entities/chunk.py b/taskingai/client/models/entities/chunk.py index 227f042..062b03d 100644 --- a/taskingai/client/models/entities/chunk.py +++ b/taskingai/client/models/entities/chunk.py @@ -19,6 +19,7 @@ class Chunk(BaseModel): + object: str = Field("Chunk") chunk_id: str = Field(..., min_length=20, max_length=30) record_id: Optional[str] = Field(..., min_length=20, max_length=30) collection_id: str = Field(..., min_length=20, max_length=30) diff --git a/taskingai/client/models/entities/collection.py b/taskingai/client/models/entities/collection.py index 58163be..c82f215 100644 --- a/taskingai/client/models/entities/collection.py +++ b/taskingai/client/models/entities/collection.py @@ -19,6 +19,7 @@ class Collection(BaseModel): + object: str = Field("Collection") collection_id: str = Field(..., min_length=24, max_length=24) name: str = Field("", min_length=0, max_length=256) description: str = Field("", min_length=0, max_length=512) diff --git a/taskingai/client/models/entities/message.py b/taskingai/client/models/entities/message.py index 0968969..cc33f6c 100644 --- a/taskingai/client/models/entities/message.py +++ b/taskingai/client/models/entities/message.py @@ -19,6 +19,7 @@ class Message(BaseModel): + object: str = Field("Message") message_id: str = Field(..., min_length=20, max_length=30) chat_id: str = Field(..., min_length=20, max_length=30) assistant_id: str = Field(None, min_length=20, max_length=30) diff --git a/taskingai/client/models/entities/message_chunk.py b/taskingai/client/models/entities/message_chunk.py index 3b21815..7d76add 100644 --- a/taskingai/client/models/entities/message_chunk.py +++ b/taskingai/client/models/entities/message_chunk.py @@ -18,6 +18,7 @@ class MessageChunk(BaseModel): + object: str = Field("MessageChunk") role: MessageRole = Field(...) index: int = Field(...) delta: str = Field(...) diff --git a/taskingai/client/models/entities/message_generation_log.py b/taskingai/client/models/entities/message_generation_log.py index cc14047..4f11e92 100644 --- a/taskingai/client/models/entities/message_generation_log.py +++ b/taskingai/client/models/entities/message_generation_log.py @@ -19,6 +19,7 @@ class MessageGenerationLog(BaseModel): + object: str = Field("MessageGenerationLog") session_id: str = Field(..., min_length=20, max_length=30) event: str = Field(...) event_id: str = Field(..., min_length=20, max_length=30) diff --git a/taskingai/client/models/entities/record.py b/taskingai/client/models/entities/record.py index ee584b1..d1228b7 100644 --- a/taskingai/client/models/entities/record.py +++ b/taskingai/client/models/entities/record.py @@ -20,6 +20,7 @@ class Record(BaseModel): + object: str = Field("Record") record_id: str = Field(..., min_length=20, max_length=30) collection_id: str = Field(..., min_length=20, max_length=30) title: str = Field("", min_length=0, max_length=255) diff --git a/taskingai/client/models/schemas/chat_completion_request.py b/taskingai/client/models/schemas/chat_completion_request.py index 84c360b..a9bbf8a 100644 --- a/taskingai/client/models/schemas/chat_completion_request.py +++ b/taskingai/client/models/schemas/chat_completion_request.py @@ -15,8 +15,8 @@ from typing import Optional, List, Dict, Union from ..entities.chat_completion_function_message import ChatCompletionFunctionMessage from ..entities.chat_completion_system_message import ChatCompletionSystemMessage -from ..entities.chat_completion_user_message import ChatCompletionUserMessage from ..entities.chat_completion_assistant_message import ChatCompletionAssistantMessage +from ..entities.chat_completion_user_message import ChatCompletionUserMessage from ..entities.chat_completion_function import ChatCompletionFunction __all__ = ["ChatCompletionRequest"] From 82858bb4a879b303c6e5c3ad6c9adbfd15af7b45 Mon Sep 17 00:00:00 2001 From: taskingaijc <150663083+taskingaijc@users.noreply.github.com> Date: Tue, 14 May 2024 11:17:49 +0800 Subject: [PATCH 2/2] test: remove chunk_keys --- test/testcase/test_async/test_async_retrieval.py | 16 ---------------- test/testcase/test_sync/test_sync_retrieval.py | 16 ---------------- 2 files changed, 32 deletions(-) diff --git a/test/testcase/test_async/test_async_retrieval.py b/test/testcase/test_async/test_async_retrieval.py index 090558d..94ae677 100644 --- a/test/testcase/test_async/test_async_retrieval.py +++ b/test/testcase/test_async/test_async_retrieval.py @@ -305,18 +305,6 @@ async def test_a_delete_record(self): @pytest.mark.test_async class TestChunk(Base): - chunk_list = [ - "chunk_id", - "record_id", - "collection_id", - "content", - "metadata", - "num_tokens", - "score", - "updated_timestamp", - "created_timestamp", - ] - chunk_keys = set(chunk_list) @pytest.mark.run(order=41) @pytest.mark.asyncio @@ -332,7 +320,6 @@ async def test_a_query_chunks(self): for chunk in res: chunk_dict = vars(chunk) assume_query_chunk_result(query_text, chunk_dict) - pytest.assume(chunk_dict.keys() == self.chunk_keys) pytest.assume(chunk_dict["score"] >= 0.04) @pytest.mark.run(order=42) @@ -345,7 +332,6 @@ async def test_create_chunk(self): } res = await a_create_chunk(**create_chunk_data) res_dict = vars(res) - pytest.assume(res_dict.keys() == self.chunk_keys) assume_chunk_result(create_chunk_data, res_dict) Base.chunk_id = res_dict["chunk_id"] @@ -386,7 +372,6 @@ async def test_get_chunk(self): res_dict = vars(res) pytest.assume(res_dict["collection_id"] == self.collection_id) pytest.assume(res_dict["chunk_id"] == chunk_id) - pytest.assume(res_dict.keys() == self.chunk_keys) @pytest.mark.run(order=45) @pytest.mark.asyncio @@ -401,7 +386,6 @@ async def test_update_chunk(self): } res = await a_update_chunk(**update_chunk_data) res_dict = vars(res) - pytest.assume(res_dict.keys() == self.chunk_keys) assume_chunk_result(update_chunk_data, res_dict) @pytest.mark.run(order=46) diff --git a/test/testcase/test_sync/test_sync_retrieval.py b/test/testcase/test_sync/test_sync_retrieval.py index 8ebe111..9255bbc 100644 --- a/test/testcase/test_sync/test_sync_retrieval.py +++ b/test/testcase/test_sync/test_sync_retrieval.py @@ -315,18 +315,6 @@ def test_delete_record(self, collection_id): @pytest.mark.test_sync class TestChunk: - chunk_list = [ - "chunk_id", - "record_id", - "collection_id", - "content", - "metadata", - "num_tokens", - "score", - "updated_timestamp", - "created_timestamp", - ] - chunk_keys = set(chunk_list) @pytest.mark.run(order=41) def test_query_chunks(self, collection_id): @@ -341,7 +329,6 @@ def test_query_chunks(self, collection_id): for chunk in res: chunk_dict = vars(chunk) assume_query_chunk_result(query_text, chunk_dict) - pytest.assume(chunk_dict.keys() == self.chunk_keys) pytest.assume(chunk_dict["score"] >= 0.04) @pytest.mark.run(order=42) @@ -353,7 +340,6 @@ def test_create_chunk(self, collection_id): } res = create_chunk(**create_chunk_data) res_dict = vars(res) - pytest.assume(res_dict.keys() == self.chunk_keys) assume_chunk_result(create_chunk_data, res_dict) @pytest.mark.run(order=43) @@ -391,7 +377,6 @@ def test_get_chunk(self, collection_id): res_dict = vars(res) pytest.assume(res_dict["collection_id"] == collection_id) pytest.assume(res_dict["chunk_id"] == chunk_id) - pytest.assume(res_dict.keys() == self.chunk_keys) @pytest.mark.run(order=45) def test_update_chunk(self, collection_id, chunk_id): @@ -405,7 +390,6 @@ def test_update_chunk(self, collection_id, chunk_id): } res = update_chunk(**update_chunk_data) res_dict = vars(res) - pytest.assume(res_dict.keys() == self.chunk_keys) assume_chunk_result(update_chunk_data, res_dict) @pytest.mark.run(order=46)