coco-framework 0.1.0__py3-none-any.whl

coco_core/utils/html.py

@@ -0,0 +1,132 @@
+"""Utilities for working with HTML."""
+
+import logging
+import re
+from collections.abc import Sequence
+from urllib.parse import urljoin, urlparse
+
+logger = logging.getLogger(__name__)
+
+PREFIXES_TO_IGNORE = ("javascript:", "mailto:", "#")
+
+SUFFIXES_TO_IGNORE = (
+    ".css",
+    ".js",
+    ".ico",
+    ".png",
+    ".jpg",
+    ".jpeg",
+    ".gif",
+    ".svg",
+    ".csv",
+    ".bz2",
+    ".zip",
+    ".epub",
+    ".webp",
+    ".pdf",
+    ".docx",
+    ".xlsx",
+    ".pptx",
+    ".pptm",
+)
+
+SUFFIXES_TO_IGNORE_REGEX = (
+    "(?!" + "|".join([re.escape(s) + r"[\#'\"]" for s in SUFFIXES_TO_IGNORE]) + ")"
+)
+
+PREFIXES_TO_IGNORE_REGEX = (
+    "(?!" + "|".join([re.escape(s) for s in PREFIXES_TO_IGNORE]) + ")"
+)
+
+DEFAULT_LINK_REGEX = (
+    rf"href=[\"']{PREFIXES_TO_IGNORE_REGEX}((?:{SUFFIXES_TO_IGNORE_REGEX}.)*?)[\#'\"]"
+)
+
+
+def find_all_links(
+    raw_html: str, *, pattern: str | re.Pattern[str] | None = None
+) -> list[str]:
+    """Extract all links from a raw HTML string.
+
+    Args:
+        raw_html: original HTML.
+        pattern: Regex to use for extracting links from raw HTML.
+
+    Returns:
+        A list of all links found in the HTML.
+    """
+    pattern = pattern or DEFAULT_LINK_REGEX
+    return list(set(re.findall(pattern, raw_html)))
+
+
+def extract_sub_links(
+    raw_html: str,
+    url: str,
+    *,
+    base_url: str | None = None,
+    pattern: str | re.Pattern[str] | None = None,
+    prevent_outside: bool = True,
+    exclude_prefixes: Sequence[str] = (),
+    continue_on_failure: bool = False,
+) -> list[str]:
+    """Extract all links from a raw HTML string and convert into absolute paths.
+
+    Args:
+        raw_html: Original HTML.
+        url: The url of the HTML.
+        base_url: the base URL to check for outside links against.
+        pattern: Regex to use for extracting links from raw HTML.
+        prevent_outside: If `True`, ignore external links which are not children
+            of the base URL.
+        exclude_prefixes: Exclude any URLs that start with one of these prefixes.
+        continue_on_failure: If `True`, continue if parsing a specific link raises an
+            exception. Otherwise, raise the exception.
+
+    Returns:
+        A list of absolute paths to sub links.
+    """
+    base_url_to_use = base_url if base_url is not None else url
+    parsed_base_url = urlparse(base_url_to_use)
+    parsed_url = urlparse(url)
+    all_links = find_all_links(raw_html, pattern=pattern)
+    absolute_paths = set()
+    for link in all_links:
+        try:
+            parsed_link = urlparse(link)
+            # Some may be absolute links like https://to/path
+            if parsed_link.scheme in {"http", "https"}:
+                absolute_path = link
+            # Some may have omitted the protocol like //to/path
+            elif link.startswith("//"):
+                absolute_path = f"{parsed_url.scheme}:{link}"
+            else:
+                absolute_path = urljoin(url, parsed_link.path)
+                if parsed_link.query:
+                    absolute_path += f"?{parsed_link.query}"
+            absolute_paths.add(absolute_path)
+        except Exception as e:
+            if continue_on_failure:
+                logger.warning(
+                    "Unable to load link %s. Raised exception:\n\n%s", link, e
+                )
+                continue
+            raise
+
+    results = []
+    for path in absolute_paths:
+        if any(path.startswith(exclude_prefix) for exclude_prefix in exclude_prefixes):
+            continue
+
+        if prevent_outside:
+            parsed_path = urlparse(path)
+
+            if parsed_base_url.netloc != parsed_path.netloc:
+                continue
+
+            # Will take care of verifying rest of path after netloc
+            # if it's more specific
+            if not path.startswith(base_url_to_use):
+                continue
+
+        results.append(path)
+    return results

