diff --git a/docs/reference/experimental/async/activity.md b/docs/reference/experimental/async/activity.md
new file mode 100644
index 000000000..59e2f0061
--- /dev/null
+++ b/docs/reference/experimental/async/activity.md
@@ -0,0 +1,24 @@
+# Activity
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API Reference
+
+::: synapseclient.models.Activity
+ options:
+ members:
+ - from_parent_async
+ - store_async
+ - delete_async
+---
+::: synapseclient.models.UsedEntity
+ options:
+ filters:
+ - "!"
+---
+::: synapseclient.models.UsedURL
+ options:
+ filters:
+ - "!"
diff --git a/docs/reference/experimental/async/agent.md b/docs/reference/experimental/async/agent.md
new file mode 100644
index 000000000..be2e74c36
--- /dev/null
+++ b/docs/reference/experimental/async/agent.md
@@ -0,0 +1,32 @@
+# Agent
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API reference
+
+::: synapseclient.models.Agent
+ options:
+ members:
+ - register_async
+ - get_async
+ - start_session_async
+ - get_session_async
+ - prompt_async
+ - get_chat_history
+---
+::: synapseclient.models.AgentSession
+ options:
+ members:
+ - start_async
+ - get_async
+ - update_async
+ - prompt_async
+---
+::: synapseclient.models.AgentPrompt
+ options:
+ inherited_members: true
+ members:
+ - send_job_and_wait_async
+---
diff --git a/docs/reference/experimental/async/file.md b/docs/reference/experimental/async/file.md
new file mode 100644
index 000000000..e2fe12300
--- /dev/null
+++ b/docs/reference/experimental/async/file.md
@@ -0,0 +1,27 @@
+# File
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API Reference
+
+::: synapseclient.models.File
+ options:
+ inherited_members: true
+ members:
+ - get_async
+ - store_async
+ - copy_async
+ - delete_async
+ - from_id_async
+ - from_path_async
+ - change_metadata_async
+ - get_permissions_async
+ - get_acl_async
+ - set_permissions_async
+---
+::: synapseclient.models.file.FileHandle
+ options:
+ filters:
+ - "!"
diff --git a/docs/reference/experimental/async/folder.md b/docs/reference/experimental/async/folder.md
new file mode 100644
index 000000000..c11983a99
--- /dev/null
+++ b/docs/reference/experimental/async/folder.md
@@ -0,0 +1,20 @@
+# Folder
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API Reference
+
+::: synapseclient.models.Folder
+ options:
+ inherited_members: true
+ members:
+ - get_async
+ - store_async
+ - delete_async
+ - copy_async
+ - sync_from_synapse_async
+ - get_permissions_async
+ - get_acl_async
+ - set_permissions_async
diff --git a/docs/reference/experimental/async/project.md b/docs/reference/experimental/async/project.md
new file mode 100644
index 000000000..b628d4e19
--- /dev/null
+++ b/docs/reference/experimental/async/project.md
@@ -0,0 +1,19 @@
+# Project
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API reference
+
+::: synapseclient.models.Project
+ options:
+ inherited_members: true
+ members:
+ - get_async
+ - store_async
+ - delete_async
+ - sync_from_synapse_async
+ - get_permissions_async
+ - get_acl_async
+ - set_permissions_async
diff --git a/docs/reference/experimental/async/table.md b/docs/reference/experimental/async/table.md
new file mode 100644
index 000000000..63f3b3a0b
--- /dev/null
+++ b/docs/reference/experimental/async/table.md
@@ -0,0 +1,21 @@
+# Table
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API Reference
+
+::: synapseclient.models.Table
+ options:
+ inherited_members: true
+ members:
+ - get_async
+ - store_schema_async
+ - store_rows_from_csv_async
+ - delete_rows_async
+ - query_async
+ - delete_async
+ - get_permissions_async
+ - get_acl_async
+ - set_permissions_async
diff --git a/docs/reference/experimental/async/team.md b/docs/reference/experimental/async/team.md
new file mode 100644
index 000000000..0dd066e35
--- /dev/null
+++ b/docs/reference/experimental/async/team.md
@@ -0,0 +1,19 @@
+# Team
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API Reference
+
+::: synapseclient.models.Team
+ options:
+ members:
+ - create_async
+ - delete_async
+ - from_id_async
+ - from_name_async
+ - members_async
+ - invite_async
+ - open_invitations_async
+---
diff --git a/docs/reference/experimental/async/user_profile.md b/docs/reference/experimental/async/user_profile.md
new file mode 100644
index 000000000..7174061d9
--- /dev/null
+++ b/docs/reference/experimental/async/user_profile.md
@@ -0,0 +1,19 @@
+# UserProfile
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API Reference
+
+::: synapseclient.models.UserProfile
+ options:
+ inherited_members: true
+ members:
+ - get_async
+ - from_id_async
+ - from_username_async
+ - is_certified_async
+---
+::: synapseclient.models.UserPreference
+---
diff --git a/docs/reference/experimental/mixins/access_controllable.md b/docs/reference/experimental/mixins/access_controllable.md
new file mode 100644
index 000000000..96e7f70b9
--- /dev/null
+++ b/docs/reference/experimental/mixins/access_controllable.md
@@ -0,0 +1,3 @@
+# AccessControllable
+
+::: synapseclient.models.mixins.AccessControllable
diff --git a/docs/reference/experimental/mixins/asynchronous_communicator.md b/docs/reference/experimental/mixins/asynchronous_communicator.md
new file mode 100644
index 000000000..bfc081057
--- /dev/null
+++ b/docs/reference/experimental/mixins/asynchronous_communicator.md
@@ -0,0 +1,3 @@
+# AsynchronousCommunicator
+
+::: synapseclient.models.mixins.AsynchronousCommunicator
diff --git a/docs/reference/experimental/mixins/failure_strategy.md b/docs/reference/experimental/mixins/failure_strategy.md
new file mode 100644
index 000000000..3809b74f5
--- /dev/null
+++ b/docs/reference/experimental/mixins/failure_strategy.md
@@ -0,0 +1,3 @@
+# FailureStrategy
+
+::: synapseclient.models.FailureStrategy
diff --git a/docs/reference/experimental/mixins/storable_container.md b/docs/reference/experimental/mixins/storable_container.md
new file mode 100644
index 000000000..49e10a5e3
--- /dev/null
+++ b/docs/reference/experimental/mixins/storable_container.md
@@ -0,0 +1,3 @@
+# StorableContainer
+
+::: synapseclient.models.mixins.StorableContainer
diff --git a/docs/reference/experimental/sync/activity.md b/docs/reference/experimental/sync/activity.md
new file mode 100644
index 000000000..f0547e13c
--- /dev/null
+++ b/docs/reference/experimental/sync/activity.md
@@ -0,0 +1,35 @@
+# Activity
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## Example Script
+
+
+ Working with activities
+
+```python
+{!docs/scripts/object_orientated_programming_poc/oop_poc_activity.py!}
+```
+
+
+## API Reference
+
+::: synapseclient.models.Activity
+ options:
+ inherited_members: true
+ members:
+ - from_parent
+ - store
+ - delete
+---
+::: synapseclient.models.UsedEntity
+ options:
+ filters:
+ - "!"
+---
+::: synapseclient.models.UsedURL
+ options:
+ filters:
+ - "!"
diff --git a/docs/reference/experimental/sync/agent.md b/docs/reference/experimental/sync/agent.md
new file mode 100644
index 000000000..3d8cb7f08
--- /dev/null
+++ b/docs/reference/experimental/sync/agent.md
@@ -0,0 +1,42 @@
+# Agent
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## Example Script:
+
+
+ Working with Synapse agents
+
+```python
+{!docs/scripts/object_orientated_programming_poc/oop_poc_agent.py!}
+```
+
+
+## API Reference
+
+::: synapseclient.models.Agent
+ options:
+ inherited_members: true
+ members:
+ - register
+ - get
+ - start_session
+ - get_session
+ - prompt
+ - get_chat_history
+---
+::: synapseclient.models.AgentSession
+ options:
+ inherited_members: true
+ members:
+ - start
+ - get
+ - update
+ - prompt
+---
+::: synapseclient.models.AgentPrompt
+ options:
+ inherited_members: true
+---
diff --git a/docs/reference/experimental/sync/file.md b/docs/reference/experimental/sync/file.md
new file mode 100644
index 000000000..9b49e7603
--- /dev/null
+++ b/docs/reference/experimental/sync/file.md
@@ -0,0 +1,37 @@
+# File
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## Example Script
+
+
+ Working with files
+
+```python
+{!docs/scripts/object_orientated_programming_poc/oop_poc_file.py!}
+```
+
+
+## API Reference
+
+::: synapseclient.models.File
+ options:
+ inherited_members: true
+ members:
+ - get
+ - store
+ - copy
+ - delete
+ - from_id
+ - from_path
+ - change_metadata
+ - get_permissions
+ - get_acl
+ - set_permissions
+---
+::: synapseclient.models.file.FileHandle
+ options:
+ filters:
+ - "!"
diff --git a/docs/reference/experimental/sync/folder.md b/docs/reference/experimental/sync/folder.md
new file mode 100644
index 000000000..5a1cb5ddb
--- /dev/null
+++ b/docs/reference/experimental/sync/folder.md
@@ -0,0 +1,30 @@
+# Folder
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## Example Script
+
+
+ Working with folders
+
+```python
+{!docs/scripts/object_orientated_programming_poc/oop_poc_folder.py!}
+```
+
+
+## API Reference
+
+::: synapseclient.models.Folder
+ options:
+ inherited_members: true
+ members:
+ - get
+ - store
+ - delete
+ - copy
+ - sync_from_synapse
+ - get_permissions
+ - get_acl
+ - set_permissions
diff --git a/docs/reference/experimental/sync/project.md b/docs/reference/experimental/sync/project.md
new file mode 100644
index 000000000..e8cebfed5
--- /dev/null
+++ b/docs/reference/experimental/sync/project.md
@@ -0,0 +1,29 @@
+# Project
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## Example Script
+
+
+ Working with a project
+
+```python
+{!docs/scripts/object_orientated_programming_poc/oop_poc_project.py!}
+```
+
+
+## API reference
+
+::: synapseclient.models.Project
+ options:
+ inherited_members: true
+ members:
+ - get
+ - store
+ - delete
+ - sync_from_synapse
+ - get_permissions
+ - get_acl
+ - set_permissions
diff --git a/docs/reference/experimental/sync/table.md b/docs/reference/experimental/sync/table.md
new file mode 100644
index 000000000..058826d0d
--- /dev/null
+++ b/docs/reference/experimental/sync/table.md
@@ -0,0 +1,31 @@
+# Table
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## Example Script
+
+
+ Working with tables
+
+```python
+{!docs/scripts/object_orientated_programming_poc/oop_poc_table.py!}
+```
+
+
+## API Reference
+
+::: synapseclient.models.Table
+ options:
+ inherited_members: true
+ members:
+ - get
+ - store_schema
+ - store_rows_from_csv
+ - delete_rows
+ - query
+ - delete
+ - get_permissions
+ - get_acl
+ - set_permissions
diff --git a/docs/reference/experimental/sync/team.md b/docs/reference/experimental/sync/team.md
new file mode 100644
index 000000000..46fc51305
--- /dev/null
+++ b/docs/reference/experimental/sync/team.md
@@ -0,0 +1,30 @@
+# Team
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## Example Script
+
+
+ Working with teams
+
+```python
+{!docs/scripts/object_orientated_programming_poc/oop_poc_team.py!}
+```
+
+
+## API Reference
+
+::: synapseclient.models.Team
+ options:
+ inherited_members: true
+ members:
+ - create
+ - delete
+ - from_id
+ - from_name
+ - members
+ - invite
+ - open_invitations
+---
diff --git a/docs/reference/experimental/sync/user_profile.md b/docs/reference/experimental/sync/user_profile.md
new file mode 100644
index 000000000..46424f4b5
--- /dev/null
+++ b/docs/reference/experimental/sync/user_profile.md
@@ -0,0 +1,19 @@
+# UserProfile
+
+Contained within this file are experimental interfaces for working with the Synapse Python
+Client. Unless otherwise noted these interfaces are subject to change at any time. Use
+at your own risk.
+
+## API Reference
+
+::: synapseclient.models.UserProfile
+ options:
+ inherited_members: true
+ members:
+ - get
+ - from_id
+ - from_username
+ - is_certified
+---
+::: synapseclient.models.UserPreference
+---
diff --git a/docs/reference/oop/models.md b/docs/reference/oop/models.md
deleted file mode 100644
index 2c7ebc153..000000000
--- a/docs/reference/oop/models.md
+++ /dev/null
@@ -1,169 +0,0 @@
-Contained within this file are experimental interfaces for working with the Synapse Python
-Client. Unless otherwise noted these interfaces are subject to change at any time. Use
-at your own risk.
-
-## Sample Scripts:
-
-
- Working with a project
-
-```python
-{!docs/scripts/object_orientated_programming_poc/oop_poc_project.py!}
-```
-
-
-
- Working with folders
-
-```python
-{!docs/scripts/object_orientated_programming_poc/oop_poc_folder.py!}
-```
-
-
-
- Working with files
-
-```python
-{!docs/scripts/object_orientated_programming_poc/oop_poc_file.py!}
-```
-
-
-
- Working with tables
-
-```python
-{!docs/scripts/object_orientated_programming_poc/oop_poc_table.py!}
-```
-
-
-
- Current Synapse interface for working with a project
-
-```python
-{!docs/scripts/object_orientated_programming_poc/synapse_project.py!}
-```
-
-
-
- Working with activities
-
-```python
-{!docs/scripts/object_orientated_programming_poc/oop_poc_activity.py!}
-```
-
-
-
- Working with teams
-
-```python
-{!docs/scripts/object_orientated_programming_poc/oop_poc_team.py!}
-```
-
-
-## API reference
-
-::: synapseclient.models.Project
- options:
- inherited_members: true
- members:
- - get
- - store
- - delete
- - sync_from_synapse
- - get_permissions
- - get_acl
- - set_permissions
----
-::: synapseclient.models.Folder
- options:
- inherited_members: true
- members:
- - get
- - store
- - delete
- - copy
- - sync_from_synapse
- - get_permissions
- - get_acl
- - set_permissions
----
-::: synapseclient.models.File
- options:
- inherited_members: true
- members:
- - get
- - store
- - copy
- - delete
- - from_id
- - from_path
- - change_metadata
- - get_permissions
- - get_acl
- - set_permissions
-::: synapseclient.models.file.FileHandle
- options:
- filters:
- - "!"
----
-::: synapseclient.models.Table
- options:
- inherited_members: true
- members:
- - get
- - store_schema
- - store_rows_from_csv
- - delete_rows
- - query
- - delete
- - get_permissions
- - get_acl
- - set_permissions
----
-::: synapseclient.models.Activity
- options:
- members:
- - from_parent
- - store
- - delete
-
-::: synapseclient.models.UsedEntity
- options:
- filters:
- - "!"
-::: synapseclient.models.UsedURL
- options:
- filters:
- - "!"
----
-::: synapseclient.models.Team
- options:
- members:
- - create
- - delete
- - from_id
- - from_name
- - members
- - invite
- - open_invitations
----
-::: synapseclient.models.UserProfile
- options:
- members:
- - get
- - from_id
- - from_username
- - is_certified
-::: synapseclient.models.UserPreference
----
-::: synapseclient.models.Annotations
- options:
- members:
- - from_dict
----
-::: synapseclient.models.mixins.AccessControllable
----
-
-::: synapseclient.models.mixins.StorableContainer
----
-::: synapseclient.models.FailureStrategy
diff --git a/docs/reference/oop/models_async.md b/docs/reference/oop/models_async.md
deleted file mode 100644
index c61ce0df6..000000000
--- a/docs/reference/oop/models_async.md
+++ /dev/null
@@ -1,100 +0,0 @@
-Contained within this file are experimental interfaces for working with the Synapse Python
-Client. Unless otherwise noted these interfaces are subject to change at any time. Use
-at your own risk.
-
-These APIs also introduce [AsyncIO](https://docs.python.org/3/library/asyncio.html) to
-the client.
-
-## Sample Scripts:
-See [this page for sample scripts](models.md#sample-scripts).
-The sample scripts are from a synchronous context,
-replace any of the method calls with the async counter-party and they will be
-functionally equivalent.
-
-## API reference
-
-::: synapseclient.models.Project
- options:
- inherited_members: true
- members:
- - get_async
- - store_async
- - delete_async
- - sync_from_synapse_async
- - get_permissions_async
- - get_acl_async
- - set_permissions_async
----
-::: synapseclient.models.Folder
- options:
- inherited_members: true
- members:
- - get_async
- - store_async
- - delete_async
- - copy_async
- - sync_from_synapse_async
- - get_permissions_async
- - get_acl_async
- - set_permissions_async
----
-::: synapseclient.models.File
- options:
- inherited_members: true
- members:
- - get_async
- - store_async
- - copy_async
- - delete_async
- - from_id_async
- - from_path_async
- - change_metadata_async
- - get_permissions_async
- - get_acl_async
- - set_permissions_async
----
-::: synapseclient.models.Table
- options:
- inherited_members: true
- members:
- - get_async
- - store_schema_async
- - store_rows_from_csv_async
- - delete_rows_async
- - query_async
- - delete_async
- - get_permissions_async
- - get_acl_async
- - set_permissions_async
----
-::: synapseclient.models.Activity
- options:
- members:
- - from_parent_async
- - store_async
- - delete_async
-
----
-::: synapseclient.models.Team
- options:
- members:
- - create_async
- - delete_async
- - from_id_async
- - from_name_async
- - members_async
- - invite_async
- - open_invitations_async
----
-::: synapseclient.models.UserProfile
- options:
- members:
- - get_async
- - from_id_async
- - from_username_async
- - is_certified_async
----
-::: synapseclient.models.Annotations
- options:
- members:
- - store_async
diff --git a/docs/scripts/object_orientated_programming_poc/oop_poc_agent.py b/docs/scripts/object_orientated_programming_poc/oop_poc_agent.py
new file mode 100644
index 000000000..2703f41a9
--- /dev/null
+++ b/docs/scripts/object_orientated_programming_poc/oop_poc_agent.py
@@ -0,0 +1,105 @@
+"""
+The purpose of this script is to demonstrate how to use the new OOP interface for Synapse AI Agents.
+
+1. Register and send a prompt to a custom agent
+2. Send a prompt to the baseline Synapse Agent
+3. Conduct more than one session with the same agent
+4. Start a new session with a custom agent and send a prompt to it
+5. Start a new session with the baseline Synapse Agent and send a prompt to it
+6. Start a new session with a custom agent and then update what the agent has access to
+"""
+
+import synapseclient
+from synapseclient.models import Agent, AgentSession, AgentSessionAccessLevel
+
+# IDs for a bedrock agent with the instructions:
+# "You are a test agent that when greeted with: 'hello' will always response with: 'world'"
+CLOUD_AGENT_ID = "QOTV3KQM1X"
+AGENT_REGISTRATION_ID = 29
+
+syn = synapseclient.Synapse(debug=True)
+syn.login()
+
+# Using the Agent class
+
+
+# Register a custom agent and send a prompt to it
+def register_and_send_prompt_to_custom_agent():
+ my_custom_agent = Agent(cloud_agent_id=CLOUD_AGENT_ID)
+ my_custom_agent.register(synapse_client=syn)
+ my_custom_agent.prompt(
+ prompt="Hello", enable_trace=True, print_response=True, synapse_client=syn
+ )
+
+
+# Create an Agent Object and prompt.
+# By default, this will send a prompt to a new session with the baseline Synapse Agent.
+def get_baseline_agent_and_send_prompt_to_it():
+ baseline_agent = Agent()
+ baseline_agent.prompt(
+ prompt="What is Synapse?",
+ enable_trace=True,
+ print_response=True,
+ synapse_client=syn,
+ )
+
+
+# Conduct more than one session with the same agent
+def conduct_multiple_sessions_with_same_agent():
+ my_agent = Agent(registration_id=AGENT_REGISTRATION_ID).get(synapse_client=syn)
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ synapse_client=syn,
+ )
+ my_second_session = my_agent.start_session(synapse_client=syn)
+ my_agent.prompt(
+ prompt="Hello again",
+ enable_trace=True,
+ print_response=True,
+ session=my_second_session,
+ synapse_client=syn,
+ )
+
+
+# Using the AgentSession class
+
+
+# Start a new session with a custom agent and send a prompt to it
+def start_new_session_with_custom_agent_and_send_prompt_to_it():
+ my_session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID).start(
+ synapse_client=syn
+ )
+ my_session.prompt(
+ prompt="Hello", enable_trace=True, print_response=True, synapse_client=syn
+ )
+
+
+# Start a new session with the baseline Synapse Agent and send a prompt to it
+def start_new_session_with_baseline_agent_and_send_prompt_to_it():
+ my_session = AgentSession().start(synapse_client=syn)
+ my_session.prompt(
+ prompt="What is Synapse?",
+ enable_trace=True,
+ print_response=True,
+ synapse_client=syn,
+ )
+
+
+# Start a new session with a custom agent and then update what the agent has access to
+def start_new_session_with_custom_agent_and_update_access_to_it():
+ my_session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID).start(
+ synapse_client=syn
+ )
+ print(f"Access level before update: {my_session.access_level}")
+ my_session.access_level = AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA
+ my_session.update(synapse_client=syn)
+ print(f"Access level after update: {my_session.access_level}")
+
+
+register_and_send_prompt_to_custom_agent()
+get_baseline_agent_and_send_prompt_to_it()
+conduct_multiple_sessions_with_same_agent()
+start_new_session_with_baseline_agent_and_send_prompt_to_it()
+start_new_session_with_custom_agent_and_update_access_to_it()
diff --git a/mkdocs.yml b/mkdocs.yml
index 768dcd0e3..68f9e0053 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -75,8 +75,28 @@ nav:
- Core: reference/core.md
- REST Apis: reference/rest_apis.md
- Experimental:
- - Object-Orientated Models: reference/oop/models.md
- - Async Object-Orientated Models: reference/oop/models_async.md
+ - Agent: reference/experimental/sync/agent.md
+ - Project: reference/experimental/sync/project.md
+ - Folder: reference/experimental/sync/folder.md
+ - File: reference/experimental/sync/file.md
+ - Table: reference/experimental/sync/table.md
+ - Activity: reference/experimental/sync/activity.md
+ - Team: reference/experimental/sync/team.md
+ - UserProfile: reference/experimental/sync/user_profile.md
+ - Asynchronous:
+ - Agent: reference/experimental/async/agent.md
+ - Project: reference/experimental/async/project.md
+ - Folder: reference/experimental/async/folder.md
+ - File: reference/experimental/async/file.md
+ - Table: reference/experimental/async/table.md
+ - Activity: reference/experimental/async/activity.md
+ - Team: reference/experimental/async/team.md
+ - UserProfile: reference/experimental/async/user_profile.md
+ - Mixins:
+ - AccessControllable: reference/experimental/mixins/access_controllable.md
+ - StorableContainer: reference/experimental/mixins/storable_container.md
+ - AsynchronousCommunicator: reference/experimental/mixins/asynchronous_communicator.md
+ - FailureStrategy: reference/experimental/mixins/failure_strategy.md
- Further Reading:
- Home: explanations/home.md
- Domain Models of Synapse: explanations/domain_models_of_synapse.md
diff --git a/synapseclient/api/__init__.py b/synapseclient/api/__init__.py
index 3211aaf38..f41f782fc 100644
--- a/synapseclient/api/__init__.py
+++ b/synapseclient/api/__init__.py
@@ -1,4 +1,12 @@
# These are all of the models that are used by the Synapse client.
+from .agent_services import (
+ get_agent,
+ get_session,
+ get_trace,
+ register_agent,
+ start_session,
+ update_session,
+)
from .annotations import set_annotations, set_annotations_async
from .configuration_services import (
get_client_authenticated_s3_profile,
@@ -78,4 +86,11 @@
"get_transfer_config",
# entity_factory
"get_from_entity_factory",
+ # agent_services
+ "register_agent",
+ "get_agent",
+ "start_session",
+ "get_session",
+ "update_session",
+ "get_trace",
]
diff --git a/synapseclient/api/agent_services.py b/synapseclient/api/agent_services.py
new file mode 100644
index 000000000..6cb65e1fd
--- /dev/null
+++ b/synapseclient/api/agent_services.py
@@ -0,0 +1,189 @@
+"""This module is responsible for exposing the services defined at:
+
+"""
+
+import json
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+if TYPE_CHECKING:
+ from synapseclient import Synapse
+
+
+async def register_agent(
+ cloud_agent_id: str,
+ cloud_alias_id: Optional[str] = None,
+ *,
+ synapse_client: Optional["Synapse"] = None,
+) -> Dict[str, Any]:
+ """
+ Registers an agent with Synapse OR gets existing agent registration.
+ Sends a request matching
+
+
+ Arguments:
+ cloud_agent_id: The cloud provider ID of the agent to register.
+ cloud_alias_id: The cloud provider alias ID of the agent to register.
+ In the Synapse API, this defaults to 'TSTALIASID'.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The registered agent matching
+
+ """
+ from synapseclient import Synapse
+
+ client = Synapse.get_client(synapse_client=synapse_client)
+
+ request = {"awsAgentId": cloud_agent_id}
+ if cloud_alias_id:
+ request["awsAliasId"] = cloud_alias_id
+ return await client.rest_put_async(
+ uri="/agent/registration", body=json.dumps(request)
+ )
+
+
+async def get_agent(
+ registration_id: str,
+ *,
+ synapse_client: Optional["Synapse"] = None,
+) -> Dict[str, Any]:
+ """
+ Gets information about an existing agent registration.
+
+ Arguments:
+ registration_id: The ID of the agent registration to get.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The requested agent registration matching
+
+ """
+ from synapseclient import Synapse
+
+ client = Synapse.get_client(synapse_client=synapse_client)
+ return await client.rest_get_async(uri=f"/agent/registration/{registration_id}")
+
+
+async def start_session(
+ access_level: str,
+ agent_registration_id: str,
+ *,
+ synapse_client: Optional["Synapse"] = None,
+) -> Dict[str, Any]:
+ """
+ Starts a new chat session with an agent.
+ Sends a request matching
+
+
+ Arguments:
+ access_level: The access level of the agent.
+ agent_registration_id: The ID of the agent registration to start the session for.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+ """
+ from synapseclient import Synapse
+
+ client = Synapse.get_client(synapse_client=synapse_client)
+
+ request = {
+ "agentAccessLevel": access_level,
+ "agentRegistrationId": agent_registration_id,
+ }
+ return await client.rest_post_async(uri="/agent/session", body=json.dumps(request))
+
+
+async def get_session(
+ id: str,
+ *,
+ synapse_client: Optional["Synapse"] = None,
+) -> Dict[str, Any]:
+ """
+ Gets information about an existing chat session.
+
+ Arguments:
+ id: The ID of the session to get.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The requested session matching
+
+ """
+ from synapseclient import Synapse
+
+ client = Synapse.get_client(synapse_client=synapse_client)
+ return await client.rest_get_async(uri=f"/agent/session/{id}")
+
+
+async def update_session(
+ id: str,
+ access_level: str,
+ *,
+ synapse_client: Optional["Synapse"] = None,
+) -> Dict[str, Any]:
+ """
+ Updates the access level for a chat session.
+ Sends a request matching
+
+
+ Arguments:
+ id: The ID of the session to update.
+ access_level: The access level of the agent.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+ """
+ from synapseclient import Synapse
+
+ client = Synapse.get_client(synapse_client=synapse_client)
+
+ request = {
+ "sessionId": id,
+ "agentAccessLevel": access_level,
+ }
+ return await client.rest_put_async(
+ uri=f"/agent/session/{id}", body=json.dumps(request)
+ )
+
+
+async def get_trace(
+ prompt_id: str,
+ *,
+ newer_than: Optional[int] = None,
+ synapse_client: Optional["Synapse"] = None,
+) -> Dict[str, Any]:
+ """
+ Gets the trace of a prompt.
+ Sends a request matching
+
+
+ Arguments:
+ prompt_id: The token of the prompt to get the trace for.
+ newer_than: The timestamp to get trace results newer than. Defaults to None (all results).
+ Timestamps should be in milliseconds since the epoch per the API documentation.
+ https://rest-docs.synapse.org/rest/org/sagebionetworks/repo/model/agent/TraceEvent.html
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The trace matching
+
+ """
+ from synapseclient import Synapse
+
+ client = Synapse.get_client(synapse_client=synapse_client)
+
+ request = {
+ "jobId": prompt_id,
+ "newerThanTimestamp": newer_than,
+ }
+ return await client.rest_post_async(
+ uri=f"/agent/chat/trace/{prompt_id}", body=json.dumps(request)
+ )
diff --git a/synapseclient/client.py b/synapseclient/client.py
index 61bcf73c5..8aba3217d 100644
--- a/synapseclient/client.py
+++ b/synapseclient/client.py
@@ -6373,20 +6373,17 @@ async def rest_get_async(
Returns:
JSON encoding of response
"""
- try:
- response = await self._rest_call_async(
- "get",
- uri,
- None,
- endpoint,
- headers,
- retry_policy,
- requests_session_async_synapse,
- **kwargs,
- )
- return self._return_rest_body(response)
- except Exception:
- self.logger.exception("Error in rest_get_async")
+ response = await self._rest_call_async(
+ "get",
+ uri,
+ None,
+ endpoint,
+ headers,
+ retry_policy,
+ requests_session_async_synapse,
+ **kwargs,
+ )
+ return self._return_rest_body(response)
async def rest_post_async(
self,
diff --git a/synapseclient/core/constants/concrete_types.py b/synapseclient/core/constants/concrete_types.py
index f8d4ee442..e2033c030 100644
--- a/synapseclient/core/constants/concrete_types.py
+++ b/synapseclient/core/constants/concrete_types.py
@@ -68,3 +68,6 @@
# Activity/Provenance
USED_URL = "org.sagebionetworks.repo.model.provenance.UsedURL"
USED_ENTITY = "org.sagebionetworks.repo.model.provenance.UsedEntity"
+
+# Agent
+AGENT_CHAT_REQUEST = "org.sagebionetworks.repo.model.agent.AgentChatRequest"
diff --git a/synapseclient/models/__init__.py b/synapseclient/models/__init__.py
index a487a3827..1e2f686ed 100644
--- a/synapseclient/models/__init__.py
+++ b/synapseclient/models/__init__.py
@@ -1,5 +1,11 @@
# These are all of the models that are used by the Synapse client.
from synapseclient.models.activity import Activity, UsedEntity, UsedURL
+from synapseclient.models.agent import (
+ Agent,
+ AgentPrompt,
+ AgentSession,
+ AgentSessionAccessLevel,
+)
from synapseclient.models.annotations import Annotations
from synapseclient.models.file import File, FileHandle
from synapseclient.models.folder import Folder
@@ -38,4 +44,8 @@
"TeamMember",
"UserProfile",
"UserPreference",
+ "Agent",
+ "AgentSession",
+ "AgentSessionAccessLevel",
+ "AgentPrompt",
]
diff --git a/synapseclient/models/agent.py b/synapseclient/models/agent.py
new file mode 100644
index 000000000..344373f8a
--- /dev/null
+++ b/synapseclient/models/agent.py
@@ -0,0 +1,944 @@
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Dict, List, Optional, Union
+
+from synapseclient import Synapse
+from synapseclient.api import (
+ get_agent,
+ get_session,
+ get_trace,
+ register_agent,
+ start_session,
+ update_session,
+)
+from synapseclient.core.async_utils import async_to_sync, otel_trace_method
+from synapseclient.core.constants.concrete_types import AGENT_CHAT_REQUEST
+from synapseclient.models.mixins import AsynchronousCommunicator
+from synapseclient.models.protocols.agent_protocol import (
+ AgentSessionSynchronousProtocol,
+ AgentSynchronousProtocol,
+)
+
+
+class AgentType(str, Enum):
+ """
+ Enum representing the type of agent as defined in
+
+
+ - BASELINE is a default agent provided by Synapse.
+ - CUSTOM is a custom agent that has been registered by a user.
+ """
+
+ BASELINE = "BASELINE"
+ CUSTOM = "CUSTOM"
+
+
+class AgentSessionAccessLevel(str, Enum):
+ """
+ Enum representing the access level of the agent session as defined in
+
+
+ - PUBLICLY_ACCESSIBLE: The agent can only access publicly accessible data.
+ - READ_YOUR_PRIVATE_DATA: The agent can read the user's private data.
+ - WRITE_YOUR_PRIVATE_DATA: The agent can write to the user's private data.
+ """
+
+ PUBLICLY_ACCESSIBLE = "PUBLICLY_ACCESSIBLE"
+ READ_YOUR_PRIVATE_DATA = "READ_YOUR_PRIVATE_DATA"
+ WRITE_YOUR_PRIVATE_DATA = "WRITE_YOUR_PRIVATE_DATA"
+
+
+@dataclass
+class AgentPrompt(AsynchronousCommunicator):
+ """Represents a prompt, response, and metadata within an AgentSession.
+
+ Attributes:
+ id: The unique ID of the agent prompt.
+ session_id: The ID of the session that the prompt is associated with.
+ prompt: The prompt to send to the agent.
+ response: The response from the agent.
+ enable_trace: Whether tracing is enabled for the prompt.
+ trace: The trace of the agent session.
+ """
+
+ concrete_type: str = AGENT_CHAT_REQUEST
+
+ id: Optional[str] = None
+ """The unique ID of the agent prompt."""
+
+ session_id: Optional[str] = None
+ """The ID of the session that the prompt is associated with."""
+
+ prompt: Optional[str] = None
+ """The prompt sent to the agent."""
+
+ response: Optional[str] = None
+ """The response from the agent."""
+
+ enable_trace: Optional[bool] = False
+ """Whether tracing is enabled for the prompt."""
+
+ trace: Optional[str] = None
+ """The trace or "thought process" of the agent when responding to the prompt."""
+
+ def to_synapse_request(self):
+ """Converts the request to a request expected of the Synapse REST API."""
+ return {
+ "concreteType": self.concrete_type,
+ "sessionId": self.session_id,
+ "chatText": self.prompt,
+ "enableTrace": self.enable_trace,
+ }
+
+ def fill_from_dict(self, synapse_response: Dict[str, str]) -> "AgentPrompt":
+ """
+ Converts a response from the REST API into this dataclass.
+
+ Arguments:
+ synapse_response: The response from the REST API.
+
+ Returns:
+ The AgentPrompt object.
+ """
+ self.id = synapse_response.get("jobId", None)
+ self.session_id = synapse_response.get("sessionId", None)
+ self.response = synapse_response.get("responseText", None)
+ return self
+
+ async def _post_exchange_async(
+ self, *, synapse_client: Optional[Synapse] = None, **kwargs
+ ) -> None:
+ """Retrieves information about the trace of this prompt with the agent.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+ """
+ if self.enable_trace:
+ trace_response = await get_trace(
+ prompt_id=self.id,
+ newer_than=kwargs.get("newer_than", None),
+ synapse_client=synapse_client,
+ )
+ self.trace = trace_response["page"][0]["message"]
+
+
+@dataclass
+@async_to_sync
+class AgentSession(AgentSessionSynchronousProtocol):
+ """Represents a [Synapse Agent Session](https://rest-docs.synapse.org/rest/org/sagebionetworks/repo/model/agent/AgentSession.html)
+
+ Attributes:
+ id: The unique ID of the agent session.
+ Can only be used by the user that created it.
+ access_level: The access level of the agent session.
+ One of PUBLICLY_ACCESSIBLE, READ_YOUR_PRIVATE_DATA,
+ or WRITE_YOUR_PRIVATE_DATA.
+ started_on: The date the agent session was started.
+ started_by: The ID of the user who started the agent session.
+ modified_on: The date the agent session was last modified.
+ agent_registration_id: The registration ID of the agent that will
+ be used for this session.
+ etag: The etag of the agent session.
+
+ Note: It is recommended to use the `Agent` class to conduct chat sessions,
+ but you are free to use AgentSession directly if you wish.
+
+ Example: Start a session and send a prompt.
+ Start a session with a custom agent by providing the agent's registration ID and calling `start()`.
+ Then, send a prompt to the agent.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession
+
+ syn = Synapse()
+ syn.login()
+
+ my_session = AgentSession(agent_registration_id="foo").start()
+ my_session.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ Example: Get an existing session and send a prompt.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession
+
+ syn = Synapse()
+ syn.login()
+
+ my_session = AgentSession(id="foo").get()
+ my_session.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ Example: Update the access level of an existing session.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, update the access level of the session and call `update()`.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession, AgentSessionAccessLevel
+
+ syn = Synapse()
+ syn.login()
+
+ my_session = AgentSession(id="foo").get()
+ my_session.access_level = AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA
+ my_session.update()
+ """
+
+ id: Optional[str] = None
+ """The unique ID of the agent session.
+ Can only be used by the user that created it."""
+
+ access_level: Optional[
+ AgentSessionAccessLevel
+ ] = AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE
+ """The access level of the agent session.
+ One of PUBLICLY_ACCESSIBLE, READ_YOUR_PRIVATE_DATA, or
+ WRITE_YOUR_PRIVATE_DATA. Defaults to PUBLICLY_ACCESSIBLE.
+ """
+
+ started_on: Optional[datetime] = None
+ """The date the agent session was started."""
+
+ started_by: Optional[int] = None
+ """The ID of the user who started the agent session."""
+
+ modified_on: Optional[datetime] = None
+ """The date the agent session was last modified."""
+
+ agent_registration_id: Optional[int] = None
+ """The registration ID of the agent that will be used for this session."""
+
+ etag: Optional[str] = None
+ """The etag of the agent session."""
+
+ chat_history: List[AgentPrompt] = field(default_factory=list)
+ """A list of AgentPrompt objects."""
+
+ def fill_from_dict(self, synapse_agent_session: Dict[str, str]) -> "AgentSession":
+ """
+ Converts a response from the REST API into this dataclass.
+
+ Arguments:
+ synapse_agent_session: The response from the REST API.
+
+ Returns:
+ The AgentSession object.
+ """
+ self.id = synapse_agent_session.get("sessionId", None)
+ self.access_level = synapse_agent_session.get("agentAccessLevel", None)
+ self.started_on = synapse_agent_session.get("startedOn", None)
+ self.started_by = synapse_agent_session.get("startedBy", None)
+ self.modified_on = synapse_agent_session.get("modifiedOn", None)
+ self.agent_registration_id = synapse_agent_session.get(
+ "agentRegistrationId", None
+ )
+ self.etag = synapse_agent_session.get("etag", None)
+ return self
+
+ @otel_trace_method(method_to_trace_name=lambda self, **kwargs: "Start_Session")
+ async def start_async(
+ self, *, synapse_client: Optional[Synapse] = None
+ ) -> "AgentSession":
+ """Starts an agent session.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The new AgentSession object.
+
+ Example: Start a session and send a prompt.
+ Start a session with a custom agent by providing the agent's registration ID and calling `start()`.
+ Then, send a prompt to the agent.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_session = await AgentSession(agent_registration_id="foo").start_async()
+ await my_session.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+ """
+ session_response = await start_session(
+ access_level=self.access_level,
+ agent_registration_id=self.agent_registration_id,
+ synapse_client=synapse_client,
+ )
+ return self.fill_from_dict(synapse_agent_session=session_response)
+
+ @otel_trace_method(
+ method_to_trace_name=lambda self, **kwargs: f"Get_Session: {self.id}"
+ )
+ async def get_async(
+ self, *, synapse_client: Optional[Synapse] = None
+ ) -> "AgentSession":
+ """Gets an agent session.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The retrieved AgentSession object.
+
+ Example: Get an existing session and send a prompt.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_session = await AgentSession(id="foo").get_async()
+ await my_session.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+ """
+ session_response = await get_session(
+ id=self.id,
+ synapse_client=synapse_client,
+ )
+ return self.fill_from_dict(synapse_agent_session=session_response)
+
+ @otel_trace_method(
+ method_to_trace_name=lambda self, **kwargs: f"Update_Session: {self.id}"
+ )
+ async def update_async(
+ self,
+ *,
+ synapse_client: Optional[Synapse] = None,
+ ) -> "AgentSession":
+ """Updates an agent session.
+ Only updates to the access level are currently supported.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The updated AgentSession object.
+
+ Example: Update the access level of an existing session.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, update the access level of the session and call `update()`.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession, AgentSessionAccessLevel
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_session = await AgentSession(id="foo").get_async()
+ my_session.access_level = AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA
+ await my_session.update_async()
+
+ asyncio.run(main())
+ """
+ session_response = await update_session(
+ id=self.id,
+ access_level=self.access_level,
+ synapse_client=synapse_client,
+ )
+ return self.fill_from_dict(synapse_agent_session=session_response)
+
+ @otel_trace_method(method_to_trace_name=lambda self, **kwargs: f"Prompt: {self.id}")
+ async def prompt_async(
+ self,
+ prompt: str,
+ enable_trace: bool = False,
+ print_response: bool = False,
+ newer_than: Optional[int] = None,
+ *,
+ synapse_client: Optional[Synapse] = None,
+ ) -> None:
+ """Sends a prompt to the agent and adds the response to the AgentSession's
+ chat history. A session must be started before sending a prompt.
+
+ Arguments:
+ prompt: The prompt to send to the agent.
+ enable_trace: Whether to enable trace for the prompt.
+ print_response: Whether to print the response to the console.
+ newer_than: The timestamp to get trace results newer than.
+ Defaults to None (all results).
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Example: Send a prompt within an existing session.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_session = await AgentSession(id="foo").get_async()
+ await my_session.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+ """
+ agent_prompt = await AgentPrompt(
+ prompt=prompt, session_id=self.id, enable_trace=enable_trace
+ ).send_job_and_wait_async(
+ synapse_client=synapse_client, post_exchange_args={"newer_than": newer_than}
+ )
+ self.chat_history.append(agent_prompt)
+ if print_response:
+ client = Synapse.get_client(synapse_client=synapse_client)
+ client.logger.info(f"PROMPT:\n{prompt}\n")
+ client.logger.info(f"RESPONSE:\n{agent_prompt.response}\n")
+ if enable_trace:
+ client.logger.info(f"TRACE:\n{agent_prompt.trace}")
+
+
+@dataclass
+@async_to_sync
+class Agent(AgentSynchronousProtocol):
+ """Represents a [Synapse Agent Registration](https://rest-docs.synapse.org/rest/org/sagebionetworks/repo/model/agent/AgentRegistration.html)
+
+ Attributes:
+ cloud_agent_id: The unique ID of the agent in the cloud provider.
+ cloud_alias_id: The alias ID of the agent in the cloud provider.
+ Defaults to 'TSTALIASID' in the Synapse API.
+ registration_id: The ID number of the agent assigned by Synapse.
+ registered_on: The date the agent was registered.
+ type: The type of agent.
+ sessions: A dictionary of AgentSession objects, keyed by session ID.
+ current_session: The current session. Prompts will be sent to this session by default.
+
+ Example: Chat with the baseline Synapse Agent
+ You can chat with the same agent which is available in the Synapse UI
+ at https://www.synapse.org/Chat:default. By default, this "baseline" agent
+ is used when a registration ID is not provided. In the background,
+ the Agent class will start a session and set that new session as the
+ current session if one is not already set.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent()
+ my_agent.prompt(
+ prompt="Can you tell me about the AD Knowledge Portal dataset?",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ Example: Register and chat with a custom agent
+ **Only available for internal users (Sage Bionetworks employees)**
+
+ Alternatively, you can register a custom agent and chat with it provided
+ you have already created it.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent(cloud_agent_id="foo")
+ my_agent.register()
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ Example: Get and chat with an existing agent
+ Retrieve an existing agent by providing the agent's registration ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent(registration_id="foo").get()
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ Advanced Example: Start and prompt multiple sessions
+ Here, we connect to a custom agent and start one session with the prompt "Hello".
+ In the background, this first session is being set as the current session
+ and future prompts will be sent to this session by default. If we want to send a
+ prompt to a different session, we can do so by starting it and calling prompt again,
+ but with our new session as an argument. We now have two sessions, both stored in the
+ `my_agent.sessions` dictionary. After the second prompt, `my_second_session` is now
+ the current session.
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent(registration_id="foo").get()
+
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ my_second_session = my_agent.start_session()
+ my_agent.prompt(
+ prompt="Hello again",
+ enable_trace=True,
+ print_response=True,
+ session=my_second_session,
+ )
+ """
+
+ cloud_agent_id: Optional[str] = None
+ """The unique ID of the agent in the cloud provider."""
+
+ cloud_alias_id: Optional[str] = None
+ """The alias ID of the agent in the cloud provider.
+ Defaults to 'TSTALIASID' in the Synapse API.
+ """
+
+ registration_id: Optional[int] = None
+ """The ID number of the agent assigned by Synapse."""
+
+ registered_on: Optional[datetime] = None
+ """The date the agent was registered."""
+
+ type: Optional[AgentType] = None
+ """The type of agent. One of either BASELINE or CUSTOM."""
+
+ sessions: Dict[str, AgentSession] = field(default_factory=dict)
+ """A dictionary of AgentSession objects, keyed by session ID."""
+
+ current_session: Optional[AgentSession] = None
+ """The current session. Prompts will be sent to this session by default."""
+
+ def fill_from_dict(self, agent_registration: Dict[str, str]) -> "Agent":
+ """
+ Converts a response from the REST API into this dataclass.
+
+ Arguments:
+ agent_registration: The response from the REST API.
+
+ Returns:
+ The Agent object.
+ """
+ self.cloud_agent_id = agent_registration.get("awsAgentId", None)
+ self.cloud_alias_id = agent_registration.get("awsAliasId", None)
+ self.registration_id = agent_registration.get("agentRegistrationId", None)
+ self.registered_on = agent_registration.get("registeredOn", None)
+ self.type = (
+ AgentType(agent_registration.get("type"))
+ if agent_registration.get("type", None)
+ else None
+ )
+ return self
+
+ @otel_trace_method(
+ method_to_trace_name=lambda self, **kwargs: f"Register_Agent: {self.registration_id}"
+ )
+ async def register_async(
+ self, *, synapse_client: Optional[Synapse] = None
+ ) -> "Agent":
+ """Registers an agent with the Synapse API.
+ If agent already exists, it will be retrieved.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The registered or existing Agent object.
+
+ Example: Register and chat with a custom agent
+ **Only available for internal users (Sage Bionetworks employees)**
+
+ Alternatively, you can register a custom agent and chat with it provided
+ you have already created it.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_agent = Agent(cloud_agent_id="foo")
+ await my_agent.register_async()
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+ """
+ agent_response = await register_agent(
+ cloud_agent_id=self.cloud_agent_id,
+ cloud_alias_id=self.cloud_alias_id,
+ synapse_client=synapse_client,
+ )
+ return self.fill_from_dict(agent_registration=agent_response)
+
+ @otel_trace_method(
+ method_to_trace_name=lambda self, **kwargs: f"Get_Agent: {self.registration_id}"
+ )
+ async def get_async(self, *, synapse_client: Optional[Synapse] = None) -> "Agent":
+ """Gets an existing custom agent. There is no need to use this method
+ if you are trying to use the baseline agent.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The existing Agent object.
+
+ Example: Get and chat with an existing agent
+ Retrieve an existing custom agent by providing the agent's registration ID and calling `get_async()`.
+ Then, send a prompt to the agent.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_agent = await Agent(registration_id="foo").get_async()
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+ """
+ if self.registration_id is None:
+ raise ValueError(
+ "Registration ID is required to retrieve a custom agent. "
+ "If you are trying to use the baseline agent, you do not need to "
+ "use `get` or `get_async`. Instead, simply create an `Agent` object "
+ "and start prompting `my_agent = Agent(); my_agent.prompt(...)`.",
+ )
+ agent_response = await get_agent(
+ registration_id=self.registration_id,
+ synapse_client=synapse_client,
+ )
+ return self.fill_from_dict(agent_registration=agent_response)
+
+ @otel_trace_method(
+ method_to_trace_name=lambda self, **kwargs: f"Start_Agent_Session: {self.registration_id}"
+ )
+ async def start_session_async(
+ self,
+ access_level: Optional[
+ AgentSessionAccessLevel
+ ] = AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ *,
+ synapse_client: Optional[Synapse] = None,
+ ) -> "AgentSession":
+ """Starts an agent session.
+ Adds the session to the Agent's sessions dictionary and sets it as the current session.
+
+ Arguments:
+ access_level: The access level of the agent session.
+ Must be one of PUBLICLY_ACCESSIBLE, READ_YOUR_PRIVATE_DATA,
+ or WRITE_YOUR_PRIVATE_DATA.
+ Defaults to PUBLICLY_ACCESSIBLE.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The new AgentSession object.
+
+ Example: Start a session and send a prompt with the baseline Synapse Agent.
+ The baseline Synapse Agent is the default agent used when a registration ID is not provided.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_agent = Agent()
+ await my_agent.start_session_async()
+ await my_agent.prompt_async(
+ prompt="Can you tell me about the AD Knowledge Portal dataset?",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+
+ Example: Start a session and send a prompt with a custom agent.
+ The baseline Synapse Agent is the default agent used when a registration ID is not provided.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_agent = Agent(cloud_agent_id="foo")
+ await my_agent.start_session_async()
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+ """
+ access_level = AgentSessionAccessLevel(access_level)
+ session = await AgentSession(
+ agent_registration_id=self.registration_id, access_level=access_level
+ ).start_async(synapse_client=synapse_client)
+ self.sessions[session.id] = session
+ self.current_session = session
+ return session
+
+ @otel_trace_method(
+ method_to_trace_name=lambda self, **kwargs: f"Get_Agent_Session: {self.registration_id}"
+ )
+ async def get_session_async(
+ self, session_id: str, *, synapse_client: Optional[Synapse] = None
+ ) -> "AgentSession":
+ """Gets an existing agent session.
+ Adds the session to the Agent's sessions dictionary and
+ sets it as the current session.
+
+ Arguments:
+ session_id: The ID of the session to get.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The existing AgentSession object.
+
+ Example: Get an existing session and send a prompt.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_session = await Agent().get_session_async(session_id="foo")
+ await my_session.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+ """
+ session = await AgentSession(id=session_id).get_async(
+ synapse_client=synapse_client
+ )
+ if session.id not in self.sessions:
+ self.sessions[session.id] = session
+ self.current_session = session
+ return session
+
+ @otel_trace_method(
+ method_to_trace_name=lambda self, **kwargs: f"Prompt_Agent_Session: {self.registration_id}"
+ )
+ async def prompt_async(
+ self,
+ prompt: str,
+ enable_trace: bool = False,
+ print_response: bool = False,
+ session: Optional[AgentSession] = None,
+ newer_than: Optional[int] = None,
+ *,
+ synapse_client: Optional[Synapse] = None,
+ ) -> None:
+ """Sends a prompt to the agent for the current session.
+ If no session is currently active, a new session will be started.
+
+ Arguments:
+ prompt: The prompt to send to the agent.
+ enable_trace: Whether to enable trace for the prompt.
+ print_response: Whether to print the response to the console.
+ session_id: The ID of the session to send the prompt to.
+ If None, the current session will be used.
+ newer_than: The timestamp to get trace results newer than. Defaults to None (all results).
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Example: Prompt the baseline Synapse Agent to add annotations to a file on Synapse
+ The baseline Synpase Agent can be used to add annotations to files.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_agent = Agent()
+ await my_agent.prompt_async(
+ prompt="Add the annotation 'test' to the file 'syn123456789'",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+
+ Example: Prompt a custom agent.
+ If you have already registered a custom agent, you can prompt it by providing the agent's registration ID.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_agent = Agent(registration_id="foo")
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ asyncio.run(main())
+
+ Advanced Example: Start and prompt multiple sessions
+ Here, we connect to a custom agent and start one session with the prompt "Hello".
+ In the background, this first session is being set as the current session
+ and future prompts will be sent to this session by default. If we want to send a
+ prompt to a different session, we can do so by starting it and calling prompt again,
+ but with our new session as an argument. We now have two sessions, both stored in the
+ `my_agent.sessions` dictionary. After the second prompt, `my_second_session` is now
+ the current session.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_agent = Agent(registration_id="foo").get()
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ my_second_session = await my_agent.start_session_async()
+ await my_agent.prompt_async(
+ prompt="Hello again",
+ enable_trace=True,
+ print_response=True,
+ session=my_second_session,
+ )
+
+ asyncio.run(main())
+ """
+ if session:
+ await self.get_session_async(
+ session_id=session.id, synapse_client=synapse_client
+ )
+ else:
+ if not self.current_session:
+ await self.start_session_async(synapse_client=synapse_client)
+
+ await self.current_session.prompt_async(
+ prompt=prompt,
+ enable_trace=enable_trace,
+ newer_than=newer_than,
+ print_response=print_response,
+ synapse_client=synapse_client,
+ )
+
+ def get_chat_history(self) -> Union[List[AgentPrompt], None]:
+ """Gets the chat history for the current session.
+
+ Example: Get the chat history for the current session.
+ First, send a prompt to the agent.
+ Then, retrieve the chat history for the current session by calling `get_chat_history()`.
+
+ import asyncio
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ async def main():
+ my_agent = Agent()
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+ print(my_agent.get_chat_history())
+
+ asyncio.run(main())
+ """
+ return self.current_session.chat_history if self.current_session else None
diff --git a/synapseclient/models/mixins/__init__.py b/synapseclient/models/mixins/__init__.py
index 0fb23dac7..93a98589c 100644
--- a/synapseclient/models/mixins/__init__.py
+++ b/synapseclient/models/mixins/__init__.py
@@ -1,9 +1,11 @@
"""References to the mixins that are used in the Synapse models."""
from synapseclient.models.mixins.access_control import AccessControllable
+from synapseclient.models.mixins.asynchronous_job import AsynchronousCommunicator
from synapseclient.models.mixins.storable_container import StorableContainer
__all__ = [
"AccessControllable",
"StorableContainer",
+ "AsynchronousCommunicator",
]
diff --git a/synapseclient/models/mixins/asynchronous_job.py b/synapseclient/models/mixins/asynchronous_job.py
new file mode 100644
index 000000000..aac481663
--- /dev/null
+++ b/synapseclient/models/mixins/asynchronous_job.py
@@ -0,0 +1,410 @@
+import asyncio
+import json
+import time
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from synapseclient import Synapse
+from synapseclient.core.constants.concrete_types import AGENT_CHAT_REQUEST
+from synapseclient.core.exceptions import SynapseError, SynapseTimeoutError
+
+ASYNC_JOB_URIS = {
+ AGENT_CHAT_REQUEST: "/agent/chat/async",
+}
+
+
+class AsynchronousCommunicator:
+ """Mixin to handle communication with the Synapse Asynchronous Job service."""
+
+ def to_synapse_request(self) -> None:
+ """Converts the request to a request expected of the Synapse REST API."""
+ raise NotImplementedError("to_synapse_request must be implemented.")
+
+ def fill_from_dict(
+ self, synapse_response: Dict[str, str]
+ ) -> "AsynchronousCommunicator":
+ """
+ Converts a response from the REST API into this dataclass.
+
+ Arguments:
+ synapse_response: The response from the REST API.
+
+ Returns:
+ An instance of this class.
+ """
+ raise NotImplementedError("fill_from_dict must be implemented.")
+
+ async def _post_exchange_async(
+ self, synapse_client: Optional[Synapse] = None, **kwargs
+ ) -> None:
+ """Any additional logic to run after the exchange with Synapse.
+
+ Arguments:
+ synapse_client: The Synapse client to use for the request.
+ **kwargs: Additional arguments to pass to the request.
+ """
+ pass
+
+ async def send_job_and_wait_async(
+ self,
+ post_exchange_args: Optional[Dict[str, Any]] = None,
+ *,
+ synapse_client: Optional[Synapse] = None,
+ ) -> "AsynchronousCommunicator":
+ """Send the job to the Asynchronous Job service and wait for it to complete.
+ Intended to be called by a class inheriting from this mixin to start a job
+ in the Synapse API and wait for it to complete. The inheriting class needs to
+ represent an asynchronous job request and response and include all necessary attributes.
+ This was initially implemented to be used in the AgentPrompt class which can be used
+ as an example.
+
+ Arguments:
+ post_exchange_args: Additional arguments to pass to the request.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ An instance of this class.
+
+ Example: Using this function
+ This function was initially implemented to be used in the AgentPrompt class
+ to send a prompt to an AI agent and wait for the response. It can also be used
+ in any other class that needs to use an Asynchronous Job.
+
+ The inheriting class (AgentPrompt) will typically not be used directly, but rather
+ through a higher level class (AgentSession), but this example shows how you would
+ use this function.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentPrompt
+
+ syn = Synapse()
+ syn.login()
+
+ agent_prompt = AgentPrompt(
+ id=None,
+ session_id="123",
+ prompt="Hello",
+ response=None,
+ enable_trace=True,
+ trace=None,
+ )
+ # This will fill the id, response, and trace
+ # attributes with the response from the API
+ agent_prompt.send_job_and_wait_async()
+ """
+ result = await send_job_and_wait_async(
+ request=self.to_synapse_request(),
+ request_type=self.concrete_type,
+ synapse_client=synapse_client,
+ )
+ self.fill_from_dict(synapse_response=result)
+ await self._post_exchange_async(
+ **post_exchange_args, synapse_client=synapse_client
+ )
+ return self
+
+
+class AsynchronousJobState(str, Enum):
+ """Enum representing the state of a Synapse Asynchronous Job:
+
+
+ - PROCESSING: The job is being processed.
+ - FAILED: The job has failed.
+ - COMPLETE: The job has been completed.
+ """
+
+ PROCESSING = "PROCESSING"
+ FAILED = "FAILED"
+ COMPLETE = "COMPLETE"
+
+
+class CallersContext(str, Enum):
+ """Enum representing information about a web service call:
+
+
+ - SESSION_ID: Each web service request is issued a unique session ID (UUID)
+ that is included in the call's access record.
+ Events that are triggered by a web service request should include the session ID
+ so that they can be linked to each other and the call's access record.
+ """
+
+ SESSION_ID = "SESSION_ID"
+
+
+@dataclass
+class AsynchronousJobStatus:
+ """Represents a Synapse Asynchronous Job Status object:
+
+
+ Attributes:
+ state: The state of the job. Either PROCESSING, FAILED, or COMPLETE.
+ canceling: Whether the job has been requested to be cancelled.
+ request_body: The body of an Asynchronous job request.
+ Will be one of the models described here:
+
+ response_body: The body of an Asynchronous job response.
+ Will be one of the models described here:
+
+ etag: The etag of the job status. Changes whenever the status changes.
+ id: The ID if the job issued when this job was started.
+ started_by_user_id: The ID of the user that started the job.
+ started_on: The date-time when the status was last changed to PROCESSING.
+ changed_on: The date-time when the status of this job was last changed.
+ progress_message: The current message of the progress tracker.
+ progress_current: A value indicating how much progress has been made.
+ I.e. a value of 50 indicates that 50% of the work has been
+ completed if progress_total is 100.
+ progress_total: A value indicating the total amount of work to complete.
+ exception: The exception that needs to be thrown if the job fails.
+ error_message: A one-line error message when the job fails.
+ error_details: Full stack trace of the error when the job fails.
+ runtime_ms: The number of milliseconds from the start to completion of this job.
+ callers_context: Contextual information about a web service call.
+ """
+
+ state: Optional["AsynchronousJobState"] = None
+ """The state of the job. Either PROCESSING, FAILED, or COMPLETE."""
+
+ canceling: Optional[bool] = False
+ """Whether the job has been requested to be cancelled."""
+
+ request_body: Optional[dict] = None
+ """The body of an Asynchronous job request. Will be one of the models described here:
+ """
+
+ response_body: Optional[dict] = None
+ """The body of an Asynchronous job response. Will be one of the models described here:
+ """
+
+ etag: Optional[str] = None
+ """The etag of the job status. Changes whenever the status changes."""
+
+ id: Optional[str] = None
+ """The ID if the job issued when this job was started."""
+
+ started_by_user_id: Optional[int] = None
+ """The ID of the user that started the job."""
+
+ started_on: Optional[str] = None
+ """The date-time when the status was last changed to PROCESSING."""
+
+ changed_on: Optional[str] = None
+ """The date-time when the status of this job was last changed."""
+
+ progress_message: Optional[str] = None
+ """The current message of the progress tracker."""
+
+ progress_current: Optional[int] = None
+ """A value indicating how much progress has been made.
+ I.e. a value of 50 indicates that 50% of the work has been
+ completed if progress_total is 100."""
+
+ progress_total: Optional[int] = None
+ """A value indicating the total amount of work to complete."""
+
+ exception: Optional[str] = None
+ """The exception that needs to be thrown if the job fails."""
+
+ error_message: Optional[str] = None
+ """A one-line error message when the job fails."""
+
+ error_details: Optional[str] = None
+ """Full stack trace of the error when the job fails."""
+
+ runtime_ms: Optional[int] = None
+ """The number of milliseconds from the start to completion of this job."""
+
+ callers_context: Optional["CallersContext"] = None
+ """Contextual information about a web service call."""
+
+ def fill_from_dict(self, async_job_status: dict) -> "AsynchronousJobStatus":
+ """Converts a response from the REST API into this dataclass.
+
+ Arguments:
+ async_job_status: The response from the REST API.
+
+ Returns:
+ A AsynchronousJobStatus object.
+ """
+ self.state = (
+ AsynchronousJobState(async_job_status.get("jobState"))
+ if async_job_status.get("jobState")
+ else None
+ )
+ self.canceling = async_job_status.get("jobCanceling", None)
+ self.request_body = async_job_status.get("requestBody", None)
+ self.response_body = async_job_status.get("responseBody", None)
+ self.etag = async_job_status.get("etag", None)
+ self.id = async_job_status.get("jobId", None)
+ self.started_by_user_id = async_job_status.get("startedByUserId", None)
+ self.started_on = async_job_status.get("startedOn", None)
+ self.changed_on = async_job_status.get("changedOn", None)
+ self.progress_message = async_job_status.get("progressMessage", None)
+ self.progress_current = async_job_status.get("progressCurrent", None)
+ self.progress_total = async_job_status.get("progressTotal", None)
+ self.exception = async_job_status.get("exception", None)
+ self.error_message = async_job_status.get("errorMessage", None)
+ self.error_details = async_job_status.get("errorDetails", None)
+ self.runtime_ms = async_job_status.get("runtimeMs", None)
+ self.callers_context = async_job_status.get("callersContext", None)
+ return self
+
+
+async def send_job_and_wait_async(
+ request: Dict[str, Any],
+ request_type: str,
+ endpoint: str = None,
+ *,
+ synapse_client: Optional["Synapse"] = None,
+) -> Dict[str, Any]:
+ """
+ Sends the job to the Synapse API and waits for the response. Request body matches:
+
+
+ Arguments:
+ request: A request matching .
+ endpoint: The endpoint to use for the request. Defaults to None.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The response body matching
+
+
+ Raises:
+ SynapseError: If the job fails.
+ SynapseTimeoutError: If the job does not complete within the timeout.
+ """
+ job_id = await send_job_async(request=request, synapse_client=synapse_client)
+ return {
+ "jobId": job_id,
+ **await get_job_async(
+ job_id=job_id,
+ request_type=request_type,
+ synapse_client=synapse_client,
+ endpoint=endpoint,
+ ),
+ }
+
+
+async def send_job_async(
+ request: Dict[str, Any],
+ *,
+ synapse_client: Optional["Synapse"] = None,
+) -> str:
+ """
+ Sends the job to the Synapse API. Request body matches:
+
+ Returns the job ID.
+
+ Arguments:
+ request: A request matching .
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The job ID retrieved from the response.
+
+ """
+ if not request:
+ raise ValueError("request must be provided.")
+
+ request_type = request.get("concreteType")
+
+ if not request_type or request_type not in ASYNC_JOB_URIS:
+ raise ValueError(f"Unsupported request type: {request_type}")
+
+ client = Synapse.get_client(synapse_client=synapse_client)
+ response = await client.rest_post_async(
+ uri=f"{ASYNC_JOB_URIS[request_type]}/start", body=json.dumps(request)
+ )
+ return response["token"]
+
+
+async def get_job_async(
+ job_id: str,
+ request_type: str,
+ endpoint: str = None,
+ sleep: int = 1,
+ timeout: int = 60,
+ *,
+ synapse_client: Optional["Synapse"] = None,
+) -> Dict[str, Any]:
+ """
+ Gets the job from the server using its ID. Handles progress tracking, failures and timeouts.
+
+ Arguments:
+ job_id: The ID of the job to get.
+ request_type: The type of the job.
+ endpoint: The endpoint to use for the request. Defaults to None.
+ sleep: The number of seconds to wait between requests. Defaults to 1.
+ timeout: The number of seconds to wait for the job to complete or progress
+ before raising a SynapseTimeoutError. Defaults to 60.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The response body matching
+
+
+ Raises:
+ SynapseError: If the job fails.
+ SynapseTimeoutError: If the job does not complete or progress within the timeout interval.
+ """
+ client = Synapse.get_client(synapse_client=synapse_client)
+ start_time = asyncio.get_event_loop().time()
+
+ last_message = ""
+ last_progress = 0
+ last_total = 1
+ progressed = False
+
+ while asyncio.get_event_loop().time() - start_time < timeout:
+ result = await client.rest_get_async(
+ uri=f"{ASYNC_JOB_URIS[request_type]}/get/{job_id}",
+ endpoint=endpoint,
+ )
+ job_status = AsynchronousJobStatus().fill_from_dict(async_job_status=result)
+ if job_status.state == AsynchronousJobState.PROCESSING:
+ progress_tracking = any(
+ [
+ job_status.progress_message,
+ job_status.progress_current,
+ job_status.progress_total,
+ ]
+ )
+ progressed = (
+ job_status.progress_message != last_message
+ or last_progress != job_status.progress_current
+ )
+ if progress_tracking and progressed:
+ last_message = job_status.progress_message
+ last_progress = job_status.progress_current
+ last_total = job_status.progress_total
+
+ client._print_transfer_progress(
+ last_progress,
+ last_total,
+ prefix=last_message,
+ isBytes=False,
+ )
+ start_time = asyncio.get_event_loop().time()
+ await asyncio.sleep(sleep)
+ elif job_status.state == AsynchronousJobState.FAILED:
+ raise SynapseError(
+ f"{job_status.error_message}\n{job_status.error_details}",
+ )
+ else:
+ break
+ else:
+ raise SynapseTimeoutError(
+ f"Timeout waiting for query results: {time.time() - start_time} seconds"
+ )
+
+ return result
diff --git a/synapseclient/models/protocols/agent_protocol.py b/synapseclient/models/protocols/agent_protocol.py
new file mode 100644
index 000000000..ef52bf3a3
--- /dev/null
+++ b/synapseclient/models/protocols/agent_protocol.py
@@ -0,0 +1,390 @@
+"""Protocol for the methods of the Agent and AgentSession classes that have
+synchronous counterparts generated at runtime."""
+
+from typing import TYPE_CHECKING, Optional, Protocol
+
+from synapseclient import Synapse
+
+if TYPE_CHECKING:
+ from synapseclient.models import Agent, AgentSession, AgentSessionAccessLevel
+
+
+class AgentSessionSynchronousProtocol(Protocol):
+ """Protocol for the methods of the AgentSession class that have synchronous counterparts
+ generated at runtime."""
+
+ def start(self, *, synapse_client: Optional[Synapse] = None) -> "AgentSession":
+ """Starts an agent session.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The new AgentSession object.
+
+ Example: Start a session and send a prompt.
+ Start a session with a custom agent by providing the agent's registration ID and calling `start()`.
+ Then, send a prompt to the agent.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession
+
+ syn = Synapse()
+ syn.login()
+
+ my_session = AgentSession(agent_registration_id="foo").start()
+ my_session.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+ """
+ return self
+
+ def get(self, *, synapse_client: Optional[Synapse] = None) -> "AgentSession":
+ """Gets an agent session.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The retrieved AgentSession object.
+
+ Example: Get an existing session and send a prompt.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession
+
+ syn = Synapse()
+ syn.login()
+
+ my_session = AgentSession(id="foo").get()
+ my_session.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+ """
+ return self
+
+ def update(self, *, synapse_client: Optional[Synapse] = None) -> "AgentSession":
+ """Updates an agent session.
+ Only updates to the access level are currently supported.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The updated AgentSession object.
+
+ Example: Update the access level of an existing session.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, update the access level of the session and call `update()`.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession, AgentSessionAccessLevel
+
+ syn = Synapse()
+ syn.login()
+
+ my_session = AgentSession(id="foo").get()
+ my_session.access_level = AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA
+ my_session.update()
+ """
+ return self
+
+ def prompt(
+ self,
+ prompt: str,
+ enable_trace: bool = False,
+ print_response: bool = False,
+ newer_than: Optional[int] = None,
+ *,
+ synapse_client: Optional[Synapse] = None,
+ ) -> None:
+ """Sends a prompt to the agent and adds the response to the AgentSession's
+ chat history. A session must be started before sending a prompt.
+
+ Arguments:
+ prompt: The prompt to send to the agent.
+ enable_trace: Whether to enable trace for the prompt.
+ print_response: Whether to print the response to the console.
+ newer_than: The timestamp to get trace results newer than.
+ Defaults to None (all results).
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Example: Send a prompt within an existing session.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import AgentSession
+
+ syn = Synapse()
+ syn.login()
+
+ my_session = AgentSession(id="foo").get()
+ my_session.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+ """
+ return None
+
+
+class AgentSynchronousProtocol(Protocol):
+ """Protocol for the methods of the Agent class that have synchronous counterparts
+ generated at runtime."""
+
+ def register(self, *, synapse_client: Optional[Synapse] = None) -> "Agent":
+ """Registers an agent with the Synapse API.
+ If agent already exists, it will be retrieved.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The registered or existing Agent object.
+
+ Example: Register and chat with a custom agent
+ **Only available for internal users (Sage Bionetworks employees)**
+
+ Alternatively, you can register a custom agent and chat with it provided
+ you have already created it.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent(cloud_agent_id="foo")
+ my_agent.register()
+
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+ """
+ return self
+
+ def get(self, *, synapse_client: Optional[Synapse] = None) -> "Agent":
+ """Gets an existing agent.
+
+ Arguments:
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The existing Agent object.
+
+ Example: Get and chat with an existing agent
+ Retrieve an existing agent by providing the agent's registration ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent(registration_id="foo").get()
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+ """
+ return self
+
+ def start_session(
+ self,
+ access_level: Optional["AgentSessionAccessLevel"] = "PUBLICLY_ACCESSIBLE",
+ *,
+ synapse_client: Optional[Synapse] = None,
+ ) -> "AgentSession":
+ """Starts an agent session.
+ Adds the session to the Agent's sessions dictionary and sets it as the current session.
+
+ Arguments:
+ access_level: The access level of the agent session.
+ Must be one of PUBLICLY_ACCESSIBLE, READ_YOUR_PRIVATE_DATA,
+ or WRITE_YOUR_PRIVATE_DATA.
+ Defaults to PUBLICLY_ACCESSIBLE.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The new AgentSession object.
+
+ Example: Start a session and send a prompt with the baseline Synapse Agent.
+ The baseline Synapse Agent is the default agent used when a registration ID is not provided.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent()
+ my_agent.start_session()
+ my_agent.prompt(
+ prompt="Can you tell me about the AD Knowledge Portal dataset?",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ Example: Start a session and send a prompt with a custom agent.
+ The baseline Synapse Agent is the default agent used when a registration ID is not provided.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent(cloud_agent_id="foo")
+ my_agent.start_session()
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+ """
+ return AgentSession()
+
+ def get_session(
+ self, session_id: str, *, synapse_client: Optional[Synapse] = None
+ ) -> "AgentSession":
+ """Gets an existing agent session.
+ Adds the session to the Agent's sessions dictionary and
+ sets it as the current session.
+
+ Arguments:
+ session_id: The ID of the session to get.
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Returns:
+ The existing AgentSession object.
+
+ Example: Get an existing session and send a prompt.
+ Retrieve an existing session by providing the session ID and calling `get()`.
+ Then, send a prompt to the agent.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_session = Agent().get_session(session_id="foo")
+ my_session.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+ """
+ return AgentSession()
+
+ def prompt(
+ self,
+ prompt: str,
+ enable_trace: bool = False,
+ print_response: bool = False,
+ session: Optional["AgentSession"] = None,
+ newer_than: Optional[int] = None,
+ *,
+ synapse_client: Optional[Synapse] = None,
+ ) -> None:
+ """Sends a prompt to the agent for the current session.
+ If no session is currently active, a new session will be started.
+
+ Arguments:
+ prompt: The prompt to send to the agent.
+ enable_trace: Whether to enable trace for the prompt.
+ print_response: Whether to print the response to the console.
+ session_id: The ID of the session to send the prompt to.
+ If None, the current session will be used.
+ newer_than: The timestamp to get trace results newer than. Defaults to None (all results).
+ synapse_client: If not passed in and caching was not disabled by
+ `Synapse.allow_client_caching(False)` this will use the last created
+ instance from the Synapse class constructor.
+
+ Example: Prompt the baseline Synapse Agent to add annotations to a file on Synapse
+ The baseline Synpase Agent can be used to add annotations to files.
+
+ from synapseclient import Synapse
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent()
+ my_agent.prompt(
+ prompt="Add the annotation 'test' to the file 'syn123456789'",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ Example: Prompt a custom agent.
+ If you have already registered a custom agent, you can prompt it by providing the agent's registration ID.
+
+ from synapseclient import Synapse
+ from synapseclient.models.agent import Agent
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent(registration_id="foo")
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ Advanced Example: Start and prompt multiple sessions
+ Here, we connect to a custom agent and start one session with the prompt "Hello".
+ In the background, this first session is being set as the current session
+ and future prompts will be sent to this session by default. If we want to send a
+ prompt to a different session, we can do so by starting it and calling prompt again,
+ but with our new session as an argument. We now have two sessions, both stored in the
+ `my_agent.sessions` dictionary. After the second prompt, `my_second_session` is now
+ the current session.
+
+ syn = Synapse()
+ syn.login()
+
+ my_agent = Agent(registration_id="foo").get()
+
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ )
+
+ my_second_session = my_agent.start_session()
+ my_agent.prompt(
+ prompt="Hello again",
+ enable_trace=True,
+ print_response=True,
+ session=my_second_session,
+ )
+ """
+ return None
diff --git a/synapseclient/models/services/storable_entity_components.py b/synapseclient/models/services/storable_entity_components.py
index 8615cb9c9..8eafa5739 100644
--- a/synapseclient/models/services/storable_entity_components.py
+++ b/synapseclient/models/services/storable_entity_components.py
@@ -4,7 +4,6 @@
from synapseclient import Synapse
from synapseclient.core.exceptions import SynapseError
-from synapseclient.models import Annotations
if TYPE_CHECKING:
from synapseclient.models import File, Folder, Project, Table
@@ -243,6 +242,8 @@ async def _store_activity_and_annotations(
or last_persistent_instance.annotations != root_resource.annotations
)
):
+ from synapseclient.models import Annotations
+
result = await Annotations(
id=root_resource.id,
etag=root_resource.etag,
diff --git a/tests/integration/synapseclient/models/async/test_agent_async.py b/tests/integration/synapseclient/models/async/test_agent_async.py
new file mode 100644
index 000000000..dd7ef53e4
--- /dev/null
+++ b/tests/integration/synapseclient/models/async/test_agent_async.py
@@ -0,0 +1,228 @@
+"""Integration tests for the asynchronous methods of the AgentPrompt, AgentSession, and Agent classes."""
+
+# These tests have been disabled until out `test` user has needed permissions
+# Context: https://sagebionetworks.jira.com/browse/SYNPY-1544?focusedCommentId=235070
+# import pytest
+
+# from synapseclient import Synapse
+# from synapseclient.core.constants.concrete_types import AGENT_CHAT_REQUEST
+# from synapseclient.models.agent import (
+# Agent,
+# AgentPrompt,
+# AgentSession,
+# AgentSessionAccessLevel,
+# )
+
+# # These are the ID values for a "Hello World" agent registered on Synapse.
+# # The Bedrock agent is hosted on Sage Bionetworks AWS infrastructure.
+# # CFN Template:
+# # https://raw.githubusercontent.com/Sage-Bionetworks-Workflows/dpe-agents/refs/heads/main/client_integration_test/template.json
+# AGENT_AWS_ID = "QOTV3KQM1X"
+# AGENT_REGISTRATION_ID = "29"
+
+
+# class TestAgentPrompt:
+# """Integration tests for the synchronous methods of the AgentPrompt class."""
+
+# @pytest.fixture(autouse=True, scope="function")
+# def init(self, syn: Synapse) -> None:
+# self.syn = syn
+
+# async def test_send_job_and_wait_async_with_post_exchange_args(self) -> None:
+# # GIVEN an AgentPrompt with a valid concrete type, prompt, and enable_trace
+# test_prompt = AgentPrompt(
+# concrete_type=AGENT_CHAT_REQUEST,
+# prompt="hello",
+# enable_trace=True,
+# )
+# # AND the ID of an existing agent session
+# test_session = await AgentSession(
+# agent_registration_id=AGENT_REGISTRATION_ID
+# ).start_async(synapse_client=self.syn)
+# test_prompt.session_id = test_session.id
+# # WHEN I send the job and wait for it to complete
+# await test_prompt.send_job_and_wait_async(
+# post_exchange_args={"newer_than": 0},
+# synapse_client=self.syn,
+# )
+# # THEN I expect the AgentPrompt to be updated with the response and trace
+# assert test_prompt.response is not None
+# assert test_prompt.trace is not None
+
+
+# class TestAgentSession:
+# """Integration tests for the synchronous methods of the AgentSession class."""
+
+# @pytest.fixture(autouse=True, scope="function")
+# def init(self, syn: Synapse) -> None:
+# self.syn = syn
+
+# async def test_start(self) -> None:
+# # GIVEN an agent session with a valid agent registration id
+# agent_session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID)
+
+# # WHEN the start method is called
+# result_session = await agent_session.start_async(synapse_client=self.syn)
+
+# # THEN the result should be an AgentSession object
+# # with expected attributes including an empty chat history
+# assert result_session.id is not None
+# assert (
+# result_session.access_level == AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE
+# )
+# assert result_session.started_on is not None
+# assert result_session.started_by is not None
+# assert result_session.modified_on is not None
+# assert result_session.agent_registration_id == AGENT_REGISTRATION_ID
+# assert result_session.etag is not None
+# assert result_session.chat_history == []
+
+# async def test_get(self) -> None:
+# # GIVEN an agent session with a valid agent registration id
+# agent_session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID)
+# # WHEN I start a session
+# await agent_session.start_async(synapse_client=self.syn)
+# # THEN I expect to be able to get the session with its id
+# new_session = await AgentSession(id=agent_session.id).get_async(
+# synapse_client=self.syn
+# )
+# assert new_session == agent_session
+
+# async def test_update(self) -> None:
+# # GIVEN an agent session with a valid agent
+# # registration id and access level set
+# agent_session = AgentSession(
+# agent_registration_id=AGENT_REGISTRATION_ID,
+# access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+# )
+# # WHEN I start a session
+# await agent_session.start_async(synapse_client=self.syn)
+# # AND I update the access level of the session
+# agent_session.access_level = AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA
+# await agent_session.update_async(synapse_client=self.syn)
+# # THEN I expect the access level to be updated
+# updated_session = await AgentSession(id=agent_session.id).get_async(
+# synapse_client=self.syn
+# )
+# assert (
+# updated_session.access_level
+# == AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA
+# )
+
+# async def test_prompt(self) -> None:
+# # GIVEN an agent session with a valid agent registration id
+# agent_session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID)
+# # WHEN I start a session
+# await agent_session.start_async(synapse_client=self.syn)
+# # THEN I expect to be able to prompt the agent
+# await agent_session.prompt_async(
+# prompt="hello",
+# enable_trace=True,
+# )
+# # AND I expect the chat history to be updated with the prompt and response
+# assert len(agent_session.chat_history) == 1
+# assert agent_session.chat_history[0].prompt == "hello"
+# assert agent_session.chat_history[0].response is not None
+# assert agent_session.chat_history[0].trace is not None
+
+
+# class TestAgent:
+# """Integration tests for the synchronous methods of the Agent class."""
+
+# def get_test_agent(self) -> Agent:
+# return Agent(
+# cloud_agent_id=AGENT_AWS_ID,
+# cloud_alias_id="TSTALIASID",
+# registration_id=AGENT_REGISTRATION_ID,
+# registered_on="2025-01-16T18:57:35.680Z",
+# type="CUSTOM",
+# sessions={},
+# current_session=None,
+# )
+
+# @pytest.fixture(autouse=True, scope="function")
+# def init(self, syn: Synapse) -> None:
+# self.syn = syn
+
+# async def test_register(self) -> None:
+# # GIVEN an Agent with a valid agent AWS id
+# agent = Agent(cloud_agent_id=AGENT_AWS_ID)
+# # WHEN I register the agent
+# await agent.register_async(synapse_client=self.syn)
+# # THEN I expect the agent to be registered
+# expected_agent = self.get_test_agent()
+# assert agent == expected_agent
+
+# async def test_get(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = Agent(registration_id=AGENT_REGISTRATION_ID)
+# # WHEN I get the agent
+# await agent.get_async(synapse_client=self.syn)
+# # THEN I expect the agent to be returned
+# expected_agent = self.get_test_agent()
+# assert agent == expected_agent
+
+# async def test_get_no_registration_id(self) -> None:
+# # GIVEN an Agent with no registration id
+# agent = Agent()
+# # WHEN I get the agent, I expect a ValueError to be raised
+# with pytest.raises(ValueError, match="Registration ID is required"):
+# await agent.get_async(synapse_client=self.syn)
+
+# async def test_start_session(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = Agent(registration_id=AGENT_REGISTRATION_ID)
+# # WHEN I start a session
+# await agent.start_session_async(synapse_client=self.syn)
+# # THEN I expect a current session to be set
+# assert agent.current_session is not None
+# # AND I expect the session to be in the sessions dictionary
+# assert agent.sessions[agent.current_session.id] == agent.current_session
+
+# async def test_get_session(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = Agent(registration_id=AGENT_REGISTRATION_ID)
+# # WHEN I start a session
+# await agent.start_session_async(synapse_client=self.syn)
+# # THEN I expect to be able to get the session with its id
+# existing_session = await agent.get_session_async(
+# session_id=agent.current_session.id
+# )
+# # AND I expect those sessions to be the same
+# assert existing_session == agent.current_session
+
+# async def test_prompt_with_session(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = await Agent(registration_id=AGENT_REGISTRATION_ID).get_async(
+# synapse_client=self.syn
+# )
+# # AND a session started separately
+# session = await AgentSession(
+# agent_registration_id=AGENT_REGISTRATION_ID
+# ).start_async(synapse_client=self.syn)
+# # WHEN I prompt the agent with a session
+# await agent.prompt_async(prompt="hello", enable_trace=True, session=session)
+# test_session = agent.sessions[session.id]
+# # THEN I expect the chat history to be updated with the prompt and response
+# assert len(test_session.chat_history) == 1
+# assert test_session.chat_history[0].prompt == "hello"
+# assert test_session.chat_history[0].response is not None
+# assert test_session.chat_history[0].trace is not None
+# # AND I expect the current session to be the session provided
+# assert agent.current_session.id == session.id
+
+# async def test_prompt_no_session(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = await Agent(registration_id=AGENT_REGISTRATION_ID).get_async(
+# synapse_client=self.syn
+# )
+# # WHEN I prompt the agent without a current session set
+# # and no session provided
+# await agent.prompt_async(prompt="hello", enable_trace=True)
+# # THEN I expect a new session to be started and set as the current session
+# assert agent.current_session is not None
+# # AND I expect the chat history to be updated with the prompt and response
+# assert len(agent.current_session.chat_history) == 1
+# assert agent.current_session.chat_history[0].prompt == "hello"
+# assert agent.current_session.chat_history[0].response is not None
+# assert agent.current_session.chat_history[0].trace is not None
diff --git a/tests/integration/synapseclient/models/synchronous/test_agent.py b/tests/integration/synapseclient/models/synchronous/test_agent.py
new file mode 100644
index 000000000..07b77291e
--- /dev/null
+++ b/tests/integration/synapseclient/models/synchronous/test_agent.py
@@ -0,0 +1,192 @@
+"""Integration tests for the synchronous methods of the AgentSession and Agent classes."""
+
+# These tests have been disabled until out `test` user has needed permissions
+# Context: https://sagebionetworks.jira.com/browse/SYNPY-1544?focusedCommentId=235070
+# import pytest
+
+# from synapseclient import Synapse
+# from synapseclient.models.agent import Agent, AgentSession, AgentSessionAccessLevel
+
+# # These are the ID values for a "Hello World" agent registered on Synapse.
+# # The Bedrock agent is hosted on Sage Bionetworks AWS infrastructure.
+# # CFN Template:
+# # https://raw.githubusercontent.com/Sage-Bionetworks-Workflows/dpe-agents/refs/heads/main/client_integration_test/template.json
+# CLOUD_AGENT_ID = "QOTV3KQM1X"
+# AGENT_REGISTRATION_ID = "29"
+
+
+# class TestAgentSession:
+# """Integration tests for the synchronous methods of the AgentSession class."""
+
+# @pytest.fixture(autouse=True, scope="function")
+# def init(self, syn: Synapse) -> None:
+# self.syn = syn
+
+# async def test_start(self) -> None:
+# # GIVEN an agent session with a valid agent registration id
+# agent_session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID)
+
+# # WHEN the start method is called
+# result_session = agent_session.start(synapse_client=self.syn)
+
+# # THEN the result should be an AgentSession object
+# # with expected attributes including an empty chat history
+# assert result_session.id is not None
+# assert (
+# result_session.access_level == AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE
+# )
+# assert result_session.started_on is not None
+# assert result_session.started_by is not None
+# assert result_session.modified_on is not None
+# assert result_session.agent_registration_id == str(AGENT_REGISTRATION_ID)
+# assert result_session.etag is not None
+# assert result_session.chat_history == []
+
+# async def test_get(self) -> None:
+# # GIVEN an agent session with a valid agent registration id
+# agent_session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID)
+# # WHEN I start a session
+# agent_session.start(synapse_client=self.syn)
+# # THEN I expect to be able to get the session with its id
+# new_session = AgentSession(id=agent_session.id).get(synapse_client=self.syn)
+# assert new_session == agent_session
+
+# async def test_update(self) -> None:
+# # GIVEN an agent session with a valid agent registration id and access level set
+# agent_session = AgentSession(
+# agent_registration_id=AGENT_REGISTRATION_ID,
+# access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+# )
+# # WHEN I start a session
+# agent_session.start(synapse_client=self.syn)
+# # AND I update the access level of the session
+# agent_session.access_level = AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA
+# agent_session.update(synapse_client=self.syn)
+# # THEN I expect the access level to be updated
+# updated_session = AgentSession(id=agent_session.id).get(synapse_client=self.syn)
+# assert (
+# updated_session.access_level
+# == AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA
+# )
+
+# async def test_prompt(self) -> None:
+# # GIVEN an agent session with a valid agent registration id
+# agent_session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID)
+# # WHEN I start a session
+# agent_session.start(synapse_client=self.syn)
+# # THEN I expect to be able to prompt the agent
+# agent_session.prompt(
+# prompt="hello",
+# enable_trace=True,
+# )
+# # AND I expect the chat history to be updated with the prompt and response
+# assert len(agent_session.chat_history) == 1
+# assert agent_session.chat_history[0].prompt == "hello"
+# assert agent_session.chat_history[0].response is not None
+# assert agent_session.chat_history[0].trace is not None
+
+
+# class TestAgent:
+# """Integration tests for the synchronous methods of the Agent class."""
+
+# def get_test_agent(self) -> Agent:
+# return Agent(
+# cloud_agent_id=CLOUD_AGENT_ID,
+# cloud_alias_id="TSTALIASID",
+# registration_id=AGENT_REGISTRATION_ID,
+# registered_on="2025-01-16T18:57:35.680Z",
+# type="CUSTOM",
+# sessions={},
+# current_session=None,
+# )
+
+# @pytest.fixture(autouse=True, scope="function")
+# def init(self, syn: Synapse) -> None:
+# self.syn = syn
+
+# async def test_register(self) -> None:
+# # GIVEN an Agent with a valid agent AWS id
+# agent = Agent(cloud_agent_id=CLOUD_AGENT_ID)
+# # WHEN I register the agent
+# agent.register(synapse_client=self.syn)
+# # THEN I expect the agent to be registered
+# expected_agent = self.get_test_agent()
+# assert agent == expected_agent
+
+# async def test_get(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = Agent(registration_id=AGENT_REGISTRATION_ID)
+# # WHEN I get the agent
+# agent.get(synapse_client=self.syn)
+# # THEN I expect the agent to be returned
+# expected_agent = self.get_test_agent()
+# assert agent == expected_agent
+
+# async def test_get_no_registration_id(self) -> None:
+# # GIVEN an Agent with no registration id
+# agent = Agent()
+# # WHEN I get the agent, I expect a ValueError to be raised
+# with pytest.raises(ValueError, match="Registration ID is required"):
+# agent.get(synapse_client=self.syn)
+
+# async def test_start_session(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = Agent(registration_id=AGENT_REGISTRATION_ID).get(
+# synapse_client=self.syn
+# )
+# # WHEN I start a session
+# agent.start_session(synapse_client=self.syn)
+# # THEN I expect a current session to be set
+# assert agent.current_session is not None
+# # AND I expect the session to be in the sessions dictionary
+# assert agent.sessions[agent.current_session.id] == agent.current_session
+
+# async def test_get_session(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = Agent(registration_id=AGENT_REGISTRATION_ID).get(
+# synapse_client=self.syn
+# )
+# # WHEN I start a session
+# session = agent.start_session(synapse_client=self.syn)
+# # THEN I expect to be able to get the session with its id
+# existing_session = agent.get_session(session_id=session.id)
+# # AND I expect those sessions to be the same
+# assert existing_session == session
+# # AND I expect it to be the current session
+# assert existing_session == agent.current_session
+
+# async def test_prompt_with_session(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = Agent(registration_id=AGENT_REGISTRATION_ID).get(
+# synapse_client=self.syn
+# )
+# # AND a session started separately
+# session = AgentSession(agent_registration_id=AGENT_REGISTRATION_ID).start(
+# synapse_client=self.syn
+# )
+# # WHEN I prompt the agent with a session
+# agent.prompt(prompt="hello", enable_trace=True, session=session)
+# test_session = agent.sessions[session.id]
+# # THEN I expect the chat history to be updated with the prompt and response
+# assert len(test_session.chat_history) == 1
+# assert test_session.chat_history[0].prompt == "hello"
+# assert test_session.chat_history[0].response is not None
+# assert test_session.chat_history[0].trace is not None
+# # AND I expect the current session to be the session provided
+# assert agent.current_session.id == session.id
+
+# async def test_prompt_no_session(self) -> None:
+# # GIVEN an Agent with a valid agent registration id
+# agent = Agent(registration_id=AGENT_REGISTRATION_ID).get(
+# synapse_client=self.syn
+# )
+# # WHEN I prompt the agent without a current session set
+# # and no session provided
+# agent.prompt(prompt="hello", enable_trace=True)
+# # THEN I expect a new session to be started and set as the current session
+# assert agent.current_session is not None
+# # AND I expect the chat history to be updated with the prompt and response
+# assert len(agent.current_session.chat_history) == 1
+# assert agent.current_session.chat_history[0].prompt == "hello"
+# assert agent.current_session.chat_history[0].response is not None
+# assert agent.current_session.chat_history[0].trace is not None
diff --git a/tests/unit/synapseclient/mixins/async/unit_test_asynchronous_job.py b/tests/unit/synapseclient/mixins/async/unit_test_asynchronous_job.py
new file mode 100644
index 000000000..056976dcc
--- /dev/null
+++ b/tests/unit/synapseclient/mixins/async/unit_test_asynchronous_job.py
@@ -0,0 +1,278 @@
+"""Unit tests for Asynchronous Job logic."""
+
+import json
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from synapseclient import Synapse
+from synapseclient.core.constants.concrete_types import AGENT_CHAT_REQUEST
+from synapseclient.core.exceptions import SynapseError, SynapseTimeoutError
+from synapseclient.models.mixins.asynchronous_job import (
+ ASYNC_JOB_URIS,
+ AsynchronousJobState,
+ AsynchronousJobStatus,
+ get_job_async,
+ send_job_and_wait_async,
+ send_job_async,
+)
+
+
+class TestSendJobAsync:
+ """Unit tests for send_job_async."""
+
+ good_request = {"concreteType": AGENT_CHAT_REQUEST}
+ bad_request_no_concrete_type = {"otherKey": "otherValue"}
+ bad_request_invalid_concrete_type = {"concreteType": "InvalidConcreteType"}
+ request_type = AGENT_CHAT_REQUEST
+
+ @pytest.fixture(autouse=True, scope="function")
+ def init_syn(self, syn: Synapse) -> None:
+ self.syn = syn
+
+ async def test_send_job_async_when_request_is_missing(self) -> None:
+ with pytest.raises(ValueError, match="request must be provided."):
+ # WHEN I call send_job_async without a request
+ # THEN I should get a ValueError
+ await send_job_async(request=None)
+
+ async def test_send_job_async_when_request_is_missing_concrete_type(self) -> None:
+ with pytest.raises(ValueError, match="Unsupported request type: None"):
+ # GIVEN a request with no concrete type
+ # WHEN I call send_job_async
+ # THEN I should get a ValueError
+ await send_job_async(request=self.bad_request_no_concrete_type)
+
+ async def test_send_job_async_when_request_is_invalid_concrete_type(self) -> None:
+ with pytest.raises(
+ ValueError, match="Unsupported request type: InvalidConcreteType"
+ ):
+ # GIVEN a request with an invalid concrete type
+ # WHEN I call send_job_async
+ # THEN I should get a ValueError
+ await send_job_async(request=self.bad_request_invalid_concrete_type)
+
+ async def test_send_job_async_when_request_is_valid(self) -> None:
+ with (
+ patch(
+ "synapseclient.Synapse.get_client",
+ return_value=self.syn,
+ ) as mock_get_client,
+ patch(
+ "synapseclient.Synapse.rest_post_async",
+ new_callable=AsyncMock,
+ return_value={"token": "123"},
+ ) as mock_rest_post_async,
+ ):
+ # WHEN I call send_job_async with a good request
+ job_id = await send_job_async(
+ request=self.good_request, synapse_client=self.syn
+ )
+ # THEN the return value should be the token
+ assert job_id == "123"
+ # AND get_client should have been called
+ mock_get_client.assert_called_once_with(synapse_client=self.syn)
+ # AND rest_post_async should have been called with the correct arguments
+ mock_rest_post_async.assert_called_once_with(
+ uri=f"{ASYNC_JOB_URIS[self.request_type]}/start",
+ body=json.dumps(self.good_request),
+ )
+
+
+class TestGetJobAsync:
+ """Unit tests for get_job_async."""
+
+ request_type = AGENT_CHAT_REQUEST
+ job_id = "123"
+
+ processing_job_status = AsynchronousJobStatus(
+ state=AsynchronousJobState.PROCESSING,
+ progress_message="Processing",
+ progress_current=1,
+ progress_total=100,
+ )
+ failed_job_status = AsynchronousJobStatus(
+ state=AsynchronousJobState.FAILED,
+ progress_message="Failed",
+ progress_current=1,
+ progress_total=100,
+ error_message="Error",
+ error_details="Details",
+ id="123",
+ )
+
+ @pytest.fixture(autouse=True, scope="function")
+ def init_syn(self, syn: Synapse) -> None:
+ self.syn = syn
+
+ async def test_get_job_async_when_job_fails(self) -> None:
+ with (
+ patch(
+ "synapseclient.Synapse.rest_get_async",
+ new_callable=AsyncMock,
+ return_value={},
+ ) as mock_rest_get_async,
+ patch.object(
+ AsynchronousJobStatus,
+ "fill_from_dict",
+ return_value=self.failed_job_status,
+ ) as mock_fill_from_dict,
+ ):
+ with pytest.raises(
+ SynapseError,
+ match=(
+ f"{self.failed_job_status.error_message}\n"
+ f"{self.failed_job_status.error_details}"
+ ),
+ ):
+ # WHEN I call get_job_async
+ # AND the job fails in the Synapse API
+ # THEN I should get a SynapseError with the error message and details
+ await get_job_async(
+ job_id="123",
+ request_type=AGENT_CHAT_REQUEST,
+ synapse_client=self.syn,
+ sleep=1,
+ timeout=60,
+ endpoint=None,
+ )
+ # AND rest_get_async should have been called once with the correct arguments
+ mock_rest_get_async.assert_called_once_with(
+ uri=f"{ASYNC_JOB_URIS[AGENT_CHAT_REQUEST]}/get/{self.job_id}",
+ endpoint=None,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ async_job_status=mock_rest_get_async.return_value,
+ )
+
+ async def test_get_job_async_when_job_times_out(self) -> None:
+ with (
+ patch(
+ "synapseclient.Synapse.rest_get_async",
+ new_callable=AsyncMock,
+ return_value={},
+ ) as mock_rest_get_async,
+ patch.object(
+ AsynchronousJobStatus,
+ "fill_from_dict",
+ return_value=self.processing_job_status,
+ ) as mock_fill_from_dict,
+ ):
+ with pytest.raises(
+ SynapseTimeoutError, match="Timeout waiting for query results:"
+ ):
+ # WHEN I call get_job_async
+ # AND the job does not complete or progress within the timeout interval
+ # THEN I should get a SynapseTimeoutError
+ await get_job_async(
+ job_id=self.job_id,
+ request_type=self.request_type,
+ synapse_client=self.syn,
+ endpoint=None,
+ timeout=0,
+ sleep=1,
+ )
+ # AND rest_get_async should not have been called
+ mock_rest_get_async.assert_not_called()
+ # AND fill_from_dict should not have been called
+ mock_fill_from_dict.assert_not_called()
+
+
+class TestSendJobAndWaitAsync:
+ """Unit tests for send_job_and_wait_async."""
+
+ good_request = {"concreteType": AGENT_CHAT_REQUEST}
+ job_id = "123"
+ request_type = AGENT_CHAT_REQUEST
+
+ @pytest.fixture(autouse=True, scope="function")
+ def init_syn(self, syn: Synapse) -> None:
+ self.syn = syn
+
+ async def test_send_job_and_wait_async(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.mixins.asynchronous_job.send_job_async",
+ new_callable=AsyncMock,
+ return_value=self.job_id,
+ ) as mock_send_job_async,
+ patch(
+ "synapseclient.models.mixins.asynchronous_job.get_job_async",
+ new_callable=AsyncMock,
+ return_value={
+ "key": "value",
+ },
+ ) as mock_get_job_async,
+ ):
+ # WHEN I call send_job_and_wait_async with a good request
+ # THEN the return value should be a dictionary with the job ID
+ # and response key value pair(s)
+ assert await send_job_and_wait_async(
+ request=self.good_request,
+ request_type=self.request_type,
+ synapse_client=self.syn,
+ endpoint=None,
+ ) == {
+ "jobId": self.job_id,
+ "key": "value",
+ }
+ # AND send_job_async should have been called once with the correct arguments
+ mock_send_job_async.assert_called_once_with(
+ request=self.good_request,
+ synapse_client=self.syn,
+ )
+ # AND get_job_async should have been called once with the correct arguments
+ mock_get_job_async.assert_called_once_with(
+ job_id=self.job_id,
+ request_type=self.request_type,
+ synapse_client=self.syn,
+ endpoint=None,
+ )
+
+
+class TestAsynchronousJobStatus:
+ """Unit tests for AsynchronousJobStatus."""
+
+ def test_fill_from_dict(self) -> None:
+ # GIVEN a dictionary with job status information
+ async_job_status_dict = {
+ "jobState": AsynchronousJobState.PROCESSING,
+ "jobCanceling": False,
+ "requestBody": {"key": "value"},
+ "responseBody": {"key": "value"},
+ "etag": "123",
+ "jobId": "123",
+ "startedByUserId": "123",
+ "startedOn": "123",
+ "changedOn": "123",
+ "progressMessage": "Processing",
+ "progressCurrent": 1,
+ "progressTotal": 100,
+ "exception": None,
+ "errorMessage": None,
+ "errorDetails": None,
+ "runtimeMs": 1000,
+ "callersContext": None,
+ }
+ # WHEN I call fill_from_dict on it
+ async_job_status = AsynchronousJobStatus().fill_from_dict(async_job_status_dict)
+ # THEN the resulting AsynchronousJobStatus object
+ # should have the correct attribute values
+ assert async_job_status.state == AsynchronousJobState.PROCESSING
+ assert async_job_status.canceling is False
+ assert async_job_status.request_body == {"key": "value"}
+ assert async_job_status.response_body == {"key": "value"}
+ assert async_job_status.etag == "123"
+ assert async_job_status.id == "123"
+ assert async_job_status.started_by_user_id == "123"
+ assert async_job_status.started_on == "123"
+ assert async_job_status.changed_on == "123"
+ assert async_job_status.progress_message == "Processing"
+ assert async_job_status.progress_current == 1
+ assert async_job_status.progress_total == 100
+ assert async_job_status.exception is None
+ assert async_job_status.error_message is None
+ assert async_job_status.error_details is None
+ assert async_job_status.runtime_ms == 1000
+ assert async_job_status.callers_context is None
diff --git a/tests/unit/synapseclient/models/async/unit_test_agent_async.py b/tests/unit/synapseclient/models/async/unit_test_agent_async.py
new file mode 100644
index 000000000..290094301
--- /dev/null
+++ b/tests/unit/synapseclient/models/async/unit_test_agent_async.py
@@ -0,0 +1,703 @@
+"""Unit tests for Asynchronous methods in Agent, AgentSession, and AgentPrompt classes."""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from synapseclient import Synapse
+from synapseclient.core.constants.concrete_types import AGENT_CHAT_REQUEST
+from synapseclient.models.agent import (
+ Agent,
+ AgentPrompt,
+ AgentSession,
+ AgentSessionAccessLevel,
+ AgentType,
+)
+
+
+class TestAgentPrompt:
+ """Unit tests for the AgentPrompt class' asynchronous methods."""
+
+ agent_prompt = AgentPrompt(
+ id="123",
+ concrete_type=AGENT_CHAT_REQUEST,
+ session_id="456",
+ prompt="Hello",
+ enable_trace=True,
+ )
+ synapse_request = {
+ "concreteType": agent_prompt.concrete_type,
+ "sessionId": agent_prompt.session_id,
+ "chatText": agent_prompt.prompt,
+ "enableTrace": agent_prompt.enable_trace,
+ }
+ synapse_response = {
+ "jobId": "123",
+ "sessionId": "456",
+ "responseText": "World",
+ }
+ trace_response = {
+ "page": [
+ {
+ "message": "I'm a robot",
+ }
+ ]
+ }
+
+ @pytest.fixture(autouse=True, scope="function")
+ def init_syn(self, syn: Synapse) -> None:
+ self.syn = syn
+
+ async def test_to_synapse_request(self):
+ # WHEN I call to_synapse_request on an initialized AgentPrompt
+ result = self.agent_prompt.to_synapse_request()
+ # THEN the result should be a dictionary with the correct keys and values
+ assert result == {
+ "concreteType": self.agent_prompt.concrete_type,
+ "sessionId": self.agent_prompt.session_id,
+ "chatText": self.agent_prompt.prompt,
+ "enableTrace": self.agent_prompt.enable_trace,
+ }
+
+ async def test_fill_from_dict(self):
+ # WHEN I call fill_from_dict on an initialized AgentPrompt with a synapse_response
+ result_agent_prompt = self.agent_prompt.fill_from_dict(self.synapse_response)
+ # THEN the result should be an AgentPrompt with the correct values
+ assert result_agent_prompt.id == self.synapse_response["jobId"]
+ assert result_agent_prompt.session_id == self.synapse_response["sessionId"]
+ assert result_agent_prompt.response == self.synapse_response["responseText"]
+
+ async def test_post_exchange_async_trace_enabled(self):
+ with patch(
+ "synapseclient.models.agent.get_trace",
+ new_callable=AsyncMock,
+ return_value=self.trace_response,
+ ) as mock_get_trace:
+ # WHEN I call _post_exchange_async on an
+ # initialized AgentPrompt with enable_trace=True
+ await self.agent_prompt._post_exchange_async(synapse_client=self.syn)
+ # THEN the mock_get_trace should have been called with the correct arguments
+ mock_get_trace.assert_called_once_with(
+ prompt_id=self.agent_prompt.id,
+ newer_than=None,
+ synapse_client=self.syn,
+ )
+ # AND the trace should be set to the response from the mock_get_trace
+ assert self.agent_prompt.trace == self.trace_response["page"][0]["message"]
+
+ async def test_post_exchange_async_trace_disabled(self):
+ with patch(
+ "synapseclient.models.agent.get_trace",
+ new_callable=AsyncMock,
+ return_value=self.trace_response,
+ ) as mock_get_trace:
+ self.agent_prompt.enable_trace = False
+ # WHEN I call _post_exchange_async on an
+ # initialized AgentPrompt with enable_trace=False
+ await self.agent_prompt._post_exchange_async(synapse_client=self.syn)
+ # THEN the mock_get_trace should not have been called
+ mock_get_trace.assert_not_called()
+
+ async def test_send_job_and_wait_async(self):
+ with (
+ patch(
+ "synapseclient.models.mixins.asynchronous_job.send_job_and_wait_async",
+ new_callable=AsyncMock,
+ return_value=self.synapse_response,
+ ) as mock_send_job_and_wait_async,
+ patch.object(
+ self.agent_prompt,
+ "to_synapse_request",
+ return_value=self.synapse_request,
+ ) as mock_to_synapse_request,
+ patch.object(
+ self.agent_prompt,
+ "fill_from_dict",
+ ) as mock_fill_from_dict,
+ patch.object(
+ self.agent_prompt,
+ "_post_exchange_async",
+ new_callable=AsyncMock,
+ ) as mock_post_exchange_async,
+ ):
+ # WHEN I call send_job_and_wait_async on an initialized AgentPrompt
+ await self.agent_prompt.send_job_and_wait_async(
+ post_exchange_args={"foo": "bar"}, synapse_client=self.syn
+ )
+ # THEN the mock_send_job_and_wait_async should
+ # have been called with the correct arguments
+ mock_send_job_and_wait_async.assert_called_once_with(
+ request=mock_to_synapse_request.return_value,
+ request_type=self.agent_prompt.concrete_type,
+ synapse_client=self.syn,
+ )
+ # THEN the mock_fill_from_dict should have been called with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ synapse_response=self.synapse_response
+ )
+ # AND the mock_post_exchange_async should have been called with the correct arguments
+ mock_post_exchange_async.assert_called_once_with(
+ synapse_client=self.syn, **{"foo": "bar"}
+ )
+
+
+class TestAgentSession:
+ """Unit tests for the AgentSession class' synchronous methods."""
+
+ test_session = AgentSession(
+ id="123",
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ started_on="2024-01-01T00:00:00Z",
+ started_by="123456789",
+ modified_on="2024-01-01T00:00:00Z",
+ agent_registration_id="0",
+ etag="11111111-1111-1111-1111-111111111111",
+ )
+
+ session_response = {
+ "sessionId": test_session.id,
+ "agentAccessLevel": test_session.access_level,
+ "startedOn": test_session.started_on,
+ "startedBy": test_session.started_by,
+ "modifiedOn": test_session.modified_on,
+ "agentRegistrationId": test_session.agent_registration_id,
+ "etag": test_session.etag,
+ }
+
+ updated_test_session = AgentSession(
+ id=test_session.id,
+ access_level=AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA,
+ started_on=test_session.started_on,
+ started_by=test_session.started_by,
+ modified_on=test_session.modified_on,
+ agent_registration_id=test_session.agent_registration_id,
+ etag=test_session.etag,
+ )
+
+ updated_session_response = {
+ "sessionId": updated_test_session.id,
+ "agentAccessLevel": updated_test_session.access_level,
+ "startedOn": updated_test_session.started_on,
+ "startedBy": updated_test_session.started_by,
+ "modifiedOn": updated_test_session.modified_on,
+ "agentRegistrationId": updated_test_session.agent_registration_id,
+ "etag": updated_test_session.etag,
+ }
+
+ test_prompt_trace_enabled = AgentPrompt(
+ concrete_type=AGENT_CHAT_REQUEST,
+ session_id="456",
+ prompt="Hello",
+ enable_trace=True,
+ response="World",
+ trace="Trace",
+ )
+
+ test_prompt_trace_disabled = AgentPrompt(
+ concrete_type=AGENT_CHAT_REQUEST,
+ session_id="456",
+ prompt="Hello",
+ enable_trace=False,
+ response="World",
+ trace=None,
+ )
+
+ @pytest.fixture(autouse=True, scope="function")
+ def init_syn(self, syn: Synapse) -> None:
+ self.syn = syn
+
+ async def test_fill_from_dict(self) -> None:
+ # WHEN I call fill_from_dict on an empty AgentSession with a synapse_response
+ result_session = AgentSession().fill_from_dict(self.session_response)
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+
+ async def test_start_async(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.start_session",
+ new_callable=AsyncMock,
+ return_value=self.session_response,
+ ) as mock_start_session,
+ patch.object(
+ AgentSession,
+ "fill_from_dict",
+ return_value=self.test_session,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an AgentSession with access_level and agent_registration_id
+ initial_session = AgentSession(
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ agent_registration_id=0,
+ )
+ # WHEN I call start
+ result_session = await initial_session.start_async(synapse_client=self.syn)
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+ # AND start_session should have been called once with the correct arguments
+ mock_start_session.assert_called_once_with(
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ agent_registration_id=0,
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ synapse_agent_session=self.session_response
+ )
+
+ async def test_get_async(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.get_session",
+ new_callable=AsyncMock,
+ return_value=self.session_response,
+ ) as mock_get_session,
+ patch.object(
+ AgentSession,
+ "fill_from_dict",
+ return_value=self.test_session,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an AgentSession with an agent_registration_id
+ initial_session = AgentSession(
+ agent_registration_id=0,
+ )
+ # WHEN I call get
+ result_session = await initial_session.get_async(synapse_client=self.syn)
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+ # AND get_session should have been called once with the correct arguments
+ mock_get_session.assert_called_once_with(
+ id=initial_session.id,
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ synapse_agent_session=self.session_response
+ )
+
+ async def test_update_async(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.update_session",
+ new_callable=AsyncMock,
+ return_value=self.updated_session_response,
+ ) as mock_update_session,
+ patch.object(
+ AgentSession,
+ "fill_from_dict",
+ return_value=self.updated_test_session,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an AgentSession with an updated access_level
+ # WHEN I call update
+ result_session = await self.updated_test_session.update_async(
+ synapse_client=self.syn
+ )
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.updated_test_session
+ # AND update_session should have been called once with the correct arguments
+ mock_update_session.assert_called_once_with(
+ id=self.updated_test_session.id,
+ access_level=AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA,
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ synapse_agent_session=self.updated_session_response
+ )
+
+ async def test_prompt_trace_enabled_print_response(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.AgentPrompt.send_job_and_wait_async",
+ new_callable=AsyncMock,
+ return_value=self.test_prompt_trace_enabled,
+ ) as mock_send_job_and_wait_async,
+ patch.object(
+ self.syn.logger,
+ "info",
+ ) as mock_logger_info,
+ ):
+ # GIVEN an existing AgentSession
+ # WHEN I call prompt with trace enabled and print_response enabled
+ await self.test_session.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ newer_than=0,
+ synapse_client=self.syn,
+ )
+ # THEN the result should be an AgentPrompt with the correct
+ # values appended to the chat history
+ assert self.test_prompt_trace_enabled in self.test_session.chat_history
+ # AND send_job_and_wait_async should have
+ # been called once with the correct arguments
+ mock_send_job_and_wait_async.assert_called_once_with(
+ synapse_client=self.syn, post_exchange_args={"newer_than": 0}
+ )
+ # AND the trace should be printed
+ mock_logger_info.assert_called_with(
+ f"TRACE:\n{self.test_prompt_trace_enabled.trace}"
+ )
+
+ async def test_prompt_trace_disabled_no_print(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.AgentPrompt.send_job_and_wait_async",
+ new_callable=AsyncMock,
+ return_value=self.test_prompt_trace_disabled,
+ ) as mock_send_job_and_wait_async,
+ patch.object(
+ self.syn.logger,
+ "info",
+ ) as mock_logger_info,
+ ):
+ # WHEN I call prompt with trace disabled and print_response disabled
+ await self.test_session.prompt_async(
+ prompt="Hello",
+ enable_trace=False,
+ print_response=False,
+ newer_than=0,
+ synapse_client=self.syn,
+ )
+ # THEN the result should be an AgentPrompt with the
+ # correct values appended to the chat history
+ assert self.test_prompt_trace_disabled in self.test_session.chat_history
+ # AND send_job_and_wait_async should have been
+ # called once with the correct arguments
+ mock_send_job_and_wait_async.assert_called_once_with(
+ synapse_client=self.syn, post_exchange_args={"newer_than": 0}
+ )
+ # AND print should not have been called
+ mock_logger_info.assert_not_called()
+
+
+class TestAgent:
+ """Unit tests for the Agent class' synchronous methods."""
+
+ def get_example_agent(self) -> Agent:
+ return Agent(
+ cloud_agent_id="123",
+ cloud_alias_id="456",
+ registration_id=0,
+ type=AgentType.BASELINE,
+ registered_on="2024-01-01T00:00:00Z",
+ sessions={},
+ current_session=None,
+ )
+
+ test_agent = Agent(
+ cloud_agent_id="123",
+ cloud_alias_id="456",
+ registration_id=0,
+ type=AgentType.BASELINE,
+ registered_on="2024-01-01T00:00:00Z",
+ sessions={},
+ current_session=None,
+ )
+
+ agent_response = {
+ "awsAgentId": test_agent.cloud_agent_id,
+ "awsAliasId": test_agent.cloud_alias_id,
+ "agentRegistrationId": test_agent.registration_id,
+ "registeredOn": test_agent.registered_on,
+ "type": test_agent.type,
+ }
+
+ test_session = AgentSession(
+ id="123",
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ started_on="2024-01-01T00:00:00Z",
+ started_by="123456789",
+ modified_on="2024-01-01T00:00:00Z",
+ agent_registration_id="0",
+ etag="11111111-1111-1111-1111-111111111111",
+ )
+
+ test_prompt = AgentPrompt(
+ concrete_type=AGENT_CHAT_REQUEST,
+ session_id="456",
+ prompt="Hello",
+ enable_trace=True,
+ response="World",
+ trace="Trace",
+ )
+
+ @pytest.fixture(autouse=True, scope="function")
+ def init_syn(self, syn: Synapse) -> None:
+ self.syn = syn
+
+ async def test_fill_from_dict(self) -> None:
+ # GIVEN an empty Agent
+ empty_agent = Agent()
+ # WHEN I call fill_from_dict on an empty Agent with a synapse_response
+ result_agent = empty_agent.fill_from_dict(
+ agent_registration=self.agent_response
+ )
+ # THEN the result should be an Agent with the correct values
+ assert result_agent == self.test_agent
+
+ async def test_register_async(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.register_agent",
+ new_callable=AsyncMock,
+ return_value=self.agent_response,
+ ) as mock_register_agent,
+ patch.object(
+ Agent,
+ "fill_from_dict",
+ return_value=self.test_agent,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an Agent with a cloud_agent_id
+ initial_agent = Agent(
+ cloud_agent_id="123",
+ cloud_alias_id="456",
+ )
+ # WHEN I call register
+ result_agent = await initial_agent.register_async(synapse_client=self.syn)
+ # THEN the result should be an Agent with the correct values
+ assert result_agent == self.test_agent
+ # AND register_agent should have been called once with the correct arguments
+ mock_register_agent.assert_called_once_with(
+ cloud_agent_id="123",
+ cloud_alias_id="456",
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ agent_registration=self.agent_response
+ )
+
+ async def test_get_async(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.get_agent",
+ new_callable=AsyncMock,
+ return_value=self.agent_response,
+ ) as mock_get_agent,
+ patch.object(
+ Agent,
+ "fill_from_dict",
+ return_value=self.test_agent,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an Agent with a registration_id
+ initial_agent = Agent(
+ registration_id=0,
+ )
+ # WHEN I call get
+ result_agent = await initial_agent.get_async(synapse_client=self.syn)
+ # THEN the result should be an Agent with the correct values
+ assert result_agent == self.test_agent
+ # AND get_agent should have been called once with the correct arguments
+ mock_get_agent.assert_called_once_with(
+ registration_id=0,
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ agent_registration=self.agent_response
+ )
+
+ async def test_start_session_async(self) -> None:
+ with (
+ patch.object(
+ AgentSession,
+ "start_async",
+ new_callable=AsyncMock,
+ return_value=self.test_session,
+ ) as mock_start_session,
+ ):
+ # GIVEN an existing Agent
+ my_agent = self.get_example_agent()
+ # WHEN I call start_session
+ result_session = await my_agent.start_session_async(
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ synapse_client=self.syn,
+ )
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+ # AND start_session should have been called once with the correct arguments
+ mock_start_session.assert_called_once_with(
+ synapse_client=self.syn,
+ )
+ # AND the current_session should be set to the new session
+ assert my_agent.current_session == self.test_session
+ # AND the sessions dictionary should have the new session
+ assert my_agent.sessions[self.test_session.id] == self.test_session
+
+ async def test_get_session_async(self) -> None:
+ with (
+ patch.object(
+ AgentSession,
+ "get_async",
+ new_callable=AsyncMock,
+ return_value=self.test_session,
+ ) as mock_get_session,
+ ):
+ # GIVEN an existing AgentSession
+ my_agent = self.get_example_agent()
+ # WHEN I call get_session
+ result_session = await my_agent.get_session_async(
+ session_id="123", synapse_client=self.syn
+ )
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+ # AND get_session should have been called once with the correct arguments
+ mock_get_session.assert_called_once_with(
+ synapse_client=self.syn,
+ )
+ # AND the current_session should be set to the session
+ assert my_agent.current_session == self.test_session
+ # AND the sessions dictionary should have the session
+ assert my_agent.sessions[self.test_session.id] == self.test_session
+
+ async def test_prompt_session_selected(self) -> None:
+ with (
+ patch.object(
+ AgentSession,
+ "get_async",
+ new_callable=AsyncMock,
+ return_value=self.test_session,
+ ) as mock_get_async,
+ patch.object(
+ Agent,
+ "start_session_async",
+ new_callable=AsyncMock,
+ ) as mock_start_session,
+ patch.object(
+ AgentSession,
+ "prompt_async",
+ ) as mock_prompt_async,
+ ):
+ # GIVEN an existing AgentSession
+ my_agent = self.get_example_agent()
+ # WHEN I call prompt with a session selected
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ session=self.test_session,
+ newer_than=0,
+ synapse_client=self.syn,
+ )
+ # AND get_session_async should have been called once with the correct arguments
+ mock_get_async.assert_called_once_with(
+ synapse_client=self.syn,
+ )
+ # AND start_session_async should not have been called
+ mock_start_session.assert_not_called()
+ # AND prompt_async should have been called once with the correct arguments
+ mock_prompt_async.assert_called_once_with(
+ prompt="Hello",
+ enable_trace=True,
+ newer_than=0,
+ print_response=True,
+ synapse_client=self.syn,
+ )
+
+ async def test_prompt_session_none_current_session_none(self) -> None:
+ with (
+ patch.object(
+ Agent,
+ "get_session_async",
+ new_callable=AsyncMock,
+ ) as mock_get_session,
+ patch.object(
+ AgentSession,
+ "start_async",
+ new_callable=AsyncMock,
+ return_value=self.test_session,
+ ) as mock_start_async,
+ patch.object(
+ AgentSession,
+ "prompt_async",
+ ) as mock_prompt_async,
+ ):
+ # GIVEN an existing Agent with no current session
+ my_agent = self.get_example_agent()
+ # WHEN I call prompt with no session selected and no current session set
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ newer_than=0,
+ synapse_client=self.syn,
+ )
+ # THEN get_session_async should not have been called
+ mock_get_session.assert_not_called()
+ # AND start_session_async should have been called once with the correct arguments
+ mock_start_async.assert_called_once_with(
+ synapse_client=self.syn,
+ )
+ # AND prompt_async should have been called once with the correct arguments
+ mock_prompt_async.assert_called_once_with(
+ prompt="Hello",
+ enable_trace=True,
+ newer_than=0,
+ print_response=True,
+ synapse_client=self.syn,
+ )
+
+ async def test_prompt_session_none_current_session_present(self) -> None:
+ with (
+ patch.object(
+ Agent,
+ "get_session_async",
+ new_callable=AsyncMock,
+ ) as mock_get_session,
+ patch.object(
+ AgentSession,
+ "start_async",
+ new_callable=AsyncMock,
+ ) as mock_start_async,
+ patch.object(
+ AgentSession,
+ "prompt_async",
+ ) as mock_prompt_async,
+ ):
+ # GIVEN an existing Agent with a current session
+ my_agent = self.get_example_agent()
+ my_agent.current_session = self.test_session
+ # WHEN I call prompt with no session selected and a current session set
+ await my_agent.prompt_async(
+ prompt="Hello",
+ enable_trace=True,
+ newer_than=0,
+ print_response=True,
+ synapse_client=self.syn,
+ )
+ # THEN get_session_async and start_session_async should not have been called
+ mock_get_session.assert_not_called()
+ mock_start_async.assert_not_called()
+ # AND prompt_async should have been called once with the correct arguments
+ mock_prompt_async.assert_called_once_with(
+ prompt="Hello",
+ enable_trace=True,
+ newer_than=0,
+ print_response=True,
+ synapse_client=self.syn,
+ )
+
+ async def test_get_chat_history_when_current_session_none(self) -> None:
+ # GIVEN an existing Agent with no current session
+ my_agent = self.get_example_agent()
+ # WHEN I call get_chat_history
+ result_chat_history = my_agent.get_chat_history()
+ # THEN the result should be None
+ assert result_chat_history is None
+
+ async def test_get_chat_history_when_current_session_and_chat_history_present(
+ self,
+ ) -> None:
+ # GIVEN an existing Agent with a current session and chat history
+ my_agent = self.get_example_agent()
+ my_agent.current_session = self.test_session
+ my_agent.current_session.chat_history = [self.test_prompt]
+ # WHEN I call get_chat_history
+ result_chat_history = my_agent.get_chat_history()
+ # THEN the result should be the chat history
+ assert self.test_prompt in result_chat_history
diff --git a/tests/unit/synapseclient/models/synchronous/unit_test_agent.py b/tests/unit/synapseclient/models/synchronous/unit_test_agent.py
new file mode 100644
index 000000000..83f33cb7b
--- /dev/null
+++ b/tests/unit/synapseclient/models/synchronous/unit_test_agent.py
@@ -0,0 +1,588 @@
+"""Unit tests for Synchronous methods in Agent, AgentSession, and AgentPrompt classes."""
+
+from unittest.mock import patch
+
+import pytest
+
+from synapseclient import Synapse
+from synapseclient.core.constants.concrete_types import AGENT_CHAT_REQUEST
+from synapseclient.models.agent import (
+ Agent,
+ AgentPrompt,
+ AgentSession,
+ AgentSessionAccessLevel,
+ AgentType,
+)
+
+
+class TestAgentPrompt:
+ """Unit tests for the AgentPrompt class' synchronous methods."""
+
+ test_prompt = AgentPrompt(
+ concrete_type=AGENT_CHAT_REQUEST,
+ session_id="456",
+ prompt="Hello",
+ enable_trace=True,
+ )
+ prompt_request = {
+ "concreteType": test_prompt.concrete_type,
+ "sessionId": test_prompt.session_id,
+ "chatText": test_prompt.prompt,
+ "enableTrace": test_prompt.enable_trace,
+ }
+ prompt_response = {
+ "jobId": "123",
+ "sessionId": "456",
+ "responseText": "World",
+ }
+
+ def test_to_synapse_request(self) -> None:
+ # GIVEN an existing AgentPrompt
+ # WHEN I call to_synapse_request
+ result_request = self.test_prompt.to_synapse_request()
+ # THEN the result should be a dictionary with the correct keys and values
+ assert result_request == self.prompt_request
+
+ def test_fill_from_dict(self) -> None:
+ # GIVEN an existing AgentPrompt
+ # WHEN I call fill_from_dict
+ result_prompt = self.test_prompt.fill_from_dict(self.prompt_response)
+ # THEN the result should be an AgentPrompt with the correct values
+ assert result_prompt == self.test_prompt
+
+
+class TestAgentSession:
+ """Unit tests for the AgentSession class' synchronous methods."""
+
+ test_session = AgentSession(
+ id="123",
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ started_on="2024-01-01T00:00:00Z",
+ started_by="123456789",
+ modified_on="2024-01-01T00:00:00Z",
+ agent_registration_id="0",
+ etag="11111111-1111-1111-1111-111111111111",
+ )
+
+ session_response = {
+ "sessionId": test_session.id,
+ "agentAccessLevel": test_session.access_level,
+ "startedOn": test_session.started_on,
+ "startedBy": test_session.started_by,
+ "modifiedOn": test_session.modified_on,
+ "agentRegistrationId": test_session.agent_registration_id,
+ "etag": test_session.etag,
+ }
+
+ updated_test_session = AgentSession(
+ id=test_session.id,
+ access_level=AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA,
+ started_on=test_session.started_on,
+ started_by=test_session.started_by,
+ modified_on=test_session.modified_on,
+ agent_registration_id=test_session.agent_registration_id,
+ etag=test_session.etag,
+ )
+
+ updated_session_response = {
+ "sessionId": updated_test_session.id,
+ "agentAccessLevel": updated_test_session.access_level,
+ "startedOn": updated_test_session.started_on,
+ "startedBy": updated_test_session.started_by,
+ "modifiedOn": updated_test_session.modified_on,
+ "agentRegistrationId": updated_test_session.agent_registration_id,
+ "etag": updated_test_session.etag,
+ }
+
+ test_prompt_trace_enabled = AgentPrompt(
+ concrete_type=AGENT_CHAT_REQUEST,
+ session_id="456",
+ prompt="Hello",
+ enable_trace=True,
+ response="World",
+ trace="Trace",
+ )
+
+ test_prompt_trace_disabled = AgentPrompt(
+ concrete_type=AGENT_CHAT_REQUEST,
+ session_id="456",
+ prompt="Hello",
+ enable_trace=False,
+ response="World",
+ trace=None,
+ )
+
+ @pytest.fixture(autouse=True, scope="function")
+ def init_syn(self, syn: Synapse) -> None:
+ self.syn = syn
+
+ def test_fill_from_dict(self) -> None:
+ # WHEN I call fill_from_dict on an empty AgentSession with a synapse_response
+ result_session = AgentSession().fill_from_dict(self.session_response)
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+
+ def test_start(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.start_session",
+ return_value=self.session_response,
+ ) as mock_start_session,
+ patch.object(
+ AgentSession,
+ "fill_from_dict",
+ return_value=self.test_session,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an AgentSession with access_level and agent_registration_id
+ initial_session = AgentSession(
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ agent_registration_id=0,
+ )
+ # WHEN I call start
+ result_session = initial_session.start(synapse_client=self.syn)
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+ # AND start_session should have been called once with the correct arguments
+ mock_start_session.assert_called_once_with(
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ agent_registration_id=0,
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ synapse_agent_session=self.session_response
+ )
+
+ def test_get(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.get_session",
+ return_value=self.session_response,
+ ) as mock_get_session,
+ patch.object(
+ AgentSession,
+ "fill_from_dict",
+ return_value=self.test_session,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an AgentSession with an agent_registration_id
+ initial_session = AgentSession(
+ agent_registration_id=0,
+ )
+ # WHEN I call get
+ result_session = initial_session.get(synapse_client=self.syn)
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+ # AND get_session should have been called once with the correct arguments
+ mock_get_session.assert_called_once_with(
+ id=initial_session.id,
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ synapse_agent_session=self.session_response
+ )
+
+ def test_update(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.update_session",
+ return_value=self.updated_session_response,
+ ) as mock_update_session,
+ patch.object(
+ AgentSession,
+ "fill_from_dict",
+ return_value=self.updated_test_session,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an AgentSession with an updated access_level
+ # WHEN I call update
+ result_session = self.updated_test_session.update(synapse_client=self.syn)
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.updated_test_session
+ # AND update_session should have been called once with the correct arguments
+ mock_update_session.assert_called_once_with(
+ id=self.updated_test_session.id,
+ access_level=AgentSessionAccessLevel.READ_YOUR_PRIVATE_DATA,
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ synapse_agent_session=self.updated_session_response
+ )
+
+ def test_prompt_trace_enabled_print_response(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.AgentPrompt.send_job_and_wait_async",
+ return_value=self.test_prompt_trace_enabled,
+ ) as mock_send_job_and_wait_async,
+ patch.object(
+ self.syn.logger,
+ "info",
+ ) as mock_logger_info,
+ ):
+ # GIVEN an existing AgentSession
+ # WHEN I call prompt with trace enabled and print_response enabled
+ self.test_session.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ newer_than=0,
+ synapse_client=self.syn,
+ )
+ # THEN the result should be an AgentPrompt with the correct values appended to the chat history
+ assert self.test_prompt_trace_enabled in self.test_session.chat_history
+ # AND send_job_and_wait_async should have been called once with the correct arguments
+ mock_send_job_and_wait_async.assert_called_once_with(
+ synapse_client=self.syn, post_exchange_args={"newer_than": 0}
+ )
+ # AND the trace should be printed
+ mock_logger_info.assert_called_with(
+ f"TRACE:\n{self.test_prompt_trace_enabled.trace}"
+ )
+
+ def test_prompt_trace_disabled_no_print(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.AgentPrompt.send_job_and_wait_async",
+ return_value=self.test_prompt_trace_disabled,
+ ) as mock_send_job_and_wait_async,
+ patch.object(
+ self.syn.logger,
+ "info",
+ ) as mock_logger_info,
+ ):
+ # WHEN I call prompt with trace disabled and print_response disabled
+ self.test_session.prompt(
+ prompt="Hello",
+ enable_trace=False,
+ print_response=False,
+ newer_than=0,
+ synapse_client=self.syn,
+ )
+ # THEN the result should be an AgentPrompt with the correct values appended to the chat history
+ assert self.test_prompt_trace_disabled in self.test_session.chat_history
+ # AND send_job_and_wait_async should have been called once with the correct arguments
+ mock_send_job_and_wait_async.assert_called_once_with(
+ synapse_client=self.syn, post_exchange_args={"newer_than": 0}
+ )
+ # AND print should not have been called
+ mock_logger_info.assert_not_called()
+
+
+class TestAgent:
+ """Unit tests for the Agent class' synchronous methods."""
+
+ def get_example_agent(self) -> Agent:
+ return Agent(
+ cloud_agent_id="123",
+ cloud_alias_id="456",
+ registration_id=0,
+ type=AgentType.BASELINE,
+ registered_on="2024-01-01T00:00:00Z",
+ sessions={},
+ current_session=None,
+ )
+
+ test_agent = Agent(
+ cloud_agent_id="123",
+ cloud_alias_id="456",
+ registration_id=0,
+ type=AgentType.BASELINE,
+ registered_on="2024-01-01T00:00:00Z",
+ sessions={},
+ current_session=None,
+ )
+
+ agent_response = {
+ "awsAgentId": test_agent.cloud_agent_id,
+ "awsAliasId": test_agent.cloud_alias_id,
+ "agentRegistrationId": test_agent.registration_id,
+ "registeredOn": test_agent.registered_on,
+ "type": test_agent.type,
+ }
+
+ test_session = AgentSession(
+ id="123",
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ started_on="2024-01-01T00:00:00Z",
+ started_by="123456789",
+ modified_on="2024-01-01T00:00:00Z",
+ agent_registration_id="0",
+ etag="11111111-1111-1111-1111-111111111111",
+ )
+
+ test_prompt = AgentPrompt(
+ concrete_type=AGENT_CHAT_REQUEST,
+ session_id="456",
+ prompt="Hello",
+ enable_trace=True,
+ response="World",
+ trace="Trace",
+ )
+
+ @pytest.fixture(autouse=True, scope="function")
+ def init_syn(self, syn: Synapse) -> None:
+ self.syn = syn
+
+ def test_fill_from_dict(self) -> None:
+ # GIVEN an empty Agent
+ empty_agent = Agent()
+ # WHEN I call fill_from_dict on an empty Agent with a synapse_response
+ result_agent = empty_agent.fill_from_dict(
+ agent_registration=self.agent_response
+ )
+ # THEN the result should be an Agent with the correct values
+ assert result_agent == self.test_agent
+
+ def test_register(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.register_agent",
+ return_value=self.agent_response,
+ ) as mock_register_agent,
+ patch.object(
+ Agent,
+ "fill_from_dict",
+ return_value=self.test_agent,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an Agent with a cloud_agent_id
+ initial_agent = Agent(
+ cloud_agent_id="123",
+ cloud_alias_id="456",
+ )
+ # WHEN I call register
+ result_agent = initial_agent.register(synapse_client=self.syn)
+ # THEN the result should be an Agent with the correct values
+ assert result_agent == self.test_agent
+ # AND register_agent should have been called once with the correct arguments
+ mock_register_agent.assert_called_once_with(
+ cloud_agent_id="123",
+ cloud_alias_id="456",
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ agent_registration=self.agent_response
+ )
+
+ def test_get(self) -> None:
+ with (
+ patch(
+ "synapseclient.models.agent.get_agent",
+ return_value=self.agent_response,
+ ) as mock_get_agent,
+ patch.object(
+ Agent,
+ "fill_from_dict",
+ return_value=self.test_agent,
+ ) as mock_fill_from_dict,
+ ):
+ # GIVEN an Agent with a registration_id
+ initial_agent = Agent(
+ registration_id=0,
+ )
+ # WHEN I call get
+ result_agent = initial_agent.get(synapse_client=self.syn)
+ # THEN the result should be an Agent with the correct values
+ assert result_agent == self.test_agent
+ # AND get_agent should have been called once with the correct arguments
+ mock_get_agent.assert_called_once_with(
+ registration_id=0,
+ synapse_client=self.syn,
+ )
+ # AND fill_from_dict should have been called once with the correct arguments
+ mock_fill_from_dict.assert_called_once_with(
+ agent_registration=self.agent_response
+ )
+
+ def test_start_session(self) -> None:
+ with patch.object(
+ AgentSession,
+ "start_async",
+ return_value=self.test_session,
+ ) as mock_start_session:
+ # GIVEN an existing Agent
+ my_agent = self.get_example_agent()
+ # WHEN I call start_session
+ result_session = my_agent.start_session(
+ access_level=AgentSessionAccessLevel.PUBLICLY_ACCESSIBLE,
+ synapse_client=self.syn,
+ )
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+ # AND start_session should have been called once with the correct arguments
+ mock_start_session.assert_called_once_with(
+ synapse_client=self.syn,
+ )
+ # AND the current_session should be set to the new session
+ assert my_agent.current_session == self.test_session
+ # AND the sessions dictionary should have the new session
+ assert my_agent.sessions[self.test_session.id] == self.test_session
+
+ def test_get_session(self) -> None:
+ with patch.object(
+ AgentSession,
+ "get_async",
+ return_value=self.test_session,
+ ) as mock_get_session:
+ # GIVEN an existing AgentSession
+ my_agent = self.get_example_agent()
+ # WHEN I call get_session
+ result_session = my_agent.get_session(
+ session_id="123", synapse_client=self.syn
+ )
+ # THEN the result should be an AgentSession with the correct values
+ assert result_session == self.test_session
+ # AND get_session should have been called once with the correct arguments
+ mock_get_session.assert_called_once_with(
+ synapse_client=self.syn,
+ )
+ # AND the current_session should be set to the session
+ assert my_agent.current_session == self.test_session
+ # AND the sessions dictionary should have the session
+ assert my_agent.sessions[self.test_session.id] == self.test_session
+
+ def test_prompt_session_selected(self) -> None:
+ with (
+ patch.object(
+ AgentSession,
+ "get_async",
+ return_value=self.test_session,
+ ) as mock_get_async,
+ patch.object(
+ Agent,
+ "start_session_async",
+ ) as mock_start_session,
+ patch.object(
+ AgentSession,
+ "prompt_async",
+ ) as mock_prompt_async,
+ ):
+ # GIVEN an existing AgentSession
+ my_agent = self.get_example_agent()
+ # WHEN I call prompt with a session selected
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ session=self.test_session,
+ newer_than=0,
+ synapse_client=self.syn,
+ )
+ # AND get_session_async should have been called once with the correct arguments
+ mock_get_async.assert_called_once_with(
+ synapse_client=self.syn,
+ )
+ # AND start_session_async should not have been called
+ mock_start_session.assert_not_called()
+ # AND prompt_async should have been called once with the correct arguments
+ mock_prompt_async.assert_called_once_with(
+ prompt="Hello",
+ enable_trace=True,
+ newer_than=0,
+ print_response=True,
+ synapse_client=self.syn,
+ )
+
+ def test_prompt_session_none_current_session_none(self) -> None:
+ with (
+ patch.object(
+ Agent,
+ "get_session_async",
+ ) as mock_get_session,
+ patch.object(
+ AgentSession,
+ "start_async",
+ return_value=self.test_session,
+ ) as mock_start_async,
+ patch.object(
+ AgentSession,
+ "prompt_async",
+ ) as mock_prompt_async,
+ ):
+ # GIVEN an existing Agent with no current session
+ my_agent = self.get_example_agent()
+ # WHEN I call prompt with no session selected and no current session set
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ print_response=True,
+ newer_than=0,
+ synapse_client=self.syn,
+ )
+ # THEN get_session_async should not have been called
+ mock_get_session.assert_not_called()
+ # AND start_session_async should have been called once with the correct arguments
+ mock_start_async.assert_called_once_with(
+ synapse_client=self.syn,
+ )
+ # AND prompt_async should have been called once with the correct arguments
+ mock_prompt_async.assert_called_once_with(
+ prompt="Hello",
+ enable_trace=True,
+ newer_than=0,
+ print_response=True,
+ synapse_client=self.syn,
+ )
+
+ def test_prompt_session_none_current_session_present(self) -> None:
+ with (
+ patch.object(
+ Agent,
+ "get_session_async",
+ ) as mock_get_session,
+ patch.object(
+ AgentSession,
+ "start_async",
+ ) as mock_start_async,
+ patch.object(
+ AgentSession,
+ "prompt_async",
+ ) as mock_prompt_async,
+ ):
+ # GIVEN an existing Agent with a current session
+ my_agent = self.get_example_agent()
+ my_agent.current_session = self.test_session
+ # WHEN I call prompt with no session selected and a current session set
+ my_agent.prompt(
+ prompt="Hello",
+ enable_trace=True,
+ newer_than=0,
+ print_response=True,
+ synapse_client=self.syn,
+ )
+ # THEN get_session_async and start_session_async should not have been called
+ mock_get_session.assert_not_called()
+ mock_start_async.assert_not_called()
+ # AND prompt_async should have been called once with the correct arguments
+ mock_prompt_async.assert_called_once_with(
+ prompt="Hello",
+ enable_trace=True,
+ newer_than=0,
+ print_response=True,
+ synapse_client=self.syn,
+ )
+
+ def test_get_chat_history_when_current_session_none(self) -> None:
+ # GIVEN an existing Agent with no current session
+ my_agent = self.get_example_agent()
+ # WHEN I call get_chat_history
+ result_chat_history = my_agent.get_chat_history()
+ # THEN the result should be None
+ assert result_chat_history is None
+
+ def test_get_chat_history_when_current_session_and_chat_history_present(
+ self,
+ ) -> None:
+ # GIVEN an existing Agent with a current session and chat history
+ my_agent = self.get_example_agent()
+ my_agent.current_session = self.test_session
+ my_agent.current_session.chat_history = [self.test_prompt]
+ # WHEN I call get_chat_history
+ result_chat_history = my_agent.get_chat_history()
+ # THEN the result should be the chat history
+ assert self.test_prompt in result_chat_history