coco-framework 0.1.0__py3-none-any.whl

coco_core/example_selectors/semantic_similarity.py

@@ -0,0 +1,366 @@
+"""Example selector that selects examples based on SemanticSimilarity."""
+
+from __future__ import annotations
+
+from abc import ABC
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, ConfigDict
+
+from coco_core.example_selectors.base import BaseExampleSelector
+from coco_core.vectorstores import VectorStore
+
+if TYPE_CHECKING:
+    from coco_core.documents import Document
+    from coco_core.embeddings import Embeddings
+
+
+def sorted_values(values: dict[str, str]) -> list[str]:
+    """Return a list of values in dict sorted by key.
+
+    Args:
+        values: A dictionary with keys as input variables
+            and values as their values.
+
+    Returns:
+        A list of values in dict sorted by key.
+    """
+    return [values[val] for val in sorted(values)]
+
+
+class _VectorStoreExampleSelector(BaseExampleSelector, BaseModel, ABC):
+    """Example selector that selects examples based on SemanticSimilarity."""
+
+    vectorstore: VectorStore
+    """VectorStore that contains information about examples."""
+
+    k: int = 4
+    """Number of examples to select."""
+
+    example_keys: list[str] | None = None
+    """Optional keys to filter examples to."""
+
+    input_keys: list[str] | None = None
+    """Optional keys to filter input to. If provided, the search is based on
+    the input variables instead of all variables."""
+
+    vectorstore_kwargs: dict[str, Any] | None = None
+    """Extra arguments passed to similarity_search function of the `VectorStore`."""
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+    )
+
+    @staticmethod
+    def _example_to_text(example: dict[str, str], input_keys: list[str] | None) -> str:
+        if input_keys:
+            return " ".join(sorted_values({key: example[key] for key in input_keys}))
+        return " ".join(sorted_values(example))
+
+    def _documents_to_examples(self, documents: list[Document]) -> list[dict[str, Any]]:
+        # Get the examples from the metadata.
+        # This assumes that examples are stored in metadata.
+        examples = [dict(e.metadata) for e in documents]
+        # If example keys are provided, filter examples to those keys.
+        if self.example_keys:
+            examples = [{k: eg[k] for k in self.example_keys} for eg in examples]
+        return examples
+
+    def add_example(self, example: dict[str, str]) -> str:
+        """Add a new example to vectorstore.
+
+        Args:
+            example: A dictionary with keys as input variables
+                and values as their values.
+
+        Returns:
+            The ID of the added example.
+        """
+        ids = self.vectorstore.add_texts(
+            [self._example_to_text(example, self.input_keys)], metadatas=[example]
+        )
+        return ids[0]
+
+    async def aadd_example(self, example: dict[str, str]) -> str:
+        """Async add new example to vectorstore.
+
+        Args:
+            example: A dictionary with keys as input variables
+                and values as their values.
+
+        Returns:
+            The ID of the added example.
+        """
+        ids = await self.vectorstore.aadd_texts(
+            [self._example_to_text(example, self.input_keys)], metadatas=[example]
+        )
+        return ids[0]
+
+
+class SemanticSimilarityExampleSelector(_VectorStoreExampleSelector):
+    """Select examples based on semantic similarity."""
+
+    def select_examples(self, input_variables: dict[str, str]) -> list[dict[str, Any]]:
+        """Select examples based on semantic similarity.
+
+        Args:
+            input_variables: The input variables to use for search.
+
+        Returns:
+            The selected examples.
+        """
+        # Get the docs with the highest similarity.
+        vectorstore_kwargs = self.vectorstore_kwargs or {}
+        example_docs = self.vectorstore.similarity_search(
+            self._example_to_text(input_variables, self.input_keys),
+            k=self.k,
+            **vectorstore_kwargs,
+        )
+        return self._documents_to_examples(example_docs)
+
+    async def aselect_examples(
+        self, input_variables: dict[str, str]
+    ) -> list[dict[str, Any]]:
+        """Asynchronously select examples based on semantic similarity.
+
+        Args:
+            input_variables: The input variables to use for search.
+
+        Returns:
+            The selected examples.
+        """
+        # Get the docs with the highest similarity.
+        vectorstore_kwargs = self.vectorstore_kwargs or {}
+        example_docs = await self.vectorstore.asimilarity_search(
+            self._example_to_text(input_variables, self.input_keys),
+            k=self.k,
+            **vectorstore_kwargs,
+        )
+        return self._documents_to_examples(example_docs)
+
+    @classmethod
+    def from_examples(
+        cls,
+        examples: list[dict[str, str]],
+        embeddings: Embeddings,
+        vectorstore_cls: type[VectorStore],
+        k: int = 4,
+        input_keys: list[str] | None = None,
+        *,
+        example_keys: list[str] | None = None,
+        vectorstore_kwargs: dict[str, Any] | None = None,
+        **vectorstore_cls_kwargs: Any,
+    ) -> SemanticSimilarityExampleSelector:
+        """Create k-shot example selector using example list and embeddings.
+
+        Reshuffles examples dynamically based on query similarity.
+
+        Args:
+            examples: List of examples to use in the prompt.
+            embeddings: An initialized embedding API interface, e.g. OpenAIEmbeddings().
+            vectorstore_cls: A vector store DB interface class, e.g. FAISS.
+            k: Number of examples to select.
+            input_keys: If provided, the search is based on the input variables
+                instead of all variables.
+            example_keys: If provided, keys to filter examples to.
+            vectorstore_kwargs: Extra arguments passed to similarity_search function
+                of the `VectorStore`.
+            vectorstore_cls_kwargs: optional kwargs containing url for vector store
+
+        Returns:
+            The ExampleSelector instantiated, backed by a vector store.
+        """
+        string_examples = [cls._example_to_text(eg, input_keys) for eg in examples]
+        vectorstore = vectorstore_cls.from_texts(
+            string_examples, embeddings, metadatas=examples, **vectorstore_cls_kwargs
+        )
+        return cls(
+            vectorstore=vectorstore,
+            k=k,
+            input_keys=input_keys,
+            example_keys=example_keys,
+            vectorstore_kwargs=vectorstore_kwargs,
+        )
+
+    @classmethod
+    async def afrom_examples(
+        cls,
+        examples: list[dict[str, str]],
+        embeddings: Embeddings,
+        vectorstore_cls: type[VectorStore],
+        k: int = 4,
+        input_keys: list[str] | None = None,
+        *,
+        example_keys: list[str] | None = None,
+        vectorstore_kwargs: dict[str, Any] | None = None,
+        **vectorstore_cls_kwargs: Any,
+    ) -> SemanticSimilarityExampleSelector:
+        """Async create k-shot example selector using example list and embeddings.
+
+        Reshuffles examples dynamically based on query similarity.
+
+        Args:
+            examples: List of examples to use in the prompt.
+            embeddings: An initialized embedding API interface, e.g. OpenAIEmbeddings().
+            vectorstore_cls: A vector store DB interface class, e.g. FAISS.
+            k: Number of examples to select.
+            input_keys: If provided, the search is based on the input variables
+                instead of all variables.
+            example_keys: If provided, keys to filter examples to.
+            vectorstore_kwargs: Extra arguments passed to similarity_search function
+                of the `VectorStore`.
+            vectorstore_cls_kwargs: optional kwargs containing url for vector store
+
+        Returns:
+            The ExampleSelector instantiated, backed by a vector store.
+        """
+        string_examples = [cls._example_to_text(eg, input_keys) for eg in examples]
+        vectorstore = await vectorstore_cls.afrom_texts(
+            string_examples, embeddings, metadatas=examples, **vectorstore_cls_kwargs
+        )
+        return cls(
+            vectorstore=vectorstore,
+            k=k,
+            input_keys=input_keys,
+            example_keys=example_keys,
+            vectorstore_kwargs=vectorstore_kwargs,
+        )
+
+
+class MaxMarginalRelevanceExampleSelector(_VectorStoreExampleSelector):
+    """Select examples based on Max Marginal Relevance.
+
+    This was shown to improve performance in this paper:
+    https://arxiv.org/pdf/2211.13892.pdf
+    """
+
+    fetch_k: int = 20
+    """Number of examples to fetch to rerank."""
+
+    def select_examples(self, input_variables: dict[str, str]) -> list[dict[str, Any]]:
+        """Select examples based on Max Marginal Relevance.
+
+        Args:
+            input_variables: The input variables to use for search.
+
+        Returns:
+            The selected examples.
+        """
+        example_docs = self.vectorstore.max_marginal_relevance_search(
+            self._example_to_text(input_variables, self.input_keys),
+            k=self.k,
+            fetch_k=self.fetch_k,
+        )
+        return self._documents_to_examples(example_docs)
+
+    async def aselect_examples(
+        self, input_variables: dict[str, str]
+    ) -> list[dict[str, Any]]:
+        """Asynchronously select examples based on Max Marginal Relevance.
+
+        Args:
+            input_variables: The input variables to use for search.
+
+        Returns:
+            The selected examples.
+        """
+        example_docs = await self.vectorstore.amax_marginal_relevance_search(
+            self._example_to_text(input_variables, self.input_keys),
+            k=self.k,
+            fetch_k=self.fetch_k,
+        )
+        return self._documents_to_examples(example_docs)
+
+    @classmethod
+    def from_examples(
+        cls,
+        examples: list[dict[str, str]],
+        embeddings: Embeddings,
+        vectorstore_cls: type[VectorStore],
+        k: int = 4,
+        input_keys: list[str] | None = None,
+        fetch_k: int = 20,
+        example_keys: list[str] | None = None,
+        vectorstore_kwargs: dict[str, Any] | None = None,
+        **vectorstore_cls_kwargs: Any,
+    ) -> MaxMarginalRelevanceExampleSelector:
+        """Create k-shot example selector using example list and embeddings.
+
+        Reshuffles examples dynamically based on Max Marginal Relevance.
+
+        Args:
+            examples: List of examples to use in the prompt.
+            embeddings: An initialized embedding API interface, e.g. OpenAIEmbeddings().
+            vectorstore_cls: A vector store DB interface class, e.g. FAISS.
+            k: Number of examples to select.
+            fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.
+            input_keys: If provided, the search is based on the input variables
+                instead of all variables.
+            example_keys: If provided, keys to filter examples to.
+            vectorstore_kwargs: Extra arguments passed to similarity_search function
+                of the `VectorStore`.
+            vectorstore_cls_kwargs: optional kwargs containing url for vector store
+
+        Returns:
+            The ExampleSelector instantiated, backed by a vector store.
+        """
+        string_examples = [cls._example_to_text(eg, input_keys) for eg in examples]
+        vectorstore = vectorstore_cls.from_texts(
+            string_examples, embeddings, metadatas=examples, **vectorstore_cls_kwargs
+        )
+        return cls(
+            vectorstore=vectorstore,
+            k=k,
+            fetch_k=fetch_k,
+            input_keys=input_keys,
+            example_keys=example_keys,
+            vectorstore_kwargs=vectorstore_kwargs,
+        )
+
+    @classmethod
+    async def afrom_examples(
+        cls,
+        examples: list[dict[str, str]],
+        embeddings: Embeddings,
+        vectorstore_cls: type[VectorStore],
+        *,
+        k: int = 4,
+        input_keys: list[str] | None = None,
+        fetch_k: int = 20,
+        example_keys: list[str] | None = None,
+        vectorstore_kwargs: dict[str, Any] | None = None,
+        **vectorstore_cls_kwargs: Any,
+    ) -> MaxMarginalRelevanceExampleSelector:
+        """Create k-shot example selector using example list and embeddings.
+
+        Reshuffles examples dynamically based on Max Marginal Relevance.
+
+        Args:
+            examples: List of examples to use in the prompt.
+            embeddings: An initialized embedding API interface, e.g. OpenAIEmbeddings().
+            vectorstore_cls: A vector store DB interface class, e.g. FAISS.
+            k: Number of examples to select.
+            fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.
+            input_keys: If provided, the search is based on the input variables
+                instead of all variables.
+            example_keys: If provided, keys to filter examples to.
+            vectorstore_kwargs: Extra arguments passed to similarity_search function
+                of the `VectorStore`.
+            vectorstore_cls_kwargs: optional kwargs containing url for vector store
+
+        Returns:
+            The ExampleSelector instantiated, backed by a vector store.
+        """
+        string_examples = [cls._example_to_text(eg, input_keys) for eg in examples]
+        vectorstore = await vectorstore_cls.afrom_texts(
+            string_examples, embeddings, metadatas=examples, **vectorstore_cls_kwargs
+        )
+        return cls(
+            vectorstore=vectorstore,
+            k=k,
+            fetch_k=fetch_k,
+            input_keys=input_keys,
+            example_keys=example_keys,
+            vectorstore_kwargs=vectorstore_kwargs,
+        )

