arize-phoenix 1.9.1rc3__py3-none-any.whl → 2.0.0__py3-none-any.whl

phoenix/experimental/evals/functions/executor.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import signal
+import traceback
+from typing import Any, Callable, Coroutine, List, Optional, Protocol, Sequence, Tuple, Union
+
+from tqdm.auto import tqdm
+
+from phoenix.exceptions import PhoenixException
+
+logger = logging.getLogger(__name__)
+
+
+class Unset:
+    pass
+
+
+_unset = Unset()
+
+
+class Executor(Protocol):
+    def run(self, inputs: Sequence[Any]) -> List[Any]:
+        ...
+
+
+class AsyncExecutor(Executor):
+    """
+    A class that provides asynchronous execution of tasks using a producer-consumer pattern.
+
+    An async interface is provided by the `execute` method, which returns a coroutine, and a sync
+    interface is provided by the `run` method.
+
+    Args:
+        generation_fn (Callable[[Any], Coroutine[Any, Any, Any]]): A coroutine function that
+        generates tasks to be executed.
+
+        concurrency (int, optional): The number of concurrent consumers. Defaults to 3.
+
+        tqdm_bar_format (Optional[str], optional): The format string for the progress bar. Defaults
+        to None.
+
+        max_retries (int, optional): The maximum number of times to retry on exceptions. Defaults to
+        10.
+
+        exit_on_error (bool, optional): Whether to exit execution on the first encountered error.
+        Defaults to True.
+
+        fallback_return_value (Union[Unset, Any], optional): The fallback return value for tasks
+        that encounter errors. Defaults to _unset.
+
+        termination_signal (signal.Signals, optional): The signal handled to terminate the executor.
+    """
+
+    def __init__(
+        self,
+        generation_fn: Callable[[Any], Coroutine[Any, Any, Any]],
+        concurrency: int = 3,
+        tqdm_bar_format: Optional[str] = None,
+        max_retries: int = 10,
+        exit_on_error: bool = True,
+        fallback_return_value: Union[Unset, Any] = _unset,
+        termination_signal: signal.Signals = signal.SIGINT,
+    ):
+        self.generate = generation_fn
+        self.fallback_return_value = fallback_return_value
+        self.concurrency = concurrency
+        self.tqdm_bar_format = tqdm_bar_format
+        self.max_retries = max_retries
+        self.exit_on_error = exit_on_error
+        self.base_priority = 0
+        self.termination_signal = termination_signal
+
+    async def producer(
+        self,
+        inputs: Sequence[Any],
+        queue: asyncio.PriorityQueue[Tuple[int, Any]],
+        max_fill: int,
+        done_producing: asyncio.Event,
+        termination_signal: asyncio.Event,
+    ) -> None:
+        try:
+            for index, input in enumerate(inputs):
+                if termination_signal.is_set():
+                    break
+                while queue.qsize() >= max_fill:
+                    # keep room in the queue for requeues
+                    await asyncio.sleep(1)
+                await queue.put((self.base_priority, (index, input)))
+        finally:
+            done_producing.set()
+
+    async def consumer(
+        self,
+        output: List[Any],
+        queue: asyncio.PriorityQueue[Tuple[int, Any]],
+        done_producing: asyncio.Event,
+        termination_event: asyncio.Event,
+        progress_bar: tqdm[Any],
+    ) -> None:
+        termination_event_watcher = None
+        while True:
+            marked_done = False
+            try:
+                priority, item = await asyncio.wait_for(queue.get(), timeout=1)
+            except asyncio.TimeoutError:
+                if done_producing.is_set() and queue.empty():
+                    break
+                continue
+            if termination_event.is_set():
+                # discard any remaining items in the queue
+                queue.task_done()
+                marked_done = True
+                continue
+
+            index, payload = item
+            try:
+                generate_task = asyncio.create_task(self.generate(payload))
+                termination_event_watcher = asyncio.create_task(termination_event.wait())
+                done, pending = await asyncio.wait(
+                    [generate_task, termination_event_watcher],
+                    timeout=120,
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                if generate_task in done:
+                    output[index] = generate_task.result()
+                    progress_bar.update()
+                elif termination_event.is_set():
+                    # discard the pending task and remaining items in the queue
+                    if not generate_task.done():
+                        generate_task.cancel()
+                        try:
+                            # allow any cleanup to finish for the cancelled task
+                            await generate_task
+                        except asyncio.CancelledError:
+                            # Handle the cancellation exception
+                            pass
+                    queue.task_done()
+                    marked_done = True
+                    continue
+                else:
+                    tqdm.write("Worker timeout, requeuing")
+                    # task timeouts are requeued at base priority
+                    await queue.put((self.base_priority, item))
+            except Exception as exc:
+                is_phoenix_exception = isinstance(exc, PhoenixException)
+                if (retry_count := abs(priority)) <= self.max_retries and not is_phoenix_exception:
+                    tqdm.write(
+                        f"Exception in worker on attempt {retry_count + 1}: raised {repr(exc)}"
+                    )
+                    tqdm.write("Requeuing...")
+                    await queue.put((priority - 1, item))
+                else:
+                    tqdm.write(f"Exception in worker: {traceback.format_exc()}")
+                    if self.exit_on_error:
+                        termination_event.set()
+                    else:
+                        progress_bar.update()
+            finally:
+                if not marked_done:
+                    queue.task_done()
+                if termination_event_watcher and not termination_event_watcher.done():
+                    termination_event_watcher.cancel()
+
+    async def execute(self, inputs: Sequence[Any]) -> List[Any]:
+        termination_event = asyncio.Event()
+
+        def termination_handler(signum: int, frame: Any) -> None:
+            termination_event.set()
+            tqdm.write("Process was interrupted. The return value will be incomplete...")
+
+        signal.signal(self.termination_signal, termination_handler)
+        outputs = [self.fallback_return_value] * len(inputs)
+        progress_bar = tqdm(total=len(inputs), bar_format=self.tqdm_bar_format)
+
+        max_queue_size = 5 * self.concurrency  # limit the queue to bound memory usage
+        max_fill = max_queue_size - (2 * self.concurrency)  # ensure there is always room to requeue
+        queue: asyncio.PriorityQueue[Tuple[int, Any]] = asyncio.PriorityQueue(
+            maxsize=max_queue_size
+        )
+        done_producing = asyncio.Event()
+
+        producer = asyncio.create_task(
+            self.producer(inputs, queue, max_fill, done_producing, termination_event)
+        )
+        consumers = [
+            asyncio.create_task(
+                self.consumer(outputs, queue, done_producing, termination_event, progress_bar)
+            )
+            for _ in range(self.concurrency)
+        ]
+
+        await asyncio.gather(producer, *consumers)
+        join_task = asyncio.create_task(queue.join())
+        termination_event_watcher = asyncio.create_task(termination_event.wait())
+        done, pending = await asyncio.wait(
+            [join_task, termination_event_watcher], return_when=asyncio.FIRST_COMPLETED
+        )
+        if termination_event_watcher in done:
+            # Cancel all tasks
+            if not join_task.done():
+                join_task.cancel()
+            if not producer.done():
+                producer.cancel()
+            for task in consumers:
+                if not task.done():
+                    task.cancel()
+
+        if not termination_event_watcher.done():
+            termination_event_watcher.cancel()
+
+        # reset the SIGTERM handler
+        signal.signal(self.termination_signal, signal.SIG_DFL)  # reset the SIGTERM handler
+        return outputs
+
+    def run(self, inputs: Sequence[Any]) -> List[Any]:
+        return asyncio.run(self.execute(inputs))
+
+
+class SyncExecutor(Executor):
+    """
+    Synchronous executor for generating outputs from inputs using a given generation function.
+
+    Args:
+        generation_fn (Callable[[Any], Any]): The generation function that takes an input and
+        returns an output.
+
+        tqdm_bar_format (Optional[str], optional): The format string for the progress bar. Defaults
+        to None.
+
+        max_retries (int, optional): The maximum number of times to retry on exceptions. Defaults to
+        10.
+
+        exit_on_error (bool, optional): Whether to exit execution on the first encountered error.
+        Defaults to True.
+
+        fallback_return_value (Union[Unset, Any], optional): The fallback return value for tasks
+        that encounter errors. Defaults to _unset.
+    """
+
+    def __init__(
+        self,
+        generation_fn: Callable[[Any], Any],
+        tqdm_bar_format: Optional[str] = None,
+        max_retries: int = 10,
+        exit_on_error: bool = True,
+        fallback_return_value: Union[Unset, Any] = _unset,
+        termination_signal: signal.Signals = signal.SIGINT,
+    ):
+        self.generate = generation_fn
+        self.fallback_return_value = fallback_return_value
+        self.tqdm_bar_format = tqdm_bar_format
+        self.max_retries = max_retries
+        self.exit_on_error = exit_on_error
+        self.termination_signal = termination_signal
+
+        self._TERMINATE = False
+
+    def _signal_handler(self, signum: int, frame: Any) -> None:
+        tqdm.write("Process was interrupted. The return value will be incomplete...")
+        self._TERMINATE = True
+
+    def run(self, inputs: Sequence[Any]) -> List[Any]:
+        signal.signal(self.termination_signal, self._signal_handler)
+        outputs = [self.fallback_return_value] * len(inputs)
+        progress_bar = tqdm(total=len(inputs), bar_format=self.tqdm_bar_format)
+
+        for index, input in enumerate(inputs):
+            try:
+                for attempt in range(self.max_retries + 1):
+                    if self._TERMINATE:
+                        return outputs
+                    try:
+                        result = self.generate(input)
+                        outputs[index] = result
+                        progress_bar.update()
+                    except Exception as exc:
+                        is_phoenix_exception = isinstance(exc, PhoenixException)
+                        if attempt >= self.max_retries or is_phoenix_exception:
+                            raise exc
+                        else:
+                            tqdm.write(f"Exception in worker on attempt {attempt + 1}: {exc}")
+                            tqdm.write("Retrying...")
+            except Exception as exc:
+                tqdm.write(f"Exception in worker: {exc}")
+                if self.exit_on_error:
+                    return outputs
+                else:
+                    progress_bar.update()
+        signal.signal(self.termination_signal, signal.SIG_DFL)  # reset the SIGTERM handler
+        return outputs
+
+
+def get_executor_on_sync_context(
+    sync_fn: Callable[[Any], Any],
+    async_fn: Callable[[Any], Coroutine[Any, Any, Any]],
+    run_sync: bool = False,
+    concurrency: int = 3,
+    tqdm_bar_format: Optional[str] = None,
+    exit_on_error: bool = True,
+    fallback_return_value: Union[Unset, Any] = _unset,
+) -> Executor:
+    if run_sync:
+        return SyncExecutor(
+            sync_fn,
+            tqdm_bar_format=tqdm_bar_format,
+            exit_on_error=exit_on_error,
+            fallback_return_value=fallback_return_value,
+        )
+
+    if _running_event_loop_exists():
+        if getattr(asyncio, "_nest_patched", False):
+            return AsyncExecutor(
+                async_fn,
+                concurrency=concurrency,
+                tqdm_bar_format=tqdm_bar_format,
+                exit_on_error=exit_on_error,
+                fallback_return_value=fallback_return_value,
+            )
+        else:
+            logger.warning(
+                "🐌!! If running llm_classify inside a notebook, patching the event loop with "
+                "nest_asyncio will allow asynchronous eval submission, and is significantly "
+                "faster. To patch the event loop, run `nest_asyncio.apply()`."
+            )
+            return SyncExecutor(
+                sync_fn,
+                tqdm_bar_format=tqdm_bar_format,
+                exit_on_error=exit_on_error,
+                fallback_return_value=fallback_return_value,
+            )
+    else:
+        return AsyncExecutor(
+            async_fn,
+            concurrency=concurrency,
+            tqdm_bar_format=tqdm_bar_format,
+            exit_on_error=exit_on_error,
+            fallback_return_value=fallback_return_value,
+        )
+
+
+def _running_event_loop_exists() -> bool:
+    """Checks for a running event loop.
+
+    Returns:
+        bool: True if a running event loop exists, False otherwise.
+    """
+    try:
+        asyncio.get_running_loop()
+        return True
+    except RuntimeError:
+        return False

phoenix/experimental/evals/functions/generate.py

@@ -1,9 +1,11 @@
 import logging
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 import pandas as pd
-from tqdm.auto import tqdm
 
+from phoenix.experimental.evals.functions.executor import (
+    get_executor_on_sync_context,
+)
 from phoenix.experimental.evals.models import BaseEvalModel, set_verbosity
 from phoenix.experimental.evals.templates import (
     PromptTemplate,
@@ -15,7 +17,7 @@ from phoenix.experimental.evals.utils import get_tqdm_progress_bar_formatter
 logger = logging.getLogger(__name__)
 
 
-def _no_op_parser(response: str) -> Dict[str, str]:
+def _no_op_parser(response: str, response_index: int) -> Dict[str, str]:
     return {"output": response}
 
 
@@ -25,7 +27,11 @@ def llm_generate(
     model: BaseEvalModel,
     system_instruction: Optional[str] = None,
     verbose: bool = False,
-    output_parser: Optional[Callable[[str], Dict[str, Any]]] = None,
+    output_parser: Optional[Callable[[str, int], Dict[str, Any]]] = None,
+    include_prompt: bool = False,
+    include_response: bool = False,
+    run_sync: bool = False,
+    concurrency: int = 20,
 ) -> pd.DataFrame:
     """
     Generates a text using a template using an LLM. This function is useful
@@ -49,10 +55,23 @@ def llm_generate(
         verbose (bool, optional): If True, prints detailed information to stdout such as model
         invocation parameters and retry info. Default False.
 
-        output_parser (Callable[[str], Dict[str, Any]], optional): An optional function
-        that takes each generated response and parses it to a dictionary. The keys of the dictionary
-        should correspond to the column names of the output dataframe. If None, the output dataframe
-        will have a single column named "output". Default None.
+        output_parser (Callable[[str, int], Dict[str, Any]], optional): An optional function
+        that takes each generated response and response index and parses it to a dictionary. The
+        keys of the dictionary should correspond to the column names of the output dataframe. If
+        None, the output dataframe will have a single column named "output". Default None.
+
+        include_prompt (bool, default=False): If True, includes a column named `prompt` in the
+        output dataframe containing the prompt used for each generation.
+
+        include_response (bool, default=False): If True, includes a column named `response` in the
+        output dataframe containing the raw response from the LLM prior to applying the output
+        parser.
+
+        run_sync (bool, default=False): If True, forces synchronous request submission. Otherwise
+        evaluations will be run asynchronously if possible.
+
+        concurrency (int, default=20): The number of concurrent evals if async submission is
+        possible.
 
     Returns:
         generations_dataframe (pandas.DataFrame): A dataframe where each row
@@ -61,28 +80,53 @@ def llm_generate(
     """
     tqdm_bar_format = get_tqdm_progress_bar_formatter("llm_generate")
     output_parser = output_parser or _no_op_parser
-    with set_verbosity(model, verbose) as verbose_model:
-        template = normalize_prompt_template(template)
-        logger.info(f"Template: \n{template.prompt()}\n")
-        logger.info(f"Template variables: {template.variables}")
-        prompts = map_template(dataframe, template)
-
-        # For each prompt, generate and parse the response
-        output = []
-        # Wrap the loop in a try / catch so that we can still return a dataframe
-        # even if the process is interrupted
-        try:
-            for prompt in tqdm(prompts, bar_format=tqdm_bar_format):
-                logger.info(f"Prompt: {prompt}")
-                response = verbose_model(prompt, instruction=system_instruction)
-                parsed_response = output_parser(response)
-                output.append(parsed_response)
-
-        except (Exception, KeyboardInterrupt) as e:
-            logger.error(e)
-            print(
-                "Process was interrupted. The return value will be incomplete",
-                e,
+    template = normalize_prompt_template(template)
+    logger.info(f"Template: \n{template.prompt()}\n")
+    logger.info(f"Template variables: {template.variables}")
+    prompts = map_template(dataframe, template)
+
+    async def _run_llm_generation_async(enumerated_prompt: Tuple[int, str]) -> Dict[str, Any]:
+        index, prompt = enumerated_prompt
+        with set_verbosity(model, verbose) as verbose_model:
+            response = await verbose_model._async_generate(
+                prompt,
+                instruction=system_instruction,
+            )
+        parsed_response = output_parser(response, index)
+        if include_prompt:
+            parsed_response["prompt"] = prompt
+        if include_response:
+            parsed_response["response"] = response
+        return parsed_response
+
+    def _run_llm_generation_sync(enumerated_prompt: Tuple[int, str]) -> Dict[str, Any]:
+        index, prompt = enumerated_prompt
+        with set_verbosity(model, verbose) as verbose_model:
+            response = verbose_model._generate(
+                prompt,
+                instruction=system_instruction,
             )
-        # Return the data as a dataframe
-        return pd.DataFrame(output)
+        parsed_response = output_parser(response, index)
+        if include_prompt:
+            parsed_response["prompt"] = prompt
+        if include_response:
+            parsed_response["response"] = response
+        return parsed_response
+
+    fallback_return_value = {
+        "output": "generation-failed",
+        **({"prompt": ""} if include_prompt else {}),
+        **({"response": ""} if include_response else {}),
+    }
+
+    executor = get_executor_on_sync_context(
+        _run_llm_generation_sync,
+        _run_llm_generation_async,
+        run_sync=run_sync,
+        concurrency=concurrency,
+        tqdm_bar_format=tqdm_bar_format,
+        exit_on_error=True,
+        fallback_return_value=fallback_return_value,
+    )
+    output = executor.run(list(enumerate(prompts.tolist())))
+    return pd.DataFrame(output)

phoenix/experimental/evals/models/rate_limiters.py

@@ -6,6 +6,7 @@ from typing import Any, Callable, Coroutine, Optional, Tuple, Type, TypeVar
 
 from typing_extensions import ParamSpec
 
+from phoenix.exceptions import PhoenixException
 from phoenix.utilities.logging import printif
 
 ParameterSpec = ParamSpec("ParameterSpec")
@@ -13,7 +14,7 @@ GenericType = TypeVar("GenericType")
 AsyncCallable = Callable[ParameterSpec, Coroutine[Any, Any, GenericType]]
 
 
-class UnavailableTokensError(Exception):
+class UnavailableTokensError(PhoenixException):
     pass
 
 
@@ -133,7 +134,7 @@ class AdaptiveTokenBucket:
                 continue
 
 
-class RateLimitError(BaseException):
+class RateLimitError(PhoenixException):
     ...
 
 
@@ -162,9 +163,9 @@ class RateLimiter:
             rate_increase_factor=rate_increase_factor,
             cooldown_seconds=cooldown_seconds,
         )
-        self._rate_limit_handling = asyncio.Event()
-        self._rate_limit_handling.set()  # allow requests to start immediately
-        self._rate_limit_handling_lock = asyncio.Lock()
+        self._rate_limit_handling: Optional[asyncio.Event] = None
+        self._rate_limit_handling_lock: Optional[asyncio.Lock] = None
+        self._current_loop: Optional[asyncio.AbstractEventLoop] = None
         self._verbose = verbose
 
     def limit(
@@ -192,11 +193,30 @@ class RateLimiter:
 
         return wrapper
 
+    def _initialize_async_primitives(self) -> None:
+        """
+        Lazily initialize async primitives to ensure they are created in the correct event loop.
+        """
+
+        loop = asyncio.get_running_loop()
+        if loop is not self._current_loop:
+            self._current_loop = loop
+            self._rate_limit_handling = asyncio.Event()
+            self._rate_limit_handling.set()
+            self._rate_limit_handling_lock = asyncio.Lock()
+
     def alimit(
         self, fn: AsyncCallable[ParameterSpec, GenericType]
     ) -> AsyncCallable[ParameterSpec, GenericType]:
         @wraps(fn)
         async def wrapper(*args: Any, **kwargs: Any) -> GenericType:
+            self._initialize_async_primitives()
+            assert self._rate_limit_handling_lock is not None and isinstance(
+                self._rate_limit_handling_lock, asyncio.Lock
+            )
+            assert self._rate_limit_handling is not None and isinstance(
+                self._rate_limit_handling, asyncio.Event
+            )
             try:
                 try:
                     await asyncio.wait_for(self._rate_limit_handling.wait(), 120)

phoenix/experimental/evals/templates/__init__.py

@@ -15,7 +15,6 @@ from .default_templates import (
     TOXICITY_PROMPT_TEMPLATE,
 )
 from .template import (
-    NOT_PARSABLE,
     ClassificationTemplate,
     PromptOptions,
     PromptTemplate,
@@ -32,7 +31,6 @@ __all__ = [
     "normalize_classification_template",
     "normalize_prompt_template",
     "map_template",
-    "NOT_PARSABLE",
     "CODE_READABILITY_PROMPT_RAILS_MAP",
     "CODE_READABILITY_PROMPT_TEMPLATE",
     "HALLUCINATION_PROMPT_RAILS_MAP",

phoenix/experimental/evals/templates/template.py

@@ -4,14 +4,11 @@ from typing import Callable, List, Mapping, Optional, Tuple, Union
 
 import pandas as pd
 
+from phoenix.experimental.evals.utils import NOT_PARSABLE
+
 DEFAULT_START_DELIM = "{"
 DEFAULT_END_DELIM = "}"
 
-# Rather than returning None, we return this string to indicate that the LLM output could not be
-# parsed.
-# This is useful for debugging as well as to just treat the output as a non-parsable category
-NOT_PARSABLE = "NOT_PARSABLE"
-
 
 @dataclass
 class PromptOptions:

phoenix/experimental/evals/utils/__init__.py

@@ -1,11 +1,24 @@
 import json
 from io import BytesIO
+from typing import List, Optional, Tuple
 from urllib.error import HTTPError
 from urllib.request import urlopen
 from zipfile import ZipFile
 
 import pandas as pd
 
+from phoenix.utilities.logging import printif
+
+# Rather than returning None, we return this string to indicate that the LLM output could not be
+# parsed.
+# This is useful for debugging as well as to just treat the output as a non-parsable category
+NOT_PARSABLE = "NOT_PARSABLE"
+
+# argument keys in the default openai function call,
+# defined here only to prevent typos
+_RESPONSE = "response"
+_EXPLANATION = "explanation"
+
 
 def download_benchmark_dataset(task: str, dataset_name: str) -> pd.DataFrame:
     """Downloads an Arize evals benchmark dataset as a pandas dataframe.
@@ -42,3 +55,56 @@ def get_tqdm_progress_bar_formatter(title: str) -> str:
         title + " |{bar}| {n_fmt}/{total_fmt} ({percentage:3.1f}%) "
         "| ⏳ {elapsed}<{remaining} | {rate_fmt}{postfix}"
     )
+
+
+def snap_to_rail(raw_string: Optional[str], rails: List[str], verbose: bool = False) -> str:
+    """
+    Snaps a string to the nearest rail, or returns None if the string cannot be
+    snapped to a rail.
+
+    Args:
+        raw_string (str): An input to be snapped to a rail.
+
+        rails (List[str]): The target set of strings to snap to.
+
+    Returns:
+        str: A string from the rails argument or "UNPARSABLE" if the input
+        string could not be snapped.
+    """
+    if not raw_string:
+        return NOT_PARSABLE
+    snap_string = raw_string.lower()
+    rails = list(set(rail.lower() for rail in rails))
+    rails.sort(key=len, reverse=True)
+    found_rails = set()
+    for rail in rails:
+        if rail in snap_string:
+            found_rails.add(rail)
+            snap_string = snap_string.replace(rail, "")
+    if len(found_rails) != 1:
+        printif(verbose, f"- Cannot snap {repr(raw_string)} to rails")
+        return NOT_PARSABLE
+    rail = list(found_rails)[0]
+    printif(verbose, f"- Snapped {repr(raw_string)} to rail: {rail}")
+    return rail
+
+
+def parse_openai_function_call(raw_output: str) -> Tuple[str, Optional[str]]:
+    """
+    Parses the output of an OpenAI function call.
+
+    Args:
+        raw_output (str): The raw output of an OpenAI function call.
+
+    Returns:
+        Tuple[str, Optional[str]]: A tuple of the unrailed label and an optional
+        explanation.
+    """
+    try:
+        function_arguments = json.loads(raw_output, strict=False)
+        unrailed_label = function_arguments.get(_RESPONSE, "")
+        explanation = function_arguments.get(_EXPLANATION)
+    except json.JSONDecodeError:
+        unrailed_label = raw_output
+        explanation = None
+    return unrailed_label, explanation