coco_core/utils/image.py

@@ -0,0 +1,15 @@
+"""Utilities for image processing."""
+
+from typing import Any
+
+
+def __getattr__(name: str) -> Any:
+    if name in {"encode_image", "image_to_data_url"}:
+        msg = (
+            f"'{name}' has been removed for security reasons.\n\n"
+            f"Usage of this utility in environments with user-input paths is a "
+            f"security vulnerability. Out of an abundance of caution, the utility "
+            f"has been removed to prevent possible misuse."
+        )
+        raise ValueError(msg)
+    raise AttributeError(name)

coco_core/utils/input.py

@@ -0,0 +1,82 @@
+"""Handle chained inputs."""
+
+from typing import TextIO
+
+_TEXT_COLOR_MAPPING = {
+    "blue": "36;1",
+    "yellow": "33;1",
+    "pink": "38;5;200",
+    "green": "32;1",
+    "red": "31;1",
+}
+
+
+def get_color_mapping(
+    items: list[str], excluded_colors: list[str] | None = None
+) -> dict[str, str]:
+    """Get mapping for items to a support color.
+
+    Args:
+        items: The items to map to colors.
+        excluded_colors: The colors to exclude.
+
+    Returns:
+        The mapping of items to colors.
+
+    Raises:
+        ValueError: If no colors are available after applying exclusions.
+    """
+    colors = list(_TEXT_COLOR_MAPPING.keys())
+    if excluded_colors is not None:
+        colors = [c for c in colors if c not in excluded_colors]
+    if not colors:
+        msg = "No colors available after applying exclusions."
+        raise ValueError(msg)
+    return {item: colors[i % len(colors)] for i, item in enumerate(items)}
+
+
+def get_colored_text(text: str, color: str) -> str:
+    """Get colored text.
+
+    Args:
+        text: The text to color.
+        color: The color to use.
+
+    Returns:
+        The colored text.
+    """
+    color_str = _TEXT_COLOR_MAPPING[color]
+    return f"\u001b[{color_str}m\033[1;3m{text}\u001b[0m"
+
+
+def get_bolded_text(text: str) -> str:
+    """Get bolded text.
+
+    Args:
+        text: The text to bold.
+
+    Returns:
+        The bolded text.
+    """
+    return f"\033[1m{text}\033[0m"
+
+
+def print_text(
+    text: str, color: str | None = None, end: str = "", file: TextIO | None = None
+) -> None:
+    """Print text with highlighting and no end characters.
+
+    If a color is provided, the text will be printed in that color.
+
+    If a file is provided, the text will be written to that file.
+
+    Args:
+        text: The text to print.
+        color: The color to use.
+        end: The end character to use.
+        file: The file to write to.
+    """
+    text_to_print = get_colored_text(text, color) if color else text
+    print(text_to_print, end=end, file=file)
+    if file:
+        file.flush()  # ensure all printed content are written to file

coco_core/utils/interactive_env.py

@@ -0,0 +1,12 @@
+"""Utilities for working with interactive environments."""
+
+import sys
+
+
+def is_interactive_env() -> bool:
+    """Determine if running within IPython or Jupyter.
+
+    Returns:
+        `True` if running in an interactive environment, `False` otherwise.
+    """
+    return hasattr(sys, "ps2")

coco_core/utils/iter.py

