File size: 1,705 Bytes
5b41acf
7cdc7d0
 
 
 
 
 
 
 
 
 
 
 
88c61d3
 
 
 
 
 
 
7cdc7d0
 
 
 
 
fe70438
 
 
b9d0035
 
 
 
 
 
 
fe70438
7cdc7d0
 
 
 
 
 
 
 
 
 
 
91ef70a
 
 
 
 
 
5b41acf
99f75f9
5b41acf
 
 
 
 
99f75f9
5b41acf
 
66630b0
99f75f9
66630b0
 
 
5b41acf
66630b0
99f75f9
66630b0
 
 
91ef70a
99f75f9
7cdc7d0
 
 
 
 
 
 
fe70438
b9d0035
 
88c61d3
91ef70a
66630b0
5b41acf
66630b0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
from typing import Any, Dict, List, Literal, NewType, Optional, TypedDict, Union

from .type_utils import register_type

Text = NewType("Text", str)
Number = NewType("Number", Union[float, int])


class Turn(TypedDict):
    role: Literal["system", "user", "agent"]
    content: Text


class RagResponse(TypedDict):
    answer: str
    contexts: List[str]
    context_ids: Union[List[int], List[str]]
    is_answerable: bool


Dialog = NewType("Dialog", List[Turn])


class Image(TypedDict):
    image: Any
    format: str


class Document(TypedDict):
    title: str
    body: str


MultiDocument = NewType("MultiDocument", List[Document])

Video = NewType("Video", List[Image])


class Audio(TypedDict):
    audio: Any


class Table(TypedDict):
    header: List[str]
    rows: List[List[Any]]


class SQLDatabase(TypedDict):
    db_id: Optional[str]
    db_type: Literal["local", "in_memory", "remote"]
    dbms: Optional[str]
    data: Optional[Dict[str, Dict]]


class JsonSchema:
    @classmethod
    def __verify_type__(cls, object):
        if not isinstance(object, dict):
            return False
        import jsonschema_rs

        jsonschema_rs.meta.validate(object)
        return True


class Tool(TypedDict):
    name: str
    description: str
    parameters: JsonSchema


class ToolCall(TypedDict):
    name: str
    arguments: Dict[str, Any]


register_type(Text)
register_type(Number)
register_type(Turn)
register_type(Dialog)
register_type(Table)
register_type(Audio)
register_type(Image)
register_type(Video)
register_type(Document)
register_type(MultiDocument)
register_type(RagResponse)
register_type(SQLDatabase)
register_type(Tool)
register_type(JsonSchema)
register_type(ToolCall)