coco-framework 0.1.0__py3-none-any.whl

coco_core/utils/usage.py

@@ -0,0 +1,61 @@
+"""Usage utilities."""
+
+from collections.abc import Callable
+from typing import Any
+
+
+def _dict_int_op(
+    left: dict[str, Any],
+    right: dict[str, Any],
+    op: Callable[[int, int], int],
+    *,
+    default: int = 0,
+    depth: int = 0,
+    max_depth: int = 100,
+) -> dict[str, Any]:
+    """Apply an integer operation to corresponding values in two dictionaries.
+
+    Recursively combines two dictionaries by applying the given operation to integer
+    values at matching keys.
+
+    Supports nested dictionaries.
+
+    Args:
+        left: First dictionary to combine.
+        right: Second dictionary to combine.
+        op: Binary operation function to apply to integer values.
+        default: Default value to use when a key is missing from a dictionary.
+        depth: Current recursion depth (used internally).
+        max_depth: Maximum recursion depth (to prevent infinite loops).
+
+    Returns:
+        A new dictionary with combined values.
+
+    Raises:
+        ValueError: If `max_depth` is exceeded or if value types are not supported.
+    """
+    if depth >= max_depth:
+        msg = f"{max_depth=} exceeded, unable to combine dicts."
+        raise ValueError(msg)
+    combined: dict[str, Any] = {}
+    for k in set(left).union(right):
+        if isinstance(left.get(k, default), int) and isinstance(
+            right.get(k, default), int
+        ):
+            combined[k] = op(left.get(k, default), right.get(k, default))
+        elif isinstance(left.get(k, {}), dict) and isinstance(right.get(k, {}), dict):
+            combined[k] = _dict_int_op(
+                left.get(k, {}),
+                right.get(k, {}),
+                op,
+                default=default,
+                depth=depth + 1,
+                max_depth=max_depth,
+            )
+        else:
+            types = [type(d[k]) for d in (left, right) if k in d]
+            msg = (
+                f"Unknown value types: {types}. Only dict and int values are supported."
+            )
+            raise ValueError(msg)  # noqa: TRY004
+    return combined

coco_core/utils/utils.py

