coco-framework 0.1.0__py3-none-any.whl

coco_core/indexing/in_memory.py

@@ -0,0 +1,104 @@
+"""In memory document index."""
+
+import operator
+import uuid
+from collections.abc import Sequence
+from typing import Any, cast
+
+from pydantic import Field
+from typing_extensions import override
+
+from coco_core._api import beta
+from coco_core.callbacks import CallbackManagerForRetrieverRun
+from coco_core.documents import Document
+from coco_core.indexing import UpsertResponse
+from coco_core.indexing.base import DeleteResponse, DocumentIndex
+
+
+@beta(message="Introduced in version 0.2.29. Underlying abstraction subject to change.")
+class InMemoryDocumentIndex(DocumentIndex):
+    """In memory document index.
+
+    This is an in-memory document index that stores documents in a dictionary.
+
+    It provides a simple search API that returns documents by the number of
+    counts the given query appears in the document.
+    """
+
+    store: dict[str, Document] = Field(default_factory=dict)
+    top_k: int = 4
+
+    @override
+    def upsert(self, items: Sequence[Document], /, **kwargs: Any) -> UpsertResponse:
+        """Upsert documents into the index.
+
+        Args:
+            items: Sequence of documents to add to the index.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            A response object that contains the list of IDs that were
+            successfully added or updated in the index and the list of IDs that
+            failed to be added or updated.
+        """
+        ok_ids = []
+
+        for item in items:
+            if item.id is None:
+                id_ = str(uuid.uuid4())
+                item_ = item.model_copy()
+                item_.id = id_
+            else:
+                item_ = item
+                id_ = item.id
+
+            self.store[id_] = item_
+            ok_ids.append(cast("str", item_.id))
+
+        return UpsertResponse(succeeded=ok_ids, failed=[])
+
+    @override
+    def delete(self, ids: list[str] | None = None, **kwargs: Any) -> DeleteResponse:
+        """Delete by IDs.
+
+        Args:
+            ids: List of IDs to delete.
+
+        Raises:
+            ValueError: If IDs is None.
+
+        Returns:
+            A response object that contains the list of IDs that were successfully
+            deleted and the list of IDs that failed to be deleted.
+        """
+        if ids is None:
+            msg = "IDs must be provided for deletion"
+            raise ValueError(msg)
+
+        ok_ids = []
+
+        for id_ in ids:
+            if id_ in self.store:
+                del self.store[id_]
+                ok_ids.append(id_)
+
+        return DeleteResponse(
+            succeeded=ok_ids, num_deleted=len(ok_ids), num_failed=0, failed=[]
+        )
+
+    @override
+    def get(self, ids: Sequence[str], /, **kwargs: Any) -> list[Document]:
+        return [self.store[id_] for id_ in ids if id_ in self.store]
+
+    @override
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> list[Document]:
+        counts_by_doc = []
+
+        for document in self.store.values():
+            count = document.page_content.count(query)
+            counts_by_doc.append((document, count))
+
+        counts_by_doc.sort(key=operator.itemgetter(1), reverse=True)
+        return [doc.model_copy() for doc, count in counts_by_doc[: self.top_k]]

coco_core/language_models/__init__.py

@@ -0,0 +1,116 @@
+"""Core language model abstractions.
+
+Coco has two main classes to work with language models: chat models and
+"old-fashioned" LLMs (string-in, string-out).
+
+**Chat models**
+
+Language models that use a sequence of messages as inputs and return chat messages
+as outputs (as opposed to using plain text).
+
+Chat models support the assignment of distinct roles to conversation messages, helping
+to distinguish messages from the AI, users, and instructions such as system messages.
+
+The key abstraction for chat models is
+[`BaseChatModel`][coco_core.language_models.BaseChatModel]. Implementations should
+inherit from this class.
+
+See existing [chat model integrations](https://docs.langchain.com/oss/python/integrations/chat).
+
+**LLMs (legacy)**
+
+Language models that takes a string as input and returns a string.
+
+These are traditionally older models (newer models generally are chat models).
+
+Although the underlying models are string in, string out, the Coco wrappers also
+allow these models to take messages as input. This gives them the same interface as
+chat models. When messages are passed in as input, they will be formatted into a string
+under the hood before being passed to the underlying model.
+"""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+from coco_core.language_models._utils import is_openai_data_block
+
+if TYPE_CHECKING:
+    from coco_core.language_models.base import (
+        BaseLanguageModel,
+        LangSmithParams,
+        LanguageModelInput,
+        LanguageModelLike,
+        LanguageModelOutput,
+        get_tokenizer,
+    )
+    from coco_core.language_models.chat_models import (
+        BaseChatModel,
+        SimpleChatModel,
+    )
+    from coco_core.language_models.fake import FakeListLLM, FakeStreamingListLLM
+    from coco_core.language_models.fake_chat_models import (
+        FakeListChatModel,
+        FakeMessagesListChatModel,
+        GenericFakeChatModel,
+        ParrotFakeChatModel,
+    )
+    from coco_core.language_models.llms import LLM, BaseLLM
+    from coco_core.language_models.model_profile import (
+        ModelProfile,
+        ModelProfileRegistry,
+    )
+
+__all__ = (
+    "LLM",
+    "BaseChatModel",
+    "BaseLLM",
+    "BaseLanguageModel",
+    "FakeListChatModel",
+    "FakeListLLM",
+    "FakeMessagesListChatModel",
+    "FakeStreamingListLLM",
+    "GenericFakeChatModel",
+    "LangSmithParams",
+    "LanguageModelInput",
+    "LanguageModelLike",
+    "LanguageModelOutput",
+    "ModelProfile",
+    "ModelProfileRegistry",
+    "ParrotFakeChatModel",
+    "SimpleChatModel",
+    "get_tokenizer",
+    "is_openai_data_block",
+)
+
+_dynamic_imports = {
+    "BaseLanguageModel": "base",
+    "LangSmithParams": "base",
+    "LanguageModelInput": "base",
+    "LanguageModelLike": "base",
+    "LanguageModelOutput": "base",
+    "get_tokenizer": "base",
+    "BaseChatModel": "chat_models",
+    "SimpleChatModel": "chat_models",
+    "FakeListLLM": "fake",
+    "FakeStreamingListLLM": "fake",
+    "FakeListChatModel": "fake_chat_models",
+    "FakeMessagesListChatModel": "fake_chat_models",
+    "GenericFakeChatModel": "fake_chat_models",
+    "ParrotFakeChatModel": "fake_chat_models",
+    "LLM": "llms",
+    "ModelProfile": "model_profile",
+    "ModelProfileRegistry": "model_profile",
+    "BaseLLM": "llms",
+    "is_openai_data_block": "_utils",
+}
+
+
+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__)