Work around type hinting problems with redis.asyncio

Typing for the asyncio redis client is messed up, and functions
return `T | Awaitable[T]` instead of `T`. Add a workaround
wrapper that asserts that we are in the second case.

See:

  redis/redis-py#3619

for an in-progress upstream fix.

Author: owtaylor

beeai/agents/backport_agent.py

@@ -30,7 +30,7 @@
 from observability import setup_observability
 from tools.commands import RunShellCommandTool
 from triage_agent import BackportData, ErrorData
-from utils import get_agent_execution_config, mcp_tools, redis_client, get_git_finalization_steps
+from utils import fixAwait, get_agent_execution_config, mcp_tools, redis_client, get_git_finalization_steps
 
 logger = logging.getLogger(__name__)
 
@@ -230,7 +230,7 @@ class Task(BaseModel):
 
             while True:
                 logger.info("Waiting for tasks from backport_queue (timeout: 30s)...")
-                element = await redis.brpop("backport_queue", timeout=30)
+                element = await fixAwait(redis.brpop(["backport_queue"], timeout=30))
                 if element is None:
                     logger.info("No tasks received, continuing to wait...")
                     continue
@@ -264,13 +264,13 @@ async def retry(task, error):
                             f"Task failed (attempt {task.attempts}/{max_retries}), "
                             f"re-queuing for retry: {backport_data.jira_issue}"
                         )
-                        await redis.lpush("backport_queue", task.model_dump_json())
+                        await fixAwait(redis.lpush("backport_queue", task.model_dump_json()))
                     else:
                         logger.error(
                             f"Task failed after {max_retries} attempts, "
                             f"moving to error list: {backport_data.jira_issue}"
                         )
-                        await redis.lpush("error_list", error)
+                        await fixAwait(redis.lpush("error_list", error))
 
                 try:
                     logger.info(f"Starting backport processing for {backport_data.jira_issue}")
@@ -287,7 +287,7 @@ async def retry(task, error):
                     rmtree(local_clone)
                     if output.success:
                         logger.info(f"Backport successful for {backport_data.jira_issue}, " f"adding to completed list")
-                        await redis.lpush("completed_backport_list", output.model_dump_json())
+                        await fixAwait(redis.lpush("completed_backport_list", output.model_dump_json()))
                     else:
                         logger.warning(f"Backport failed for {backport_data.jira_issue}: {output.error}")
                         await retry(task, output.error)

beeai/agents/rebase_agent.py

@@ -24,7 +24,7 @@
 from observability import setup_observability
 from tools.commands import RunShellCommandTool
 from triage_agent import RebaseData, ErrorData
-from utils import get_agent_execution_config, mcp_tools, redis_client, get_git_finalization_steps
+from utils import fixAwait, get_agent_execution_config, mcp_tools, redis_client, get_git_finalization_steps
 
 logger = logging.getLogger(__name__)
 
@@ -195,7 +195,7 @@ class Task(BaseModel):
 
             while True:
                 logger.info("Waiting for tasks from rebase_queue (timeout: 30s)...")
-                element = await redis.brpop("rebase_queue", timeout=30)
+                element = await fixAwait(redis.brpop(["rebase_queue"], timeout=30))
                 if element is None:
                     logger.info("No tasks received, continuing to wait...")
                     continue
@@ -225,13 +225,13 @@ async def retry(task, error):
                             f"Task failed (attempt {task.attempts}/{max_retries}), "
                             f"re-queuing for retry: {rebase_data.jira_issue}"
                         )
-                        await redis.lpush("rebase_queue", task.model_dump_json())
+                        await fixAwait(redis.lpush("rebase_queue", task.model_dump_json()))
                     else:
                         logger.error(
                             f"Task failed after {max_retries} attempts, "
                             f"moving to error list: {rebase_data.jira_issue}"
                         )
-                        await redis.lpush("error_list", error)
+                        await fixAwait(redis.lpush("error_list", error))
 
                 try:
                     logger.info(f"Starting rebase processing for {rebase_data.jira_issue}")
@@ -246,7 +246,7 @@ async def retry(task, error):
                 else:
                     if output.success:
                         logger.info(f"Rebase successful for {rebase_data.jira_issue}, " f"adding to completed list")
-                        await redis.lpush("completed_rebase_list", output.model_dump_json())
+                        await fixAwait(redis.lpush("completed_rebase_list", output.model_dump_json()))
                     else:
                         logger.warning(f"Rebase failed for {rebase_data.jira_issue}: {output.error}")
                         await retry(task, output.error)