@@ -0,0 +1,229 @@
+"""Utilities for working with iterators."""
+
+from collections import deque
+from collections.abc import Generator, Iterable, Iterator
+from contextlib import AbstractContextManager
+from itertools import islice
+from types import TracebackType
+from typing import (
+    Any,
+    Generic,
+    Literal,
+    TypeVar,
+    overload,
+)
+
+T = TypeVar("T")
+
+
+class NoLock:
+    """Dummy lock that provides the proper interface but no protection."""
+
+    def __enter__(self) -> None:
+        """Do nothing."""
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> Literal[False]:
+        """Return False (exception not suppressed)."""
+        return False
+
+
+def tee_peer(
+    iterator: Iterator[T],
+    # the buffer specific to this peer
+    buffer: deque[T],
+    # the buffers of all peers, including our own
+    peers: list[deque[T]],
+    lock: AbstractContextManager[Any],
+) -> Generator[T, None, None]:
+    """An individual iterator of a `.tee`.
+
+    This function is a generator that yields items from the shared iterator `iterator`.
+    It buffers items until the least advanced iterator has yielded them as well. The
+    buffer is shared with all other peers.
+
+    Args:
+        iterator: The shared iterator.
+        buffer: The buffer for this peer.
+        peers: The buffers of all peers.
+        lock: The lock to synchronise access to the shared buffers.
+
+    Yields:
+        The next item from the shared iterator.
+    """
+    try:
+        while True:
+            if not buffer:
+                with lock:
+                    # Another peer produced an item while we were waiting for the lock.
+                    # Proceed with the next loop iteration to yield the item.
+                    if buffer:
+                        continue
+                    try:
+                        item = next(iterator)
+                    except StopIteration:
+                        break
+                    else:
+                        # Append to all buffers, including our own. We'll fetch our
+                        # item from the buffer again, instead of yielding it directly.
+                        # This ensures the proper item ordering if any of our peers
+                        # are fetching items concurrently. They may have buffered their
+                        # item already.
+                        for peer_buffer in peers:
+                            peer_buffer.append(item)
+            yield buffer.popleft()
+    finally:
+        with lock:
+            # this peer is done - remove its buffer
+            for idx, peer_buffer in enumerate(peers):  # pragma: no branch
+                if peer_buffer is buffer:
+                    peers.pop(idx)
+                    break
+            # if we are the last peer, try and close the iterator
+            if not peers and hasattr(iterator, "close"):
+                iterator.close()
+
+
+class Tee(Generic[T]):
+    """Create `n` separate asynchronous iterators over `iterable`.
+
+    This splits a single `iterable` into multiple iterators, each providing the same
+    items in the same order.
+
+    All child iterators may advance separately but share the same items from `iterable`
+    -- when the most advanced iterator retrieves an item, it is buffered until the least
+    advanced iterator has yielded it as well. A `tee` works lazily and can handle an
+    infinite `iterable`, provided that all iterators advance.
+
+    ```python
+    async def derivative(sensor_data):
+        previous, current = a.tee(sensor_data, n=2)
+        await a.anext(previous)  # advance one iterator
+        return a.map(operator.sub, previous, current)
+    ```
+
+    Unlike `itertools.tee`, `.tee` returns a custom type instead of a `tuple`. Like a
+    tuple, it can be indexed, iterated and unpacked to get the child iterators. In
+    addition, its `.tee.aclose` method immediately closes all children, and it can be
+    used in an `async with` context for the same effect.
+
+    If `iterable` is an iterator and read elsewhere, `tee` will *not* provide these
+    items. Also, `tee` must internally buffer each item until the last iterator has
+    yielded it; if the most and least advanced iterator differ by most data, using a
+    `list` is more efficient (but not lazy).
+
+    If the underlying iterable is concurrency safe (`anext` may be awaited concurrently)
+    the resulting iterators are concurrency safe as well. Otherwise, the iterators are
+    safe if there is only ever one single "most advanced" iterator. To enforce
+    sequential use of `anext`, provide a `lock`
+
+    - e.g., an `asyncio.Lock` instance in an `asyncio` application - and access is
+        automatically synchronised.
+
+    """
+
+    def __init__(
+        self,
+        iterable: Iterator[T],
+        n: int = 2,
+        *,
+        lock: AbstractContextManager[Any] | None = None,
+    ):
+        """Create a `tee`.
+
+        Args:
+            iterable: The iterable to split.
+            n: The number of iterators to create.
+            lock: The lock to synchronise access to the shared buffers.
+
+        """
+        self._iterator = iter(iterable)
+        self._buffers: list[deque[T]] = [deque() for _ in range(n)]
+        self._children = tuple(
+            tee_peer(
+                iterator=self._iterator,
+                buffer=buffer,
+                peers=self._buffers,
+                lock=lock if lock is not None else NoLock(),
+            )
+            for buffer in self._buffers
+        )
+
+    def __len__(self) -> int:
+        """Return the number of child iterators."""
+        return len(self._children)
+
+    @overload
+    def __getitem__(self, item: int) -> Iterator[T]: ...
+
+    @overload
+    def __getitem__(self, item: slice) -> tuple[Iterator[T], ...]: ...
+
+    def __getitem__(self, item: int | slice) -> Iterator[T] | tuple[Iterator[T], ...]:
+        """Return the child iterator(s) at the given index or slice."""
+        return self._children[item]
+
+    def __iter__(self) -> Iterator[Iterator[T]]:
+        """Return an iterator over the child iterators.
+
+        Yields:
+            The child iterators.
+        """
+        yield from self._children
+
+    def __enter__(self) -> "Tee[T]":
+        """Return `Tee` instance."""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> Literal[False]:
+        """Close all child iterators.
+
+        Returns:
+            `False` (exception not suppressed).
+        """
+        self.close()
+        return False
+
+    def close(self) -> None:
+        """Close all child iterators."""
+        for child in self._children:
+            child.close()
+
+
+# Why this is needed https://stackoverflow.com/a/44638570
+safetee = Tee
+
+
+def batch_iterate(size: int | None, iterable: Iterable[T]) -> Iterator[list[T]]:
+    """Utility batching function.
+
+    Args:
+        size: The size of the batch.
+
+            If `None`, returns a single batch.
+        iterable: The iterable to batch.
+
+    Yields:
+        The batches of the iterable.
+
+    Raises:
+        ValueError: If `size` is not `None` and is not a positive integer.
+    """
+    if size is not None and size <= 0:
+        msg = f"Batch size must be a positive integer, got {size}."
+        raise ValueError(msg)
+    it = iter(iterable)
+    while True:
+        chunk = list(islice(it, size))
+        if not chunk:
+            return
+        yield chunk