coco_core/exceptions.py

@@ -0,0 +1,166 @@
+"""Custom **exceptions** for Coco."""
+
+from enum import Enum
+from typing import Any
+
+
+class CocoException(Exception):  # noqa: N818
+    """General Coco exception."""
+
+
+class TracerException(CocoException):
+    """Base class for exceptions in tracers module."""
+
+
+class OutputParserException(ValueError, CocoException):  # noqa: N818
+    """Exception that output parsers should raise to signify a parsing error.
+
+    This exists to differentiate parsing errors from other code or execution errors
+    that also may arise inside the output parser.
+
+    `OutputParserException` will be available to catch and handle in ways to fix the
+    parsing error, while other errors will be raised.
+    """
+
+    def __init__(
+        self,
+        error: Any,
+        observation: str | None = None,
+        llm_output: str | None = None,
+        send_to_llm: bool = False,  # noqa: FBT001,FBT002
+    ):
+        """Create an `OutputParserException`.
+
+        Args:
+            error: The error that's being re-raised or an error message.
+            observation: String explanation of error which can be passed to a model to
+                try and remediate the issue.
+            llm_output: String model output which is error-ing.
+
+            send_to_llm: Whether to send the observation and llm_output back to an Agent
+                after an `OutputParserException` has been raised.
+
+                This gives the underlying model driving the agent the context that the
+                previous output was improperly structured, in the hopes that it will
+                update the output to the correct format.
+
+        Raises:
+            ValueError: If `send_to_llm` is `True` but either observation or
+                `llm_output` are not provided.
+        """
+        if isinstance(error, str):
+            error = create_message(
+                message=error, error_code=ErrorCode.OUTPUT_PARSING_FAILURE
+            )
+
+        super().__init__(error)
+        if send_to_llm and (observation is None or llm_output is None):
+            msg = (
+                "Arguments 'observation' & 'llm_output'"
+                " are required if 'send_to_llm' is True"
+            )
+            raise ValueError(msg)
+        self.observation = observation
+        self.llm_output = llm_output
+        self.send_to_llm = send_to_llm
+
+
+class ModelError(CocoException):
+    """Base exception for failures related to model invocation.
+
+    Subclasses correspond to conditions that model providers report consistently,
+    keyed to the HTTP status they surface it with, so the same condition maps to
+    the same exception type regardless of provider.
+
+    Provider integrations raise subclasses that also inherit from the provider
+    SDK's own exception type, so code catching either continues to work.
+    """
+
+    is_retryable = False
+    """Whether retrying the same model request may succeed."""
+
+
+class ModelAuthenticationError(ModelError):
+    """Exception raised when model provider authentication fails (HTTP 401)."""
+
+
+class ModelPermissionDeniedError(ModelError):
+    """Exception raised when credentials lack permission for a request (HTTP 403)."""
+
+
+class ModelInvalidRequestError(ModelError):
+    """Exception raised when a provider rejects a request as invalid (e.g. HTTP 400)."""
+
+
+class ModelNotFoundError(ModelError):
+    """Exception raised when the requested model cannot be found (HTTP 404)."""
+
+
+class ModelRateLimitError(ModelError):
+    """Exception raised when a model provider rate limit is exceeded (HTTP 429)."""
+
+    is_retryable = True
+
+
+class ModelAPIError(ModelError):
+    """Exception raised when a model provider reports a server failure (HTTP 5xx)."""
+
+    is_retryable = True
+
+
+class ModelConnectionError(ModelError):
+    """Exception raised when a model provider cannot be reached."""
+
+    is_retryable = True
+
+
+class ModelTimeoutError(ModelError):
+    """Exception raised when a model request times out."""
+
+    is_retryable = True
+
+
+class ContextOverflowError(ModelError):
+    """Exception raised when input exceeds the model's context limit.
+
+    This exception is raised by chat models when the input tokens exceed
+    the maximum context window supported by the model.
+    """
+
+
+class ErrorCode(Enum):
+    """Error codes."""
+
+    INVALID_PROMPT_INPUT = "INVALID_PROMPT_INPUT"
+    INVALID_TOOL_RESULTS = "INVALID_TOOL_RESULTS"  # Used in JS; not Py (yet)
+    MESSAGE_COERCION_FAILURE = "MESSAGE_COERCION_FAILURE"
+    MODEL_AUTHENTICATION = "MODEL_AUTHENTICATION"  # Used in JS; not Py (yet)
+    MODEL_NOT_FOUND = "MODEL_NOT_FOUND"  # Used in JS; not Py (yet)
+    MODEL_RATE_LIMIT = "MODEL_RATE_LIMIT"  # Used in JS; not Py (yet)
+    OUTPUT_PARSING_FAILURE = "OUTPUT_PARSING_FAILURE"
+
+
+def create_message(*, message: str, error_code: ErrorCode) -> str:
+    """Create a message with a link to the Coco troubleshooting guide.
+
+    Args:
+        message: The message to display.
+        error_code: The error code to display.
+
+    Returns:
+        The full message with the troubleshooting link.
+
+    Example:
+        ```python
+        create_message(
+            message="Failed to parse output",
+            error_code=ErrorCode.OUTPUT_PARSING_FAILURE,
+        )
+        "Failed to parse output. For troubleshooting, visit: ..."
+        ```
+    """
+    return (
+        f"{message}\n"
+        "For troubleshooting, visit: https://docs.langchain.com/oss/python/langchain"
+        f"/errors/{error_code.value} "
+    )

