coco-framework 0.1.0__py3-none-any.whl

coco_core/rate_limiters.py

@@ -0,0 +1,254 @@
+"""Interface for a rate limiter and an in-memory rate limiter."""
+
+from __future__ import annotations
+
+import abc
+import asyncio
+import threading
+import time
+
+
+class BaseRateLimiter(abc.ABC):
+    """Base class for rate limiters.
+
+    Usage of the base limiter is through the acquire and aacquire methods depending
+    on whether running in a sync or async context.
+
+    Implementations are free to add a timeout parameter to their initialize method
+    to allow users to specify a timeout for acquiring the necessary tokens when
+    using a blocking call.
+
+    Current limitations:
+
+    - Rate limiting information is not surfaced in tracing or callbacks. This means
+        that the total time it takes to invoke a chat model will encompass both
+        the time spent waiting for tokens and the time spent making the request.
+    """
+
+    @abc.abstractmethod
+    def acquire(self, *, blocking: bool = True) -> bool:
+        """Attempt to acquire the necessary tokens for the rate limiter.
+
+        This method blocks until the required tokens are available if `blocking`
+        is set to `True`.
+
+        If `blocking` is set to `False`, the method will immediately return the result
+        of the attempt to acquire the tokens.
+
+        Args:
+            blocking: If `True`, the method will block until the tokens are available.
+                If `False`, the method will return immediately with the result of
+                the attempt.
+
+        Returns:
+            `True` if the tokens were successfully acquired, `False` otherwise.
+        """
+
+    @abc.abstractmethod
+    async def aacquire(self, *, blocking: bool = True) -> bool:
+        """Attempt to acquire the necessary tokens for the rate limiter.
+
+        This method blocks until the required tokens are available if `blocking`
+        is set to `True`.
+
+        If `blocking` is set to `False`, the method will immediately return the result
+        of the attempt to acquire the tokens.
+
+        Args:
+            blocking: If `True`, the method will block until the tokens are available.
+                If `False`, the method will return immediately with the result of
+                the attempt.
+
+        Returns:
+            `True` if the tokens were successfully acquired, `False` otherwise.
+        """
+
+
+class InMemoryRateLimiter(BaseRateLimiter):
+    """An in memory rate limiter based on a token bucket algorithm.
+
+    This is an in memory rate limiter, so it cannot rate limit across
+    different processes.
+
+    The rate limiter only allows time-based rate limiting and does not
+    take into account any information about the input or the output, so it
+    cannot be used to rate limit based on the size of the request.
+
+    It is thread safe and can be used in either a sync or async context.
+
+    The in memory rate limiter is based on a token bucket. The bucket is filled
+    with tokens at a given rate. Each request consumes a token. If there are
+    not enough tokens in the bucket, the request is blocked until there are
+    enough tokens.
+
+    These tokens have nothing to do with LLM tokens. They are just
+    a way to keep track of how many requests can be made at a given time.
+
+    Current limitations:
+
+    - The rate limiter is not designed to work across different processes. It is
+        an in-memory rate limiter, but it is thread safe.
+    - The rate limiter only supports time-based rate limiting. It does not take
+        into account the size of the request or any other factors.
+
+    Example:
+        ```python
+        import time
+
+        from coco_core.rate_limiters import InMemoryRateLimiter
+
+        rate_limiter = InMemoryRateLimiter(
+            requests_per_second=0.1,  # <-- Can only make a request once every 10 seconds!!
+            check_every_n_seconds=0.1,  # Wake up every 100 ms to check whether allowed to make a request,
+            max_bucket_size=10,  # Controls the maximum burst size.
+        )
+
+        from coco_anthropic import ChatAnthropic
+
+        model = ChatAnthropic(model_name="claude-sonnet-5", rate_limiter=rate_limiter)
+
+        for _ in range(5):
+            tic = time.time()
+            model.invoke("hello")
+            toc = time.time()
+            print(toc - tic)
+        ```
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        *,
+        requests_per_second: float = 1,
+        check_every_n_seconds: float = 0.1,
+        max_bucket_size: float = 1,
+    ) -> None:
+        """A rate limiter based on a token bucket.
+
+        These tokens have nothing to do with LLM tokens. They are just
+        a way to keep track of how many requests can be made at a given time.
+
+        This rate limiter is designed to work in a threaded environment.
+
+        It works by filling up a bucket with tokens at a given rate. Each
+        request consumes a given number of tokens. If there are not enough
+        tokens in the bucket, the request is blocked until there are enough
+        tokens.
+
+        Args:
+            requests_per_second: The number of tokens to add per second to the bucket.
+                The tokens represent "credit" that can be used to make requests.
+            check_every_n_seconds: Check whether the tokens are available
+                every this many seconds. Can be a float to represent
+                fractions of a second.
+            max_bucket_size: The maximum number of tokens that can be in the bucket.
+                Must be at least `1`. Used to prevent bursts of requests.
+        """
+        # Number of requests that we can make per second.
+        self.requests_per_second = requests_per_second
+
+        # Number of tokens in the bucket.
+        self.available_tokens = 0.0
+
+        self.max_bucket_size = max_bucket_size
+
+        # A lock to ensure that tokens can only be consumed by one thread
+        # at a given time.
+        self._consume_lock = threading.Lock()
+
+        # The last time we tried to consume tokens.
+        self.last: float | None = None
+
+        self.check_every_n_seconds = check_every_n_seconds
+
+    def _consume(self) -> bool:
+        """Try to consume a token.
+
+        Returns:
+            True means that the tokens were consumed, and the caller can proceed to
+            make the request. A False means that the tokens were not consumed, and
+            the caller should try again later.
+        """
+        with self._consume_lock:
+            now = time.monotonic()
+
+            # initialize on first call to avoid a burst
+            if self.last is None:
+                self.last = now
+
+            elapsed = now - self.last
+
+            if elapsed * self.requests_per_second >= 1:
+                self.available_tokens += elapsed * self.requests_per_second
+                self.last = now
+
+            # Make sure that we don't exceed the bucket size.
+            # This is used to prevent bursts of requests.
+            self.available_tokens = min(self.available_tokens, self.max_bucket_size)
+
+            # As long as we have at least one token, we can proceed.
+            if self.available_tokens >= 1:
+                self.available_tokens -= 1
+                return True
+
+            return False
+
+    def acquire(self, *, blocking: bool = True) -> bool:
+        """Attempt to acquire a token from the rate limiter.
+
+        This method blocks until the required tokens are available if `blocking`
+        is set to `True`.
+
+        If `blocking` is set to `False`, the method will immediately return the result
+        of the attempt to acquire the tokens.
+
+        Args:
+            blocking: If `True`, the method will block until the tokens are available.
+                If `False`, the method will return immediately with the result of
+                the attempt.
+
+        Returns:
+            `True` if the tokens were successfully acquired, `False` otherwise.
+        """
+        if not blocking:
+            return self._consume()
+
+        while not self._consume():
+            time.sleep(self.check_every_n_seconds)
+
+        return True
+
+    async def aacquire(self, *, blocking: bool = True) -> bool:
+        """Attempt to acquire a token from the rate limiter. Async version.
+
+        This method blocks until the required tokens are available if `blocking`
+        is set to `True`.
+
+        If `blocking` is set to `False`, the method will immediately return the result
+        of the attempt to acquire the tokens.
+
+        Args:
+            blocking: If `True`, the method will block until the tokens are available.
+                If `False`, the method will return immediately with the result of
+                the attempt.
+
+        Returns:
+            `True` if the tokens were successfully acquired, `False` otherwise.
+        """
+        if not blocking:
+            return self._consume()
+
+        while not self._consume():  # noqa: ASYNC110
+            # This code ignores the ASYNC110 warning which is a false positive in this
+            # case.
+            # There is no external actor that can mark that the Event is done
+            # since the tokens are managed by the rate limiter itself.
+            # It needs to wake up to re-fill the tokens.
+            # https://docs.astral.sh/ruff/rules/async-busy-wait/
+            await asyncio.sleep(self.check_every_n_seconds)
+        return True
+
+
+__all__ = [
+    "BaseRateLimiter",
+    "InMemoryRateLimiter",
+]

coco_core/retrievers.py

@@ -0,0 +1,328 @@
+"""**Retriever** class returns `Document` objects given a text **query**.
+
+It is more general than a vector store. A retriever does not need to be able to
+store documents, only to return (or retrieve) it. Vector stores can be used as
+the backbone of a retriever, but there are other types of retrievers as well.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from inspect import signature
+from typing import TYPE_CHECKING, Any
+
+from pydantic import ConfigDict
+from typing_extensions import Self, TypedDict, override
+
+from coco_core.callbacks.manager import AsyncCallbackManager, CallbackManager
+from coco_core.documents import Document
+from coco_core.runnables import (
+    Runnable,
+    RunnableConfig,
+    RunnableSerializable,
+    ensure_config,
+)
+from coco_core.runnables.config import run_in_executor
+
+if TYPE_CHECKING:
+    from coco_core.callbacks.manager import (
+        AsyncCallbackManagerForRetrieverRun,
+        CallbackManagerForRetrieverRun,
+    )
+
+RetrieverInput = str
+RetrieverOutput = list[Document]
+RetrieverLike = Runnable[RetrieverInput, RetrieverOutput]
+RetrieverOutputLike = Runnable[Any, RetrieverOutput]
+
+
+class LangSmithRetrieverParams(TypedDict, total=False):
+    """LangSmith parameters for tracing."""
+
+    ls_retriever_name: str
+    """Retriever name."""
+
+    ls_vector_store_provider: str | None
+    """Vector store provider."""
+
+    ls_embedding_provider: str | None
+    """Embedding provider."""
+
+    ls_embedding_model: str | None
+    """Embedding model."""
+
+
+class BaseRetriever(RunnableSerializable[RetrieverInput, RetrieverOutput], ABC):
+    """Abstract base class for a document retrieval system.
+
+    A retrieval system is defined as something that can take string queries and return
+    the most 'relevant' documents from some source.
+
+    Usage:
+
+    A retriever follows the standard `Runnable` interface, and should be used via the
+    standard `Runnable` methods of `invoke`, `ainvoke`, `batch`, `abatch`.
+
+    Implementation:
+
+    When implementing a custom retriever, the class should implement the
+    `_get_relevant_documents` method to define the logic for retrieving documents.
+
+    Optionally, an async native implementations can be provided by overriding the
+    `_aget_relevant_documents` method.
+
+    !!! example "Retriever that returns the first 5 documents from a list of documents"
+
+        ```python
+        from coco_core.documents import Document
+        from coco_core.retrievers import BaseRetriever
+
+        class SimpleRetriever(BaseRetriever):
+            docs: list[Document]
+            k: int = 5
+
+            def _get_relevant_documents(self, query: str) -> list[Document]:
+                \"\"\"Return the first k documents from the list of documents\"\"\"
+                return self.docs[:self.k]
+
+            async def _aget_relevant_documents(self, query: str) -> list[Document]:
+                \"\"\"(Optional) async native implementation.\"\"\"
+                return self.docs[:self.k]
+        ```
+
+    !!! example "Simple retriever based on a scikit-learn vectorizer"
+
+        ```python
+        from sklearn.metrics.pairwise import cosine_similarity
+
+
+        class TFIDFRetriever(BaseRetriever, BaseModel):
+            vectorizer: Any
+            docs: list[Document]
+            tfidf_array: Any
+            k: int = 4
+
+            class Config:
+                arbitrary_types_allowed = True
+
+            def _get_relevant_documents(self, query: str) -> list[Document]:
+                # Ip -- (n_docs,x), Op -- (n_docs,n_Feats)
+                query_vec = self.vectorizer.transform([query])
+                # Op -- (n_docs,1) -- Cosine Sim with each doc
+                results = cosine_similarity(self.tfidf_array, query_vec).reshape((-1,))
+                return [self.docs[i] for i in results.argsort()[-self.k :][::-1]]
+        ```
+    """
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+    _new_arg_supported: bool = False
+
+    _expects_other_args: bool = False
+
+    tags: list[str] | None = None
+    """Optional list of tags associated with the retriever.
+
+    These tags will be associated with each call to this retriever,
+    and passed as arguments to the handlers defined in `callbacks`.
+
+    You can use these to eg identify a specific instance of a retriever with its
+    use case.
+    """
+
+    metadata: dict[str, Any] | None = None
+    """Optional metadata associated with the retriever.
+
+    This metadata will be associated with each call to this retriever,
+    and passed as arguments to the handlers defined in `callbacks`.
+
+    You can use these to eg identify a specific instance of a retriever with its
+    use case.
+    """
+
+    @override
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        parameters = signature(cls._get_relevant_documents).parameters
+        cls._new_arg_supported = parameters.get("run_manager") is not None
+        if (
+            not cls._new_arg_supported
+            and cls._aget_relevant_documents == BaseRetriever._aget_relevant_documents
+        ):
+            # we need to tolerate no run_manager in _aget_relevant_documents signature
+            async def _aget_relevant_documents(
+                self: Self, query: str
+            ) -> list[Document]:
+                return await run_in_executor(None, self._get_relevant_documents, query)  # type: ignore[call-arg]
+
+            cls._aget_relevant_documents = _aget_relevant_documents  # type: ignore[assignment]
+
+        # If a V1 retriever broke the interface and expects additional arguments
+        cls._expects_other_args = (
+            len(set(parameters.keys()) - {"self", "query", "run_manager"}) > 0
+        )
+
+    def _get_ls_params(self, **_kwargs: Any) -> LangSmithRetrieverParams:
+        """Get standard params for tracing."""
+        default_retriever_name = self.get_name()
+        if default_retriever_name.startswith("Retriever"):
+            default_retriever_name = default_retriever_name[9:]
+        elif default_retriever_name.endswith("Retriever"):
+            default_retriever_name = default_retriever_name[:-9]
+        default_retriever_name = default_retriever_name.lower()
+
+        return LangSmithRetrieverParams(ls_retriever_name=default_retriever_name)
+
+    @override
+    def invoke(
+        self, input: str, config: RunnableConfig | None = None, **kwargs: Any
+    ) -> list[Document]:
+        """Invoke the retriever to get relevant documents.
+
+        Main entry point for synchronous retriever invocations.
+
+        Args:
+            input: The query string.
+            config: Configuration for the retriever.
+            **kwargs: Additional arguments to pass to the retriever.
+
+        Returns:
+            List of relevant documents.
+
+        Examples:
+        ```python
+        retriever.invoke("query")
+        ```
+        """
+        config = ensure_config(config)
+        inheritable_metadata = {
+            **(config.get("metadata") or {}),
+            **self._get_ls_params(**kwargs),
+        }
+        callback_manager = CallbackManager.configure(
+            config.get("callbacks"),
+            None,
+            verbose=kwargs.get("verbose", False),
+            inheritable_tags=config.get("tags"),
+            local_tags=self.tags,
+            inheritable_metadata=inheritable_metadata,
+            local_metadata=self.metadata,
+        )
+        run_manager = callback_manager.on_retriever_start(
+            None,
+            input,
+            name=config.get("run_name") or self.get_name(),
+            run_id=kwargs.pop("run_id", None),
+        )
+        try:
+            kwargs_ = kwargs if self._expects_other_args else {}
+            if self._new_arg_supported:
+                result = self._get_relevant_documents(
+                    input, run_manager=run_manager, **kwargs_
+                )
+            else:
+                result = self._get_relevant_documents(input, **kwargs_)
+        except Exception as e:
+            run_manager.on_retriever_error(e)
+            raise
+        else:
+            run_manager.on_retriever_end(
+                result,
+            )
+            return result
+
+    @override
+    async def ainvoke(
+        self,
+        input: str,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> list[Document]:
+        """Asynchronously invoke the retriever to get relevant documents.
+
+        Main entry point for asynchronous retriever invocations.
+
+        Args:
+            input: The query string.
+            config: Configuration for the retriever.
+            **kwargs: Additional arguments to pass to the retriever.
+
+        Returns:
+            List of relevant documents.
+
+        Examples:
+        ```python
+        await retriever.ainvoke("query")
+        ```
+        """
+        config = ensure_config(config)
+        inheritable_metadata = {
+            **(config.get("metadata") or {}),
+            **self._get_ls_params(**kwargs),
+        }
+        callback_manager = AsyncCallbackManager.configure(
+            config.get("callbacks"),
+            None,
+            verbose=kwargs.get("verbose", False),
+            inheritable_tags=config.get("tags"),
+            local_tags=self.tags,
+            inheritable_metadata=inheritable_metadata,
+            local_metadata=self.metadata,
+        )
+        run_manager = await callback_manager.on_retriever_start(
+            None,
+            input,
+            name=config.get("run_name") or self.get_name(),
+            run_id=kwargs.pop("run_id", None),
+        )
+        try:
+            kwargs_ = kwargs if self._expects_other_args else {}
+            if self._new_arg_supported:
+                result = await self._aget_relevant_documents(
+                    input, run_manager=run_manager, **kwargs_
+                )
+            else:
+                result = await self._aget_relevant_documents(input, **kwargs_)
+        except Exception as e:
+            await run_manager.on_retriever_error(e)
+            raise
+        else:
+            await run_manager.on_retriever_end(
+                result,
+            )
+            return result
+
+    @abstractmethod
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> list[Document]:
+        """Get documents relevant to a query.
+
+        Args:
+            query: String to find relevant documents for.
+            run_manager: The callback handler to use.
+
+        Returns:
+            List of relevant documents.
+        """
+
+    async def _aget_relevant_documents(
+        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
+    ) -> list[Document]:
+        """Asynchronously get documents relevant to a query.
+
+        Args:
+            query: String to find relevant documents for
+            run_manager: The callback handler to use
+
+        Returns:
+            List of relevant documents
+        """
+        return await run_in_executor(
+            None,
+            self._get_relevant_documents,
+            query,
+            run_manager=run_manager.get_sync(),
+        )