11import os
22from abc import ABC
3+ from collections .abc import AsyncGenerator , Awaitable
34from dataclasses import dataclass
45from typing import (
56 Any ,
6- AsyncGenerator ,
7- Awaitable ,
87 Callable ,
9- Dict ,
10- List ,
118 Optional ,
129 TypedDict ,
1310 Union ,
4138class Document :
4239 id : Optional [str ]
4340 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 ]]
4643 category : Optional [str ]
4744 sourcepage : Optional [str ]
4845 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 ]
5249 score : Optional [float ] = None
5350 reranker_score : Optional [float ] = None
5451
@@ -80,7 +77,7 @@ def serialize_for_results(self) -> dict[str, Any]:
8077 }
8178
8279 @classmethod
83- def trim_embedding (cls , embedding : Optional [List [float ]]) -> Optional [str ]:
80+ def trim_embedding (cls , embedding : Optional [list [float ]]) -> Optional [str ]:
8481 """Returns a trimmed list of floats from the vector embedding."""
8582 if embedding :
8683 if len (embedding ) > 2 :
@@ -105,15 +102,15 @@ def update_token_usage(self, usage: CompletionUsage) -> None:
105102
106103@dataclass
107104class DataPoints :
108- text : Optional [List [str ]] = None
109- images : Optional [List ] = None
105+ text : Optional [list [str ]] = None
106+ images : Optional [list ] = None
110107
111108
112109@dataclass
113110class ExtraInfo :
114111 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
117114
118115
119116@dataclass
@@ -201,15 +198,15 @@ async def search(
201198 top : int ,
202199 query_text : Optional [str ],
203200 filter : Optional [str ],
204- vectors : List [VectorQuery ],
201+ vectors : list [VectorQuery ],
205202 use_text_search : bool ,
206203 use_vector_search : bool ,
207204 use_semantic_ranker : bool ,
208205 use_semantic_captions : bool ,
209206 minimum_search_score : Optional [float ] = None ,
210207 minimum_reranker_score : Optional [float ] = None ,
211208 use_query_rewriting : Optional [bool ] = None ,
212- ) -> List [Document ]:
209+ ) -> list [Document ]:
213210 search_text = query_text if use_text_search else ""
214211 search_vectors = vectors if use_vector_search else []
215212 if use_semantic_ranker :
@@ -248,7 +245,7 @@ async def search(
248245 sourcefile = document .get ("sourcefile" ),
249246 oids = document .get ("oids" ),
250247 groups = document .get ("groups" ),
251- captions = cast (List [QueryCaptionResult ], document .get ("@search.captions" )),
248+ captions = cast (list [QueryCaptionResult ], document .get ("@search.captions" )),
252249 score = document .get ("@search.score" ),
253250 reranker_score = document .get ("@search.reranker_score" ),
254251 )
@@ -266,7 +263,7 @@ async def search(
266263 return qualified_documents
267264
268265 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
270267 ) -> list [str ]:
271268
272269 def nonewlines (s : str ) -> str :
@@ -358,13 +355,13 @@ def create_chat_completion(
358355 overrides : dict [str , Any ],
359356 response_token_limit : int ,
360357 should_stream : bool = False ,
361- tools : Optional [List [ChatCompletionToolParam ]] = None ,
358+ tools : Optional [list [ChatCompletionToolParam ]] = None ,
362359 temperature : Optional [float ] = None ,
363360 n : Optional [int ] = None ,
364361 reasoning_effort : Optional [ChatCompletionReasoningEffort ] = None ,
365362 ) -> Union [Awaitable [ChatCompletion ], Awaitable [AsyncStream [ChatCompletionChunk ]]]:
366363 if chatgpt_model in self .GPT_REASONING_MODELS :
367- params : Dict [str , Any ] = {
364+ params : dict [str , Any ] = {
368365 # max_tokens is not supported
369366 "max_completion_tokens" : response_token_limit
370367 }
@@ -400,14 +397,14 @@ def create_chat_completion(
400397 def format_thought_step_for_chatcompletion (
401398 self ,
402399 title : str ,
403- messages : List [ChatCompletionMessageParam ],
400+ messages : list [ChatCompletionMessageParam ],
404401 overrides : dict [str , Any ],
405402 model : str ,
406403 deployment : Optional [str ],
407404 usage : Optional [CompletionUsage ] = None ,
408405 reasoning_effort : Optional [ChatCompletionReasoningEffort ] = None ,
409406 ) -> ThoughtStep :
410- properties : Dict [str , Any ] = {"model" : model }
407+ properties : dict [str , Any ] = {"model" : model }
411408 if deployment :
412409 properties ["deployment" ] = deployment
413410 # Only add reasoning_effort setting if the model supports it
0 commit comments