@@ -0,0 +1,528 @@
+"""Generic utility functions."""
+
+import contextlib
+import datetime
+import functools
+import importlib
+import os
+import warnings
+from collections.abc import Callable, Iterator, Sequence
+from importlib.metadata import version
+from typing import Any, overload
+from uuid import uuid4
+
+from packaging.version import parse
+from pydantic import SecretStr
+from requests import HTTPError, Response
+from typing_extensions import override
+
+from coco_core.utils.pydantic import (
+    is_pydantic_v1_subclass,
+)
+
+
+def xor_args(*arg_groups: tuple[str, ...]) -> Callable[..., Any]:
+    """Validate specified keyword args are mutually exclusive.
+
+    Args:
+        *arg_groups: Groups of mutually exclusive keyword args.
+
+    Returns:
+        Decorator that validates the specified keyword args are mutually exclusive.
+    """
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            """Validate exactly one arg in each group is not None."""
+            counts = [
+                sum(1 for arg in arg_group if kwargs.get(arg) is not None)
+                for arg_group in arg_groups
+            ]
+            invalid_groups = [i for i, count in enumerate(counts) if count != 1]
+            if invalid_groups:
+                invalid_group_names = [", ".join(arg_groups[i]) for i in invalid_groups]
+                msg = (
+                    "Exactly one argument in each of the following"
+                    " groups must be defined:"
+                    f" {', '.join(invalid_group_names)}"
+                )
+                raise ValueError(msg)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def raise_for_status_with_text(response: Response) -> None:
+    """Raise an error with the response text.
+
+    Args:
+        response: The response to check for errors.
+
+    Raises:
+        ValueError: If the response has an error status code.
+    """
+    try:
+        response.raise_for_status()
+    except HTTPError as e:
+        raise ValueError(response.text) from e
+
+
+@contextlib.contextmanager
+def mock_now(dt_value: datetime.datetime) -> Iterator[type]:
+    """Context manager for mocking out datetime.now() in unit tests.
+
+    Args:
+        dt_value: The datetime value to use for datetime.now().
+
+    Yields:
+        The mocked datetime class.
+
+    Example:
+        ```python
+        with mock_now(datetime.datetime(2011, 2, 3, 10, 11)):
+            assert datetime.datetime.now() == datetime.datetime(2011, 2, 3, 10, 11)
+        ```
+    """
+
+    class MockDateTime(datetime.datetime):
+        """Mock datetime.datetime.now() with a fixed datetime."""
+
+        @classmethod
+        @override
+        def now(cls, tz: datetime.tzinfo | None = None) -> "MockDateTime":
+            # Create a copy of dt_value.
+            return MockDateTime(
+                dt_value.year,
+                dt_value.month,
+                dt_value.day,
+                dt_value.hour,
+                dt_value.minute,
+                dt_value.second,
+                dt_value.microsecond,
+                dt_value.tzinfo,
+            )
+
+    real_datetime = datetime.datetime
+    datetime.datetime = MockDateTime  # type: ignore[misc]
+    try:
+        yield datetime.datetime
+    finally:
+        datetime.datetime = real_datetime  # type: ignore[misc]
+
+
+def guard_import(
+    module_name: str, *, pip_name: str | None = None, package: str | None = None
+) -> Any:
+    """Dynamically import a module.
+
+    Raise an exception if the module is not installed.
+
+    Args:
+        module_name: The name of the module to import.
+        pip_name: The name of the module to install with pip.
+        package: The package to import the module from.
+
+    Returns:
+        The imported module.
+
+    Raises:
+        ImportError: If the module is not installed.
+    """
+    try:
+        module = importlib.import_module(module_name, package)
+    except (ImportError, ModuleNotFoundError) as e:
+        top_level = module_name.split(".", maxsplit=1)[0]
+        # Coco modules ship in trustcoco distributions: coco_core in
+        # trustcoco-core, coco in trustcoco. Any other module keeps the
+        # underscore to hyphen rule.
+        default = top_level.replace("_", "-")
+        if top_level == "coco" or top_level.startswith("coco_"):
+            default = "trust" + default
+        pip_name = pip_name or default
+        msg = (
+            f"Could not import {module_name} python package. "
+            f"Please install it with `pip install {pip_name}`."
+        )
+        raise ImportError(msg) from e
+    return module
+
+
+def check_package_version(
+    package: str,
+    lt_version: str | None = None,
+    lte_version: str | None = None,
+    gt_version: str | None = None,
+    gte_version: str | None = None,
+) -> None:
+    """Check the version of a package.
+
+    Args:
+        package: The name of the package.
+        lt_version: The version must be less than this.
+        lte_version: The version must be less than or equal to this.
+        gt_version: The version must be greater than this.
+        gte_version: The version must be greater than or equal to this.
+
+
+    Raises:
+        ValueError: If the package version does not meet the requirements.
+    """
+    imported_version = parse(version(package))
+    if lt_version is not None and imported_version >= parse(lt_version):
+        msg = (
+            f"Expected {package} version to be < {lt_version}. Received "
+            f"{imported_version}."
+        )
+        raise ValueError(msg)
+    if lte_version is not None and imported_version > parse(lte_version):
+        msg = (
+            f"Expected {package} version to be <= {lte_version}. Received "
+            f"{imported_version}."
+        )
+        raise ValueError(msg)
+    if gt_version is not None and imported_version <= parse(gt_version):
+        msg = (
+            f"Expected {package} version to be > {gt_version}. Received "
+            f"{imported_version}."
+        )
+        raise ValueError(msg)
+    if gte_version is not None and imported_version < parse(gte_version):
+        msg = (
+            f"Expected {package} version to be >= {gte_version}. Received "
+            f"{imported_version}."
+        )
+        raise ValueError(msg)
+
+
+def get_pydantic_field_names(pydantic_cls: Any) -> set[str]:
+    """Get field names, including aliases, for a pydantic class.
+
+    Args:
+        pydantic_cls: Pydantic class.
+
+    Returns:
+        Field names.
+    """
+    all_required_field_names = set()
+    if is_pydantic_v1_subclass(pydantic_cls):
+        for field in pydantic_cls.__fields__.values():
+            all_required_field_names.add(field.name)
+            if field.has_alias:
+                all_required_field_names.add(field.alias)
+    else:  # Assuming pydantic 2 for now
+        for name, field in pydantic_cls.model_fields.items():
+            all_required_field_names.add(name)
+            if field.alias:
+                all_required_field_names.add(field.alias)
+    return all_required_field_names
+
+
+def _build_model_kwargs(
+    values: dict[str, Any],
+    all_required_field_names: set[str],
+) -> dict[str, Any]:
+    """Build `model_kwargs` param from Pydantic constructor values.
+
+    Args:
+        values: All init args passed in by user.
+        all_required_field_names: All required field names for the pydantic class.
+
+    Returns:
+        Extra kwargs.
+
+    Raises:
+        ValueError: If a field is specified in both `values` and `extra_kwargs`.
+        ValueError: If a field is specified in `model_kwargs`.
+    """
+    extra_kwargs = values.get("model_kwargs", {})
+    for field_name in list(values):
+        if field_name in extra_kwargs:
+            msg = f"Found {field_name} supplied twice."
+            raise ValueError(msg)
+        if field_name not in all_required_field_names:
+            warnings.warn(
+                f"""WARNING! {field_name} is not default parameter.
+                {field_name} was transferred to model_kwargs.
+                Please confirm that {field_name} is what you intended.""",
+                stacklevel=7,
+            )
+            extra_kwargs[field_name] = values.pop(field_name)
+
+    invalid_model_kwargs = all_required_field_names.intersection(extra_kwargs.keys())
+    if invalid_model_kwargs:
+        warnings.warn(
+            f"Parameters {invalid_model_kwargs} should be specified explicitly. "
+            f"Instead they were passed in as part of `model_kwargs` parameter.",
+            stacklevel=7,
+        )
+        for k in invalid_model_kwargs:
+            values[k] = extra_kwargs.pop(k)
+
+    values["model_kwargs"] = extra_kwargs
+    return values
+
+
+# DON'T USE! Kept for backwards-compatibility but should never have been public.
+def build_extra_kwargs(
+    extra_kwargs: dict[str, Any],
+    values: dict[str, Any],
+    all_required_field_names: set[str],
+) -> dict[str, Any]:
+    """Build extra kwargs from values and extra_kwargs.
+
+    !!! danger "DON'T USE"
+
+        Kept for backwards-compatibility but should never have been public. Use the
+        internal `_build_model_kwargs` function instead.
+
+    Args:
+        extra_kwargs: Extra kwargs passed in by user.
+        values: Values passed in by user.
+        all_required_field_names: All required field names for the pydantic class.
+
+    Returns:
+        Extra kwargs.
+
+    Raises:
+        ValueError: If a field is specified in both `values` and `extra_kwargs`.
+        ValueError: If a field is specified in `model_kwargs`.
+    """
+    # DON'T USE! Kept for backwards-compatibility but should never have been public.
+    for field_name in list(values):
+        if field_name in extra_kwargs:
+            msg = f"Found {field_name} supplied twice."
+            raise ValueError(msg)
+        if field_name not in all_required_field_names:
+            warnings.warn(
+                f"""WARNING! {field_name} is not default parameter.
+                {field_name} was transferred to model_kwargs.
+                Please confirm that {field_name} is what you intended.""",
+                stacklevel=7,
+            )
+            extra_kwargs[field_name] = values.pop(field_name)
+
+    # DON'T USE! Kept for backwards-compatibility but should never have been public.
+    invalid_model_kwargs = all_required_field_names.intersection(extra_kwargs.keys())
+    if invalid_model_kwargs:
+        msg = (
+            f"Parameters {invalid_model_kwargs} should be specified explicitly. "
+            f"Instead they were passed in as part of `model_kwargs` parameter."
+        )
+        raise ValueError(msg)
+
+    # DON'T USE! Kept for backwards-compatibility but should never have been public.
+    return extra_kwargs
+
+
+def convert_to_secret_str(value: SecretStr | str) -> SecretStr:
+    """Convert a string to a `SecretStr` if needed.
+
+    Args:
+        value: The value to convert.
+
+    Returns:
+        The `SecretStr` value.
+    """
+    if isinstance(value, SecretStr):
+        return value
+    return SecretStr(value)
+
+
+class _NoDefaultType:
+    """Type to indicate no default value is provided."""
+
+
+_NoDefault = _NoDefaultType()
+
+
+@overload
+def from_env(key: str, /) -> Callable[[], str]: ...
+
+
+@overload
+def from_env(key: str, /, *, default: str) -> Callable[[], str]: ...
+
+
+@overload
+def from_env(key: Sequence[str], /, *, default: str) -> Callable[[], str]: ...
+
+
+@overload
+def from_env(key: str, /, *, error_message: str) -> Callable[[], str]: ...
+
+
+@overload
+def from_env(
+    key: str | Sequence[str], /, *, default: str, error_message: str | None
+) -> Callable[[], str]: ...
+
+
+@overload
+def from_env(
+    key: str, /, *, default: None, error_message: str | None
+) -> Callable[[], str | None]: ...
+
+
+@overload
+def from_env(
+    key: str | Sequence[str], /, *, default: None
+) -> Callable[[], str | None]: ...
+
+
+def from_env(
+    key: str | Sequence[str],
+    /,
+    *,
+    default: str | _NoDefaultType | None = _NoDefault,
+    error_message: str | None = None,
+) -> Callable[[], str] | Callable[[], str | None]:
+    """Create a factory method that gets a value from an environment variable.
+
+    Args:
+        key: The environment variable to look up.
+
+            If a list of keys is provided, the first key found in the environment will
+            be used. If no key is found, the default value will be used if set,
+            otherwise an error will be raised.
+        default: The default value to return if the environment variable is not set.
+        error_message: The error message which will be raised if the key is not found
+            and no default value is provided.
+
+            This will be raised as a ValueError.
+
+    Returns:
+        Factory method that will look up the value from the environment.
+    """
+
+    def get_from_env_fn() -> str | None:
+        """Get a value from an environment variable.
+
+        Raises:
+            ValueError: If the environment variable is not set and no default is
+                provided.
+
+        Returns:
+            The value from the environment.
+        """
+        if isinstance(key, (list, tuple)):
+            for k in key:
+                if k in os.environ:
+                    return os.environ[k]
+        if isinstance(key, str) and key in os.environ:
+            return os.environ[key]
+
+        if isinstance(default, (str, type(None))):
+            return default
+        if error_message:
+            raise ValueError(error_message)
+        msg = (
+            f"Did not find {key}, please add an environment variable"
+            f" `{key}` which contains it, or pass"
+            f" `{key}` as a named parameter."
+        )
+        raise ValueError(msg)
+
+    return get_from_env_fn
+
+
+@overload
+def secret_from_env(key: str | Sequence[str], /) -> Callable[[], SecretStr]: ...
+
+
+@overload
+def secret_from_env(key: str, /, *, default: str) -> Callable[[], SecretStr]: ...
+
+
+@overload
+def secret_from_env(
+    key: str | Sequence[str], /, *, default: None
+) -> Callable[[], SecretStr | None]: ...
+
+
+@overload
+def secret_from_env(key: str, /, *, error_message: str) -> Callable[[], SecretStr]: ...
+
+
+def secret_from_env(
+    key: str | Sequence[str],
+    /,
+    *,
+    default: str | _NoDefaultType | None = _NoDefault,
+    error_message: str | None = None,
+) -> Callable[[], SecretStr | None] | Callable[[], SecretStr]:
+    """Secret from env.
+
+    Args:
+        key: The environment variable to look up.
+        default: The default value to return if the environment variable is not set.
+        error_message: The error message which will be raised if the key is not found
+            and no default value is provided.
+
+            This will be raised as a `ValueError`.
+
+    Returns:
+        Factory method that will look up the secret from the environment.
+    """
+
+    def get_secret_from_env() -> SecretStr | None:
+        """Get a value from an environment variable.
+
+        Raises:
+            ValueError: If the environment variable is not set and no default is
+                provided.
+
+        Returns:
+            The secret from the environment.
+        """
+        if isinstance(key, (list, tuple)):
+            for k in key:
+                if k in os.environ:
+                    return SecretStr(os.environ[k])
+        if isinstance(key, str) and key in os.environ:
+            return SecretStr(os.environ[key])
+        if isinstance(default, str):
+            return SecretStr(default)
+        if default is None:
+            return None
+        if error_message:
+            raise ValueError(error_message)
+        msg = (
+            f"Did not find {key}, please add an environment variable"
+            f" `{key}` which contains it, or pass"
+            f" `{key}` as a named parameter."
+        )
+        raise ValueError(msg)
+
+    return get_secret_from_env
+
+
+LC_AUTO_PREFIX = "lc_"
+"""Coco auto-generated ID prefix for messages and content blocks."""
+
+LC_ID_PREFIX = "lc_run-"
+"""Internal tracing/callback system identifier.
+
+Used for:
+
+- Tracing. Every Coco operation (LLM call, chain execution, tool use, etc.)
+    gets a unique run_id (UUID)
+- Enables tracking parent-child relationships between operations
+"""
+
+
+def ensure_id(id_val: str | None) -> str:
+    """Ensure the ID is a valid string, generating a new UUID if not provided.
+
+    Auto-generated UUIDs are prefixed by `'lc_'` to indicate they are
+    Coco-generated IDs.
+
+    Args:
+        id_val: Optional string ID value to validate.
+
+    Returns:
+        A string ID, either the validated provided value or a newly generated UUID4.
+    """
+    return id_val or f"{LC_AUTO_PREFIX}{uuid4()}"

