-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy pathembedding.py
More file actions
123 lines (103 loc) · 3.31 KB
/
embedding.py
File metadata and controls
123 lines (103 loc) · 3.31 KB
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
from typing import Any, Dict, List, Literal, Union, cast, overload
from typing_extensions import NotRequired, TypedDict
from ._config import ClientConfig
from ._types import BaseResponse
from .async_request import AsyncRequest
from .request import Request, RequestConfig
class EmbeddingParams(TypedDict):
text: NotRequired[str]
file_content: NotRequired[Any]
type: Literal["text", "text-other", "image", "audio", "pdf"]
url: NotRequired[str]
file_store_key: NotRequired[str]
token_overflow_mode: NotRequired[Literal["truncate", "error"]]
class Chunk(TypedDict):
text: str
timestamp: List[int]
class EmbeddingResponse(BaseResponse):
embeddings: List[List[float]]
chunks: Union[List[Chunk], List[str]]
class Embedding(ClientConfig):
config: RequestConfig
def __init__(
self,
api_key: str,
base_url: str,
headers: Union[Dict[str, str], None] = None,
):
super().__init__(api_key, base_url, headers)
self.config = RequestConfig(
base_url=base_url,
api_key=api_key,
headers=headers,
)
@overload
def execute(self, params: EmbeddingParams) -> EmbeddingResponse: ...
@overload
def execute(self, blob: bytes, options: EmbeddingParams = None) -> EmbeddingResponse: ...
def execute(
self,
blob: Union[EmbeddingParams, bytes],
options: EmbeddingParams = None,
) -> EmbeddingResponse:
path = "/embedding"
options = options or {}
if isinstance(blob, dict):
resp = Request(
config=self.config,
path=path,
params=cast(Dict[Any, Any], blob),
verb="post",
).perform_with_content()
return resp
files = {"file": blob}
resp = Request(
config=self.config,
path=path,
params=options,
files=files,
verb="post",
).perform_with_content()
return resp
class AsyncEmbedding(ClientConfig):
config: RequestConfig
def __init__(
self,
api_key: str,
base_url: str,
headers: Union[Dict[str, str], None] = None,
):
super().__init__(api_key, base_url, headers)
self.config = RequestConfig(
base_url=base_url,
api_key=api_key,
headers=headers,
)
@overload
async def execute(self, params: EmbeddingParams) -> EmbeddingResponse: ...
@overload
async def execute(self, blob: bytes, options: EmbeddingParams = None) -> EmbeddingResponse: ...
async def execute(
self,
blob: Union[EmbeddingParams, bytes],
options: EmbeddingParams = None,
) -> EmbeddingResponse:
path = "/embedding"
options = options or {}
if isinstance(blob, dict):
resp = await AsyncRequest(
config=self.config,
path=path,
params=cast(Dict[Any, Any], blob),
verb="post",
).perform_with_content()
return resp
files = {"file": blob}
resp = await AsyncRequest(
config=self.config,
path=path,
params=options,
files=files,
verb="post",
).perform_with_content()
return resp