From aaa9c506d83b98c6dda7e31685e64d950fe6d609 Mon Sep 17 00:00:00 2001 From: Brad Macdonald <52762200+BWMac@users.noreply.github.com> Date: Tue, 28 Jan 2025 13:23:37 -0500 Subject: [PATCH 1/8] [SYNPY-1544] Synapse Agent OOP Model (#1152) * Adds async convenience functions * expose convenience functions * updates convenience functions * updates agent_services * removes rest_get_async exception handling * pre-commit fixes * delete accidentally committed script * adds initial agent implementation * clean up agent * adds missing docstrings * pre-commit * updates agent_services * updates agent.py * Updates alias ID handling * adds syncronous interface * prevent cicular import in storable_entity_components * remove promt sending and receiving from agent_service * adds initial (dirty) async job mixin * pre-commit run * [SYNPY-1544] potential changes to mixin (#1153) * Changes for async mixin * Remove arg * bug fix * generalizes send_job_and_wait_async * removes typing.Self --------- Co-authored-by: bwmac * cleans up agent logic * adds async job unit tests * updates async job tests * adds agent unit tests * adds integration tests * pre-commit * adds examples to agent.py * removes todos * adds POC script * add to mixins * adds agent docs * updates agent docs * reorganize documentation * updates poc script * clean up * add docstring * removes unused imports * split too long lines * force synapse_client kwarg * updates agent.py * updates synapse_client docstring description * updates asynchronous_job * updates integration tests * pre-commit * agent inherited members * updates docs for inherited members * missing inherited members * updates doc formatting * try team formatting change * updates script description * adds Annotation lazy import * try team formatting change * more formatting changes * address review comments in agent.py * move synchronous docs up a layer * adds syn login * adds warning message to docs * updates docstring examples * updates docstrings * adds error handling for agent.get * async integration tests * fix conditional * disables integration tests * updates docstring for clarity --------- Co-authored-by: BryanFauble <17128019+BryanFauble@users.noreply.github.com> --- docs/reference/experimental/async/activity.md | 24 + docs/reference/experimental/async/agent.md | 32 + docs/reference/experimental/async/file.md | 27 + docs/reference/experimental/async/folder.md | 20 + docs/reference/experimental/async/project.md | 19 + docs/reference/experimental/async/table.md | 21 + docs/reference/experimental/async/team.md | 19 + .../experimental/async/user_profile.md | 19 + .../mixins/access_controllable.md | 3 + .../mixins/asynchronous_communicator.md | 3 + .../experimental/mixins/failure_strategy.md | 3 + .../experimental/mixins/storable_container.md | 3 + docs/reference/experimental/sync/activity.md | 35 + docs/reference/experimental/sync/agent.md | 42 + docs/reference/experimental/sync/file.md | 37 + docs/reference/experimental/sync/folder.md | 30 + docs/reference/experimental/sync/project.md | 29 + docs/reference/experimental/sync/table.md | 31 + docs/reference/experimental/sync/team.md | 30 + .../experimental/sync/user_profile.md | 19 + docs/reference/oop/models.md | 169 ---- docs/reference/oop/models_async.md | 100 -- .../oop_poc_agent.py | 105 ++ mkdocs.yml | 24 +- synapseclient/api/__init__.py | 15 + synapseclient/api/agent_services.py | 189 ++++ synapseclient/client.py | 25 +- .../core/constants/concrete_types.py | 3 + synapseclient/models/__init__.py | 10 + synapseclient/models/agent.py | 944 ++++++++++++++++++ synapseclient/models/mixins/__init__.py | 2 + .../models/mixins/asynchronous_job.py | 410 ++++++++ .../models/protocols/agent_protocol.py | 390 ++++++++ .../services/storable_entity_components.py | 3 +- .../models/async/test_agent_async.py | 228 +++++ .../models/synchronous/test_agent.py | 192 ++++ .../async/unit_test_asynchronous_job.py | 278 ++++++ .../models/async/unit_test_agent_async.py | 703 +++++++++++++ .../models/synchronous/unit_test_agent.py | 588 +++++++++++ 39 files changed, 4538 insertions(+), 286 deletions(-) create mode 100644 docs/reference/experimental/async/activity.md create mode 100644 docs/reference/experimental/async/agent.md create mode 100644 docs/reference/experimental/async/file.md create mode 100644 docs/reference/experimental/async/folder.md create mode 100644 docs/reference/experimental/async/project.md create mode 100644 docs/reference/experimental/async/table.md create mode 100644 docs/reference/experimental/async/team.md create mode 100644 docs/reference/experimental/async/user_profile.md create mode 100644 docs/reference/experimental/mixins/access_controllable.md create mode 100644 docs/reference/experimental/mixins/asynchronous_communicator.md create mode 100644 docs/reference/experimental/mixins/failure_strategy.md create mode 100644 docs/reference/experimental/mixins/storable_container.md create mode 100644 docs/reference/experimental/sync/activity.md create mode 100644 docs/reference/experimental/sync/agent.md create mode 100644 docs/reference/experimental/sync/file.md create mode 100644 docs/reference/experimental/sync/folder.md create mode 100644 docs/reference/experimental/sync/project.md create mode 100644 docs/reference/experimental/sync/table.md create mode 100644 docs/reference/experimental/sync/team.md create mode 100644 docs/reference/experimental/sync/user_profile.md delete mode 100644 docs/reference/oop/models.md delete mode 100644 docs/reference/oop/models_async.md create mode 100644 docs/scripts/object_orientated_programming_poc/oop_poc_agent.py create mode 100644 synapseclient/api/agent_services.py create mode 100644 synapseclient/models/agent.py create mode 100644 synapseclient/models/mixins/asynchronous_job.py create mode 100644 synapseclient/models/protocols/agent_protocol.py create mode 100644 tests/integration/synapseclient/models/async/test_agent_async.py create mode 100644 tests/integration/synapseclient/models/synchronous/test_agent.py create mode 100644 tests/unit/synapseclient/mixins/async/unit_test_asynchronous_job.py create mode 100644 tests/unit/synapseclient/models/async/unit_test_agent_async.py create mode 100644 tests/unit/synapseclient/models/synchronous/unit_test_agent.py 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 From d085b02bb6aafbd4a872923650ef976b013b6873 Mon Sep 17 00:00:00 2001 From: Brad Macdonald <52762200+BWMac@users.noreply.github.com> Date: Tue, 28 Jan 2025 15:33:54 -0500 Subject: [PATCH 2/8] [SYNPY-1544] Fixes docstring (#1155) * fixes docstring * protocol docstring * fix imports --- synapseclient/models/agent.py | 5 ++++- synapseclient/models/protocols/agent_protocol.py | 4 ++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/synapseclient/models/agent.py b/synapseclient/models/agent.py index 344373f8a..18d737e0d 100644 --- a/synapseclient/models/agent.py +++ b/synapseclient/models/agent.py @@ -647,7 +647,7 @@ async def get_async(self, *, synapse_client: Optional[Synapse] = None) -> "Agent import asyncio from synapseclient import Synapse - from synapseclient.models.agent import Agent + from synapseclient.models import Agent, AgentSessionAccessLevel syn = Synapse() syn.login() @@ -838,6 +838,9 @@ async def prompt_async( async def main(): my_agent = Agent() + await my_agent.start_session_async( + access_level=AgentSessionAccessLevel.WRITE_YOUR_PRIVATE_DATA + ) await my_agent.prompt_async( prompt="Add the annotation 'test' to the file 'syn123456789'", enable_trace=True, diff --git a/synapseclient/models/protocols/agent_protocol.py b/synapseclient/models/protocols/agent_protocol.py index ef52bf3a3..19df5e0ab 100644 --- a/synapseclient/models/protocols/agent_protocol.py +++ b/synapseclient/models/protocols/agent_protocol.py @@ -332,11 +332,15 @@ def prompt( The baseline Synpase Agent can be used to add annotations to files. from synapseclient import Synapse + from synapseclient.models import Agent, AgentSessionAccessLevel syn = Synapse() syn.login() my_agent = Agent() + my_agent.start_session( + access_level=AgentSessionAccessLevel.WRITE_YOUR_PRIVATE_DATA + ) my_agent.prompt( prompt="Add the annotation 'test' to the file 'syn123456789'", enable_trace=True, From cd8820a378d8a2add7edb0ed552b2f883cb926aa Mon Sep 17 00:00:00 2001 From: Brad Macdonald <52762200+BWMac@users.noreply.github.com> Date: Wed, 29 Jan 2025 09:22:24 -0500 Subject: [PATCH 3/8] Removes example setting annotations with Agent class (#1156) * removes annotation example * pre-commit --- synapseclient/models/agent.py | 9 +++------ synapseclient/models/protocols/agent_protocol.py | 11 ++++------- 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/synapseclient/models/agent.py b/synapseclient/models/agent.py index 18d737e0d..943af6287 100644 --- a/synapseclient/models/agent.py +++ b/synapseclient/models/agent.py @@ -826,8 +826,8 @@ async def prompt_async( `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. + Example: Prompt the baseline Synapse Agent. + The baseline Synapse Agent is equivilent to the Agent available in the Synapse UI. import asyncio from synapseclient import Synapse @@ -838,11 +838,8 @@ async def prompt_async( async def main(): my_agent = Agent() - await my_agent.start_session_async( - access_level=AgentSessionAccessLevel.WRITE_YOUR_PRIVATE_DATA - ) await my_agent.prompt_async( - prompt="Add the annotation 'test' to the file 'syn123456789'", + prompt="Can you tell me about the AD Knowledge Portal dataset?", enable_trace=True, print_response=True, ) diff --git a/synapseclient/models/protocols/agent_protocol.py b/synapseclient/models/protocols/agent_protocol.py index 19df5e0ab..e6b05fac7 100644 --- a/synapseclient/models/protocols/agent_protocol.py +++ b/synapseclient/models/protocols/agent_protocol.py @@ -328,21 +328,18 @@ def prompt( `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. + Example: Prompt the baseline Synapse Agent. + The baseline Synapse Agent is equivilent to the Agent available in the Synapse UI. from synapseclient import Synapse - from synapseclient.models import Agent, AgentSessionAccessLevel + from synapseclient.models import Agent syn = Synapse() syn.login() my_agent = Agent() - my_agent.start_session( - access_level=AgentSessionAccessLevel.WRITE_YOUR_PRIVATE_DATA - ) my_agent.prompt( - prompt="Add the annotation 'test' to the file 'syn123456789'", + prompt="Can you tell me about the AD Knowledge Portal dataset?", enable_trace=True, print_response=True, ) From 65422cf717b03eac176a6122f0b14575a2c93bec Mon Sep 17 00:00:00 2001 From: Brad Macdonald <52762200+BWMac@users.noreply.github.com> Date: Thu, 30 Jan 2025 12:09:11 -0500 Subject: [PATCH 4/8] [SYNPY-1557] Sync a Linked Folder Bug (#1157) * fixes docstring * protocol docstring * fix imports * adds integration test for expected behavior * adds fix * merge weirdness * fix test docstring --- .../models/mixins/storable_container.py | 2 + .../synapseutils/test_synapseutils_sync.py | 91 ++++++++++++++++++- 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/synapseclient/models/mixins/storable_container.py b/synapseclient/models/mixins/storable_container.py index e1815aedb..667766263 100644 --- a/synapseclient/models/mixins/storable_container.py +++ b/synapseclient/models/mixins/storable_container.py @@ -686,6 +686,7 @@ async def _follow_link( or not (entity := entity_bundle.get("entity", None)) or not (links_to := entity.get("linksTo", None)) or not (link_class_name := entity.get("linksToClassName", None)) + or not (link_target_name := entity.get("name", None)) or not (link_target_id := links_to.get("targetId", None)) ): return @@ -693,6 +694,7 @@ async def _follow_link( pending_tasks = self._create_task_for_child( child={ "id": link_target_id, + "name": link_target_name, "type": link_class_name, }, recursive=recursive, diff --git a/tests/integration/synapseutils/test_synapseutils_sync.py b/tests/integration/synapseutils/test_synapseutils_sync.py index e985ba37a..635d69761 100644 --- a/tests/integration/synapseutils/test_synapseutils_sync.py +++ b/tests/integration/synapseutils/test_synapseutils_sync.py @@ -1992,7 +1992,7 @@ async def test_folder_sync_from_synapse_files_spread_across_folders( assert pd.isna(matching_row[ACTIVITY_DESCRIPTION_COLUMN].values[0]) assert found_matching_file - async def test_sync_from_synapse_follow_links( + async def test_sync_from_synapse_follow_links_files( self, syn: Synapse, schedule_for_cleanup: Callable[..., None], @@ -2082,6 +2082,95 @@ async def test_sync_from_synapse_follow_links( assert pd.isna(matching_row[ACTIVITY_NAME_COLUMN].values[0]) assert pd.isna(matching_row[ACTIVITY_DESCRIPTION_COLUMN].values[0]) + async def test_sync_from_synapse_follow_links_folder( + self, + syn: Synapse, + schedule_for_cleanup: Callable[..., None], + project_model: Project, + ) -> None: + """ + Testing for this case: + + project_model (root) + ├── folder_with_files + │ ├── file1 (uploaded) + │ └── file2 (uploaded) + └── folder_with_links - This is the folder we are syncing from + └── link_to_folder_with_files -> ../folder_with_files + """ + # GIVEN a folder + folder_with_files = await Folder( + name=str(uuid.uuid4()), parent_id=project_model.id + ).store_async() + schedule_for_cleanup(folder_with_files.id) + + # AND two files in the folder + temp_files = [utils.make_bogus_uuid_file() for _ in range(2)] + file_entities = [] + for file in temp_files: + schedule_for_cleanup(file) + file_entity = syn.store(SynapseFile(path=file, parent=folder_with_files.id)) + schedule_for_cleanup(file_entity["id"]) + file_entities.append(file_entity) + + # AND a second folder to sync from + folder_with_links = await Folder( + name=str(uuid.uuid4()), parent_id=project_model.id + ).store_async() + schedule_for_cleanup(folder_with_links.id) + + # AND a link to folder_with_files in folder_with_links + syn.store(obj=Link(targetId=folder_with_files.id, parent=folder_with_links.id)) + + # AND a temp directory to write the manifest file to + temp_dir = tempfile.mkdtemp() + + # WHEN I sync the parent folder from Synapse + sync_result = synapseutils.syncFromSynapse( + syn=syn, entity=folder_with_links.id, path=temp_dir, followLink=True + ) + + # THEN I expect that the result has all of the files + assert len(sync_result) == 2 + + # AND each of the files are the ones we uploaded + for file in sync_result: + assert file in file_entities + + # AND the manifest that is created matches the expected values + manifest_df = pd.read_csv(os.path.join(temp_dir, MANIFEST_FILE), sep="\t") + assert manifest_df.shape[0] == 2 + assert PATH_COLUMN in manifest_df.columns + assert PARENT_COLUMN in manifest_df.columns + assert USED_COLUMN in manifest_df.columns + assert EXECUTED_COLUMN in manifest_df.columns + assert ACTIVITY_NAME_COLUMN in manifest_df.columns + assert ACTIVITY_DESCRIPTION_COLUMN in manifest_df.columns + assert CONTENT_TYPE_COLUMN in manifest_df.columns + assert ID_COLUMN in manifest_df.columns + assert SYNAPSE_STORE_COLUMN in manifest_df.columns + assert NAME_COLUMN in manifest_df.columns + assert manifest_df.shape[1] == 10 + + for file in sync_result: + matching_row = manifest_df[manifest_df[PATH_COLUMN] == file[PATH_COLUMN]] + assert not matching_row.empty + assert matching_row[PARENT_COLUMN].values[0] == file[PARENT_ATTRIBUTE] + assert ( + matching_row[CONTENT_TYPE_COLUMN].values[0] == file[CONTENT_TYPE_COLUMN] + ) + assert matching_row[ID_COLUMN].values[0] == file[ID_COLUMN] + assert ( + matching_row[SYNAPSE_STORE_COLUMN].values[0] + == file[SYNAPSE_STORE_COLUMN] + ) + assert matching_row[NAME_COLUMN].values[0] == file[NAME_COLUMN] + + assert pd.isna(matching_row[USED_COLUMN].values[0]) + assert pd.isna(matching_row[EXECUTED_COLUMN].values[0]) + assert pd.isna(matching_row[ACTIVITY_NAME_COLUMN].values[0]) + assert pd.isna(matching_row[ACTIVITY_DESCRIPTION_COLUMN].values[0]) + async def test_sync_from_synapse_follow_links_sync_contains_all_folders( self, syn: Synapse, From b950c6d5be264ce8a915379df174e57f4f2ea00f Mon Sep 17 00:00:00 2001 From: BryanFauble <17128019+BryanFauble@users.noreply.github.com> Date: Thu, 30 Jan 2025 15:52:06 -0700 Subject: [PATCH 5/8] [SYNPY-1544] Return the AgentPrompt when calling the prompt function (#1158) --- synapseclient/models/agent.py | 7 ++++--- synapseclient/models/protocols/agent_protocol.py | 15 ++++++++++----- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/synapseclient/models/agent.py b/synapseclient/models/agent.py index 943af6287..3fe1306ac 100644 --- a/synapseclient/models/agent.py +++ b/synapseclient/models/agent.py @@ -384,7 +384,7 @@ async def prompt_async( newer_than: Optional[int] = None, *, synapse_client: Optional[Synapse] = None, - ) -> None: + ) -> AgentPrompt: """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. @@ -431,6 +431,7 @@ async def main(): client.logger.info(f"RESPONSE:\n{agent_prompt.response}\n") if enable_trace: client.logger.info(f"TRACE:\n{agent_prompt.trace}") + return agent_prompt @dataclass @@ -811,7 +812,7 @@ async def prompt_async( newer_than: Optional[int] = None, *, synapse_client: Optional[Synapse] = None, - ) -> None: + ) -> AgentPrompt: """Sends a prompt to the agent for the current session. If no session is currently active, a new session will be started. @@ -908,7 +909,7 @@ async def main(): if not self.current_session: await self.start_session_async(synapse_client=synapse_client) - await self.current_session.prompt_async( + return await self.current_session.prompt_async( prompt=prompt, enable_trace=enable_trace, newer_than=newer_than, diff --git a/synapseclient/models/protocols/agent_protocol.py b/synapseclient/models/protocols/agent_protocol.py index e6b05fac7..bc729e5f9 100644 --- a/synapseclient/models/protocols/agent_protocol.py +++ b/synapseclient/models/protocols/agent_protocol.py @@ -6,7 +6,12 @@ from synapseclient import Synapse if TYPE_CHECKING: - from synapseclient.models import Agent, AgentSession, AgentSessionAccessLevel + from synapseclient.models import ( + Agent, + AgentPrompt, + AgentSession, + AgentSessionAccessLevel, + ) class AgentSessionSynchronousProtocol(Protocol): @@ -109,7 +114,7 @@ def prompt( newer_than: Optional[int] = None, *, synapse_client: Optional[Synapse] = None, - ) -> None: + ) -> "AgentPrompt": """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. @@ -140,7 +145,7 @@ def prompt( print_response=True, ) """ - return None + return AgentPrompt() class AgentSynchronousProtocol(Protocol): @@ -313,7 +318,7 @@ def prompt( newer_than: Optional[int] = None, *, synapse_client: Optional[Synapse] = None, - ) -> None: + ) -> "AgentPrompt": """Sends a prompt to the agent for the current session. If no session is currently active, a new session will be started. @@ -388,4 +393,4 @@ def prompt( session=my_second_session, ) """ - return None + return AgentPrompt() From b1a53fed81e6cffbe7f4504029324bd9e90ddcb7 Mon Sep 17 00:00:00 2001 From: bwmac Date: Fri, 31 Jan 2025 10:06:51 -0500 Subject: [PATCH 6/8] update version v4.7.0 --- synapseclient/synapsePythonClient | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/synapseclient/synapsePythonClient b/synapseclient/synapsePythonClient index 3ccb1602e..5aeb673c5 100644 --- a/synapseclient/synapsePythonClient +++ b/synapseclient/synapsePythonClient @@ -1,6 +1,6 @@ { "client": "synapsePythonClient", - "latestVersion": "4.6.1", + "latestVersion": "4.7.0", "blacklist": [ "0.0.0", "0.4.1", From 155f4ba6fece2aa60e0bc6f7561f2a3e75f47916 Mon Sep 17 00:00:00 2001 From: bwmac Date: Fri, 31 Jan 2025 10:20:00 -0500 Subject: [PATCH 7/8] version v4.7.0 docs --- docs/news.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/news.md b/docs/news.md index 4aff7f747..028c9289c 100644 --- a/docs/news.md +++ b/docs/news.md @@ -9,6 +9,21 @@ detailing some of the changes. the 4.x.x versions hidden behind optional feature flags or different import paths. Any breaking changes will not be included until v5.0. +## 4.7.0 (2025-01-31) + +### Highlights +- **Added functionality for interacting with Synapse Agents:** + - The new `Agent` OOP model allows you to chat with the baseline Synapse Agent, + register and chat with custom Synapse Agents, manage multiple chat sessions and more. + - See the `Agent` documentation for more details and example code to get started. + +### Bug Fixes +- \[[SYNPY-1557](https://sagebionetworks.jira.com/browse/SYNPY-1557)\] - Synapse get recursive link download issue + +### Stories +- \[[SYNPY-1544](https://sagebionetworks.jira.com/browse/SYNPY-1544)\] - Create Synapse Agent OOP Model +- \[[SYNPY-1566](https://sagebionetworks.jira.com/browse/SYNPY-1566)\] - Release python client v4.7.0 + ## 4.6.1 (2024-12-17) ### Highlights From 98f56fc2e57c598efe458983ebda4d218650b635 Mon Sep 17 00:00:00 2001 From: bwmac Date: Fri, 31 Jan 2025 10:24:42 -0500 Subject: [PATCH 8/8] pre-commit --- docs/news.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/news.md b/docs/news.md index 028c9289c..9b16bf59e 100644 --- a/docs/news.md +++ b/docs/news.md @@ -12,11 +12,11 @@ breaking changes will not be included until v5.0. ## 4.7.0 (2025-01-31) ### Highlights -- **Added functionality for interacting with Synapse Agents:** +- **Added functionality for interacting with Synapse Agents:** - The new `Agent` OOP model allows you to chat with the baseline Synapse Agent, register and chat with custom Synapse Agents, manage multiple chat sessions and more. - See the `Agent` documentation for more details and example code to get started. - + ### Bug Fixes - \[[SYNPY-1557](https://sagebionetworks.jira.com/browse/SYNPY-1557)\] - Synapse get recursive link download issue