coco_core/utils/json.py

@@ -0,0 +1,230 @@
+"""Utilities for JSON."""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import TYPE_CHECKING, Any
+
+from coco_core.exceptions import OutputParserException
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+def _replace_new_line(match: re.Match[str]) -> str:
+    """Replace newline characters in a regex match with escaped sequences.
+
+    Args:
+        match: Regex match object containing the string to process.
+
+    Returns:
+        String with newlines, carriage returns, tabs, and quotes properly escaped.
+    """
+    value = match.group(2)
+    value = re.sub(r"\n", r"\\n", value)
+    value = re.sub(r"\r", r"\\r", value)
+    value = re.sub(r"\t", r"\\t", value)
+    value = re.sub(r'(?<!\\)"', r"\"", value)
+
+    return match.group(1) + value + match.group(3)
+
+
+def _custom_parser(multiline_string: str | bytes | bytearray) -> str:
+    r"""Custom parser for multiline strings.
+
+    The LLM response for `action_input` may be a multiline string containing unescaped
+    newlines, tabs or quotes. This function replaces those characters with their escaped
+    counterparts. (newlines in JSON must be double-escaped: `\\n`).
+
+    Returns:
+        The modified string with escaped newlines, tabs and quotes.
+    """
+    if isinstance(multiline_string, (bytes, bytearray)):
+        multiline_string = multiline_string.decode()
+
+    return re.sub(
+        r'("action_input"\:\s*")(.*?)(")',
+        _replace_new_line,
+        multiline_string,
+        flags=re.DOTALL,
+    )
+
+
+# Adapted from https://github.com/KillianLucas/open-interpreter/blob/5b6080fae1f8c68938a1e4fa8667e3744084ee21/interpreter/utils/parse_partial_json.py
+# MIT License
+
+
+def parse_partial_json(s: str, *, strict: bool = False) -> Any:
+    """Parse a JSON string that may be missing closing braces.
+
+    Args:
+        s: The JSON string to parse.
+        strict: Whether to use strict parsing.
+
+    Returns:
+        The parsed JSON object as a Python dictionary.
+    """
+    # Attempt to parse the string as-is.
+    try:
+        return json.loads(s, strict=strict)
+    except json.JSONDecodeError:
+        pass
+
+    # Initialize variables.
+    new_chars = []
+    stack = []
+    is_inside_string = False
+    escaped = False
+
+    # Process each character in the string one at a time.
+    for char in s:
+        new_char = char
+        if is_inside_string:
+            if char == '"' and not escaped:
+                is_inside_string = False
+            elif char == "\n" and not escaped:
+                new_char = (
+                    "\\n"  # Replace the newline character with the escape sequence.
+                )
+            elif char == "\\":
+                escaped = not escaped
+            else:
+                escaped = False
+        elif char == '"':
+            is_inside_string = True
+            escaped = False
+        elif char == "{":
+            stack.append("}")
+        elif char == "[":
+            stack.append("]")
+        elif char in {"}", "]"}:
+            if stack and stack[-1] == char:
+                stack.pop()
+            else:
+                # Mismatched closing character; the input is malformed.
+                return None
+
+        # Append the processed character to the new string.
+        new_chars.append(new_char)
+
+    # If we're still inside a string at the end of processing,
+    # we need to close the string.
+    if is_inside_string:
+        if escaped:  # Remove unterminated escape character
+            new_chars.pop()
+        new_chars.append('"')
+
+    # Reverse the stack to get the closing characters.
+    stack.reverse()
+
+    # Try to parse mods of string until we succeed or run out of characters.
+    while new_chars:
+        # Close any remaining open structures in the reverse
+        # order that they were opened.
+        # Attempt to parse the modified string as JSON.
+        try:
+            return json.loads("".join(new_chars + stack), strict=strict)
+        except json.JSONDecodeError:
+            # If we still can't parse the string as JSON,
+            # try removing the last character
+            new_chars.pop()
+
+    # If we got here, we ran out of characters to remove
+    # and still couldn't parse the string as JSON, so return the parse error
+    # for the original string.
+    return json.loads(s, strict=strict)
+
+
+_json_markdown_re = re.compile(r"```(json)?(.*)", re.DOTALL)
+
+
+def parse_json_markdown(
+    json_string: str, *, parser: Callable[[str], Any] = parse_partial_json
+) -> Any:
+    """Parse a JSON string from a Markdown string.
+
+    Args:
+        json_string: The Markdown string.
+        parser: The parser to use.
+
+    Returns:
+        The parsed JSON object as a Python dictionary.
+    """
+    try:
+        return _parse_json(json_string, parser=parser)
+    except json.JSONDecodeError:
+        # Try to find JSON string within triple backticks
+        match = _json_markdown_re.search(json_string)
+
+        # If no match found, assume the entire string is a JSON string
+        # Else, use the content within the backticks
+        json_str = json_string if match is None else match.group(2)
+    return _parse_json(json_str, parser=parser)
+
+
+_json_strip_chars = " \n\r\t`"
+
+
+def _parse_json(
+    json_str: str, *, parser: Callable[[str], Any] = parse_partial_json
+) -> Any:
+    """Parse a JSON string, handling special characters and whitespace.
+
+    Strips whitespace, newlines, and backticks from the start and end of the string,
+    then processes special characters before parsing.
+
+    Args:
+        json_str: The JSON string to parse.
+        parser: Optional custom parser function.
+
+    Returns:
+        Parsed JSON object.
+    """
+    # Strip whitespace,newlines,backtick from the start and end
+    json_str = json_str.strip(_json_strip_chars)
+
+    # handle newlines and other special characters inside the returned value
+    json_str = _custom_parser(json_str)
+
+    # Parse the JSON string into a Python dictionary
+    return parser(json_str)
+
+
+def parse_and_check_json_markdown(
+    text: str, expected_keys: list[str]
+) -> dict[str, Any]:
+    """Parse and check a JSON string from a Markdown string.
+
+    Checks that it contains the expected keys.
+
+    Args:
+        text: The Markdown string.
+        expected_keys: The expected keys in the JSON string.
+
+    Returns:
+        The parsed JSON object as a Python dictionary.
+
+    Raises:
+        OutputParserException: If the JSON string is invalid or does not contain
+            the expected keys.
+    """
+    try:
+        json_obj = parse_json_markdown(text)
+    except json.JSONDecodeError as e:
+        msg = f"Got invalid JSON object. Error: {e}"
+        raise OutputParserException(msg) from e
+    if not isinstance(json_obj, dict):
+        error_message = (
+            f"Expected JSON object (dict), but got: {type(json_obj).__name__}. "
+        )
+        raise OutputParserException(error_message, llm_output=text)
+
+    for key in expected_keys:
+        if key not in json_obj:
+            msg = (
+                f"Got invalid return object. Expected key `{key}` "
+                f"to be present, but got {json_obj}"
+            )
+            raise OutputParserException(msg)
+    return json_obj