From 1d5a746f1ad6ac46aeb643e42b2981b96f23cad4 Mon Sep 17 00:00:00 2001 From: Dttbd Date: Fri, 10 May 2024 18:11:53 +0800 Subject: [PATCH] feat: update model --- taskingai/client/models/entities/__init__.py | 3 +++ .../client/models/entities/chat_completion.py | 2 ++ ...completion_function_parameters_property.py | 7 +++--- ...tion_function_parameters_property_items.py | 21 ++++++++++++++++++ .../models/entities/chat_completion_usage.py | 22 +++++++++++++++++++ .../models/entities/text_embedding_usage.py | 21 ++++++++++++++++++ .../client/models/entities/text_splitter.py | 2 +- .../models/schemas/chat_completion_request.py | 7 +++--- .../models/schemas/text_embedding_request.py | 2 +- .../models/schemas/text_embedding_response.py | 2 ++ taskingai/retrieval/text_splitter.py | 11 ++++++++++ 11 files changed, 92 insertions(+), 8 deletions(-) create mode 100644 taskingai/client/models/entities/chat_completion_function_parameters_property_items.py create mode 100644 taskingai/client/models/entities/chat_completion_usage.py create mode 100644 taskingai/client/models/entities/text_embedding_usage.py diff --git a/taskingai/client/models/entities/__init__.py b/taskingai/client/models/entities/__init__.py index 68316a1..e2020dd 100644 --- a/taskingai/client/models/entities/__init__.py +++ b/taskingai/client/models/entities/__init__.py @@ -30,9 +30,11 @@ from .chat_completion_function_message import * from .chat_completion_function_parameters import * from .chat_completion_function_parameters_property import * +from .chat_completion_function_parameters_property_items import * from .chat_completion_message import * from .chat_completion_role import * from .chat_completion_system_message import * +from .chat_completion_usage import * from .chat_completion_user_message import * from .chat_memory import * from .chat_memory_message import * @@ -54,6 +56,7 @@ from .status import * from .text_embedding_input_type import * from .text_embedding_output import * +from .text_embedding_usage import * from .text_splitter import * from .text_splitter_type import * from .tool_ref import * diff --git a/taskingai/client/models/entities/chat_completion.py b/taskingai/client/models/entities/chat_completion.py index 2229f6e..27483d1 100644 --- a/taskingai/client/models/entities/chat_completion.py +++ b/taskingai/client/models/entities/chat_completion.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, Field from .chat_completion_finish_reason import ChatCompletionFinishReason from .chat_completion_assistant_message import ChatCompletionAssistantMessage +from .chat_completion_usage import ChatCompletionUsage __all__ = ["ChatCompletion"] @@ -22,3 +23,4 @@ class ChatCompletion(BaseModel): finish_reason: ChatCompletionFinishReason = Field(...) message: ChatCompletionAssistantMessage = Field(...) created_timestamp: int = Field(...) + usage: ChatCompletionUsage = Field(...) diff --git a/taskingai/client/models/entities/chat_completion_function_parameters_property.py b/taskingai/client/models/entities/chat_completion_function_parameters_property.py index 056009d..6b81707 100644 --- a/taskingai/client/models/entities/chat_completion_function_parameters_property.py +++ b/taskingai/client/models/entities/chat_completion_function_parameters_property.py @@ -13,12 +13,13 @@ from pydantic import BaseModel, Field from typing import Optional, List - +from .chat_completion_function_parameters_property_items import ChatCompletionFunctionParametersPropertyItems __all__ = ["ChatCompletionFunctionParametersProperty"] class ChatCompletionFunctionParametersProperty(BaseModel): - type: str = Field(..., pattern="^(string|number|integer|boolean)$") - description: str = Field("", max_length=256) + type: str = Field(..., pattern="^(string|number|integer|boolean|array)$") + description: str = Field("", max_length=512) enum: Optional[List[str]] = Field(None) + items: Optional[ChatCompletionFunctionParametersPropertyItems] = Field(None) diff --git a/taskingai/client/models/entities/chat_completion_function_parameters_property_items.py b/taskingai/client/models/entities/chat_completion_function_parameters_property_items.py new file mode 100644 index 0000000..4f72b86 --- /dev/null +++ b/taskingai/client/models/entities/chat_completion_function_parameters_property_items.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +# chat_completion_function_parameters_property_items.py + +""" +This script is automatically generated for TaskingAI python client +Do not modify the file manually + +Author: James Yao +Organization: TaskingAI +License: Apache 2.0 +""" + +from pydantic import BaseModel, Field + + +__all__ = ["ChatCompletionFunctionParametersPropertyItems"] + + +class ChatCompletionFunctionParametersPropertyItems(BaseModel): + type: str = Field(..., pattern="^(string|number|integer|boolean)$") diff --git a/taskingai/client/models/entities/chat_completion_usage.py b/taskingai/client/models/entities/chat_completion_usage.py new file mode 100644 index 0000000..a0eea53 --- /dev/null +++ b/taskingai/client/models/entities/chat_completion_usage.py @@ -0,0 +1,22 @@ +# -*- coding: utf-8 -*- + +# chat_completion_usage.py + +""" +This script is automatically generated for TaskingAI python client +Do not modify the file manually + +Author: James Yao +Organization: TaskingAI +License: Apache 2.0 +""" + +from pydantic import BaseModel, Field + + +__all__ = ["ChatCompletionUsage"] + + +class ChatCompletionUsage(BaseModel): + input_tokens: int = Field(...) + output_tokens: int = Field(...) diff --git a/taskingai/client/models/entities/text_embedding_usage.py b/taskingai/client/models/entities/text_embedding_usage.py new file mode 100644 index 0000000..fa40007 --- /dev/null +++ b/taskingai/client/models/entities/text_embedding_usage.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +# text_embedding_usage.py + +""" +This script is automatically generated for TaskingAI python client +Do not modify the file manually + +Author: James Yao +Organization: TaskingAI +License: Apache 2.0 +""" + +from pydantic import BaseModel, Field + + +__all__ = ["TextEmbeddingUsage"] + + +class TextEmbeddingUsage(BaseModel): + input_tokens: int = Field(...) diff --git a/taskingai/client/models/entities/text_splitter.py b/taskingai/client/models/entities/text_splitter.py index cd1e5fa..eb3031f 100644 --- a/taskingai/client/models/entities/text_splitter.py +++ b/taskingai/client/models/entities/text_splitter.py @@ -19,7 +19,7 @@ class TextSplitter(BaseModel): - type: TextSplitterType = Field(...) + type: TextSplitterType = Field("token") chunk_size: Optional[int] = Field(None, ge=50, le=1000) chunk_overlap: Optional[int] = Field(None, ge=0, le=200) separators: Optional[List[str]] = Field(None, min_length=1, max_length=16) diff --git a/taskingai/client/models/schemas/chat_completion_request.py b/taskingai/client/models/schemas/chat_completion_request.py index 1341c52..84c360b 100644 --- a/taskingai/client/models/schemas/chat_completion_request.py +++ b/taskingai/client/models/schemas/chat_completion_request.py @@ -14,16 +14,16 @@ from pydantic import BaseModel, Field from typing import Optional, List, Dict, Union from ..entities.chat_completion_function_message import ChatCompletionFunctionMessage -from ..entities.chat_completion_assistant_message import ChatCompletionAssistantMessage -from ..entities.chat_completion_user_message import ChatCompletionUserMessage 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_function import ChatCompletionFunction __all__ = ["ChatCompletionRequest"] class ChatCompletionRequest(BaseModel): - model_id: str = Field(..., min_length=8, max_length=8) + model_id: str = Field(..., min_length=1, max_length=255) configs: Optional[Dict] = Field(None) stream: bool = Field(False) messages: List[ @@ -36,3 +36,4 @@ class ChatCompletionRequest(BaseModel): ] = Field(...) function_call: Optional[str] = Field(None) functions: Optional[List[ChatCompletionFunction]] = Field(None) + save_logs: bool = Field(False) diff --git a/taskingai/client/models/schemas/text_embedding_request.py b/taskingai/client/models/schemas/text_embedding_request.py index b5b0892..a38ad9b 100644 --- a/taskingai/client/models/schemas/text_embedding_request.py +++ b/taskingai/client/models/schemas/text_embedding_request.py @@ -19,6 +19,6 @@ class TextEmbeddingRequest(BaseModel): - model_id: str = Field(..., min_length=8, max_length=8) + model_id: str = Field(..., min_length=1, max_length=255) input: Union[str, List[str]] = Field(...) input_type: Optional[TextEmbeddingInputType] = Field(None) diff --git a/taskingai/client/models/schemas/text_embedding_response.py b/taskingai/client/models/schemas/text_embedding_response.py index 9c43730..1cb349a 100644 --- a/taskingai/client/models/schemas/text_embedding_response.py +++ b/taskingai/client/models/schemas/text_embedding_response.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, Field from typing import List from ..entities.text_embedding_output import TextEmbeddingOutput +from ..entities.text_embedding_usage import TextEmbeddingUsage __all__ = ["TextEmbeddingResponse"] @@ -21,3 +22,4 @@ class TextEmbeddingResponse(BaseModel): status: str = Field("success") data: List[TextEmbeddingOutput] = Field(...) + usage: TextEmbeddingUsage = Field(...) diff --git a/taskingai/retrieval/text_splitter.py b/taskingai/retrieval/text_splitter.py index 7b9812f..e5ae97f 100644 --- a/taskingai/retrieval/text_splitter.py +++ b/taskingai/retrieval/text_splitter.py @@ -4,6 +4,7 @@ "TextSplitter", "TextSplitterType", "TokenTextSplitter", + "SeparatorTextSplitter", ] @@ -14,3 +15,13 @@ def __init__(self, chunk_size: int, chunk_overlap: int): chunk_size=chunk_size, chunk_overlap=chunk_overlap, ) + + +class SeparatorTextSplitter(TextSplitter): + def __init__(self, chunk_size: int, chunk_overlap: int, separators: list[str]): + super().__init__( + type=TextSplitterType.SEPARATOR, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + separators=separators, + )