coco-framework 0.1.0__py3-none-any.whl

coco_core/output_parsers/transform.py

@@ -0,0 +1,175 @@
+"""Base classes for output parsers that can handle streaming input."""
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+from typing_extensions import override
+
+from coco_core.messages import BaseMessage, BaseMessageChunk
+from coco_core.output_parsers.base import BaseOutputParser, T
+from coco_core.outputs import (
+    ChatGeneration,
+    ChatGenerationChunk,
+    Generation,
+    GenerationChunk,
+)
+from coco_core.runnables.config import run_in_executor
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Iterator
+
+    from coco_core.runnables import RunnableConfig
+
+
+class BaseTransformOutputParser(BaseOutputParser[T]):
+    """Base class for an output parser that can handle streaming input."""
+
+    def _transform(
+        self,
+        input: Iterator[str | BaseMessage],
+    ) -> Iterator[T]:
+        for chunk in input:
+            if isinstance(chunk, BaseMessage):
+                yield self.parse_result([ChatGeneration(message=chunk)])
+            else:
+                yield self.parse_result([Generation(text=chunk)])
+
+    async def _atransform(
+        self,
+        input: AsyncIterator[str | BaseMessage],
+    ) -> AsyncIterator[T]:
+        async for chunk in input:
+            if isinstance(chunk, BaseMessage):
+                yield await run_in_executor(
+                    None, self.parse_result, [ChatGeneration(message=chunk)]
+                )
+            else:
+                yield await run_in_executor(
+                    None, self.parse_result, [Generation(text=chunk)]
+                )
+
+    @override
+    def transform(
+        self,
+        input: Iterator[str | BaseMessage],
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> Iterator[T]:
+        """Transform the input into the output format.
+
+        Args:
+            input: The input to transform.
+            config: The configuration to use for the transformation.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            The transformed output.
+        """
+        yield from self._transform_stream_with_config(
+            input, self._transform, config, run_type="parser"
+        )
+
+    @override
+    async def atransform(
+        self,
+        input: AsyncIterator[str | BaseMessage],
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[T]:
+        """Async transform the input into the output format.
+
+        Args:
+            input: The input to transform.
+            config: The configuration to use for the transformation.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            The transformed output.
+        """
+        async for chunk in self._atransform_stream_with_config(
+            input, self._atransform, config, run_type="parser"
+        ):
+            yield chunk
+
+
+class BaseCumulativeTransformOutputParser(BaseTransformOutputParser[T]):
+    """Base class for an output parser that can handle streaming input."""
+
+    diff: bool = False
+    """In streaming mode, whether to yield diffs between the previous and current parsed
+    output, or just the current parsed output.
+    """
+
+    def _diff(
+        self,
+        prev: T | None,
+        next: T,  # noqa: A002
+    ) -> T:
+        """Convert parsed outputs into a diff format.
+
+        The semantics of this are up to the output parser.
+
+        Args:
+            prev: The previous parsed output.
+            next: The current parsed output.
+
+        Returns:
+            The diff between the previous and current parsed output.
+        """
+        raise NotImplementedError
+
+    @override
+    def _transform(self, input: Iterator[str | BaseMessage]) -> Iterator[Any]:
+        prev_parsed = None
+        acc_gen: GenerationChunk | ChatGenerationChunk | None = None
+        for chunk in input:
+            chunk_gen: GenerationChunk | ChatGenerationChunk
+            if isinstance(chunk, BaseMessageChunk):
+                chunk_gen = ChatGenerationChunk(message=chunk)
+            elif isinstance(chunk, BaseMessage):
+                chunk_gen = ChatGenerationChunk(
+                    message=BaseMessageChunk(**chunk.model_dump())
+                )
+            else:
+                chunk_gen = GenerationChunk(text=chunk)
+
+            acc_gen = chunk_gen if acc_gen is None else acc_gen + chunk_gen  # type: ignore[operator]
+
+            parsed = self.parse_result([acc_gen], partial=True)
+            if parsed is not None and parsed != prev_parsed:
+                if self.diff:
+                    yield self._diff(prev_parsed, parsed)
+                else:
+                    yield parsed
+                prev_parsed = parsed
+
+    @override
+    async def _atransform(
+        self, input: AsyncIterator[str | BaseMessage]
+    ) -> AsyncIterator[T]:
+        prev_parsed = None
+        acc_gen: GenerationChunk | ChatGenerationChunk | None = None
+        async for chunk in input:
+            chunk_gen: GenerationChunk | ChatGenerationChunk
+            if isinstance(chunk, BaseMessageChunk):
+                chunk_gen = ChatGenerationChunk(message=chunk)
+            elif isinstance(chunk, BaseMessage):
+                chunk_gen = ChatGenerationChunk(
+                    message=BaseMessageChunk(**chunk.model_dump())
+                )
+            else:
+                chunk_gen = GenerationChunk(text=chunk)
+
+            acc_gen = chunk_gen if acc_gen is None else acc_gen + chunk_gen  # type: ignore[operator]
+
+            parsed = await self.aparse_result([acc_gen], partial=True)
+            if parsed is not None and parsed != prev_parsed:
+                if self.diff:
+                    yield await run_in_executor(None, self._diff, prev_parsed, parsed)
+                else:
+                    yield parsed
+                prev_parsed = parsed

coco_core/output_parsers/xml.py

@@ -0,0 +1,301 @@
+"""Output parser for XML format."""
+
+import contextlib
+import re
+import xml
+import xml.etree.ElementTree as ET
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Literal
+from xml.etree.ElementTree import TreeBuilder
+
+from typing_extensions import override
+
+from coco_core.exceptions import OutputParserException
+from coco_core.messages import BaseMessage
+from coco_core.output_parsers.transform import BaseTransformOutputParser
+from coco_core.runnables.utils import AddableDict
+
+try:
+    from defusedxml import ElementTree  # type: ignore[import-untyped]
+    from defusedxml.ElementTree import XMLParser  # type: ignore[import-untyped]
+
+    _HAS_DEFUSEDXML = True
+except ImportError:
+    _HAS_DEFUSEDXML = False
+
+XML_FORMAT_INSTRUCTIONS = """The output should be formatted as a XML file.
+1. Output should conform to the tags below.
+2. If tags are not given, make them on your own.
+3. Remember to always open and close all the tags.
+
+As an example, for the tags ["foo", "bar", "baz"]:
+1. String "<foo>\n   <bar>\n      <baz></baz>\n   </bar>\n</foo>" is a well-formatted instance of the schema.
+2. String "<foo>\n   <bar>\n   </foo>" is a badly-formatted instance.
+3. String "<foo>\n   <tag>\n   </tag>\n</foo>" is a badly-formatted instance.
+
+Here are the output tags:
+```
+{tags}
+```"""  # noqa: E501
+
+
+class _StreamingParser:
+    """Streaming parser for XML.
+
+    This implementation is pulled into a class to avoid implementation drift between
+    `transform` and `atransform` of the `XMLOutputParser`.
+    """
+
+    def __init__(self, parser: Literal["defusedxml", "xml"]) -> None:
+        """Initialize the streaming parser.
+
+        Args:
+            parser: Parser to use for XML parsing.
+
+                Can be either `'defusedxml'` or `'xml'`. See documentation in
+                `XMLOutputParser` for more information.
+
+        Raises:
+            ImportError: If `defusedxml` is not installed and the `defusedxml` parser is
+                requested.
+        """
+        if parser == "defusedxml":
+            if not _HAS_DEFUSEDXML:
+                msg = (
+                    "defusedxml is not installed. "
+                    "Please install it to use the defusedxml parser. "
+                    "You can install it with `pip install defusedxml`. "
+                    "See https://github.com/tiran/defusedxml for more details"
+                )
+                raise ImportError(msg)
+            parser_ = XMLParser(target=TreeBuilder())
+        else:
+            parser_ = None
+        self.pull_parser = ET.XMLPullParser(["start", "end"], _parser=parser_)
+        self.xml_start_re = re.compile(r"<[a-zA-Z:_]")
+        self.current_path: list[str] = []
+        self.current_path_has_children = False
+        self.buffer = ""
+        self.xml_started = False
+
+    def parse(self, chunk: str | BaseMessage) -> Iterator[AddableDict]:
+        """Parse a chunk of text.
+
+        Args:
+            chunk: A chunk of text to parse. This can be a `str` or a `BaseMessage`.
+
+        Yields:
+            A `dict` representing the parsed XML element.
+
+        Raises:
+            xml.etree.ElementTree.ParseError: If the XML is not well-formed.
+        """
+        if isinstance(chunk, BaseMessage):
+            # extract text
+            chunk_content = chunk.content
+            if not isinstance(chunk_content, str):
+                # ignore non-string messages (e.g., function calls)
+                return
+            chunk = chunk_content
+        # add chunk to buffer of unprocessed text
+        self.buffer += chunk
+        # if xml string hasn't started yet, continue to next chunk
+        if not self.xml_started:
+            if match := self.xml_start_re.search(self.buffer):
+                # if xml string has started, remove all text before it
+                self.buffer = self.buffer[match.start() :]
+                self.xml_started = True
+            else:
+                return
+        # feed buffer to parser
+        self.pull_parser.feed(self.buffer)
+        self.buffer = ""
+        # yield all events
+        try:
+            events = self.pull_parser.read_events()
+            for event, elem in events:  # type: ignore[misc]
+                if event == "start":
+                    # update current path
+                    self.current_path.append(elem.tag)  # type: ignore[union-attr]
+                    self.current_path_has_children = False
+                elif event == "end":
+                    # remove last element from current path
+                    #
+                    self.current_path.pop()
+                    # yield element
+                    if not self.current_path_has_children:
+                        yield nested_element(self.current_path, elem)  # type: ignore[arg-type]
+                    # prevent yielding of parent element
+                    if self.current_path:
+                        self.current_path_has_children = True
+                    else:
+                        self.xml_started = False
+        except xml.etree.ElementTree.ParseError:
+            # This might be junk at the end of the XML input.
+            # Let's check whether the current path is empty.
+            if not self.current_path:
+                # If it is empty, we can ignore this error.
+                return
+            else:
+                raise
+
+    def close(self) -> None:
+        """Close the parser.
+
+        This should be called after all chunks have been parsed.
+        """
+        # Ignore ParseError. This will ignore any incomplete XML at the end of the input
+        with contextlib.suppress(xml.etree.ElementTree.ParseError):
+            self.pull_parser.close()
+
+
+class XMLOutputParser(BaseTransformOutputParser[dict[str, Any]]):
+    """Parse an output using xml format.
+
+    Returns a dictionary of tags.
+    """
+
+    tags: list[str] | None = None
+    """Tags to tell the LLM to expect in the XML output.
+
+    Note this may not be perfect depending on the LLM implementation.
+
+    For example, with `tags=["foo", "bar", "baz"]`:
+
+    1. A well-formatted XML instance:
+        `'<foo>\n   <bar>\n      <baz></baz>\n   </bar>\n</foo>'`
+
+    2. A badly-formatted XML instance (missing closing tag for 'bar'):
+        `'<foo>\n   <bar>\n   </foo>'`
+
+    3. A badly-formatted XML instance (unexpected 'tag' element):
+        `'<foo>\n   <tag>\n   </tag>\n</foo>'`
+    """
+    encoding_matcher: re.Pattern[str] = re.compile(
+        r"<([^>]*encoding[^>]*)>\n(.*)", re.MULTILINE | re.DOTALL
+    )
+
+    parser: Literal["defusedxml", "xml"] = "defusedxml"
+    """Parser to use for XML parsing.
+
+    Can be either `'defusedxml'` or `'xml'`.
+
+    - `'defusedxml'` is the default parser and is used to prevent XML vulnerabilities
+        present in some distributions of Python's standard library xml. `defusedxml` is
+        a wrapper around the standard library parser that sets up the parser with secure
+        defaults.
+    - `'xml'` is the standard library parser.
+
+    !!! warning
+
+        Use `xml` only if you are sure that your distribution of the standard library is
+        not vulnerable to XML vulnerabilities.
+
+    Review the following resources for more information:
+
+    * https://docs.python.org/3/library/xml.html#xml-vulnerabilities
+    * https://github.com/tiran/defusedxml
+
+    The standard library relies on [`libexpat`](https://github.com/libexpat/libexpat)
+    for parsing XML.
+    """
+
+    def get_format_instructions(self) -> str:
+        """Return the format instructions for the XML output."""
+        return XML_FORMAT_INSTRUCTIONS.format(tags=self.tags)
+
+    def parse(self, text: str) -> dict[str, str | list[Any]]:
+        """Parse the output of an LLM call.
+
+        Args:
+            text: The output of an LLM call.
+
+        Returns:
+            A `dict` representing the parsed XML.
+
+        Raises:
+            OutputParserException: If the XML is not well-formed.
+            ImportError: If `defusedxml` is not installed and the `defusedxml` parser
+                is requested.
+        """
+        # Try to find XML string within triple backticks
+        # Imports are temporarily placed here to avoid issue with caching on CI
+        # likely if you're reading this you can move them to the top of the file
+        if self.parser == "defusedxml":
+            if not _HAS_DEFUSEDXML:
+                msg = (
+                    "defusedxml is not installed. "
+                    "Please install it to use the defusedxml parser. "
+                    "You can install it with `pip install defusedxml`. "
+                    "See https://github.com/tiran/defusedxml for more details"
+                )
+                raise ImportError(msg)
+            et = ElementTree  # Use the defusedxml parser
+        else:
+            et = ET  # Use the standard library parser
+
+        match = re.search(r"```(xml)?(.*)```", text, re.DOTALL)
+        if match is not None:
+            # If match found, use the content within the backticks
+            text = match.group(2)
+        encoding_match = self.encoding_matcher.search(text)
+        if encoding_match:
+            text = encoding_match.group(2)
+
+        text = text.strip()
+        try:
+            root = et.fromstring(text)
+            return self._root_to_dict(root)
+        except et.ParseError as e:
+            msg = f"Failed to parse XML format from completion {text}. Got: {e}"
+            raise OutputParserException(msg, llm_output=text) from e
+
+    @override
+    def _transform(self, input: Iterator[str | BaseMessage]) -> Iterator[AddableDict]:
+        streaming_parser = _StreamingParser(self.parser)
+        for chunk in input:
+            yield from streaming_parser.parse(chunk)
+        streaming_parser.close()
+
+    @override
+    async def _atransform(
+        self, input: AsyncIterator[str | BaseMessage]
+    ) -> AsyncIterator[AddableDict]:
+        streaming_parser = _StreamingParser(self.parser)
+        async for chunk in input:
+            for output in streaming_parser.parse(chunk):
+                yield output
+        streaming_parser.close()
+
+    def _root_to_dict(self, root: ET.Element) -> dict[str, str | list[Any]]:
+        """Converts xml tree to python dictionary."""
+        if root.text and bool(re.search(r"\S", root.text)):
+            # If root text contains any non-whitespace character it
+            # returns {root.tag: root.text}
+            return {root.tag: root.text}
+        root_tag: list[Any] = []
+        for child in root:
+            if len(child) == 0:
+                root_tag.append({child.tag: child.text})
+            else:
+                root_tag.append(self._root_to_dict(child))
+        return {root.tag: root_tag}
+
+    @property
+    def _type(self) -> str:
+        return "xml"
+
+
+def nested_element(path: list[str], elem: ET.Element) -> Any:
+    """Get nested element from path.
+
+    Args:
+        path: The path to the element.
+        elem: The element to extract.
+
+    Returns:
+        The nested element.
+    """
+    if len(path) == 0:
+        return AddableDict({elem.tag: elem.text})
+    return AddableDict({path[0]: [nested_element(path[1:], elem)]})

coco_core/outputs/__init__.py

@@ -0,0 +1,66 @@
+"""Output classes.
+
+Used to represent the output of a language model call and the output of a chat.
+
+The top container for information is the `LLMResult` object. `LLMResult` is used by both
+chat models and LLMs. This object contains the output of the language model and any
+additional information that the model provider wants to return.
+
+When invoking models via the standard runnable methods (e.g. invoke, batch, etc.):
+
+- Chat models will return `AIMessage` objects.
+- LLMs will return regular text strings.
+
+In addition, users can access the raw output of either LLMs or chat models via
+callbacks. The `on_chat_model_end` and `on_llm_end` callbacks will return an `LLMResult`
+object containing the generated outputs and any additional information returned by the
+model provider.
+
+In general, if information is already available in the AIMessage object, it is
+recommended to access it from there rather than from the `LLMResult` object.
+"""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.outputs.chat_generation import (
+        ChatGeneration,
+        ChatGenerationChunk,
+    )
+    from coco_core.outputs.chat_result import ChatResult
+    from coco_core.outputs.generation import Generation, GenerationChunk
+    from coco_core.outputs.llm_result import LLMResult
+    from coco_core.outputs.run_info import RunInfo
+
+__all__ = (
+    "ChatGeneration",
+    "ChatGenerationChunk",
+    "ChatResult",
+    "Generation",
+    "GenerationChunk",
+    "LLMResult",
+    "RunInfo",
+)
+
+_dynamic_imports = {
+    "ChatGeneration": "chat_generation",
+    "ChatGenerationChunk": "chat_generation",
+    "ChatResult": "chat_result",
+    "Generation": "generation",
+    "GenerationChunk": "generation",
+    "LLMResult": "llm_result",
+    "RunInfo": "run_info",
+}
+
+
+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__)

coco_core/outputs/chat_generation.py

@@ -0,0 +1,147 @@
+"""Chat generation output classes."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from pydantic import model_validator
+
+from coco_core.messages import BaseMessage, BaseMessageChunk
+from coco_core.outputs.generation import Generation
+from coco_core.utils._merge import merge_dicts
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+
+class ChatGeneration(Generation):
+    """A single chat generation output.
+
+    A subclass of `Generation` that represents the response from a chat model that
+    generates chat messages.
+
+    The `message` attribute is a structured representation of the chat message. Most of
+    the time, the message will be of type `AIMessage`.
+
+    Users working with chat models will usually access information via either
+    `AIMessage` (returned from runnable interfaces) or `LLMResult` (available via
+    callbacks).
+    """
+
+    text: str = ""
+    """The text contents of the output message.
+
+    !!! warning "SHOULD NOT BE SET DIRECTLY!"
+
+    """
+    message: BaseMessage
+    """The message output by the chat model."""
+
+    # Override type to be ChatGeneration, ignore mypy error as this is intentional
+    type: Literal["ChatGeneration"] = "ChatGeneration"  # type: ignore[assignment]
+    """Type is used exclusively for serialization purposes."""
+
+    @model_validator(mode="after")
+    def set_text(self) -> Self:
+        """Set the text attribute to be the contents of the message."""
+        # Check for legacy blocks with "text" key but no "type" field.
+        # Otherwise, delegate to `message.text`.
+        if isinstance(self.message.content, list):
+            has_legacy_blocks = any(
+                isinstance(block, dict)
+                and "text" in block
+                and block.get("type") is None
+                for block in self.message.content
+            )
+
+            if has_legacy_blocks:
+                blocks = []
+                for block in self.message.content:
+                    if isinstance(block, str):
+                        blocks.append(block)
+                    elif isinstance(block, dict):
+                        block_type = block.get("type")
+                        if block_type == "text" or (
+                            block_type is None and "text" in block
+                        ):
+                            blocks.append(block.get("text", ""))
+                self.text = "".join(blocks)
+            else:
+                self.text = self.message.text
+        else:
+            self.text = self.message.text
+
+        return self
+
+
+class ChatGenerationChunk(ChatGeneration):
+    """`ChatGeneration` chunk.
+
+    `ChatGeneration` chunks can be concatenated with other `ChatGeneration` chunks.
+    """
+
+    message: BaseMessageChunk
+    """The message chunk output by the chat model."""
+    # Override type to be ChatGeneration, ignore mypy error as this is intentional
+
+    type: Literal["ChatGenerationChunk"] = "ChatGenerationChunk"  # type: ignore[assignment]
+    """Type is used exclusively for serialization purposes."""
+
+    def __add__(
+        self, other: ChatGenerationChunk | list[ChatGenerationChunk]
+    ) -> ChatGenerationChunk:
+        """Concatenate two `ChatGenerationChunk`s.
+
+        Args:
+            other: The other `ChatGenerationChunk` or list of `ChatGenerationChunk` to
+                concatenate.
+
+        Raises:
+            TypeError: If other is not a `ChatGenerationChunk` or list of
+                `ChatGenerationChunk`.
+
+        Returns:
+            A new `ChatGenerationChunk` concatenated from self and other.
+        """
+        if isinstance(other, ChatGenerationChunk):
+            generation_info = merge_dicts(
+                self.generation_info or {},
+                other.generation_info or {},
+            )
+            return ChatGenerationChunk(
+                message=self.message + other.message,
+                generation_info=generation_info or None,
+            )
+        if isinstance(other, list) and all(
+            isinstance(x, ChatGenerationChunk) for x in other
+        ):
+            generation_info = merge_dicts(
+                self.generation_info or {},
+                *[chunk.generation_info for chunk in other if chunk.generation_info],
+            )
+            return ChatGenerationChunk(
+                message=self.message + [chunk.message for chunk in other],
+                generation_info=generation_info or None,
+            )
+        msg = f"unsupported operand type(s) for +: '{type(self)}' and '{type(other)}'"
+        raise TypeError(msg)
+
+
+def merge_chat_generation_chunks(
+    chunks: list[ChatGenerationChunk],
+) -> ChatGenerationChunk | None:
+    """Merge a list of `ChatGenerationChunk`s into a single `ChatGenerationChunk`.
+
+    Args:
+        chunks: A list of `ChatGenerationChunk` to merge.
+
+    Returns:
+        A merged `ChatGenerationChunk`, or `None` if the input list is empty.
+    """
+    if not chunks:
+        return None
+
+    if len(chunks) == 1:
+        return chunks[0]
+
+    return chunks[0] + chunks[1:]