21
21
from typing_extensions import TypedDict
22
22
23
23
import google .ai .generativelanguage as glm
24
+ from google .generativeai import protos
24
25
25
26
from google .generativeai .client import (
26
27
get_default_generative_client ,
35
36
36
37
DEFAULT_ANSWER_MODEL = "models/aqa"
37
38
38
- AnswerStyle = glm .GenerateAnswerRequest .AnswerStyle
39
+ AnswerStyle = protos .GenerateAnswerRequest .AnswerStyle
39
40
40
41
AnswerStyleOptions = Union [int , str , AnswerStyle ]
41
42
@@ -66,28 +67,30 @@ def to_answer_style(x: AnswerStyleOptions) -> AnswerStyle:
66
67
67
68
68
69
GroundingPassageOptions = (
69
- Union [glm .GroundingPassage , tuple [str , content_types .ContentType ], content_types .ContentType ],
70
+ Union [
71
+ protos .GroundingPassage , tuple [str , content_types .ContentType ], content_types .ContentType
72
+ ],
70
73
)
71
74
72
75
GroundingPassagesOptions = Union [
73
- glm .GroundingPassages ,
76
+ protos .GroundingPassages ,
74
77
Iterable [GroundingPassageOptions ],
75
78
Mapping [str , content_types .ContentType ],
76
79
]
77
80
78
81
79
- def _make_grounding_passages (source : GroundingPassagesOptions ) -> glm .GroundingPassages :
82
+ def _make_grounding_passages (source : GroundingPassagesOptions ) -> protos .GroundingPassages :
80
83
"""
81
- Converts the `source` into a `glm .GroundingPassage`. A `GroundingPassages` contains a list of
82
- `glm .GroundingPassage` objects, which each contain a `glm .Contant` and a string `id`.
84
+ Converts the `source` into a `protos .GroundingPassage`. A `GroundingPassages` contains a list of
85
+ `protos .GroundingPassage` objects, which each contain a `protos .Contant` and a string `id`.
83
86
84
87
Args:
85
- source: `Content` or a `GroundingPassagesOptions` that will be converted to glm .GroundingPassages.
88
+ source: `Content` or a `GroundingPassagesOptions` that will be converted to protos .GroundingPassages.
86
89
87
90
Return:
88
- `glm .GroundingPassages` to be passed into `glm .GenerateAnswer`.
91
+ `protos .GroundingPassages` to be passed into `protos .GenerateAnswer`.
89
92
"""
90
- if isinstance (source , glm .GroundingPassages ):
93
+ if isinstance (source , protos .GroundingPassages ):
91
94
return source
92
95
93
96
if not isinstance (source , Iterable ):
@@ -100,19 +103,19 @@ def _make_grounding_passages(source: GroundingPassagesOptions) -> glm.GroundingP
100
103
source = source .items ()
101
104
102
105
for n , data in enumerate (source ):
103
- if isinstance (data , glm .GroundingPassage ):
106
+ if isinstance (data , protos .GroundingPassage ):
104
107
passages .append (data )
105
108
elif isinstance (data , tuple ):
106
109
id , content = data # tuple must have exactly 2 items.
107
110
passages .append ({"id" : id , "content" : content_types .to_content (content )})
108
111
else :
109
112
passages .append ({"id" : str (n ), "content" : content_types .to_content (data )})
110
113
111
- return glm .GroundingPassages (passages = passages )
114
+ return protos .GroundingPassages (passages = passages )
112
115
113
116
114
117
SourceNameType = Union [
115
- str , retriever_types .Corpus , glm .Corpus , retriever_types .Document , glm .Document
118
+ str , retriever_types .Corpus , protos .Corpus , retriever_types .Document , protos .Document
116
119
]
117
120
118
121
@@ -127,15 +130,15 @@ class SemanticRetrieverConfigDict(TypedDict):
127
130
SemanticRetrieverConfigOptions = Union [
128
131
SourceNameType ,
129
132
SemanticRetrieverConfigDict ,
130
- glm .SemanticRetrieverConfig ,
133
+ protos .SemanticRetrieverConfig ,
131
134
]
132
135
133
136
134
137
def _maybe_get_source_name (source ) -> str | None :
135
138
if isinstance (source , str ):
136
139
return source
137
140
elif isinstance (
138
- source , (retriever_types .Corpus , glm .Corpus , retriever_types .Document , glm .Document )
141
+ source , (retriever_types .Corpus , protos .Corpus , retriever_types .Document , protos .Document )
139
142
):
140
143
return source .name
141
144
else :
@@ -145,8 +148,8 @@ def _maybe_get_source_name(source) -> str | None:
145
148
def _make_semantic_retriever_config (
146
149
source : SemanticRetrieverConfigOptions ,
147
150
query : content_types .ContentsType ,
148
- ) -> glm .SemanticRetrieverConfig :
149
- if isinstance (source , glm .SemanticRetrieverConfig ):
151
+ ) -> protos .SemanticRetrieverConfig :
152
+ if isinstance (source , protos .SemanticRetrieverConfig ):
150
153
return source
151
154
152
155
name = _maybe_get_source_name (source )
@@ -156,7 +159,7 @@ def _make_semantic_retriever_config(
156
159
source ["source" ] = _maybe_get_source_name (source ["source" ])
157
160
else :
158
161
raise TypeError (
159
- f"Invalid input: Failed to create a 'glm .SemanticRetrieverConfig' from the provided source. "
162
+ f"Invalid input: Failed to create a 'protos .SemanticRetrieverConfig' from the provided source. "
160
163
f"Received type: { type (source ).__name__ } , "
161
164
f"Received value: { source } "
162
165
)
@@ -166,7 +169,7 @@ def _make_semantic_retriever_config(
166
169
elif isinstance (source ["query" ], str ):
167
170
source ["query" ] = content_types .to_content (source ["query" ])
168
171
169
- return glm .SemanticRetrieverConfig (source )
172
+ return protos .SemanticRetrieverConfig (source )
170
173
171
174
172
175
def _make_generate_answer_request (
@@ -178,26 +181,26 @@ def _make_generate_answer_request(
178
181
answer_style : AnswerStyle | None = None ,
179
182
safety_settings : safety_types .SafetySettingOptions | None = None ,
180
183
temperature : float | None = None ,
181
- ) -> glm .GenerateAnswerRequest :
184
+ ) -> protos .GenerateAnswerRequest :
182
185
"""
183
- constructs a glm .GenerateAnswerRequest object by organizing the input parameters for the API call to generate a grounded answer from the model.
186
+ constructs a protos .GenerateAnswerRequest object by organizing the input parameters for the API call to generate a grounded answer from the model.
184
187
185
188
Args:
186
189
model: Name of the model used to generate the grounded response.
187
190
contents: Content of the current conversation with the model. For single-turn query, this is a
188
191
single question to answer. For multi-turn queries, this is a repeated field that contains
189
192
conversation history and the last `Content` in the list containing the question.
190
193
inline_passages: Grounding passages (a list of `Content`-like objects or `(id, content)` pairs,
191
- or a `glm .GroundingPassages`) to send inline with the request. Exclusive with `semantic_retreiver`,
194
+ or a `protos .GroundingPassages`) to send inline with the request. Exclusive with `semantic_retreiver`,
192
195
one must be set, but not both.
193
- semantic_retriever: A Corpus, Document, or `glm .SemanticRetrieverConfig` to use for grounding. Exclusive with
196
+ semantic_retriever: A Corpus, Document, or `protos .SemanticRetrieverConfig` to use for grounding. Exclusive with
194
197
`inline_passages`, one must be set, but not both.
195
198
answer_style: Style for grounded answers.
196
199
safety_settings: Safety settings for generated output.
197
200
temperature: The temperature for randomness in the output.
198
201
199
202
Returns:
200
- Call for glm .GenerateAnswerRequest().
203
+ Call for protos .GenerateAnswerRequest().
201
204
"""
202
205
model = model_types .make_model_name (model )
203
206
@@ -224,7 +227,7 @@ def _make_generate_answer_request(
224
227
if answer_style :
225
228
answer_style = to_answer_style (answer_style )
226
229
227
- return glm .GenerateAnswerRequest (
230
+ return protos .GenerateAnswerRequest (
228
231
model = model ,
229
232
contents = contents ,
230
233
inline_passages = inline_passages ,
@@ -273,9 +276,9 @@ def generate_answer(
273
276
contents: The question to be answered by the model, grounded in the
274
277
provided source.
275
278
inline_passages: Grounding passages (a list of `Content`-like objects or (id, content) pairs,
276
- or a `glm .GroundingPassages`) to send inline with the request. Exclusive with `semantic_retreiver`,
279
+ or a `protos .GroundingPassages`) to send inline with the request. Exclusive with `semantic_retreiver`,
277
280
one must be set, but not both.
278
- semantic_retriever: A Corpus, Document, or `glm .SemanticRetrieverConfig` to use for grounding. Exclusive with
281
+ semantic_retriever: A Corpus, Document, or `protos .SemanticRetrieverConfig` to use for grounding. Exclusive with
279
282
`inline_passages`, one must be set, but not both.
280
283
answer_style: Style in which the grounded answer should be returned.
281
284
safety_settings: Safety settings for generated output. Defaults to None.
@@ -327,9 +330,9 @@ async def generate_answer_async(
327
330
contents: The question to be answered by the model, grounded in the
328
331
provided source.
329
332
inline_passages: Grounding passages (a list of `Content`-like objects or (id, content) pairs,
330
- or a `glm .GroundingPassages`) to send inline with the request. Exclusive with `semantic_retreiver`,
333
+ or a `protos .GroundingPassages`) to send inline with the request. Exclusive with `semantic_retreiver`,
331
334
one must be set, but not both.
332
- semantic_retriever: A Corpus, Document, or `glm .SemanticRetrieverConfig` to use for grounding. Exclusive with
335
+ semantic_retriever: A Corpus, Document, or `protos .SemanticRetrieverConfig` to use for grounding. Exclusive with
333
336
`inline_passages`, one must be set, but not both.
334
337
answer_style: Style in which the grounded answer should be returned.
335
338
safety_settings: Safety settings for generated output. Defaults to None.
0 commit comments