-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #202 from AI21Labs/chat-model-with-tools-docs-and-…
…response-format feat: Chat model with tools docs and response format
- Loading branch information
Showing
30 changed files
with
673 additions
and
51 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,8 +1,28 @@ | ||
from __future__ import annotations | ||
from typing import Literal, List, Optional | ||
|
||
from ai21.models.ai21_base_model import AI21BaseModel | ||
from ai21.models.chat.tool_call import ToolCall | ||
|
||
|
||
class ChatMessage(AI21BaseModel): | ||
role: str | ||
content: str | ||
|
||
|
||
class AssistantMessage(ChatMessage): | ||
role: Literal["assistant"] = "assistant" | ||
tool_calls: Optional[List[ToolCall]] = None | ||
content: Optional[str] = None | ||
|
||
|
||
class ToolMessage(ChatMessage): | ||
role: Literal["tool"] = "tool" | ||
tool_call_id: str | ||
|
||
|
||
class UserMessage(ChatMessage): | ||
role: Literal["user"] = "user" | ||
|
||
|
||
class SystemMessage(ChatMessage): | ||
role: Literal["system"] = "system" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,9 @@ | ||
from typing import Optional, Dict | ||
|
||
from ai21.models.ai21_base_model import AI21BaseModel | ||
|
||
|
||
class DocumentSchema(AI21BaseModel): | ||
content: str | ||
id: Optional[str] = None | ||
metadata: Optional[Dict[str, str]] = None |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,9 @@ | ||
from typing_extensions import TypedDict, Required | ||
|
||
from ai21.models.chat.tool_parameters import ToolParameters | ||
|
||
|
||
class FunctionToolDefinition(TypedDict, total=False): | ||
name: Required[str] | ||
description: str | ||
parameters: ToolParameters |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,7 @@ | ||
from typing import Literal | ||
|
||
from ai21.models.ai21_base_model import AI21BaseModel | ||
|
||
|
||
class ResponseFormat(AI21BaseModel): | ||
type: Literal["text", "json_object"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -4,3 +4,5 @@ | |
class RoleType(str, Enum): | ||
USER = "user" | ||
ASSISTANT = "assistant" | ||
TOOL = "tool" | ||
SYSTEM = "system" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,10 @@ | ||
from typing import Literal | ||
|
||
from ai21.models.ai21_base_model import AI21BaseModel | ||
from ai21.models.chat.tool_function import ToolFunction | ||
|
||
|
||
class ToolCall(AI21BaseModel): | ||
id: str | ||
function: ToolFunction | ||
type: Literal["function"] = "function" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,8 @@ | ||
from typing_extensions import Literal, TypedDict, Required | ||
|
||
from ai21.models.chat import FunctionToolDefinition | ||
|
||
|
||
class ToolDefinition(TypedDict, total=False): | ||
type: Required[Literal["function"]] | ||
function: Required[FunctionToolDefinition] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,6 @@ | ||
from ai21.models.ai21_base_model import AI21BaseModel | ||
|
||
|
||
class ToolFunction(AI21BaseModel): | ||
name: str | ||
arguments: str |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,7 @@ | ||
from typing_extensions import Literal, Any, Dict, List, TypedDict, Required | ||
|
||
|
||
class ToolParameters(TypedDict, total=False): | ||
type: Literal["object"] | ||
properties: Required[Dict[str, Any]] | ||
required: List[str] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.