coco_core/utils/uuid.py

@@ -0,0 +1,57 @@
+"""UUID utility functions.
+
+This module exports a uuid7 function to generate monotonic, time-ordered UUIDs
+for tracing and similar operations.
+"""
+
+from __future__ import annotations
+
+import typing
+from uuid import UUID
+
+from uuid_utils.compat import uuid7 as _uuid_utils_uuid7
+
+if typing.TYPE_CHECKING:
+    from uuid import UUID
+
+_NANOS_PER_SECOND: typing.Final = 1_000_000_000
+
+
+def _to_timestamp_and_nanos(nanoseconds: int) -> tuple[int, int]:
+    """Split a nanosecond timestamp into seconds and remaining nanoseconds."""
+    seconds, nanos = divmod(nanoseconds, _NANOS_PER_SECOND)
+    return seconds, nanos
+
+
+def uuid7(nanoseconds: int | None = None) -> UUID:
+    """Generate a UUID from a Unix timestamp in nanoseconds and random bits.
+
+    UUIDv7 objects feature monotonicity within a millisecond.
+
+    Args:
+        nanoseconds: Optional ns timestamp. If not provided, uses current time.
+
+    Returns:
+        A UUIDv7 object.
+    """
+    # --- 48 ---   -- 4 --   --- 12 ---   -- 2 --   --- 30 ---   - 32 -
+    # unix_ts_ms | version | counter_hi | variant | counter_lo | random
+    #
+    # 'counter = counter_hi | counter_lo' is a 42-bit counter constructed
+    # with Method 1 of RFC 9562, §6.2, and its MSB is set to 0.
+    #
+    # 'random' is a 32-bit random value regenerated for every new UUID.
+    #
+    # If multiple UUIDs are generated within the same millisecond, the LSB
+    # of 'counter' is incremented by 1. When overflowing, the timestamp is
+    # advanced and the counter is reset to a random 42-bit integer with MSB
+    # set to 0.
+
+    # For now, just delegate to the uuid_utils implementation
+    if nanoseconds is None:
+        return _uuid_utils_uuid7()
+    seconds, nanos = _to_timestamp_and_nanos(nanoseconds)
+    return _uuid_utils_uuid7(timestamp=seconds, nanos=nanos)
+
+
+__all__ = ["uuid7"]

