Coverage for qreward / client / patch_openai.py: 100.00%

1from typing import Callable, List, Union (empty)

2 

3from typing_extensions import Literal (empty)

4 

5from openai._base_client import make_request_options (empty)

6from openai._models import BaseModel (empty)

7from openai._types import Omit (empty)

8from openai._utils import is_given, maybe_transform (empty)

9from openai.types import embedding_create_params (empty)

10from openai.types.create_embedding_response import Embedding, Usage (empty)

11from openai.types.embedding_model import EmbeddingModel (empty)

12 

13 

14class HackCreateEmbeddingResponse(BaseModel): (empty)

15 embeddings: List[Embedding] (empty)

16 

17 model: str (empty)

18 """The name of the model used to generate the embedding.""" (empty)

19 

20 object: Literal["list"] (empty)

21 """The object type, which is always "list".""" (empty)

22 

23 usage: Usage (empty)

24 """The usage information for the request.""" (empty)

25 

26 

27def hack_parser( (empty)

28 obj: HackCreateEmbeddingResponse, 

29) -> HackCreateEmbeddingResponse: 

30 if not obj.embeddings: (empty)

31 raise ValueError("No embedding data received") (empty)

32 

33 return obj (empty)

34 

35 

36# Store the original method reference for unpatch support 

37_original_async_embeddings_create = None (empty)

38 

39 

40# ===== PATCH 函数 ===== 

41def patch_openai_embeddings( (empty)

42 custom_return_cls: type = HackCreateEmbeddingResponse, 

43 custom_parser: Callable = hack_parser, 

44) -> None: 

45 """Replace OpenAI AsyncEmbeddings.create with a custom implementation. 

46 

47 Replaces the method at runtime to return a custom Response class 

48 and use a custom parser. Can be reversed with unpatch_openai_embeddings(). 

49 

50 Args: 

51 custom_return_cls: Custom return class (BaseModel subclass) 

52 custom_parser: Custom parser function that receives and returns 

53 a custom_return_cls instance 

54 """ 

55 global _original_async_embeddings_create 

56 from openai.resources.embeddings import AsyncEmbeddings (empty)

57 

58 # Save original method only on first patch to support idempotent unpatch. 

59 # Use getattr to handle test scenarios where AsyncEmbeddings may be mocked 

60 # with a dummy class that lacks a 'create' attribute. 

61 if _original_async_embeddings_create is None: (empty)

62 _original_async_embeddings_create = getattr( (empty)

63 AsyncEmbeddings, "create", None 

64 ) 

65 

66 async def patched_create( (empty)

67 self, 

68 *, 

69 input, 

70 model: Union[str, EmbeddingModel], 

71 dimensions: int | Omit = Omit(), 

72 encoding_format: Literal["float", "base64"] | Omit = Omit(), 

73 user: str | Omit = Omit(), 

74 extra_headers=None, 

75 extra_query=None, 

76 extra_body=None, 

77 timeout=None, 

78 ): 

79 params = { (empty)

80 "input": input, 

81 "model": model, 

82 "user": user, 

83 "dimensions": dimensions, 

84 "encoding_format": encoding_format, 

85 } 

86 if not is_given(encoding_format): (empty)

87 params["encoding_format"] = "base64" (empty)

88 

89 return await self._post( (empty)

90 "/embeddings", 

91 body=maybe_transform( 

92 params, 

93 embedding_create_params.EmbeddingCreateParams, 

94 ), 

95 options=make_request_options( 

96 extra_headers=extra_headers, 

97 extra_query=extra_query, 

98 extra_body=extra_body, 

99 timeout=timeout, 

100 post_parser=custom_parser, 

101 ), 

102 cast_to=custom_return_cls, 

103 ) 

104 

105 AsyncEmbeddings.create = patched_create (empty)

106 

107 

108def unpatch_openai_embeddings() -> None: (empty)

109 """Restore OpenAI AsyncEmbeddings.create to its original implementation. 

110 

111 This reverses the effect of patch_openai_embeddings(). If patch has not 

112 been applied, this function is a no-op (idempotent). 

113 """ 

114 global _original_async_embeddings_create 

115 if _original_async_embeddings_create is None: (empty)

116 return (empty)

117 

118 from openai.resources.embeddings import AsyncEmbeddings (empty)

119 

120 AsyncEmbeddings.create = _original_async_embeddings_create (empty)

121 _original_async_embeddings_create = None (empty)