From 65736884568bc99d780c564a1993401870f34d7d Mon Sep 17 00:00:00 2001
From: Hamada Salhab <hamada.a.salhab@gmail.com>
Date: Wed, 6 Nov 2024 02:07:47 +0300
Subject: [PATCH] fix(agents-api): Add `mmr_strength` in search tools &
 increase timeout for `transition_step` (#808)

<!-- ELLIPSIS_HIDDEN -->



> [!IMPORTANT]
> Add `mmr_strength` to search requests and increase `transition_step`
timeout based on environment in `agents-api`.
>
>   - **Behavior**:
> - Add `mmr_strength` parameter to `_create_search_request()` in
`execute_system.py` for `HybridDocSearchRequest`,
`TextOnlyDocSearchRequest`, and `VectorDocSearchRequest`.
> - Increase `schedule_to_close_timeout` for `transition_step` in
`transition.py` to 600 seconds unless in debug or testing mode.
>   - **Misc**:
>     - Import `debug` and `testing` from `env` in `transition.py`.
>
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=julep-ai%2Fjulep&utm_source=github&utm_medium=referral)<sup>
for bd64aa27254a7a19a7aa8320746599e67337ed85. It will automatically
update as commits are pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->
---
 agents-api/agents_api/activities/execute_system.py         | 3 +++
 .../agents_api/workflows/task_execution/transition.py      | 7 ++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/agents-api/agents_api/activities/execute_system.py b/agents-api/agents_api/activities/execute_system.py
index 2365c8261..02623c3a6 100644
--- a/agents-api/agents_api/activities/execute_system.py
+++ b/agents-api/agents_api/activities/execute_system.py
@@ -106,6 +106,7 @@ def _create_search_request(arguments: dict) -> Any:
     if "text" in arguments and "vector" in arguments:
         return HybridDocSearchRequest(
             text=arguments.pop("text"),
+            mmr_strength=arguments.pop("mmr_strength", 0),
             vector=arguments.pop("vector"),
             alpha=arguments.pop("alpha", 0.75),
             confidence=arguments.pop("confidence", 0.5),
@@ -114,11 +115,13 @@ def _create_search_request(arguments: dict) -> Any:
     elif "text" in arguments:
         return TextOnlyDocSearchRequest(
             text=arguments.pop("text"),
+            mmr_strength=arguments.pop("mmr_strength", 0),
             limit=arguments.get("limit", 10),
         )
     elif "vector" in arguments:
         return VectorDocSearchRequest(
             vector=arguments.pop("vector"),
+            mmr_strength=arguments.pop("mmr_strength", 0),
             confidence=arguments.pop("confidence", 0.7),
             limit=arguments.get("limit", 10),
         )
diff --git a/agents-api/agents_api/workflows/task_execution/transition.py b/agents-api/agents_api/workflows/task_execution/transition.py
index 150be6720..5ba726281 100644
--- a/agents-api/agents_api/workflows/task_execution/transition.py
+++ b/agents-api/agents_api/workflows/task_execution/transition.py
@@ -12,6 +12,9 @@
 from ...common.protocol.tasks import PartialTransition, StepContext
 from ...common.retry_policies import DEFAULT_RETRY_POLICY
 
+with workflow.unsafe.imports_passed_through():
+    from ...env import debug, testing
+
 
 async def transition(
     context: StepContext, state: PartialTransition | None = None, **kwargs
@@ -46,7 +49,9 @@ async def transition(
         return await workflow.execute_activity(
             task_steps.transition_step,
             args=[context, transition_request],
-            schedule_to_close_timeout=timedelta(seconds=30),
+            schedule_to_close_timeout=timedelta(
+                seconds=30 if debug or testing else 600
+            ),
             retry_policy=DEFAULT_RETRY_POLICY,
         )