core[patch]: fix xml output parser transform

libs/core/langchain_core/output_parsers/xml.py

@@ -1,10 +1,9 @@
 import re
 from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
 from xml.etree import ElementTree as ET
-from xml.etree.ElementTree import TreeBuilder
 
 from langchain_core.exceptions import OutputParserException
 from langchain_core.messages import BaseMessage
 from langchain_core.output_parsers.transform import BaseTransformOutputParser
 from langchain_core.runnables.utils import AddableDict
 
@@ -58,116 +57,6 @@
             msg = f"Failed to parse XML format from completion {text}. Got: {e}"
             raise OutputParserException(msg, llm_output=text) from e
 
-    def _transform(
-        self, input: Iterator[Union[str, BaseMessage]]
-    ) -> Iterator[AddableDict]:
-        # 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
-        from defusedxml.ElementTree import DefusedXMLParser  # type: ignore[import]
-
-        parser = ET.XMLPullParser(
-            ["start", "end"], _parser=DefusedXMLParser(target=TreeBuilder())
-        )
-        xml_start_re = re.compile(r"<[a-zA-Z:_]")
-        xml_started = False
-        current_path: List[str] = []
-        current_path_has_children = False
-        buffer = ""
-        for chunk in input:
-            if isinstance(chunk, BaseMessage):
-                # extract text
-                chunk_content = chunk.content
-                if not isinstance(chunk_content, str):
-                    continue
-                chunk = chunk_content
-            # add chunk to buffer of unprocessed text
-            buffer += chunk
-            # if xml string hasn't started yet, continue to next chunk
-            if not xml_started:
-                if match := xml_start_re.search(buffer):
-                    # if xml string has started, remove all text before it
-                    buffer = buffer[match.start() :]
-                    xml_started = True
-                else:
-                    continue
-            # feed buffer to parser
-            parser.feed(buffer)
-            buffer = ""
-            # yield all events
-
-            for event, elem in parser.read_events():
-                if event == "start":
-                    # update current path
-                    current_path.append(elem.tag)
-                    current_path_has_children = False
-                elif event == "end":
-                    # remove last element from current path
-                    current_path.pop()
-                    # yield element
-                    if not current_path_has_children:
-                        yield nested_element(current_path, elem)
-                    # prevent yielding of parent element
-                    if current_path:
-                        current_path_has_children = True
-                    else:
-                        xml_started = False
-        # close parser
-        parser.close()
-
-    async def _atransform(
-        self, input: AsyncIterator[Union[str, BaseMessage]]
-    ) -> AsyncIterator[AddableDict]:
-        # 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
-        from defusedxml.ElementTree import DefusedXMLParser  # type: ignore[import]
-
-        _parser = DefusedXMLParser(target=TreeBuilder())
-        parser = ET.XMLPullParser(["start", "end"], _parser=_parser)
-        xml_start_re = re.compile(r"<[a-zA-Z:_]")
-        xml_started = False
-        current_path: List[str] = []
-        current_path_has_children = False
-        buffer = ""
-        async for chunk in input:
-            if isinstance(chunk, BaseMessage):
-                # extract text
-                chunk_content = chunk.content
-                if not isinstance(chunk_content, str):
-                    continue
-                chunk = chunk_content
-            # add chunk to buffer of unprocessed text
-            buffer += chunk
-            # if xml string hasn't started yet, continue to next chunk
-            if not xml_started:
-                if match := xml_start_re.search(buffer):
-                    # if xml string has started, remove all text before it
-                    buffer = buffer[match.start() :]
-                    xml_started = True
-                else:
-                    continue
-            # feed buffer to parser
-            parser.feed(buffer)
-            buffer = ""
-            # yield all events
-            for event, elem in parser.read_events():
-                if event == "start":
-                    # update current path
-                    current_path.append(elem.tag)
-                    current_path_has_children = False
-                elif event == "end":
-                    # remove last element from current path
-                    current_path.pop()
-                    # yield element
-                    if not current_path_has_children:
-                        yield nested_element(current_path, elem)
-                    # prevent yielding of parent element
-                    if current_path:
-                        current_path_has_children = True
-                    else:
-                        xml_started = False
-        # close parser
-        parser.close()
-
     def _root_to_dict(self, root: ET.Element) -> Dict[str, List[Any]]:
         """Converts xml tree to python dictionary."""
         result: Dict[str, List[Any]] = {root.tag: []}

libs/core/tests/unit_tests/output_parsers/test_xml_parser.py

@@ -1,6 +1,6 @@
 """Test XMLOutputParser"""
 from typing import AsyncIterator
 from xml.etree.ElementTree import ParseError
 
 import pytest
 
@@ -49,22 +49,6 @@
     xml_parser = XMLOutputParser()
     assert DEF_RESULT_EXPECTED == xml_parser.parse(result)
     assert DEF_RESULT_EXPECTED == (await xml_parser.aparse(result))
-    assert list(xml_parser.transform(iter(result))) == [
-        {"foo": [{"bar": [{"baz": None}]}]},
-        {"foo": [{"bar": [{"baz": "slim.shady"}]}]},
-        {"foo": [{"baz": "tag"}]},
-    ]
-
-    async def _as_iter(string: str) -> AsyncIterator[str]:
-        for c in string:
-            yield c
-
-    chunks = [chunk async for chunk in xml_parser.atransform(_as_iter(result))]
-    assert chunks == [
-        {"foo": [{"bar": [{"baz": None}]}]},
-        {"foo": [{"bar": [{"baz": "slim.shady"}]}]},
-        {"foo": [{"baz": "tag"}]},
-    ]
 
 
 @pytest.mark.parametrize("result", ["foo></foo>", "<foo></foo", "foo></foo", "foofoo"])
@@ -100,17 +84,3 @@
 
     with pytest.raises(OutputParserException):
         await parser.aparse(MALICIOUS_XML)
-
-    with pytest.raises(ParseError):
-        # Right now raises undefined entity error
-        assert list(parser.transform(iter(MALICIOUS_XML))) == [
-            {"foo": [{"bar": [{"baz": None}]}]}
-        ]
-
-    async def _as_iter(string: str) -> AsyncIterator[str]:
-        for c in string:
-            yield c
-
-    with pytest.raises(ParseError):
-        chunks = [chunk async for chunk in parser.atransform(_as_iter(MALICIOUS_XML))]
-        assert chunks == [{"foo": [{"bar": [{"baz": None}]}]}]