diff --git a/tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py b/tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py index 3f3dc0e142..a1a8b8930f 100644 --- a/tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py +++ b/tests/unit/vertexai/genai/replays/test_create_agent_engine_session.py @@ -85,6 +85,35 @@ def test_create_session_with_expire_time(client): client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) +def test_create_session_with_custom_session_id(client): + agent_engine = client.agent_engines.create() + try: + assert isinstance(agent_engine, types.AgentEngine) + assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + + operation = client.agent_engines.sessions.create( + name=agent_engine.api_resource.name, + user_id="test-user-123", + config=types.CreateAgentEngineSessionConfig( + display_name="my_session", + session_state={"foo": "bar"}, + session_id="my-session-id", + ), + ) + assert isinstance(operation, types.AgentEngineSessionOperation) + assert operation.response.display_name == "my_session" + assert operation.response.session_state == {"foo": "bar"} + assert operation.response.user_id == "test-user-123" + assert ( + operation.response.name + == f"{agent_engine.api_resource.name}/sessions/my-session-id" + ) + assert operation.done + finally: + # Clean up resources. + client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), diff --git a/vertexai/_genai/sessions.py b/vertexai/_genai/sessions.py index 994fdf9587..68f4876dc6 100644 --- a/vertexai/_genai/sessions.py +++ b/vertexai/_genai/sessions.py @@ -64,6 +64,9 @@ def _CreateAgentEngineSessionConfig_to_vertex( if getv(from_object, ["labels"]) is not None: setv(parent_object, ["labels"], getv(from_object, ["labels"])) + if getv(from_object, ["session_id"]) is not None: + setv(parent_object, ["_query", "sessionId"], getv(from_object, ["session_id"])) + return to_object @@ -176,6 +179,9 @@ def _UpdateAgentEngineSessionConfig_to_vertex( if getv(from_object, ["labels"]) is not None: setv(parent_object, ["labels"], getv(from_object, ["labels"])) + if getv(from_object, ["session_id"]) is not None: + setv(parent_object, ["_query", "sessionId"], getv(from_object, ["session_id"])) + if getv(from_object, ["update_mask"]) is not None: setv( parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index e6a7ea4b3c..6db4dfbab7 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -10790,6 +10790,10 @@ class CreateAgentEngineSessionConfig(_common.BaseModel): default=None, description="""Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", ) + session_id: Optional[str] = Field( + default=None, + description="""Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", + ) class CreateAgentEngineSessionConfigDict(TypedDict, total=False): @@ -10818,6 +10822,9 @@ class CreateAgentEngineSessionConfigDict(TypedDict, total=False): labels: Optional[dict[str, str]] """Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" + session_id: Optional[str] + """Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" + CreateAgentEngineSessionConfigOrDict = Union[ CreateAgentEngineSessionConfig, CreateAgentEngineSessionConfigDict @@ -11267,6 +11274,10 @@ class UpdateAgentEngineSessionConfig(_common.BaseModel): default=None, description="""Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", ) + session_id: Optional[str] = Field( + default=None, + description="""Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", + ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -11303,6 +11314,9 @@ class UpdateAgentEngineSessionConfigDict(TypedDict, total=False): labels: Optional[dict[str, str]] """Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" + session_id: Optional[str] + """Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" + update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""