Skip to content

feat: GenAI Client(evals) - add import_evaluation_set #7061

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: GenAI Client(evals) - add import_evaluation_set #7061

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
277 changes: 277 additions & 0 deletions agentplatform/_genai/evals.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -851,6 +851,32 @@ def _GetEvaluationSetParameters_to_vertex(
return to_object


def _ImportEvaluationSetParameters_to_vertex(
from_object: Union[dict[str, Any], object],
parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
to_object: dict[str, Any] = {}
if getv(from_object, ["evaluation_set"]) is not None:
setv(to_object, ["evaluationSet"], getv(from_object, ["evaluation_set"]))

if getv(from_object, ["gcs_destination"]) is not None:
setv(to_object, ["gcsDestination"], getv(from_object, ["gcs_destination"]))

if getv(from_object, ["gcs_source"]) is not None:
setv(to_object, ["gcsSource"], getv(from_object, ["gcs_source"]))

if getv(from_object, ["inline_source"]) is not None:
setv(to_object, ["inlineSource"], getv(from_object, ["inline_source"]))

if getv(from_object, ["cloud_trace_source"]) is not None:
setv(to_object, ["cloudTraceSource"], getv(from_object, ["cloud_trace_source"]))

if getv(from_object, ["config"]) is not None:
setv(to_object, ["config"], getv(from_object, ["config"]))

return to_object


def _ListEvaluationExperimentsConfig_to_vertex(
from_object: Union[dict[str, Any], object],
parent_object: Optional[dict[str, Any]] = None,
Expand Down Expand Up @@ -2449,6 +2475,86 @@ def _get_evaluation_item(
self._api_client._verify_response(return_value)
return return_value

def _import_evaluation_set(
self,
*,
evaluation_set: types.EvaluationSetOrDict,
gcs_destination: genai_types.GcsDestinationOrDict,
gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None,
inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None,
cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None,
config: Optional[types.ImportEvaluationSetConfigOrDict] = None,
) -> types.ImportEvaluationSetOperation:
"""
Imports data into an EvaluationSet.
"""

parameter_model = types._ImportEvaluationSetParameters(
evaluation_set=evaluation_set,
gcs_destination=gcs_destination,
gcs_source=gcs_source,
inline_source=inline_source,
cloud_trace_source=cloud_trace_source,
config=config,
)

request_url_dict: Optional[dict[str, str]]
if not self._api_client.vertexai:
raise ValueError(
"This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
)
else:
request_dict = _ImportEvaluationSetParameters_to_vertex(parameter_model)
request_url_dict = request_dict.get("_url")
if request_url_dict:
path = "evaluationSets:import".format_map(request_url_dict)
else:
path = "evaluationSets:import"

query_params = request_dict.get("_query")
if query_params:
path = f"{path}?{urlencode(query_params)}"
# TODO: remove the hack that pops config.
request_dict.pop("config", None)

http_options: Optional[types.HttpOptions] = None
if (
parameter_model.config is not None
and parameter_model.config.http_options is not None
):
http_options = parameter_model.config.http_options

request_dict = _common.convert_to_dict(request_dict)
request_dict = _common.encode_unserializable_types(request_dict)

response = self._api_client.request("post", path, request_dict, http_options)

response_dict = {} if not response.body else json.loads(response.body)

return_value = types.ImportEvaluationSetOperation._from_response(
response=response_dict,
kwargs=(
{
"config": {
"response_schema": getattr(
parameter_model.config, "response_schema", None
),
"response_json_schema": getattr(
parameter_model.config, "response_json_schema", None
),
"include_all_fields": getattr(
parameter_model.config, "include_all_fields", None
),
}
}
if getattr(parameter_model, "config", None)
else {}
),
)

self._api_client._verify_response(return_value)
return return_value

def list_evaluation_experiments(
self, *, config: Optional[types.ListEvaluationExperimentsConfigOrDict] = None
) -> types.ListEvaluationExperimentsResponse:
Expand Down Expand Up @@ -3551,6 +3657,50 @@ def create_evaluation_set(
config=config,
)

def import_evaluation_set(
self,
*,
evaluation_set: types.EvaluationSetOrDict,
gcs_destination: genai_types.GcsDestinationOrDict,
gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None,
inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None,
cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None,
config: Optional[types.ImportEvaluationSetConfigOrDict] = None,
) -> types.ImportEvaluationSetOperation:
"""Imports data into an EvaluationSet.

This is a long-running operation that imports data from a source (such as
OpenTelemetry traces) into a managed EvaluationSet.

Args:
evaluation_set: The EvaluationSet to create. Used to specify
``display_name`` and ``metadata``. The ``evaluation_items`` field is
ignored and populated by the import process.
gcs_destination: The Cloud Storage location where the resulting
EvaluationItem payloads will be stored.
gcs_source: The Cloud Storage source of the input data. Exactly one of
``gcs_source``, ``inline_source``, or ``cloud_trace_source`` must be
provided.
inline_source: The inline source for small payloads (< 4MB). Exactly one
of ``gcs_source``, ``inline_source``, or ``cloud_trace_source`` must be
provided.
cloud_trace_source: The Cloud Trace source for loading traces directly
from Cloud Trace. Exactly one of ``gcs_source``, ``inline_source``, or
``cloud_trace_source`` must be provided.
config: The optional configuration for the import operation.

Returns:
The long-running operation for the import.
"""
return self._import_evaluation_set(
evaluation_set=evaluation_set,
gcs_destination=gcs_destination,
gcs_source=gcs_source,
inline_source=inline_source,
cloud_trace_source=cloud_trace_source,
config=config,
)

def generate_conversation_scenarios(
self,
*,
Expand Down Expand Up @@ -5119,6 +5269,88 @@ async def _get_evaluation_item(
self._api_client._verify_response(return_value)
return return_value

async def _import_evaluation_set(
self,
*,
evaluation_set: types.EvaluationSetOrDict,
gcs_destination: genai_types.GcsDestinationOrDict,
gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None,
inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None,
cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None,
config: Optional[types.ImportEvaluationSetConfigOrDict] = None,
) -> types.ImportEvaluationSetOperation:
"""
Imports data into an EvaluationSet.
"""

parameter_model = types._ImportEvaluationSetParameters(
evaluation_set=evaluation_set,
gcs_destination=gcs_destination,
gcs_source=gcs_source,
inline_source=inline_source,
cloud_trace_source=cloud_trace_source,
config=config,
)

request_url_dict: Optional[dict[str, str]]
if not self._api_client.vertexai:
raise ValueError(
"This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
)
else:
request_dict = _ImportEvaluationSetParameters_to_vertex(parameter_model)
request_url_dict = request_dict.get("_url")
if request_url_dict:
path = "evaluationSets:import".format_map(request_url_dict)
else:
path = "evaluationSets:import"

query_params = request_dict.get("_query")
if query_params:
path = f"{path}?{urlencode(query_params)}"
# TODO: remove the hack that pops config.
request_dict.pop("config", None)

http_options: Optional[types.HttpOptions] = None
if (
parameter_model.config is not None
and parameter_model.config.http_options is not None
):
http_options = parameter_model.config.http_options

request_dict = _common.convert_to_dict(request_dict)
request_dict = _common.encode_unserializable_types(request_dict)

response = await self._api_client.async_request(
"post", path, request_dict, http_options
)

response_dict = {} if not response.body else json.loads(response.body)

return_value = types.ImportEvaluationSetOperation._from_response(
response=response_dict,
kwargs=(
{
"config": {
"response_schema": getattr(
parameter_model.config, "response_schema", None
),
"response_json_schema": getattr(
parameter_model.config, "response_json_schema", None
),
"include_all_fields": getattr(
parameter_model.config, "include_all_fields", None
),
}
}
if getattr(parameter_model, "config", None)
else {}
),
)

self._api_client._verify_response(return_value)
return return_value

async def list_evaluation_experiments(
self, *, config: Optional[types.ListEvaluationExperimentsConfigOrDict] = None
) -> types.ListEvaluationExperimentsResponse:
Expand Down Expand Up @@ -5837,6 +6069,51 @@ async def create_evaluation_set(
)
return result

async def import_evaluation_set(
self,
*,
evaluation_set: types.EvaluationSetOrDict,
gcs_destination: genai_types.GcsDestinationOrDict,
gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None,
inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None,
cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None,
config: Optional[types.ImportEvaluationSetConfigOrDict] = None,
) -> types.ImportEvaluationSetOperation:
"""Imports data into an EvaluationSet.

This is a long-running operation that imports data from a source (such as
OpenTelemetry traces) into a managed EvaluationSet.

Args:
evaluation_set: The EvaluationSet to create. Used to specify
``display_name`` and ``metadata``. The ``evaluation_items`` field is
ignored and populated by the import process.
gcs_destination: The Cloud Storage location where the resulting
EvaluationItem payloads will be stored.
gcs_source: The Cloud Storage source of the input data. Exactly one of
``gcs_source``, ``inline_source``, or ``cloud_trace_source`` must be
provided.
inline_source: The inline source for small payloads (< 4MB). Exactly one
of ``gcs_source``, ``inline_source``, or ``cloud_trace_source`` must be
provided.
cloud_trace_source: The Cloud Trace source for loading traces directly
from Cloud Trace. Exactly one of ``gcs_source``, ``inline_source``, or
``cloud_trace_source`` must be provided.
config: The optional configuration for the import operation.

Returns:
The long-running operation for the import.
"""
result = await self._import_evaluation_set(
evaluation_set=evaluation_set,
gcs_destination=gcs_destination,
gcs_source=gcs_source,
inline_source=inline_source,
cloud_trace_source=cloud_trace_source,
config=config,
)
return result

async def generate_conversation_scenarios(
self,
*,
Expand Down
Loading
Loading