genesis-flow 1.0.0__py3-none-any.whl

mlflow/langchain/output_parsers.py

@@ -0,0 +1,142 @@
+from dataclasses import asdict
+from typing import Any, Iterator
+from uuid import uuid4
+
+from langchain_core.messages.base import BaseMessage
+from langchain_core.output_parsers.transform import BaseTransformOutputParser
+
+from mlflow.models.rag_signatures import (
+    ChainCompletionChoice,
+    Message,
+    StringResponse,
+)
+from mlflow.models.rag_signatures import (
+    ChatCompletionResponse as RagChatCompletionResponse,
+)
+from mlflow.types.agent import ChatAgentChunk, ChatAgentMessage, ChatAgentResponse
+from mlflow.types.llm import (
+    ChatChoice,
+    ChatChoiceDelta,
+    ChatChunkChoice,
+    ChatCompletionChunk,
+    ChatCompletionResponse,
+    ChatMessage,
+)
+from mlflow.utils.annotations import deprecated, experimental
+
+
+@deprecated("mlflow.langchain.output_parser.ChatCompletionOutputParser")
+class ChatCompletionsOutputParser(BaseTransformOutputParser[dict[str, Any]]):
+    """
+    OutputParser that wraps the string output into a dictionary representation of a
+    :py:class:`ChatCompletionResponse`
+    """
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @property
+    def _type(self) -> str:
+        """Return the output parser type for serialization."""
+        return "mlflow_simplified_chat_completions"
+
+    def parse(self, text: str) -> dict[str, Any]:
+        return asdict(
+            RagChatCompletionResponse(
+                choices=[ChainCompletionChoice(message=Message(role="assistant", content=text))],
+                object="chat.completion",
+            )
+        )
+
+
+class ChatCompletionOutputParser(BaseTransformOutputParser[str]):
+    """
+    OutputParser that wraps the string output into a dictionary representation of a
+    :py:class:`ChatCompletionResponse` or :py:class:`ChatCompletionChunk`
+    when streaming
+    """
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @property
+    def _type(self) -> str:
+        """Return the output parser type for serialization."""
+        return "mlflow_chat_completion"
+
+    def parse(self, text: str) -> dict[str, Any]:
+        """Returns the input text as a ChatCompletionResponse with no changes."""
+        return ChatCompletionResponse(
+            choices=[ChatChoice(message=ChatMessage(role="assistant", content=text))]
+        ).to_dict()
+
+    def transform(self, input: Iterator[BaseMessage], config, **kwargs) -> Iterator[dict[str, Any]]:
+        """Returns a generator of ChatCompletionChunk objects"""
+        for chunk in input:
+            yield ChatCompletionChunk(
+                choices=[ChatChunkChoice(delta=ChatChoiceDelta(content=chunk.content))]
+            ).to_dict()
+
+
+@deprecated("mlflow.langchain.output_parser.ChatCompletionOutputParser")
+class StringResponseOutputParser(BaseTransformOutputParser[dict[str, Any]]):
+    """
+    OutputParser that wraps the string output into an dictionary representation of a
+    :py:class:`StringResponse`
+    """
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @property
+    def _type(self) -> str:
+        """Return the output parser type for serialization."""
+        return "mlflow_simplified_str_object"
+
+    def parse(self, text: str) -> dict[str, Any]:
+        return asdict(StringResponse(content=text))
+
+
+@experimental(version="2.21.0")
+class ChatAgentOutputParser(BaseTransformOutputParser[str]):
+    """
+    OutputParser that wraps the string output into a dictionary representation of a
+    :py:class:`ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>` or a
+    :py:class:`ChatAgentChunk <mlflow.types.agent.ChatAgentChunk>` for easy interoperability.
+    """
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @property
+    def _type(self) -> str:
+        """Return the output parser type for serialization."""
+        return "mlflow_chat_agent"
+
+    def parse(self, text: str) -> dict[str, Any]:
+        """
+        Returns the output text as a dictionary representation of a
+        :py:class:`ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>`.
+        """
+        return ChatAgentResponse(
+            messages=[ChatAgentMessage(content=text, role="assistant", id=str(uuid4()))]
+        ).model_dump_compat(exclude_none=True)
+
+    def transform(self, input: Iterator[BaseMessage], config, **kwargs) -> Iterator[dict[str, Any]]:
+        """
+        Returns a generator of
+        :py:class:`ChatAgentChunk <mlflow.types.agent.ChatAgentChunk>` objects
+        """
+        for chunk in input:
+            if chunk.content:
+                yield ChatAgentChunk(
+                    delta=ChatAgentMessage(content=chunk.content, role="assistant", id=chunk.id)
+                ).model_dump_compat(exclude_none=True)