coco_core/globals.py

@@ -0,0 +1,72 @@
+"""Global values and configuration that apply to all of Coco."""
+
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from coco_core.caches import BaseCache
+
+
+# DO NOT USE THESE VALUES DIRECTLY!
+# Use them only via `get_<X>()` and `set_<X>()` below,
+# or else your code may behave unexpectedly with other uses of these global settings:
+# https://github.com/langchain-ai/langchain/pull/11311#issuecomment-1743780004
+_verbose: bool = False
+_debug: bool = False
+_llm_cache: Optional["BaseCache"] = None
+
+
+def set_verbose(value: bool) -> None:  # noqa: FBT001
+    """Set a new value for the `verbose` global setting.
+
+    Args:
+        value: The new value for the `verbose` global setting.
+    """
+    global _verbose  # noqa: PLW0603
+    _verbose = value
+
+
+def get_verbose() -> bool:
+    """Get the value of the `verbose` global setting.
+
+    Returns:
+        The value of the `verbose` global setting.
+    """
+    return _verbose
+
+
+def set_debug(value: bool) -> None:  # noqa: FBT001
+    """Set a new value for the `debug` global setting.
+
+    Args:
+        value: The new value for the `debug` global setting.
+    """
+    global _debug  # noqa: PLW0603
+    _debug = value
+
+
+def get_debug() -> bool:
+    """Get the value of the `debug` global setting.
+
+    Returns:
+        The value of the `debug` global setting.
+    """
+    return _debug
+
+
+def set_llm_cache(value: Optional["BaseCache"]) -> None:
+    """Set a new LLM cache, overwriting the previous value, if any.
+
+    Args:
+        value: The new LLM cache to use. If `None`, the LLM cache is disabled.
+    """
+    global _llm_cache  # noqa: PLW0603
+    _llm_cache = value
+
+
+def get_llm_cache() -> Optional["BaseCache"]:
+    """Get the value of the `llm_cache` global setting.
+
+    Returns:
+        The value of the `llm_cache` global setting.
+    """
+    return _llm_cache

coco_core/indexing/__init__.py

@@ -0,0 +1,53 @@
+"""Code to help indexing data into a vectorstore.
+
+This package contains helper logic to help deal with indexing data into
+a `VectorStore` while avoiding duplicated content and over-writing content
+if it's unchanged.
+"""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.indexing.api import IndexingResult, aindex, index
+    from coco_core.indexing.base import (
+        DeleteResponse,
+        DocumentIndex,
+        InMemoryRecordManager,
+        RecordManager,
+        UpsertResponse,
+    )
+
+__all__ = (
+    "DeleteResponse",
+    "DocumentIndex",
+    "InMemoryRecordManager",
+    "IndexingResult",
+    "RecordManager",
+    "UpsertResponse",
+    "aindex",
+    "index",
+)
+
+_dynamic_imports = {
+    "aindex": "api",
+    "index": "api",
+    "IndexingResult": "api",
+    "DeleteResponse": "base",
+    "DocumentIndex": "base",
+    "InMemoryRecordManager": "base",
+    "RecordManager": "base",
+    "UpsertResponse": "base",
+}
+
+
+def __getattr__(attr_name: str) -> object:
+    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 list(__all__)