Skip to content

Commit b90e084

Browse files
sararobcopybara-github
sararob
authored and
copybara-github
committed
feat: GenAI client - Add retrieve_contexts and ask_contexts methods to RAG module
PiperOrigin-RevId: 924738166
1 parent 94cd20c commit b90e084

5 files changed

Lines changed: 1162 additions & 0 deletions

File tree

‎agentplatform/_genai/rag.py‎

Lines changed: 339 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
from google.genai import _api_module
2424
from google.genai import _common
25+
from google.genai import types as genai_types
2526
from google.genai._common import get_value_by_path as getv
2627
from google.genai._common import set_value_by_path as setv
2728

@@ -30,6 +31,23 @@
3031
logger = logging.getLogger("agentplatform_genai.rag")
3132

3233

34+
def _AskContextsRequestParameters_to_vertex(
35+
from_object: Union[dict[str, Any], object],
36+
parent_object: Optional[dict[str, Any]] = None,
37+
) -> dict[str, Any]:
38+
to_object: dict[str, Any] = {}
39+
if getv(from_object, ["query"]) is not None:
40+
setv(to_object, ["query"], getv(from_object, ["query"]))
41+
42+
if getv(from_object, ["config"]) is not None:
43+
setv(to_object, ["config"], getv(from_object, ["config"]))
44+
45+
if getv(from_object, ["tools"]) is not None:
46+
setv(to_object, ["tools"], getv(from_object, ["tools"]))
47+
48+
return to_object
49+
50+
3351
def _CreateRagCorpusRequestParameters_to_vertex(
3452
from_object: Union[dict[str, Any], object],
3553
parent_object: Optional[dict[str, Any]] = None,
@@ -313,8 +331,99 @@ def _RagCorpus_to_vertex(
313331
return to_object
314332

315333

334+
def _RetrieveRagContextsRequestParameters_to_vertex(
335+
from_object: Union[dict[str, Any], object],
336+
parent_object: Optional[dict[str, Any]] = None,
337+
) -> dict[str, Any]:
338+
to_object: dict[str, Any] = {}
339+
if getv(from_object, ["vertex_rag_store"]) is not None:
340+
setv(to_object, ["vertexRagStore"], getv(from_object, ["vertex_rag_store"]))
341+
342+
if getv(from_object, ["query"]) is not None:
343+
setv(to_object, ["query"], getv(from_object, ["query"]))
344+
345+
if getv(from_object, ["config"]) is not None:
346+
setv(to_object, ["config"], getv(from_object, ["config"]))
347+
348+
return to_object
349+
350+
316351
class Rag(_api_module.BaseModule):
317352

353+
def ask_contexts(
354+
self,
355+
*,
356+
query: types.RagQueryOrDict,
357+
config: Optional[types.AskContextsConfigOrDict] = None,
358+
tools: Optional[list[genai_types.ToolOrDict]] = None,
359+
) -> types.AskContextsResponse:
360+
"""
361+
Asks a RAG Contexts.
362+
"""
363+
364+
parameter_model = types._AskContextsRequestParameters(
365+
query=query,
366+
config=config,
367+
tools=tools,
368+
)
369+
370+
request_url_dict: Optional[dict[str, str]]
371+
if not self._api_client.vertexai:
372+
raise ValueError(
373+
"This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
374+
)
375+
else:
376+
request_dict = _AskContextsRequestParameters_to_vertex(parameter_model)
377+
request_url_dict = request_dict.get("_url")
378+
if request_url_dict:
379+
path = ":askContexts".format_map(request_url_dict)
380+
else:
381+
path = ":askContexts"
382+
383+
query_params = request_dict.get("_query")
384+
if query_params:
385+
path = f"{path}?{urlencode(query_params)}"
386+
# TODO: remove the hack that pops config.
387+
request_dict.pop("config", None)
388+
389+
http_options: Optional[types.HttpOptions] = None
390+
if (
391+
parameter_model.config is not None
392+
and parameter_model.config.http_options is not None
393+
):
394+
http_options = parameter_model.config.http_options
395+
396+
request_dict = _common.convert_to_dict(request_dict)
397+
request_dict = _common.encode_unserializable_types(request_dict)
398+
399+
response = self._api_client.request("post", path, request_dict, http_options)
400+
401+
response_dict = {} if not response.body else json.loads(response.body)
402+
403+
return_value = types.AskContextsResponse._from_response(
404+
response=response_dict,
405+
kwargs=(
406+
{
407+
"config": {
408+
"response_schema": getattr(
409+
parameter_model.config, "response_schema", None
410+
),
411+
"response_json_schema": getattr(
412+
parameter_model.config, "response_json_schema", None
413+
),
414+
"include_all_fields": getattr(
415+
parameter_model.config, "include_all_fields", None
416+
),
417+
}
418+
}
419+
if getattr(parameter_model, "config", None)
420+
else {}
421+
),
422+
)
423+
424+
self._api_client._verify_response(return_value)
425+
return return_value
426+
318427
def _create_corpus(
319428
self,
320429
*,
@@ -743,9 +852,161 @@ def get_config(
743852
self._api_client._verify_response(return_value)
744853
return return_value
745854

855+
def retrieve_contexts(
856+
self,
857+
*,
858+
vertex_rag_store: Optional[types.VertexRagStoreOrDict] = None,
859+
query: types.RagQueryOrDict,
860+
config: Optional[types.RetrieveContextsConfigOrDict] = None,
861+
) -> types.RetrieveContextsResponse:
862+
"""
863+
Retrieves a RAG Contexts.
864+
"""
865+
866+
parameter_model = types._RetrieveRagContextsRequestParameters(
867+
vertex_rag_store=vertex_rag_store,
868+
query=query,
869+
config=config,
870+
)
871+
872+
request_url_dict: Optional[dict[str, str]]
873+
if not self._api_client.vertexai:
874+
raise ValueError(
875+
"This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
876+
)
877+
else:
878+
request_dict = _RetrieveRagContextsRequestParameters_to_vertex(
879+
parameter_model
880+
)
881+
request_url_dict = request_dict.get("_url")
882+
if request_url_dict:
883+
path = ":retrieveContexts".format_map(request_url_dict)
884+
else:
885+
path = ":retrieveContexts"
886+
887+
query_params = request_dict.get("_query")
888+
if query_params:
889+
path = f"{path}?{urlencode(query_params)}"
890+
# TODO: remove the hack that pops config.
891+
request_dict.pop("config", None)
892+
893+
http_options: Optional[types.HttpOptions] = None
894+
if (
895+
parameter_model.config is not None
896+
and parameter_model.config.http_options is not None
897+
):
898+
http_options = parameter_model.config.http_options
899+
900+
request_dict = _common.convert_to_dict(request_dict)
901+
request_dict = _common.encode_unserializable_types(request_dict)
902+
903+
response = self._api_client.request("post", path, request_dict, http_options)
904+
905+
response_dict = {} if not response.body else json.loads(response.body)
906+
907+
return_value = types.RetrieveContextsResponse._from_response(
908+
response=response_dict,
909+
kwargs=(
910+
{
911+
"config": {
912+
"response_schema": getattr(
913+
parameter_model.config, "response_schema", None
914+
),
915+
"response_json_schema": getattr(
916+
parameter_model.config, "response_json_schema", None
917+
),
918+
"include_all_fields": getattr(
919+
parameter_model.config, "include_all_fields", None
920+
),
921+
}
922+
}
923+
if getattr(parameter_model, "config", None)
924+
else {}
925+
),
926+
)
927+
928+
self._api_client._verify_response(return_value)
929+
return return_value
930+
746931

747932
class AsyncRag(_api_module.BaseModule):
748933

934+
async def ask_contexts(
935+
self,
936+
*,
937+
query: types.RagQueryOrDict,
938+
config: Optional[types.AskContextsConfigOrDict] = None,
939+
tools: Optional[list[genai_types.ToolOrDict]] = None,
940+
) -> types.AskContextsResponse:
941+
"""
942+
Asks a RAG Contexts.
943+
"""
944+
945+
parameter_model = types._AskContextsRequestParameters(
946+
query=query,
947+
config=config,
948+
tools=tools,
949+
)
950+
951+
request_url_dict: Optional[dict[str, str]]
952+
if not self._api_client.vertexai:
953+
raise ValueError(
954+
"This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
955+
)
956+
else:
957+
request_dict = _AskContextsRequestParameters_to_vertex(parameter_model)
958+
request_url_dict = request_dict.get("_url")
959+
if request_url_dict:
960+
path = ":askContexts".format_map(request_url_dict)
961+
else:
962+
path = ":askContexts"
963+
964+
query_params = request_dict.get("_query")
965+
if query_params:
966+
path = f"{path}?{urlencode(query_params)}"
967+
# TODO: remove the hack that pops config.
968+
request_dict.pop("config", None)
969+
970+
http_options: Optional[types.HttpOptions] = None
971+
if (
972+
parameter_model.config is not None
973+
and parameter_model.config.http_options is not None
974+
):
975+
http_options = parameter_model.config.http_options
976+
977+
request_dict = _common.convert_to_dict(request_dict)
978+
request_dict = _common.encode_unserializable_types(request_dict)
979+
980+
response = await self._api_client.async_request(
981+
"post", path, request_dict, http_options
982+
)
983+
984+
response_dict = {} if not response.body else json.loads(response.body)
985+
986+
return_value = types.AskContextsResponse._from_response(
987+
response=response_dict,
988+
kwargs=(
989+
{
990+
"config": {
991+
"response_schema": getattr(
992+
parameter_model.config, "response_schema", None
993+
),
994+
"response_json_schema": getattr(
995+
parameter_model.config, "response_json_schema", None
996+
),
997+
"include_all_fields": getattr(
998+
parameter_model.config, "include_all_fields", None
999+
),
1000+
}
1001+
}
1002+
if getattr(parameter_model, "config", None)
1003+
else {}
1004+
),
1005+
)
1006+
1007+
self._api_client._verify_response(return_value)
1008+
return return_value
1009+
7491010
async def _create_corpus(
7501011
self,
7511012
*,
@@ -1185,3 +1446,81 @@ async def get_config(
11851446

11861447
self._api_client._verify_response(return_value)
11871448
return return_value
1449+
1450+
async def retrieve_contexts(
1451+
self,
1452+
*,
1453+
vertex_rag_store: Optional[types.VertexRagStoreOrDict] = None,
1454+
query: types.RagQueryOrDict,
1455+
config: Optional[types.RetrieveContextsConfigOrDict] = None,
1456+
) -> types.RetrieveContextsResponse:
1457+
"""
1458+
Retrieves a RAG Contexts.
1459+
"""
1460+
1461+
parameter_model = types._RetrieveRagContextsRequestParameters(
1462+
vertex_rag_store=vertex_rag_store,
1463+
query=query,
1464+
config=config,
1465+
)
1466+
1467+
request_url_dict: Optional[dict[str, str]]
1468+
if not self._api_client.vertexai:
1469+
raise ValueError(
1470+
"This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
1471+
)
1472+
else:
1473+
request_dict = _RetrieveRagContextsRequestParameters_to_vertex(
1474+
parameter_model
1475+
)
1476+
request_url_dict = request_dict.get("_url")
1477+
if request_url_dict:
1478+
path = ":retrieveContexts".format_map(request_url_dict)
1479+
else:
1480+
path = ":retrieveContexts"
1481+
1482+
query_params = request_dict.get("_query")
1483+
if query_params:
1484+
path = f"{path}?{urlencode(query_params)}"
1485+
# TODO: remove the hack that pops config.
1486+
request_dict.pop("config", None)
1487+
1488+
http_options: Optional[types.HttpOptions] = None
1489+
if (
1490+
parameter_model.config is not None
1491+
and parameter_model.config.http_options is not None
1492+
):
1493+
http_options = parameter_model.config.http_options
1494+
1495+
request_dict = _common.convert_to_dict(request_dict)
1496+
request_dict = _common.encode_unserializable_types(request_dict)
1497+
1498+
response = await self._api_client.async_request(
1499+
"post", path, request_dict, http_options
1500+
)
1501+
1502+
response_dict = {} if not response.body else json.loads(response.body)
1503+
1504+
return_value = types.RetrieveContextsResponse._from_response(
1505+
response=response_dict,
1506+
kwargs=(
1507+
{
1508+
"config": {
1509+
"response_schema": getattr(
1510+
parameter_model.config, "response_schema", None
1511+
),
1512+
"response_json_schema": getattr(
1513+
parameter_model.config, "response_json_schema", None
1514+
),
1515+
"include_all_fields": getattr(
1516+
parameter_model.config, "include_all_fields", None
1517+
),
1518+
}
1519+
}
1520+
if getattr(parameter_model, "config", None)
1521+
else {}
1522+
),
1523+
)
1524+
1525+
self._api_client._verify_response(return_value)
1526+
return return_value

0 commit comments

Comments
 (0)