coco_core/vectorstores/__init__.py

@@ -0,0 +1,53 @@
+"""Vector stores."""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.vectorstores.base import VST, VectorStore, VectorStoreRetriever
+    from coco_core.vectorstores.in_memory import InMemoryVectorStore
+
+__all__ = (
+    "VST",
+    "InMemoryVectorStore",
+    "VectorStore",
+    "VectorStoreRetriever",
+)
+
+_dynamic_imports = {
+    "VectorStore": "base",
+    "VST": "base",
+    "VectorStoreRetriever": "base",
+    "InMemoryVectorStore": "in_memory",
+}
+
+
+def __getattr__(attr_name: str) -> object:
+    """Dynamically import and return an attribute from a submodule.
+
+    This function enables lazy loading of vectorstore classes from submodules, reducing
+    initial import time and circular dependency issues.
+
+    Args:
+        attr_name: Name of the attribute to import.
+
+    Returns:
+        The imported attribute object.
+
+    Raises:
+        AttributeError: If the attribute is not found in `_dynamic_imports`.
+    """
+    module_name = _dynamic_imports.get(attr_name)
+    result = import_attr(attr_name, module_name, __spec__.parent)
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    """Return a list of available attributes for this module.
+
+    Returns:
+        List of attribute names that can be imported from this module.
+    """
+    return list(__all__)