Skip to content

feat: Renaming. #6979

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: Renaming. #6979

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
22 changes: 11 additions & 11 deletions agentplatform/_genai/_evals_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ def _get_api_client_with_location(

def _get_agent_engine_instance(
agent_name: str, api_client: BaseApiClient
) -> Union[types.AgentEngine, Any]:
) -> Union[types.Runtime, Any]:
"""Gets or creates an agent engine instance for the current thread."""
if not hasattr(_thread_local_data, "agent_engine_instances"):
_thread_local_data.agent_engine_instances = {}
Expand All @@ -110,7 +110,7 @@ def _get_agent_engine_instance(
location=api_client.location,
)
_thread_local_data.agent_engine_instances[agent_name] = (
client.agent_engines.get(name=agent_name)
client.runtimes.get(name=agent_name)
)
return _thread_local_data.agent_engine_instances[agent_name]

Expand Down Expand Up @@ -614,7 +614,7 @@ def _execute_inference_concurrently(
model_or_fn: Optional[Union[str, Callable[[Any], Any]]] = None,
gemini_config: Optional[genai_types.GenerateContentConfig] = None,
inference_fn: Optional[Callable[..., Any]] = None,
agent_engine: Optional[Union[str, types.AgentEngine]] = None,
agent_engine: Optional[Union[str, types.Runtime]] = None,
agent: Optional["LlmAgent"] = None, # type: ignore # noqa: F821
user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None,
) -> list[
Expand Down Expand Up @@ -1404,7 +1404,7 @@ def _execute_inference(
api_client: BaseApiClient,
src: Union[str, pd.DataFrame],
model: Optional[Union[Callable[[Any], Any], str]] = None,
agent_engine: Optional[Union[str, types.AgentEngine]] = None,
agent_engine: Optional[Union[str, types.Runtime]] = None,
agent: Optional["LlmAgent"] = None, # type: ignore # noqa: F821
dest: Optional[str] = None,
config: Optional[genai_types.GenerateContentConfig] = None,
Expand Down Expand Up @@ -1499,7 +1499,7 @@ def _execute_inference(
f"Unsupported agent_engine type: {type(agent_engine)}. Expecting a"
" string (agent engine resource name in"
" 'projects/{project_id}/locations/{location_id}/reasoningEngines/{reasoning_engine_id}'"
" format) or a types.AgentEngine instance."
" format) or a types.Runtime instance."
)
if (
_evals_constant.INTERMEDIATE_EVENTS in prompt_dataset.columns
Expand Down Expand Up @@ -2102,7 +2102,7 @@ def _create_agent_results_dataframe(

def _run_agent_internal(
api_client: BaseApiClient,
agent_engine: Optional[Union[str, types.AgentEngine]],
agent_engine: Optional[Union[str, types.Runtime]],
agent: Optional["LlmAgent"], # type: ignore # noqa: F821
prompt_dataset: pd.DataFrame,
user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None,
Expand Down Expand Up @@ -2153,7 +2153,7 @@ def _run_agent_internal(

def _run_agent(
api_client: BaseApiClient,
agent_engine: Optional[Union[str, types.AgentEngine]],
agent_engine: Optional[Union[str, types.Runtime]],
agent: Optional["LlmAgent"], # type: ignore # noqa: F821
prompt_dataset: pd.DataFrame,
user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None,
Expand Down Expand Up @@ -2224,7 +2224,7 @@ def _run_agent(

def _create_agent_engine_session(
*,
agent_engine: types.AgentEngine,
agent_engine: types.Runtime,
user_id: str,
session_state: Optional[dict[str, Any]] = None,
) -> Any:
Expand Down Expand Up @@ -2271,7 +2271,7 @@ def _create_agent_engine_session(
operation = agent_engine.api_client.sessions.create(
name=agent_engine.api_resource.name,
user_id=user_id,
config=types.CreateAgentEngineSessionConfig(
config=types.CreateRuntimeSessionConfig(
session_state=session_state,
),
)
Expand All @@ -2293,7 +2293,7 @@ def _create_agent_engine_session(
def _execute_agent_run_with_retry(
row: pd.Series,
contents: Union[genai_types.ContentListUnion, genai_types.ContentListUnionDict],
agent_engine: types.AgentEngine,
agent_engine: types.Runtime,
max_retries: int = 3,
) -> Union[list[dict[str, Any]], dict[str, Any]]:
"""Executes agent run over agent engine for a single prompt."""
Expand Down Expand Up @@ -2340,7 +2340,7 @@ def _execute_agent_run_with_retry(
author=ag_event.author or "user",
invocation_id="history",
timestamp=base_ts + datetime.timedelta(seconds=i),
config=types.AppendAgentEngineSessionEventConfig(
config=types.AppendRuntimeSessionEventConfig(
content=ag_event.content,
),
)
Expand Down
84 changes: 39 additions & 45 deletions agentplatform/_genai/_agent_engines_utils.py → agentplatform/_genai/_runtimes_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Utility functions for agent engines."""
"""Utility functions for runtimes."""

import abc
import asyncio
Expand Down Expand Up @@ -312,7 +312,7 @@ def register_operations(self, **kwargs: Any) -> dict[str, list[str]]:
except (ImportError, AttributeError):
ADKAgent = None # type: ignore[no-redef]

_AgentEngineInterface = Union[
_RuntimeInterface = Union[
ADKAgent,
AsyncQueryable,
AsyncStreamQueryable,
Expand All @@ -328,7 +328,7 @@ class _ModuleAgentAttributes(TypedDict, total=False):
agent_name: str
register_operations: Dict[str, list[str]]
sys_paths: Optional[Sequence[str]]
agent: _AgentEngineInterface
agent: _RuntimeInterface


class ModuleAgent(Cloneable, OperationRegistrable):
Expand Down Expand Up @@ -441,24 +441,22 @@ class _RequirementsValidationResult(TypedDict):
actions: _RequirementsValidationActions


AgentEngineOperationUnion = Union[
genai_types.AgentEngineOperation,
genai_types.AgentEngineMemoryOperation,
genai_types.AgentEngineGenerateMemoriesOperation,
RuntimeOperationUnion = Union[
genai_types.RuntimeOperation,
genai_types.RuntimeMemoryOperation,
genai_types.RuntimeGenerateMemoriesOperation,
]


class GetOperationFunction(Protocol):
def __call__(
self, *, operation_name: str, **kwargs: Any
) -> AgentEngineOperationUnion:
def __call__(self, *, operation_name: str, **kwargs: Any) -> RuntimeOperationUnion:
pass


class GetAsyncOperationFunction(Protocol):
async def __call__(
self, *, operation_name: str, **kwargs: Any
) -> AgentEngineOperationUnion:
) -> RuntimeOperationUnion:
pass


Expand Down Expand Up @@ -490,8 +488,7 @@ def _get_reasoning_engine_id(operation_name: str = "", resource_name: str = "")
if match:
return match.group(1)
raise ValueError(
"Failed to parse reasoning engine ID from operation name: "
f"`{operation_name}`"
f"Failed to parse reasoning engine ID from operation name: `{operation_name}`"
)


Expand Down Expand Up @@ -611,7 +608,7 @@ def _compare_requirements(

def _generate_class_methods_spec_or_raise(
*,
agent: _AgentEngineInterface,
agent: _RuntimeInterface,
operations: Dict[str, List[str]],
) -> List[proto.Message]:
"""Generates a ReasoningEngineSpec based on the registered operations.
Expand Down Expand Up @@ -755,7 +752,8 @@ def _generate_schema(
)
# For a bidi endpoint, it requires an asyncio.Queue as the input, but
# it is not JSON serializable. We hence exclude it from the schema.
and param.annotation != asyncio.Queue and _is_pydantic_serializable(param)
and param.annotation != asyncio.Queue
and _is_pydantic_serializable(param)
}
parameters = pydantic.create_model(f.__name__, **fields_dict).schema()
# Postprocessing
Expand Down Expand Up @@ -811,7 +809,7 @@ def _generate_schema(
def _get_agent_framework(
*,
agent_framework: Optional[str],
agent: _AgentEngineInterface,
agent: _RuntimeInterface,
) -> Union[str, Any]:
"""Gets the agent framework to use.

Expand All @@ -823,7 +821,7 @@ def _get_agent_framework(
Args:
agent_framework (str):
The agent framework provided by the user.
agent (_AgentEngineInterface):
agent (_RuntimeInterface):
The agent engine instance.

Returns:
Expand Down Expand Up @@ -871,7 +869,7 @@ def _get_gcs_bucket(

def _get_registered_operations(
*,
agent: _AgentEngineInterface,
agent: _RuntimeInterface,
) -> dict[str, list[str]]:
"""Retrieves registered operations for a AgentEngine."""
if isinstance(agent, OperationRegistrable):
Expand Down Expand Up @@ -985,7 +983,7 @@ def _parse_constraints(

def _prepare(
*,
agent: Optional[_AgentEngineInterface],
agent: Optional[_RuntimeInterface],
requirements: Optional[Sequence[str]],
extra_packages: Optional[Sequence[str]],
project: str,
Expand Down Expand Up @@ -1042,7 +1040,7 @@ def _prepare(

def _register_api_methods_or_raise(
*,
agent_engine: genai_types.AgentEngine | genai_types.AgentEngineRuntimeRevision,
agent_engine: genai_types.Runtime | genai_types.RuntimeRevision,
wrap_operation_fn: Optional[
dict[str, Callable[[str, str], Callable[..., Any]]]
] = None,
Expand Down Expand Up @@ -1126,9 +1124,7 @@ def _register_api_methods_or_raise(
# Bind the method to the object.
if api_mode == _A2A_EXTENSION_MODE:
agent_card = operation_schema.get(_A2A_AGENT_CARD)
method = _wrap_operation(
method_name=method_name, agent_card=agent_card
) # type: ignore[call-arg]
method = _wrap_operation(method_name=method_name, agent_card=agent_card) # type: ignore[call-arg]
else:
method = _wrap_operation(method_name=method_name) # type: ignore[call-arg]
method.__name__ = method_name
Expand Down Expand Up @@ -1243,7 +1239,7 @@ def _to_proto(

def _upload_agent_engine(
*,
agent: _AgentEngineInterface,
agent: _RuntimeInterface,
gcs_bucket: _StorageBucket,
gcs_dir_name: str,
) -> None:
Expand Down Expand Up @@ -1412,7 +1408,7 @@ def _validate_staging_bucket_or_raise(*, staging_bucket: str) -> str:
"""Tries to validate the staging bucket."""
if not staging_bucket:
raise ValueError(
"Please provide a `staging_bucket` in `client.agent_engines.create(...)`."
"Please provide a `staging_bucket` in `client.runtimes.create(...)`."
)
if not staging_bucket.startswith("gs://"):
raise ValueError(f"{staging_bucket=} must start with `gs://`")
Expand Down Expand Up @@ -1475,8 +1471,8 @@ def _validate_requirements_or_raise(

def _validate_agent_or_raise(
*,
agent: _AgentEngineInterface,
) -> _AgentEngineInterface:
agent: _RuntimeInterface,
) -> _RuntimeInterface:
"""Tries to validate the agent engine.

The agent engine must have one of the following:
Expand All @@ -1503,9 +1499,9 @@ def _validate_agent_or_raise(

if isinstance(agent, BaseAgent):
logger.info("Deploying google.adk.agents.Agent as an application.")
from agentplatform import agent_engines
from agentplatform import frameworks

agent = agent_engines.AdkApp(agent=agent)
agent = frameworks.AdkApp(agent=agent)
except Exception:
pass
is_queryable = isinstance(agent, Queryable) and callable(agent.query)
Expand Down Expand Up @@ -1640,7 +1636,7 @@ def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]:
the `query` API.
"""

def _method(self: genai_types.AgentEngine, **kwargs) -> Any: # type: ignore[no-untyped-def]
def _method(self: genai_types.Runtime, **kwargs) -> Any: # type: ignore[no-untyped-def]
if not self.api_client:
raise ValueError("api_client is not initialized.")
if not self.api_resource:
Expand Down Expand Up @@ -1683,7 +1679,7 @@ def _wrap_async_query_operation(
"""

async def _method(
self: genai_types.AgentEngine, **kwargs: Any
self: genai_types.Runtime, **kwargs: Any
) -> Union[Coroutine[Any, Any, Any], Any]:
if not self.api_async_client:
raise ValueError("api_async_client is not initialized.")
Expand Down Expand Up @@ -1724,7 +1720,7 @@ def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterator[
the `stream_query` API.
"""

def _method(self: genai_types.AgentEngine, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-def]
def _method(self: genai_types.Runtime, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-def]
if not self.api_client:
raise ValueError("api_client is not initialized.")
if not self.api_resource:
Expand Down Expand Up @@ -1768,7 +1764,7 @@ def _wrap_async_stream_query_operation(
the `stream_query` API.
"""

async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]: # type: ignore[no-untyped-def]
async def _method(self: genai_types.Runtime, **kwargs) -> AsyncIterator[Any]: # type: ignore[no-untyped-def]
if not self.api_client:
raise ValueError("api_client is not initialized.")
if not self.api_resource:
Expand Down Expand Up @@ -1836,9 +1832,9 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def]

base_url = self.api_client._api_client._http_options.base_url.rstrip("/")
api_version = self.api_client._api_client._http_options.api_version
a2a_agent_card.supported_interfaces[0].url = (
f"{base_url}/{api_version}/{self.api_resource.name}/a2a"
)
a2a_agent_card.supported_interfaces[
0
].url = f"{base_url}/{api_version}/{self.api_resource.name}/a2a"

config = ClientConfig(
supported_protocol_bindings=[
Expand Down Expand Up @@ -1991,7 +1987,7 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None:

if cpu not in [1, 2, 4, 6, 8]:
raise ValueError(
"resource_limits['cpu'] must be one of 1, 2, 4, 6, 8. Got" f" {cpu}"
f"resource_limits['cpu'] must be one of 1, 2, 4, 6, 8. Got {cpu}"
)

if not isinstance(memory_str, str) or not memory_str.endswith("Gi"):
Expand All @@ -2004,8 +2000,7 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None:
memory_gb = int(memory_str[:-2])
except ValueError:
raise ValueError(
f"Invalid memory value: {memory_str}. Must be an integer"
" followed by 'Gi'."
f"Invalid memory value: {memory_str}. Must be an integer followed by 'Gi'."
)

# https://cloud.google.com/run/docs/configuring/memory-limits
Expand All @@ -2026,12 +2021,11 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None:

if cpu < min_cpu:
raise ValueError(
f"Memory size of {memory_str} requires at least {min_cpu} CPUs."
f" Got {cpu}"
f"Memory size of {memory_str} requires at least {min_cpu} CPUs. Got {cpu}"
)


def _is_adk_agent(agent_engine: _AgentEngineInterface) -> bool:
def _is_adk_agent(agent_engine: _RuntimeInterface) -> bool:
"""Checks if the agent engine is an ADK agent.

Args:
Expand All @@ -2041,13 +2035,13 @@ def _is_adk_agent(agent_engine: _AgentEngineInterface) -> bool:
True if the agent engine is an ADK agent, False otherwise.
"""

from agentplatform.agent_engines.templates import adk
from agentplatform.frameworks import AdkApp

return isinstance(agent_engine, adk.AdkApp)
return isinstance(agent_engine, AdkApp)


def _add_telemetry_enablement_env(
env_vars: Optional[Dict[str, Union[str, Any]]]
env_vars: Optional[Dict[str, Union[str, Any]]],
) -> Optional[Dict[str, Union[str, Any]]]:
"""Adds telemetry enablement env var to the env vars.

Expand Down
Loading
Loading