feat: Add custom session id functionality to vertex ai session service

PiperOrigin-RevId: 891919371

Author: vertex-sdk-bot

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(),

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"])

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."""