coco-framework 0.1.0__py3-none-any.whl

coco_core/stores.py

@@ -0,0 +1,291 @@
+"""**Store** implements the key-value stores and storage helpers.
+
+Module provides implementations of various key-value stores that conform
+to a simple key-value interface.
+
+The primary goal of these storages is to support implementation of caching.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Iterator, Sequence
+from typing import (
+    Any,
+    Generic,
+    TypeVar,
+)
+
+from typing_extensions import override
+
+from coco_core.exceptions import CocoException
+from coco_core.runnables import run_in_executor
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class BaseStore(ABC, Generic[K, V]):
+    """Abstract interface for a key-value store.
+
+    This is an interface that's meant to abstract away the details of different
+    key-value stores. It provides a simple interface for getting, setting, and deleting
+    key-value pairs.
+
+    The basic methods are `mget`, `mset`, and `mdelete` for getting, setting, and
+    deleting multiple key-value pairs at once. The `yield_keys` method is used to
+    iterate over keys that match a given prefix.
+
+    The async versions of these methods are also provided, which are meant to be used in
+    async contexts. The async methods are named with an `a` prefix, e.g., `amget`,
+    `amset`, `amdelete`, and `ayield_keys`.
+
+    By default, the `amget`, `amset`, `amdelete`, and `ayield_keys` methods are
+    implemented using the synchronous methods. If the store can natively support async
+    operations, it should override these methods.
+
+    By design the methods only accept batches of keys and values, and not single keys or
+    values. This is done to force user code to work with batches which will usually be
+    more efficient by saving on round trips to the store.
+
+    Examples:
+        ```python
+        from coco.storage import BaseStore
+
+
+        class MyInMemoryStore(BaseStore[str, int]):
+            def __init__(self) -> None:
+                self.store: dict[str, int] = {}
+
+            def mget(self, keys: Sequence[str]) -> list[int | None]:
+                return [self.store.get(key) for key in keys]
+
+            def mset(self, key_value_pairs: Sequence[tuple[str, int]]) -> None:
+                for key, value in key_value_pairs:
+                    self.store[key] = value
+
+            def mdelete(self, keys: Sequence[str]) -> None:
+                for key in keys:
+                    if key in self.store:
+                        del self.store[key]
+
+            def yield_keys(self, prefix: str | None = None) -> Iterator[str]:
+                if prefix is None:
+                    yield from self.store.keys()
+                else:
+                    for key in self.store.keys():
+                        if key.startswith(prefix):
+                            yield key
+        ```
+    """
+
+    @abstractmethod
+    def mget(self, keys: Sequence[K]) -> list[V | None]:
+        """Get the values associated with the given keys.
+
+        Args:
+            keys: A sequence of keys.
+
+        Returns:
+            A sequence of optional values associated with the keys.
+                If a key is not found, the corresponding value will be `None`.
+        """
+
+    async def amget(self, keys: Sequence[K]) -> list[V | None]:
+        """Async get the values associated with the given keys.
+
+        Args:
+            keys: A sequence of keys.
+
+        Returns:
+            A sequence of optional values associated with the keys.
+                If a key is not found, the corresponding value will be `None`.
+        """
+        return await run_in_executor(None, self.mget, keys)
+
+    @abstractmethod
+    def mset(self, key_value_pairs: Sequence[tuple[K, V]]) -> None:
+        """Set the values for the given keys.
+
+        Args:
+            key_value_pairs: A sequence of key-value pairs.
+        """
+
+    async def amset(self, key_value_pairs: Sequence[tuple[K, V]]) -> None:
+        """Async set the values for the given keys.
+
+        Args:
+            key_value_pairs: A sequence of key-value pairs.
+        """
+        return await run_in_executor(None, self.mset, key_value_pairs)
+
+    @abstractmethod
+    def mdelete(self, keys: Sequence[K]) -> None:
+        """Delete the given keys and their associated values.
+
+        Args:
+            keys: A sequence of keys to delete.
+        """
+
+    async def amdelete(self, keys: Sequence[K]) -> None:
+        """Async delete the given keys and their associated values.
+
+        Args:
+            keys: A sequence of keys to delete.
+        """
+        return await run_in_executor(None, self.mdelete, keys)
+
+    @abstractmethod
+    def yield_keys(self, *, prefix: str | None = None) -> Iterator[K] | Iterator[str]:
+        """Get an iterator over keys that match the given prefix.
+
+        Args:
+            prefix: The prefix to match.
+
+        Yields:
+            An iterator over keys that match the given prefix.
+
+                This method is allowed to return an iterator over either K or str
+                depending on what makes more sense for the given store.
+        """
+
+    async def ayield_keys(
+        self, *, prefix: str | None = None
+    ) -> AsyncIterator[K] | AsyncIterator[str]:
+        """Async get an iterator over keys that match the given prefix.
+
+        Args:
+            prefix: The prefix to match.
+
+        Yields:
+            The keys that match the given prefix.
+
+                This method is allowed to return an iterator over either K or str
+                depending on what makes more sense for the given store.
+        """
+        iterator = await run_in_executor(None, self.yield_keys, prefix=prefix)
+        done = object()
+        while True:
+            item = await run_in_executor(None, lambda it: next(it, done), iterator)
+            if item is done:
+                break
+            yield item  # type: ignore[misc]
+
+
+ByteStore = BaseStore[str, bytes]
+
+
+class InMemoryBaseStore(BaseStore[str, V], Generic[V]):
+    """In-memory implementation of the `BaseStore` using a dictionary."""
+
+    def __init__(self) -> None:
+        """Initialize an empty store."""
+        self.store: dict[str, V] = {}
+
+    @override
+    def mget(self, keys: Sequence[str]) -> list[V | None]:
+        return [self.store.get(key) for key in keys]
+
+    @override
+    async def amget(self, keys: Sequence[str]) -> list[V | None]:
+        return self.mget(keys)
+
+    @override
+    def mset(self, key_value_pairs: Sequence[tuple[str, V]]) -> None:
+        for key, value in key_value_pairs:
+            self.store[key] = value
+
+    @override
+    async def amset(self, key_value_pairs: Sequence[tuple[str, V]]) -> None:
+        return self.mset(key_value_pairs)
+
+    @override
+    def mdelete(self, keys: Sequence[str]) -> None:
+        for key in keys:
+            if key in self.store:
+                del self.store[key]
+
+    @override
+    async def amdelete(self, keys: Sequence[str]) -> None:
+        self.mdelete(keys)
+
+    def yield_keys(self, *, prefix: str | None = None) -> Iterator[str]:
+        """Get an iterator over keys that match the given prefix.
+
+        Args:
+            prefix: The prefix to match.
+
+        Yields:
+            The keys that match the given prefix.
+        """
+        if prefix is None:
+            yield from self.store.keys()
+        else:
+            for key in self.store:
+                if key.startswith(prefix):
+                    yield key
+
+    async def ayield_keys(self, *, prefix: str | None = None) -> AsyncIterator[str]:
+        """Async get an async iterator over keys that match the given prefix.
+
+        Args:
+            prefix: The prefix to match.
+
+        Yields:
+            The keys that match the given prefix.
+        """
+        if prefix is None:
+            for key in self.store:
+                yield key
+        else:
+            for key in self.store:
+                if key.startswith(prefix):
+                    yield key
+
+
+class InMemoryStore(InMemoryBaseStore[Any]):
+    """In-memory store for any type of data.
+
+    Attributes:
+        store: The underlying dictionary that stores the key-value pairs.
+
+    Examples:
+        ```python
+        from coco.storage import InMemoryStore
+
+        store = InMemoryStore()
+        store.mset([("key1", "value1"), ("key2", "value2")])
+        store.mget(["key1", "key2"])
+        # ['value1', 'value2']
+        store.mdelete(["key1"])
+        list(store.yield_keys())
+        # ['key2']
+        list(store.yield_keys(prefix="k"))
+        # ['key2']
+        ```
+    """
+
+
+class InMemoryByteStore(InMemoryBaseStore[bytes]):
+    """In-memory store for bytes.
+
+    Attributes:
+        store: The underlying dictionary that stores the key-value pairs.
+
+    Examples:
+        ```python
+        from coco.storage import InMemoryByteStore
+
+        store = InMemoryByteStore()
+        store.mset([("key1", b"value1"), ("key2", b"value2")])
+        store.mget(["key1", "key2"])
+        # [b'value1', b'value2']
+        store.mdelete(["key1"])
+        list(store.yield_keys())
+        # ['key2']
+        list(store.yield_keys(prefix="k"))
+        # ['key2']
+        ```
+    """
+
+
+class InvalidKeyException(CocoException):
+    """Raised when a key is invalid; e.g., uses incorrect characters."""

coco_core/structured_query.py

@@ -0,0 +1,203 @@
+"""Internal representation of a structured query language."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+class Visitor(ABC):
+    """Defines interface for IR translation using a visitor pattern."""
+
+    allowed_comparators: Sequence[Comparator] | None = None
+    """Allowed comparators for the visitor."""
+
+    allowed_operators: Sequence[Operator] | None = None
+    """Allowed operators for the visitor."""
+
+    def _validate_func(self, func: Operator | Comparator) -> None:
+        if (
+            isinstance(func, Operator)
+            and self.allowed_operators is not None
+            and func not in self.allowed_operators
+        ):
+            msg = (
+                f"Received disallowed operator {func}. Allowed "
+                f"comparators are {self.allowed_operators}"
+            )
+            raise ValueError(msg)
+        if (
+            isinstance(func, Comparator)
+            and self.allowed_comparators is not None
+            and func not in self.allowed_comparators
+        ):
+            msg = (
+                f"Received disallowed comparator {func}. Allowed "
+                f"comparators are {self.allowed_comparators}"
+            )
+            raise ValueError(msg)
+
+    @abstractmethod
+    def visit_operation(self, operation: Operation) -> Any:
+        """Translate an Operation.
+
+        Args:
+            operation: Operation to translate.
+        """
+
+    @abstractmethod
+    def visit_comparison(self, comparison: Comparison) -> Any:
+        """Translate a Comparison.
+
+        Args:
+            comparison: Comparison to translate.
+        """
+
+    @abstractmethod
+    def visit_structured_query(self, structured_query: StructuredQuery) -> Any:
+        """Translate a StructuredQuery.
+
+        Args:
+            structured_query: StructuredQuery to translate.
+        """
+
+
+def _to_snake_case(name: str) -> str:
+    """Convert a name into snake_case."""
+    snake_case = ""
+    for i, char in enumerate(name):
+        if char.isupper() and i != 0:
+            snake_case += "_" + char.lower()
+        else:
+            snake_case += char.lower()
+    return snake_case
+
+
+class Expr(BaseModel):
+    """Base class for all expressions."""
+
+    def accept(self, visitor: Visitor) -> Any:
+        """Accept a visitor.
+
+        Args:
+            visitor: visitor to accept.
+
+        Returns:
+            result of visiting.
+        """
+        return getattr(visitor, f"visit_{_to_snake_case(self.__class__.__name__)}")(
+            self
+        )
+
+
+class Operator(str, Enum):
+    """Enumerator of the operations."""
+
+    AND = "and"
+    OR = "or"
+    NOT = "not"
+
+
+class Comparator(str, Enum):
+    """Enumerator of the comparison operators."""
+
+    EQ = "eq"
+    NE = "ne"
+    GT = "gt"
+    GTE = "gte"
+    LT = "lt"
+    LTE = "lte"
+    CONTAIN = "contain"
+    LIKE = "like"
+    IN = "in"
+    NIN = "nin"
+
+
+class FilterDirective(Expr, ABC):
+    """Filtering expression."""
+
+
+class Comparison(FilterDirective):
+    """Comparison to a value."""
+
+    comparator: Comparator
+    """The comparator to use."""
+
+    attribute: str
+    """The attribute to compare."""
+
+    value: Any
+    """The value to compare to."""
+
+    def __init__(
+        self, comparator: Comparator, attribute: str, value: Any, **kwargs: Any
+    ) -> None:
+        """Create a Comparison.
+
+        Args:
+            comparator: The comparator to use.
+            attribute: The attribute to compare.
+            value: The value to compare to.
+        """
+        # super exists from BaseModel
+        super().__init__(
+            comparator=comparator, attribute=attribute, value=value, **kwargs
+        )
+
+
+class Operation(FilterDirective):
+    """Logical operation over other directives."""
+
+    operator: Operator
+    """The operator to use."""
+
+    arguments: list[FilterDirective]
+    """The arguments to the operator."""
+
+    def __init__(
+        self, operator: Operator, arguments: list[FilterDirective], **kwargs: Any
+    ) -> None:
+        """Create an Operation.
+
+        Args:
+            operator: The operator to use.
+            arguments: The arguments to the operator.
+        """
+        # super exists from BaseModel
+        super().__init__(operator=operator, arguments=arguments, **kwargs)
+
+
+class StructuredQuery(Expr):
+    """Structured query."""
+
+    query: str
+    """Query string."""
+
+    filter: FilterDirective | None
+    """Filtering expression."""
+
+    limit: int | None
+    """Limit on the number of results."""
+
+    def __init__(
+        self,
+        query: str,
+        filter: FilterDirective | None,  # noqa: A002
+        limit: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Create a StructuredQuery.
+
+        Args:
+            query: The query string.
+            filter: The filtering expression.
+            limit: The limit on the number of results.
+        """
+        # super exists from BaseModel
+        super().__init__(query=query, filter=filter, limit=limit, **kwargs)

coco_core/sys_info.py

@@ -0,0 +1,154 @@
+"""Print information about the system and coco packages for debugging purposes."""
+
+import pkgutil
+import platform
+import re
+import sys
+from collections.abc import Sequence
+from importlib import metadata, util
+
+
+def _dist_name(pkg: str) -> str:
+    """Map a Coco module name (`coco`, `coco_core`) to its `trustcoco` distribution."""
+    dist = pkg.replace("_", "-")
+    if pkg == "coco" or pkg.startswith("coco_"):
+        return f"trust{dist}"
+    return dist
+
+
+def _get_sub_deps(packages: Sequence[str]) -> list[str]:
+    """Get any specified sub-dependencies."""
+    sub_deps = set()
+    underscored_packages = {pkg.replace("-", "_") for pkg in packages} | {
+        _dist_name(pkg).replace("-", "_") for pkg in packages
+    }
+
+    for pkg in packages:
+        try:
+            required = metadata.requires(_dist_name(pkg))
+        except metadata.PackageNotFoundError:
+            try:
+                required = metadata.requires(pkg)
+            except metadata.PackageNotFoundError:
+                continue
+
+        if not required:
+            continue
+
+        for req in required:
+            # Extract package name (e.g., "httpx<1,>=0.23.0" -> "httpx")
+            match = re.match(r"^([a-zA-Z0-9_.-]+)", req)
+            if match:
+                pkg_name = match.group(1)
+                if pkg_name.replace("-", "_") not in underscored_packages:
+                    sub_deps.add(pkg_name)
+
+    return sorted(sub_deps, key=lambda x: x.lower())
+
+
+def print_sys_info(*, additional_pkgs: Sequence[str] = ()) -> None:
+    """Print information about the environment for debugging purposes.
+
+    Args:
+        additional_pkgs: Additional packages to include in the output.
+    """
+    # Packages that do not start with "coco" prefix or are namespace packages.
+    other_coco_packages = [
+        "coco_graph",
+        "langsmith",
+        "deepagents",
+        "deepagents-cli",
+    ]
+
+    coco_pkgs = [
+        name for _, name, _ in pkgutil.iter_modules() if name.startswith("coco")
+    ]
+
+    coco_graph_pkgs = [
+        name for _, name, _ in pkgutil.iter_modules() if name.startswith("coco_graph")
+    ]
+
+    all_packages = sorted(
+        set(
+            coco_pkgs
+            + coco_graph_pkgs
+            + other_coco_packages
+            + list(additional_pkgs)
+        )
+    )
+
+    # Always surface these packages to the top
+    order_by = ["coco_core", "coco", "coco_community", "langsmith"]
+
+    for pkg in reversed(order_by):
+        if pkg in all_packages:
+            all_packages.remove(pkg)
+            all_packages = [pkg, *list(all_packages)]
+
+    system_info = {
+        "OS": platform.system(),
+        "OS Version": platform.version(),
+        "Python Version": sys.version,
+    }
+    print()
+    print("System Information")
+    print("------------------")
+    print("> OS: ", system_info["OS"])
+    print("> OS Version: ", system_info["OS Version"])
+    print("> Python Version: ", system_info["Python Version"])
+
+    # Print out only coco packages
+    print()
+    print("Package Information")
+    print("-------------------")
+
+    not_installed = []
+
+    for pkg in all_packages:
+        try:
+            found_package = util.find_spec(pkg)
+        except Exception:
+            found_package = None
+        if found_package is None:
+            not_installed.append(pkg)
+            continue
+
+        # Package version
+        try:
+            package_version = metadata.version(_dist_name(pkg))
+        except Exception:
+            try:
+                package_version = metadata.version(pkg)
+            except Exception:
+                package_version = None
+
+        # Print package with version
+        if package_version is not None:
+            print(f"> {pkg}: {package_version}")
+
+    if not_installed:
+        print()
+        print("Optional packages not installed")
+        print("-------------------------------")
+        for pkg in not_installed:
+            print(f"> {pkg}")
+
+    sub_dependencies = _get_sub_deps(all_packages)
+
+    if sub_dependencies:
+        print()
+        print("Other Dependencies")
+        print("------------------")
+
+        for dep in sub_dependencies:
+            try:
+                dep_version = metadata.version(dep)
+            except Exception:
+                dep_version = None
+
+            if dep_version is not None:
+                print(f"> {dep}: {dep_version}")
+
+
+if __name__ == "__main__":
+    print_sys_info()