coco-framework 0.1.0__py3-none-any.whl

coco_core/runnables/history.py

@@ -0,0 +1,631 @@
+"""`Runnable` that manages chat message history for another `Runnable`."""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable, Sequence
+from types import GenericAlias
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+from pydantic import BaseModel
+from typing_extensions import override
+
+from coco_core._api.deprecation import warn_deprecated
+from coco_core.chat_history import BaseChatMessageHistory
+from coco_core.load.load import load
+from coco_core.messages import AIMessage, BaseMessage, HumanMessage
+from coco_core.runnables.base import Runnable, RunnableBindingBase, RunnableLambda
+from coco_core.runnables.passthrough import RunnablePassthrough
+from coco_core.runnables.utils import (
+    ConfigurableFieldSpec,
+    Output,
+    get_unique_config_specs,
+)
+from coco_core.utils.pydantic import create_model_v2
+
+if TYPE_CHECKING:
+    from coco_core.language_models.base import LanguageModelLike
+    from coco_core.runnables.config import RunnableConfig
+    from coco_core.tracers.schemas import Run
+
+
+MessagesOrDictWithMessages = Sequence["BaseMessage"] | dict[str, Any]
+GetSessionHistoryCallable = Callable[..., BaseChatMessageHistory]
+
+
+class RunnableWithMessageHistory(RunnableBindingBase[Any, Any]):  # type: ignore[no-redef]
+    """`Runnable` that manages chat message history for another `Runnable`.
+
+    A chat message history is a sequence of messages that represent a conversation.
+
+    `RunnableWithMessageHistory` wraps another `Runnable` and manages the chat message
+    history for it; it is responsible for reading and updating the chat message
+    history.
+
+    The formats supported for the inputs and outputs of the wrapped `Runnable`
+    are described below.
+
+    `RunnableWithMessageHistory` must always be called with a config that contains
+    the appropriate parameters for the chat message history factory.
+
+    By default, the `Runnable` is expected to take a single configuration parameter
+    called `session_id` which is a string. This parameter is used to create a new
+    or look up an existing chat message history that matches the given `session_id`.
+
+    In this case, the invocation would look like this:
+
+    `with_history.invoke(..., config={"configurable": {"session_id": "bar"}})`
+    ; e.g., `{"configurable": {"session_id": "<SESSION_ID>"}}`.
+
+    The configuration can be customized by passing in a list of
+    `ConfigurableFieldSpec` objects to the `history_factory_config` parameter (see
+    example below).
+
+    In the examples, we will use a chat message history with an in-memory
+    implementation to make it easy to experiment and see the results.
+
+    For production use cases, you will want to use a persistent implementation
+    of chat message history, such as `RedisChatMessageHistory`.
+
+    Example: Chat message history with an in-memory implementation for testing.
+
+        ```python
+        from operator import itemgetter
+
+        from coco_openai.chat_models import ChatOpenAI
+
+        from coco_core.chat_history import BaseChatMessageHistory
+        from coco_core.documents import Document
+        from coco_core.messages import BaseMessage, AIMessage
+        from coco_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+        from pydantic import BaseModel, Field
+        from coco_core.runnables import (
+            RunnableLambda,
+            ConfigurableFieldSpec,
+            RunnablePassthrough,
+        )
+        from coco_core.runnables.history import RunnableWithMessageHistory
+
+
+        class InMemoryHistory(BaseChatMessageHistory, BaseModel):
+            \"\"\"In memory implementation of chat message history.\"\"\"
+
+            messages: list[BaseMessage] = Field(default_factory=list)
+
+            def add_messages(self, messages: list[BaseMessage]) -> None:
+                \"\"\"Add a list of messages to the store\"\"\"
+                self.messages.extend(messages)
+
+            def clear(self) -> None:
+                self.messages = []
+
+        # Here we use a global variable to store the chat message history.
+        # This will make it easier to inspect it to see the underlying results.
+        store = {}
+
+        def get_by_session_id(session_id: str) -> BaseChatMessageHistory:
+            if session_id not in store:
+                store[session_id] = InMemoryHistory()
+            return store[session_id]
+
+
+        history = get_by_session_id("1")
+        history.add_message(AIMessage(content="hello"))
+        print(store)  # noqa: T201
+
+        ```
+
+    Example where the wrapped `Runnable` takes a dictionary input:
+
+        ```python
+        from typing import Optional
+
+        from coco_anthropic import ChatAnthropic
+        from coco_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+        from coco_core.runnables.history import RunnableWithMessageHistory
+
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                ("system", "You're an assistant who's good at {ability}"),
+                MessagesPlaceholder(variable_name="history"),
+                ("human", "{question}"),
+            ]
+        )
+
+        chain = prompt | ChatAnthropic(model="claude-sonnet-5")
+
+        chain_with_history = RunnableWithMessageHistory(
+            chain,
+            # Uses the get_by_session_id function defined in the example
+            # above.
+            get_by_session_id,
+            input_messages_key="question",
+            history_messages_key="history",
+        )
+
+        print(
+            chain_with_history.invoke(  # noqa: T201
+                {"ability": "math", "question": "What does cosine mean?"},
+                config={"configurable": {"session_id": "foo"}},
+            )
+        )
+
+        # Uses the store defined in the example above.
+        print(store)  # noqa: T201
+
+        print(
+            chain_with_history.invoke(  # noqa: T201
+                {"ability": "math", "question": "What's its inverse"},
+                config={"configurable": {"session_id": "foo"}},
+            )
+        )
+
+        print(store)  # noqa: T201
+        ```
+
+    Example where the session factory takes two keys (`user_id` and `conversation_id`):
+
+        ```python
+        store = {}
+
+
+        def get_session_history(
+            user_id: str, conversation_id: str
+        ) -> BaseChatMessageHistory:
+            if (user_id, conversation_id) not in store:
+                store[(user_id, conversation_id)] = InMemoryHistory()
+            return store[(user_id, conversation_id)]
+
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                ("system", "You're an assistant who's good at {ability}"),
+                MessagesPlaceholder(variable_name="history"),
+                ("human", "{question}"),
+            ]
+        )
+
+        chain = prompt | ChatAnthropic(model="claude-sonnet-5")
+
+        with_message_history = RunnableWithMessageHistory(
+            chain,
+            get_session_history=get_session_history,
+            input_messages_key="question",
+            history_messages_key="history",
+            history_factory_config=[
+                ConfigurableFieldSpec(
+                    id="user_id",
+                    annotation=str,
+                    name="User ID",
+                    description="Unique identifier for the user.",
+                    default="",
+                    is_shared=True,
+                ),
+                ConfigurableFieldSpec(
+                    id="conversation_id",
+                    annotation=str,
+                    name="Conversation ID",
+                    description="Unique identifier for the conversation.",
+                    default="",
+                    is_shared=True,
+                ),
+            ],
+        )
+
+        with_message_history.invoke(
+            {"ability": "math", "question": "What does cosine mean?"},
+            config={"configurable": {"user_id": "123", "conversation_id": "1"}},
+        )
+        ```
+    """
+
+    get_session_history: GetSessionHistoryCallable
+    """Function that returns a new `BaseChatMessageHistory`.
+
+    This function should either take a single positional argument `session_id` of type
+    string and return a corresponding chat message history instance
+    """
+    input_messages_key: str | None = None
+    """Must be specified if the base `Runnable` accepts a `dict` as input.
+    The key in the input `dict` that contains the messages.
+    """
+    output_messages_key: str | None = None
+    """Must be specified if the base `Runnable` returns a `dict` as output.
+    The key in the output `dict` that contains the messages.
+    """
+    history_messages_key: str | None = None
+    """Must be specified if the base `Runnable` accepts a `dict` as input and expects a
+    separate key for historical messages.
+    """
+    history_factory_config: Sequence[ConfigurableFieldSpec]
+    """Configure fields that should be passed to the chat history factory.
+
+    See `ConfigurableFieldSpec` for more details.
+    """
+
+    def __init__(
+        self,
+        runnable: Runnable[
+            list[BaseMessage], str | BaseMessage | MessagesOrDictWithMessages
+        ]
+        | Runnable[dict[str, Any], str | BaseMessage | MessagesOrDictWithMessages]
+        | LanguageModelLike,
+        get_session_history: GetSessionHistoryCallable,
+        *,
+        input_messages_key: str | None = None,
+        output_messages_key: str | None = None,
+        history_messages_key: str | None = None,
+        history_factory_config: Sequence[ConfigurableFieldSpec] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize `RunnableWithMessageHistory`.
+
+        Args:
+            runnable: The base `Runnable` to be wrapped.
+
+                Must take as input one of:
+
+                1. A list of `BaseMessage`
+                2. A `dict` with one key for all messages
+                3. A `dict` with one key for the current input string/message(s) and
+                    a separate key for historical messages. If the input key points
+                    to a string, it will be treated as a `HumanMessage` in history.
+
+                Must return as output one of:
+
+                1. A string which can be treated as an `AIMessage`
+                2. A `BaseMessage` or sequence of `BaseMessage`
+                3. A `dict` with a key for a `BaseMessage` or sequence of
+                    `BaseMessage`
+
+            get_session_history: Function that returns a new `BaseChatMessageHistory`.
+
+                This function should either take a single positional argument
+                `session_id` of type string and return a corresponding
+                chat message history instance.
+
+                ```python
+                def get_session_history(
+                    session_id: str, *, user_id: str | None = None
+                ) -> BaseChatMessageHistory: ...
+                ```
+
+                Or it should take keyword arguments that match the keys of
+                `session_history_config_specs` and return a corresponding
+                chat message history instance.
+
+                ```python
+                def get_session_history(
+                    *,
+                    user_id: str,
+                    thread_id: str,
+                ) -> BaseChatMessageHistory: ...
+                ```
+
+            input_messages_key: Must be specified if the base runnable accepts a `dict`
+                as input.
+            output_messages_key: Must be specified if the base runnable returns a `dict`
+                as output.
+            history_messages_key: Must be specified if the base runnable accepts a
+                `dict` as input and expects a separate key for historical messages.
+            history_factory_config: Configure fields that should be passed to the
+                chat history factory. See `ConfigurableFieldSpec` for more details.
+
+                Specifying these allows you to pass multiple config keys into the
+                `get_session_history` factory.
+            **kwargs: Arbitrary additional kwargs to pass to parent class
+                `RunnableBindingBase` init.
+
+        """
+        warn_deprecated(
+            since="1.3.3",
+            message=(
+                "RunnableWithMessageHistory is deprecated. "
+                "Use CocoGraph's built-in persistence instead."
+            ),
+            removal="2.0.0",
+        )
+        history_chain: Runnable[Any, Any] = RunnableLambda(
+            self._enter_history, self._aenter_history
+        ).with_config(run_name="load_history")
+        messages_key = history_messages_key or input_messages_key
+        if messages_key:
+            history_chain = RunnablePassthrough.assign(
+                **{messages_key: history_chain}
+            ).with_config(run_name="insert_history")
+
+        runnable_sync = runnable.with_listeners(on_end=self._exit_history)
+        runnable_async = runnable.with_alisteners(on_end=self._aexit_history)
+
+        def _call_runnable_sync(_input: Any) -> Runnable[Any, Any]:
+            return runnable_sync
+
+        async def _call_runnable_async(_input: Any) -> Runnable[Any, Any]:
+            return runnable_async
+
+        bound = (
+            history_chain
+            | RunnableLambda(
+                _call_runnable_sync,
+                _call_runnable_async,
+            ).with_config(run_name="check_sync_or_async")
+        ).with_config(run_name="RunnableWithMessageHistory")
+
+        if history_factory_config:
+            config_specs = history_factory_config
+        else:
+            # If not provided, then we'll use the default session_id field
+            config_specs = [
+                ConfigurableFieldSpec(
+                    id="session_id",
+                    annotation=str,
+                    name="Session ID",
+                    description="Unique identifier for a session.",
+                    default="",
+                    is_shared=True,
+                ),
+            ]
+
+        super().__init__(
+            get_session_history=get_session_history,
+            input_messages_key=input_messages_key,
+            output_messages_key=output_messages_key,
+            bound=bound,
+            history_messages_key=history_messages_key,
+            history_factory_config=config_specs,
+            **kwargs,
+        )
+        self._history_chain = history_chain
+
+    @property
+    @override
+    def config_specs(self) -> list[ConfigurableFieldSpec]:
+        """Get the configuration specs for the `RunnableWithMessageHistory`."""
+        return get_unique_config_specs(
+            super().config_specs + list(self.history_factory_config)
+        )
+
+    @override
+    def get_input_schema(self, config: RunnableConfig | None = None) -> type[BaseModel]:
+        fields: dict[str, Any] = {}
+        if self.input_messages_key and self.history_messages_key:
+            fields[self.input_messages_key] = (
+                str | BaseMessage | Sequence[BaseMessage],
+                ...,
+            )
+        elif self.input_messages_key:
+            fields[self.input_messages_key] = (Sequence[BaseMessage], ...)
+        else:
+            return create_model_v2(
+                "RunnableWithChatHistoryInput",
+                module_name=self.__class__.__module__,
+                root=(Sequence[BaseMessage], ...),
+            )
+        return create_model_v2(
+            "RunnableWithChatHistoryInput",
+            field_definitions=fields,
+            module_name=self.__class__.__module__,
+        )
+
+    @property
+    @override
+    def OutputType(self) -> type[Output]:
+        return self._history_chain.OutputType
+
+    @override
+    def get_output_schema(
+        self, config: RunnableConfig | None = None
+    ) -> type[BaseModel]:
+        """Get a Pydantic model that can be used to validate output to the `Runnable`.
+
+        `Runnable` objects that leverage the `configurable_fields` and
+        `configurable_alternatives` methods will have a dynamic output schema that
+        depends on which configuration the `Runnable` is invoked with.
+
+        This method allows to get an output schema for a specific configuration.
+
+        Args:
+            config: A config to use when generating the schema.
+
+        Returns:
+            A Pydantic model that can be used to validate output.
+        """
+        root_type = self.OutputType
+
+        if (
+            inspect.isclass(root_type)
+            and not isinstance(root_type, GenericAlias)
+            and issubclass(root_type, BaseModel)
+        ):
+            return root_type
+
+        return create_model_v2(
+            "RunnableWithChatHistoryOutput",
+            root=root_type,
+            module_name=self.__class__.__module__,
+        )
+
+    def _get_input_messages(
+        self, input_val: str | BaseMessage | Sequence[BaseMessage] | dict[str, Any]
+    ) -> list[BaseMessage]:
+        # If dictionary, try to pluck the single key representing messages
+        if isinstance(input_val, dict):
+            if self.input_messages_key:
+                key = self.input_messages_key
+            elif len(input_val) == 1:
+                key = next(iter(input_val.keys()))
+            else:
+                key = "input"
+            input_val = input_val[key]
+
+        # If value is a string, convert to a human message
+        if isinstance(input_val, str):
+            return [HumanMessage(content=input_val)]
+        # If value is a single message, convert to a list
+        if isinstance(input_val, BaseMessage):
+            return [input_val]
+        # If value is a list or tuple...
+        if isinstance(input_val, (list, tuple)):
+            # Handle empty case
+            if len(input_val) == 0:
+                return list(input_val)
+            # If is a list of list, then return the first value
+            # This occurs for chat models - since we batch inputs
+            if isinstance(input_val[0], list):
+                if len(input_val) != 1:
+                    msg = f"Expected a single list of messages. Got {input_val}."
+                    raise ValueError(msg)
+                return input_val[0]
+            return list(input_val)
+        msg = (
+            f"Expected str, BaseMessage, list[BaseMessage], or tuple[BaseMessage]. "
+            f"Got {input_val}."
+        )
+        raise ValueError(msg)
+
+    def _get_output_messages(
+        self, output_val: str | BaseMessage | Sequence[BaseMessage] | dict[str, Any]
+    ) -> list[BaseMessage]:
+        # If dictionary, try to pluck the single key representing messages
+        if isinstance(output_val, dict):
+            if self.output_messages_key:
+                key = self.output_messages_key
+            elif len(output_val) == 1:
+                key = next(iter(output_val.keys()))
+            else:
+                key = "output"
+            # If you are wrapping a chat model directly
+            # The output is actually this weird generations object
+            if key not in output_val and "generations" in output_val:
+                output_val = output_val["generations"][0][0]["message"]
+            else:
+                output_val = output_val[key]
+
+        if isinstance(output_val, str):
+            return [AIMessage(content=output_val)]
+        # If value is a single message, convert to a list
+        if isinstance(output_val, BaseMessage):
+            return [output_val]
+        if isinstance(output_val, (list, tuple)):
+            return list(output_val)
+        msg = (
+            f"Expected str, BaseMessage, list[BaseMessage], or tuple[BaseMessage]. "
+            f"Got {output_val}."
+        )
+        raise ValueError(msg)
+
+    def _enter_history(self, value: Any, config: RunnableConfig) -> list[BaseMessage]:
+        hist: BaseChatMessageHistory = config["configurable"]["message_history"]
+        messages = hist.messages.copy()
+
+        if not self.history_messages_key:
+            # return all messages
+            input_val = (
+                value if not self.input_messages_key else value[self.input_messages_key]
+            )
+            messages += self._get_input_messages(input_val)
+        return messages
+
+    async def _aenter_history(
+        self, value: dict[str, Any], config: RunnableConfig
+    ) -> list[BaseMessage]:
+        hist: BaseChatMessageHistory = config["configurable"]["message_history"]
+        messages = (await hist.aget_messages()).copy()
+
+        if not self.history_messages_key:
+            # return all messages
+            input_val = (
+                value if not self.input_messages_key else value[self.input_messages_key]
+            )
+            messages += self._get_input_messages(input_val)
+        return messages
+
+    def _exit_history(self, run: Run, config: RunnableConfig) -> None:
+        hist: BaseChatMessageHistory = config["configurable"]["message_history"]
+
+        # Get the input messages
+        inputs = load(run.inputs, allowed_objects="messages")
+        input_messages = self._get_input_messages(inputs)
+        # If historic messages were prepended to the input messages, remove them to
+        # avoid adding duplicate messages to history.
+        if not self.history_messages_key:
+            historic_messages = config["configurable"]["message_history"].messages
+            input_messages = input_messages[len(historic_messages) :]
+
+        # Get the output messages
+        output_val = load(run.outputs, allowed_objects="messages")
+        output_messages = self._get_output_messages(output_val)
+        hist.add_messages(input_messages + output_messages)
+
+    async def _aexit_history(self, run: Run, config: RunnableConfig) -> None:
+        hist: BaseChatMessageHistory = config["configurable"]["message_history"]
+
+        # Get the input messages
+        inputs = load(run.inputs, allowed_objects="messages")
+        input_messages = self._get_input_messages(inputs)
+        # If historic messages were prepended to the input messages, remove them to
+        # avoid adding duplicate messages to history.
+        if not self.history_messages_key:
+            historic_messages = await hist.aget_messages()
+            input_messages = input_messages[len(historic_messages) :]
+
+        # Get the output messages
+        output_val = load(run.outputs, allowed_objects="messages")
+        output_messages = self._get_output_messages(output_val)
+        await hist.aadd_messages(input_messages + output_messages)
+
+    def _merge_configs(self, *configs: RunnableConfig | None) -> RunnableConfig:
+        config = super()._merge_configs(*configs)
+        expected_keys = [field_spec.id for field_spec in self.history_factory_config]
+
+        configurable = config.get("configurable", {})
+
+        missing_keys = set(expected_keys) - set(configurable.keys())
+        parameter_names = _get_parameter_names(self.get_session_history)
+
+        if missing_keys and parameter_names:
+            example_input = {self.input_messages_key: "foo"}
+            example_configurable = dict.fromkeys(missing_keys, "[your-value-here]")
+            example_config = {"configurable": example_configurable}
+            msg = (
+                f"Missing keys {sorted(missing_keys)} in config['configurable'] "
+                f"Expected keys are {sorted(expected_keys)}."
+                f"When using via .invoke() or .stream(), pass in a config; "
+                f"e.g., chain.invoke({example_input}, {example_config})"
+            )
+            raise ValueError(msg)
+
+        if len(expected_keys) == 1:
+            if parameter_names:
+                # If arity = 1, then invoke function by positional arguments
+                message_history = self.get_session_history(
+                    configurable[expected_keys[0]]
+                )
+            else:
+                if not config:
+                    config["configurable"] = {}
+                message_history = self.get_session_history()
+        else:
+            # otherwise verify that names of keys patch and invoke by named arguments
+            if set(expected_keys) != set(parameter_names):
+                msg = (
+                    f"Expected keys {sorted(expected_keys)} do not match parameter "
+                    f"names {sorted(parameter_names)} of get_session_history."
+                )
+                raise ValueError(msg)
+
+            message_history = self.get_session_history(
+                **{key: configurable[key] for key in expected_keys}
+            )
+        config["configurable"]["message_history"] = message_history
+        return config
+
+
+def _get_parameter_names(callable_: GetSessionHistoryCallable) -> list[str]:
+    """Get the parameter names of the `Callable`."""
+    sig = inspect.signature(callable_)
+    return list(sig.parameters.keys())