feat: update text splitter and retrieval config

Author: Dttbd

taskingai/client/models/entities/action.py

@@ -12,13 +12,9 @@
 """
 
 from pydantic import BaseModel, Field
-from typing import Optional, Any, Dict
+from typing import Any, Dict
 from .action_method import ActionMethod
-from .action_param import ActionParam
-from .action_param import ActionParam
 from .action_body_type import ActionBodyType
-from .action_param import ActionParam
-from .chat_completion_function import ChatCompletionFunction
 from .action_authentication import ActionAuthentication
 
 __all__ = ["Action"]

taskingai/client/models/entities/retrieval_config.py

@@ -21,4 +21,6 @@
 class RetrievalConfig(BaseModel):
     top_k: int = Field(3, ge=1, le=20)
     max_tokens: Optional[int] = Field(None, ge=1, le=8192)
+    score_threshold: Optional[float] = Field(None, ge=0.0, le=1.0)
     method: RetrievalMethod = Field(...)
+    function_description: Optional[str] = Field(None, min_length=0, max_length=1024)

taskingai/client/models/entities/text_splitter.py

@@ -12,7 +12,7 @@
 """
 
 from pydantic import BaseModel, Field
-from typing import Optional
+from typing import Optional, List
 from .text_splitter_type import TextSplitterType
 
 __all__ = ["TextSplitter"]
@@ -22,3 +22,4 @@ class TextSplitter(BaseModel):
     type: TextSplitterType = Field(...)
     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)

taskingai/client/models/entities/text_splitter_type.py

@@ -18,3 +18,4 @@
 
 class TextSplitterType(str, Enum):
     TOKEN = "token"
+    SEPARATOR = "separator"

taskingai/client/models/schemas/chat_completion_request.py

@@ -14,9 +14,9 @@
 from pydantic import BaseModel, Field
 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_system_message import ChatCompletionSystemMessage
 from ..entities.chat_completion_function import ChatCompletionFunction
 
 __all__ = ["ChatCompletionRequest"]

taskingai/client/models/schemas/chunk_query_request.py

@@ -21,4 +21,5 @@
 class ChunkQueryRequest(BaseModel):
     top_k: int = Field(..., ge=1, le=20)
     max_tokens: Optional[int] = Field(None, ge=1)
+    score_threshold: Optional[float] = Field(None, ge=0.0, le=1.0)
     query_text: str = Field(..., min_length=1, max_length=32768)

taskingai/retrieval/chunk.py

@@ -248,6 +248,7 @@ def query_chunks(
     *,
     query_text: str,
     top_k: int = 3,
+    score_threshold: Optional[float] = None,
     max_tokens: Optional[int] = None,
 ) -> List[Chunk]:
     """
@@ -262,6 +263,7 @@ def query_chunks(
     body = ChunkQueryRequest(
         top_k=top_k,
         query_text=query_text,
+        score_threshold=score_threshold,
         max_tokens=max_tokens,
     )
     response: ChunkQueryResponse = api_query_collection_chunks(
@@ -276,6 +278,7 @@ async def a_query_chunks(
     *,
     query_text: str,
     top_k: int = 3,
+    score_threshold: Optional[float] = None,
     max_tokens: Optional[int] = None,
 ) -> List[Chunk]:
     """
@@ -290,6 +293,7 @@ async def a_query_chunks(
     body = ChunkQueryRequest(
         top_k=top_k,
         query_text=query_text,
+        score_threshold=score_threshold,
         max_tokens=max_tokens,
     )
     response: ChunkQueryResponse = await async_api_query_collection_chunks(