mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2025-09-14 17:19:42 +00:00
54 lines
1.2 KiB
Python
54 lines
1.2 KiB
Python
from enum import Enum
|
|
from typing import List, Optional, Union
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from ktransformers.server.schemas.base import ObjectID
|
|
|
|
|
|
class ToolType(str, Enum):
|
|
CODE_INTERPRETER = "code_interpreter"
|
|
FILE_SEARCH = "file_search"
|
|
RELATED_THREADS = "related_threads"
|
|
FUNCTION = "function"
|
|
|
|
|
|
class ToolBase(BaseModel):
|
|
type: ToolType
|
|
|
|
|
|
class CodeInterpreter(ToolBase):
|
|
pass
|
|
|
|
|
|
class FileSearch(ToolBase):
|
|
pass
|
|
|
|
|
|
class RelatedThreads(ToolBase):
|
|
pass
|
|
|
|
|
|
class FuntionTool(ToolBase):
|
|
description: str
|
|
name: str
|
|
parameters: List[str]
|
|
|
|
|
|
Tool = Union[CodeInterpreter, FileSearch, RelatedThreads, FuntionTool]
|
|
|
|
|
|
class CodeInterpreterResource(BaseModel):
|
|
file_ids: Optional[List[str]] = Field(default_factory=list, max_length=20)
|
|
|
|
|
|
class FileSearchResource(BaseModel):
|
|
vector_store_ids: Optional[List[str]] = Field(default_factory=list, max_length=1)
|
|
vector_stores: Optional[List[str]] = Field(default_factory=list, max_length=1)
|
|
|
|
|
|
class RelatedThreadsResource(BaseModel):
|
|
thread_ids: List[ObjectID] = Field(default=[])
|
|
|
|
|
|
ToolResource = Union[CodeInterpreterResource,FileSearchResource,RelatedThreadsResource]
|