langroid 0.1.139__py3-none-any.whl → 0.1.219__py3-none-any.whl

langroid/agent/batch.py

@@ -1,73 +1,159 @@
 import asyncio
 import copy
 import inspect
-from typing import Any, Callable, Coroutine, List
+from typing import Any, Callable, Coroutine, Iterable, List, Optional, TypeVar
 
 from dotenv import load_dotenv
-from rich.console import Console
 
 from langroid.agent.base import Agent
 from langroid.agent.chat_document import ChatDocument
 from langroid.agent.task import Task
-from langroid.utils.configuration import quiet_mode, settings
+from langroid.parsing.utils import batched
+from langroid.utils.configuration import quiet_mode
 from langroid.utils.logging import setup_colored_logging
-from langroid.utils.output.printing import SuppressLoggerWarnings
-
-console = Console(quiet=settings.quiet)
+from langroid.utils.output import SuppressLoggerWarnings, status
 
 setup_colored_logging()
 
 load_dotenv()
 
+T = TypeVar("T")
+U = TypeVar("U")
 
-def run_batch_tasks(
-    task: Task,
-    items: List[Any],
-    input_map: Callable[[Any], str | ChatDocument] = lambda x: str(x),
-    output_map: Callable[[ChatDocument | None], Any] = lambda x: x,
-) -> List[Any]:
+
+def run_batch_task_gen(
+    gen_task: Callable[[int], Task],
+    items: list[T],
+    input_map: Callable[[T], str | ChatDocument] = lambda x: str(x),
+    output_map: Callable[[ChatDocument | None], U] = lambda x: x,  # type: ignore
+    sequential: bool = True,
+    batch_size: Optional[int] = None,
+    turns: int = -1,
+    message: Optional[str] = None,
+    handle_exceptions: bool = False,
+) -> list[U]:
     """
-    Run copies of `task` async/concurrently one per item in `items` list.
+    Generate and run copies of a task async/concurrently one per item in `items` list.
     For each item, apply `input_map` to get the initial message to process.
     For each result, apply `output_map` to get the final result.
     Args:
-        task (Task): task to run
-        items (List[Any]): list of items to process
-        input_map (Callable[[Any], str|ChatDocument]): function to map item to
+        gen_task (Callable[[int], Task]): generates the tasks to run
+        items (list[T]): list of items to process
+        input_map (Callable[[T], str|ChatDocument]): function to map item to
             initial message to process
-        output_map (Callable[[ChatDocument|str], Any]): function to map result
+        output_map (Callable[[ChatDocument|str], U]): function to map result
             to final result
+        sequential (bool): whether to run sequentially
+            (e.g. some APIs such as ooba don't support concurrent requests)
+        batch_size (Optional[int]): The number of tasks to run at a time,
+            if None, unbatched
+        turns (int): number of turns to run, -1 for infinite
+        message (Optional[str]): optionally overrides the console status messages
+        handle_exceptions: bool: Whether to replace exceptions with outputs of None
 
     Returns:
-        List[Any]: list of final results
+        list[Any]: list of final results
     """
-
     inputs = [input_map(item) for item in items]
 
-    async def _do_task(input: str | ChatDocument, i: int) -> Any:
-        task_i = task.clone(i)
-        if task_i.agent.config.llm is not None:
-            task_i.agent.config.llm.stream = False
-            task_i.agent.config.show_stats = False
+    async def _do_task(input: str | ChatDocument, i: int) -> Optional[ChatDocument]:
+        task_i = gen_task(i)
+        if task_i.agent.llm is not None:
+            task_i.agent.llm.set_stream(False)
+        task_i.agent.config.show_stats = False
 
-        result = await task_i.run_async(input)
-        return output_map(result)
+        result = await task_i.run_async(input, turns=turns)
+        return result
 
