finish agent

Former-commit-id: d8d9d3afe32725fe79120fcd1a0970fdcdc45625

src/llmtuner/api/protocol.py

@@ -11,12 +11,15 @@ class Role(str, Enum):
     USER = "user"
     ASSISTANT = "assistant"
     SYSTEM = "system"
+    FUNCTION = "function"
+    TOOL = "tool"
 
 
 @unique
 class Finish(str, Enum):
     STOP = "stop"
     LENGTH = "length"
+    TOOL = "tool_calls"
 
 
 class ModelCard(BaseModel):
@@ -31,19 +34,32 @@ class ModelList(BaseModel):
     data: List[ModelCard] = []
 
 
+class Function(BaseModel):
+    name: str
+    arguments: str
+
+
+class FunctionCall(BaseModel):
+    id: Literal["call_default"] = "call_default"
+    type: Literal["function"] = "function"
+    function: Function
+
+
 class ChatMessage(BaseModel):
     role: Role
     content: str
 
 
-class DeltaMessage(BaseModel):
+class ChatCompletionMessage(BaseModel):
     role: Optional[Role] = None
     content: Optional[str] = None
+    tool_calls: Optional[List[FunctionCall]] = None
 
 
 class ChatCompletionRequest(BaseModel):
     model: str
     messages: List[ChatMessage]
+    tools: Optional[list] = []
     do_sample: bool = True
     temperature: Optional[float] = None
     top_p: Optional[float] = None
@@ -54,13 +70,13 @@ class ChatCompletionRequest(BaseModel):
 
 class ChatCompletionResponseChoice(BaseModel):
     index: int
-    message: ChatMessage
+    message: ChatCompletionMessage
     finish_reason: Finish
 
 
 class ChatCompletionResponseStreamChoice(BaseModel):
     index: int
-    delta: DeltaMessage
+    delta: ChatCompletionMessage
     finish_reason: Optional[Finish] = None