mlflow/langchain/retriever_chain.py

@@ -0,0 +1,153 @@
+"""Chain for wrapping a retriever."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import yaml
+from langchain.callbacks.manager import AsyncCallbackManagerForChainRun, CallbackManagerForChainRun
+from langchain.chains.base import Chain
+from langchain.schema import BaseRetriever, Document
+from pydantic import Extra, Field
+
+
+class _RetrieverChain(Chain):
+    """
+    Chain that wraps a retriever for use with MLflow.
+
+    The MLflow ``langchain`` flavor provides the functionality to log a retriever object and
+    evaluate it individually. This is useful if you want to evaluate the quality of the
+    relevant documents returned by a retriever object without directing these documents
+    through a large language model (LLM) to yield a summarized response.
+
+    In order to log the retriever object in the ``langchain`` flavor, the retriever object
+    needs to be wrapped within a ``_RetrieverChain``.
+
+    See ``examples/langchain/retriever_chain.py`` for how to log the ``_RetrieverChain``.
+
+    Args:
+        retriever: The retriever to wrap.
+    """
+
+    input_key: str = "query"
+    output_key: str = "source_documents"
+    retriever: BaseRetriever = Field(exclude=True)
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+    @property
+    def input_keys(self) -> list[str]:
+        """Return the input keys."""
+        return [self.input_key]
+
+    @property
+    def output_keys(self) -> list[str]:
+        """Return the output keys."""
+        return [self.output_key]
+
+    def _get_docs(self, question: str) -> list[Document]:
+        """Get documents from the retriever."""
+        return self.retriever.get_relevant_documents(question)
+
+    def _call(
+        self,
+        inputs: dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> dict[str, Any]:
+        """Run _get_docs on input query.
+        Returns the retrieved documents under the key 'source_documents'.
+
+        Example:
+
+        .. code-block:: python
+
+            chain = _RetrieverChain(retriever=...)
+            res = chain({"query": "This is my query"})
+            docs = res["source_documents"]
+        """
+        question = inputs[self.input_key]
+        docs = self._get_docs(question)
+        list_of_str_page_content = [doc.page_content for doc in docs]
+        return {self.output_key: json.dumps(list_of_str_page_content)}
+
+    async def _aget_docs(self, question: str) -> list[Document]:
+        """Get documents from the retriever."""
+        return await self.retriever.aget_relevant_documents(question)
+
+    async def _acall(
+        self,
+        inputs: dict[str, Any],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> dict[str, Any]:
+        """Run _get_docs on input query.
+        Returns the retrieved documents under the key 'source_documents'.
+
+        Example:
+
+        .. code-block:: python
+
+            chain = _RetrieverChain(retriever=...)
+            res = chain({"query": "This is my query"})
+            docs = res["source_documents"]
+        """
+        question = inputs[self.input_key]
+        docs = await self._aget_docs(question)
+        list_of_str_page_content = [doc.page_content for doc in docs]
+        return {self.output_key: json.dumps(list_of_str_page_content)}
+
+    @property
+    def _chain_type(self) -> str:
+        """Return the chain type."""
+        return "retriever_chain"
+
+    @classmethod
+    def load(cls, file: Union[str, Path], **kwargs: Any) -> _RetrieverChain:
+        """Load a _RetrieverChain from a file."""
+        # Convert file to Path object.
+        file_path = Path(file) if isinstance(file, str) else file
+        # Load from either json or yaml.
+        if file_path.suffix == ".json":
+            with open(file_path) as f:
+                config = json.load(f)
+        elif file_path.suffix in (".yaml", ".yml"):
+            with open(file_path) as f:
+                # This is to ignore certain tags that are not supported
+                # with pydantic >= 2.0
+                yaml.add_multi_constructor(
+                    "tag:yaml.org,2002:python/object",
+                    lambda loader, suffix, node: None,
+                    Loader=yaml.SafeLoader,
+                )
+                config = yaml.load(f, yaml.SafeLoader)
+        else:
+            raise ValueError("File type must be json or yaml")
+
+        # Override default 'verbose' and 'memory' for the chain
+        if verbose := kwargs.pop("verbose", None):
+            config["verbose"] = verbose
+        if memory := kwargs.pop("memory", None):
+            config["memory"] = memory
+
+        if "_type" not in config:
+            raise ValueError("Must specify a chain Type in config")
+        config_type = config.pop("_type")
+
+        if config_type != "retriever_chain":
+            raise ValueError(f"Loading {config_type} chain not supported")
+
+        retriever = kwargs.pop("retriever", None)
+        if retriever is None:
+            raise ValueError("`retriever` must be present.")
+
+        config.pop("retriever", None)
+
+        return cls(
+            retriever=retriever,
+            **config,
+        )