-    async def _do_all() -> List[Any]:
-        with quiet_mode(not settings.debug), SuppressLoggerWarnings():
-            return await asyncio.gather(
-                *(_do_task(input, i) for i, input in enumerate(inputs))
+    async def _do_all(
+        inputs: Iterable[str | ChatDocument], start_idx: int = 0
+    ) -> list[U]:
+        results: list[Optional[ChatDocument]] = []
+        if sequential:
+            for i, input in enumerate(inputs):
+                try:
+                    result = await _do_task(input, i + start_idx)
+                except BaseException as e:
+                    if handle_exceptions:
+                        result = None
+                    else:
+                        raise e
+                results.append(result)
+        else:
+            results_with_exceptions = await asyncio.gather(
+                *(_do_task(input, i + start_idx) for i, input in enumerate(inputs)),
+                return_exceptions=handle_exceptions,
             )
 
-    # show rich console spinner
+            results = [
+                r if not isinstance(r, BaseException) else None
+                for r in results_with_exceptions
+            ]
 
-    n = len(items)
-    with console.status(f"[bold green]Running {n} copies of {task.name}..."):
-        results = asyncio.run(_do_all())
+        return list(map(output_map, results))
+
+    results: List[U] = []
+    if batch_size is None:
+        msg = message or f"[bold green]Running {len(items)} tasks:"
+
+        with status(msg), SuppressLoggerWarnings():
+            results = asyncio.run(_do_all(inputs))
+    else:
+        batches = batched(inputs, batch_size)
+
+        for batch in batches:
+            start_idx = len(results)
+            complete_str = f", {start_idx} complete" if start_idx > 0 else ""
+            msg = message or f"[bold green]Running {len(items)} tasks{complete_str}:"
+
+            with status(msg), SuppressLoggerWarnings():
+                results.extend(asyncio.run(_do_all(batch, start_idx=start_idx)))
 
     return results
 
 
+def run_batch_tasks(
+    task: Task,
+    items: list[T],
+    input_map: Callable[[T], str | ChatDocument] = lambda x: str(x),
+    output_map: Callable[[ChatDocument | None], U] = lambda x: x,  # type: ignore
+    sequential: bool = True,
+    batch_size: Optional[int] = None,
+    turns: int = -1,
+) -> List[U]:
+    """
+    Run copies of `task` async/concurrently one per item in `items` list.
+    For each item, apply `input_map` to get the initial message to process.
+    For each result, apply `output_map` to get the final result.
+    Args:
+        task (Task): task to run
+        items (list[T]): list of items to process
+        input_map (Callable[[T], str|ChatDocument]): function to map item to
+            initial message to process
+        output_map (Callable[[ChatDocument|str], U]): function to map result
+            to final result
+        sequential (bool): whether to run sequentially
+            (e.g. some APIs such as ooba don't support concurrent requests)
+        batch_size (Optional[int]): The number of tasks to run at a time,
+            if None, unbatched
+        turns (int): number of turns to run, -1 for infinite
+
+    Returns:
+        list[Any]: list of final results
+    """
+    message = f"[bold green]Running {len(items)} copies of {task.name}..."
+    return run_batch_task_gen(
+        lambda i: task.clone(i),
+        items,
+        input_map,
+        output_map,
+        sequential,
+        batch_size,
+        turns,
+        message,
+    )
+
+
 def run_batch_agent_method(
     agent: Agent,
     method: Callable[
@@ -76,6 +162,7 @@ def run_batch_agent_method(
     items: List[Any],
     input_map: Callable[[Any], str | ChatDocument] = lambda x: str(x),
     output_map: Callable[[ChatDocument | None], Any] = lambda x: x,
+    sequential: bool = True,
 ) -> List[Any]:
     """
     Run the `method` on copies of `agent`, async/concurrently one per
@@ -97,6 +184,8 @@ def run_batch_agent_method(
             initial message to process
         output_map (Callable[[ChatDocument|str], Any]): function to map result
             to final result
+        sequential (bool): whether to run sequentially
+            (e.g. some APIs such as ooba don't support concurrent requests)
     Returns:
         List[Any]: list of final results
     """
@@ -123,13 +212,19 @@ def run_batch_agent_method(
         return output_map(result)
 
     async def _do_all() -> List[Any]:
+        if sequential:
+            results = []
+            for i, input in enumerate(inputs):
+                result = await _do_task(input, i)
+                results.append(result)
+            return results
         with quiet_mode(), SuppressLoggerWarnings():
             return await asyncio.gather(
                 *(_do_task(input, i) for i, input in enumerate(inputs))
             )
 
     n = len(items)
-    with console.status(f"[bold green]Running {n} copies of {agent_name}..."):
+    with status(f"[bold green]Running {n} copies of {agent_name}..."):
         results = asyncio.run(_do_all())
 
     return results
@@ -140,6 +235,7 @@ def llm_response_batch(
     items: List[Any],
     input_map: Callable[[Any], str | ChatDocument] = lambda x: str(x),
     output_map: Callable[[ChatDocument | None], Any] = lambda x: x,
+    sequential: bool = True,
 ) -> List[Any]:
     return run_batch_agent_method(
         agent,
@@ -147,6 +243,7 @@ def llm_response_batch(
         items,
         input_map=input_map,
         output_map=output_map,
+        sequential=sequential,
     )
 
 
@@ -155,6 +252,7 @@ def agent_response_batch(
     items: List[Any],
     input_map: Callable[[Any], str | ChatDocument] = lambda x: str(x),
     output_map: Callable[[ChatDocument | None], Any] = lambda x: x,
+    sequential: bool = True,
 ) -> List[Any]:
     return run_batch_agent_method(
         agent,
@@ -162,4 +260,5 @@ def agent_response_batch(
         items,
         input_map=input_map,
         output_map=output_map,
+        sequential=sequential,
     )