1
1
import os
2
2
from abc import ABC
3
+ from collections .abc import AsyncGenerator , Awaitable
3
4
from dataclasses import dataclass
4
5
from typing import (
5
6
Any ,
6
- AsyncGenerator ,
7
- Awaitable ,
8
7
Callable ,
9
- Dict ,
10
- List ,
11
8
Optional ,
12
9
TypedDict ,
13
10
Union ,
41
38
class Document :
42
39
id : Optional [str ]
43
40
content : Optional [str ]
44
- embedding : Optional [List [float ]]
45
- image_embedding : Optional [List [float ]]
41
+ embedding : Optional [list [float ]]
42
+ image_embedding : Optional [list [float ]]
46
43
category : Optional [str ]
47
44
sourcepage : Optional [str ]
48
45
sourcefile : Optional [str ]
49
- oids : Optional [List [str ]]
50
- groups : Optional [List [str ]]
51
- captions : List [QueryCaptionResult ]
46
+ oids : Optional [list [str ]]
47
+ groups : Optional [list [str ]]
48
+ captions : list [QueryCaptionResult ]
52
49
score : Optional [float ] = None
53
50
reranker_score : Optional [float ] = None
54
51
@@ -80,7 +77,7 @@ def serialize_for_results(self) -> dict[str, Any]:
80
77
}
81
78
82
79
@classmethod
83
- def trim_embedding (cls , embedding : Optional [List [float ]]) -> Optional [str ]:
80
+ def trim_embedding (cls , embedding : Optional [list [float ]]) -> Optional [str ]:
84
81
"""Returns a trimmed list of floats from the vector embedding."""
85
82
if embedding :
86
83
if len (embedding ) > 2 :
@@ -105,15 +102,15 @@ def update_token_usage(self, usage: CompletionUsage) -> None:
105
102
106
103
@dataclass
107
104
class DataPoints :
108
- text : Optional [List [str ]] = None
109
- images : Optional [List ] = None
105
+ text : Optional [list [str ]] = None
106
+ images : Optional [list ] = None
110
107
111
108
112
109
@dataclass
113
110
class ExtraInfo :
114
111
data_points : DataPoints
115
- thoughts : Optional [List [ThoughtStep ]] = None
116
- followup_questions : Optional [List [Any ]] = None
112
+ thoughts : Optional [list [ThoughtStep ]] = None
113
+ followup_questions : Optional [list [Any ]] = None
117
114
118
115
119
116
@dataclass
@@ -201,15 +198,15 @@ async def search(
201
198
top : int ,
202
199
query_text : Optional [str ],
203
200
filter : Optional [str ],
204
- vectors : List [VectorQuery ],
201
+ vectors : list [VectorQuery ],
205
202
use_text_search : bool ,
206
203
use_vector_search : bool ,
207
204
use_semantic_ranker : bool ,
208
205
use_semantic_captions : bool ,
209
206
minimum_search_score : Optional [float ] = None ,
210
207
minimum_reranker_score : Optional [float ] = None ,
211
208
use_query_rewriting : Optional [bool ] = None ,
212
- ) -> List [Document ]:
209
+ ) -> list [Document ]:
213
210
search_text = query_text if use_text_search else ""
214
211
search_vectors = vectors if use_vector_search else []
215
212
if use_semantic_ranker :
@@ -248,7 +245,7 @@ async def search(
248
245
sourcefile = document .get ("sourcefile" ),
249
246
oids = document .get ("oids" ),
250
247
groups = document .get ("groups" ),
251
- captions = cast (List [QueryCaptionResult ], document .get ("@search.captions" )),
248
+ captions = cast (list [QueryCaptionResult ], document .get ("@search.captions" )),
252
249
score = document .get ("@search.score" ),
253
250
reranker_score = document .get ("@search.reranker_score" ),
254
251
)
@@ -266,7 +263,7 @@ async def search(
266
263
return qualified_documents
267
264
268
265
def get_sources_content (
269
- self , results : List [Document ], use_semantic_captions : bool , use_image_citation : bool
266
+ self , results : list [Document ], use_semantic_captions : bool , use_image_citation : bool
270
267
) -> list [str ]:
271
268
272
269
def nonewlines (s : str ) -> str :
@@ -358,13 +355,13 @@ def create_chat_completion(
358
355
overrides : dict [str , Any ],
359
356
response_token_limit : int ,
360
357
should_stream : bool = False ,
361
- tools : Optional [List [ChatCompletionToolParam ]] = None ,
358
+ tools : Optional [list [ChatCompletionToolParam ]] = None ,
362
359
temperature : Optional [float ] = None ,
363
360
n : Optional [int ] = None ,
364
361
reasoning_effort : Optional [ChatCompletionReasoningEffort ] = None ,
365
362
) -> Union [Awaitable [ChatCompletion ], Awaitable [AsyncStream [ChatCompletionChunk ]]]:
366
363
if chatgpt_model in self .GPT_REASONING_MODELS :
367
- params : Dict [str , Any ] = {
364
+ params : dict [str , Any ] = {
368
365
# max_tokens is not supported
369
366
"max_completion_tokens" : response_token_limit
370
367
}
@@ -400,14 +397,14 @@ def create_chat_completion(
400
397
def format_thought_step_for_chatcompletion (
401
398
self ,
402
399
title : str ,
403
- messages : List [ChatCompletionMessageParam ],
400
+ messages : list [ChatCompletionMessageParam ],
404
401
overrides : dict [str , Any ],
405
402
model : str ,
406
403
deployment : Optional [str ],
407
404
usage : Optional [CompletionUsage ] = None ,
408
405
reasoning_effort : Optional [ChatCompletionReasoningEffort ] = None ,
409
406
) -> ThoughtStep :
410
- properties : Dict [str , Any ] = {"model" : model }
407
+ properties : dict [str , Any ] = {"model" : model }
411
408
if deployment :
412
409
properties ["deployment" ] = deployment
413
410
# Only add reasoning_effort setting if the model supports it
0 commit comments