beeai/agents/triage_agent.py

@@ -23,7 +23,7 @@
 from observability import setup_observability
 from tools.commands import RunShellCommandTool
 from tools.patch_validator import PatchValidatorTool
-from utils import get_agent_execution_config, mcp_tools, redis_client
+from utils import fixAwait, get_agent_execution_config, mcp_tools, redis_client
 
 logger = logging.getLogger(__name__)
 
@@ -285,7 +285,7 @@ class Task(BaseModel):
 
             while True:
                 logger.info("Waiting for tasks from triage_queue (timeout: 30s)...")
-                element = await redis.brpop("triage_queue", timeout=30)
+                element = await fixAwait(redis.brpop(["triage_queue"], timeout=30))
                 if element is None:
                     logger.info("No tasks received, continuing to wait...")
                     continue
@@ -304,12 +304,12 @@ async def retry(task, error):
                             f"Task failed (attempt {task.attempts}/{max_retries}), "
                             f"re-queuing for retry: {input.issue}"
                         )
-                        await redis.lpush("triage_queue", task.model_dump_json())
+                        await fixAwait(redis.lpush("triage_queue", task.model_dump_json()))
                     else:
                         logger.error(
                             f"Task failed after {max_retries} attempts, " f"moving to error list: {input.issue}"
                         )
-                        await redis.lpush("error_list", error)
+                        await fixAwait(redis.lpush("error_list", error))
 
                 try:
                     logger.info(f"Starting triage processing for {input.issue}")
@@ -329,21 +329,21 @@ async def retry(task, error):
                     if output.resolution == Resolution.REBASE:
                         logger.info(f"Triage resolved as REBASE for {input.issue}, " f"adding to rebase queue")
                         task = Task(metadata=output.data.model_dump())
-                        await redis.lpush("rebase_queue", task.model_dump_json())
+                        await fixAwait(redis.lpush("rebase_queue", task.model_dump_json()))
                     elif output.resolution == Resolution.BACKPORT:
                         logger.info(f"Triage resolved as BACKPORT for {input.issue}, " f"adding to backport queue")
                         task = Task(metadata=output.data.model_dump())
-                        await redis.lpush("backport_queue", task.model_dump_json())
+                        await fixAwait(redis.lpush("backport_queue", task.model_dump_json()))
                     elif output.resolution == Resolution.CLARIFICATION_NEEDED:
                         logger.info(
                             f"Triage resolved as CLARIFICATION_NEEDED for {input.issue}, "
                             f"adding to clarification needed queue"
                         )
                         task = Task(metadata=output.data.model_dump())
-                        await redis.lpush("clarification_needed_queue", task.model_dump_json())
+                        await fixAwait(redis.lpush("clarification_needed_queue", task.model_dump_json()))
                     elif output.resolution == Resolution.NO_ACTION:
                         logger.info(f"Triage resolved as NO_ACTION for {input.issue}, " f"adding to no action list")
-                        await redis.lpush("no_action_list", output.data.model_dump_json())
+                        await fixAwait(redis.lpush("no_action_list", output.data.model_dump_json()))
                     elif output.resolution == Resolution.ERROR:
                         logger.warning(f"Triage resolved as ERROR for {input.issue}, retrying")
                         await retry(task, output.data.model_dump_json())

beeai/agents/utils.py

@@ -1,7 +1,8 @@
+import inspect
 import os
 
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, Callable
+from typing import AsyncGenerator, Awaitable, Callable, TypeVar
 
 import redis.asyncio as redis
 from mcp import ClientSession
@@ -29,6 +30,26 @@ async def redis_client(redis_url: str) -> AsyncGenerator[redis.Redis, None]:
         await client.aclose()
 
 
+T = TypeVar("T")
+
+async def fixAwait(v: T | Awaitable[T]) -> T:
+    """
+    Work around typing problems in the asyncio redis client.
+
+    Typing for the asyncio redis client is messed up, and functions
+    return `T | Awaitable[T]` instead of `T`. This function
+    fixes the type error by asserting that the value is awaitable
+    before awaiting it.
+
+    For a proper fix, see: https://github.com/redis/redis-py/pull/3619
+
+
+    Usage: `await fixAwait(redis.get("key"))`
+    """
+    assert inspect.isawaitable(v)
+    return await v
+
+
 @asynccontextmanager
 async def mcp_tools(
     sse_url: str, filter: Callable[[str], bool] | None = None