xstructured 0.1.0__py3-none-any.whl

xstructured/langchain/runnable.py

@@ -0,0 +1,416 @@
+"""LangChain Runnable integration for the xstructured protocol."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
+from time import perf_counter
+from types import MappingProxyType
+from typing import Any, Generic, Literal, TypeVar, cast, overload
+
+from langchain_core.messages import BaseMessage, SystemMessage
+from langchain_core.prompt_values import PromptValue
+from langchain_core.runnables import Runnable, RunnableConfig
+
+from xstructured.core import ParseError, ParserConfig, RepairConfig
+from xstructured.envelope import EnvelopeSpec
+from xstructured.parser import StructuredParser
+from xstructured.schema import (
+    NamedSchemas,
+    NamedSchemaTargets,
+    SchemaTarget,
+    fingerprint_schema,
+    inspect_named_schemas,
+    schema_instructions,
+)
+from xstructured.streaming import StreamDecoder, StreamEvent, StreamEventKind
+
+from ._support import output_text
+from .repair import Repairer
+from .result import XStructuredResult
+
+Input = TypeVar("Input")
+T = TypeVar("T")
+
+
+class XStructuredRunnable(
+    Runnable[Input, XStructuredResult[T] | StreamEvent[T]],
+    Generic[Input, T],
+):
+    """A Runnable that applies and decodes the xstructured protocol."""
+
+    def __init__(
+        self,
+        runnable: Runnable[Input, Any],
+        schema: SchemaTarget | NamedSchemaTargets | NamedSchemas,
+        *,
+        envelope: EnvelopeSpec | None = None,
+        parser_config: ParserConfig | None = None,
+        multiple: bool = False,
+        multiple_envelopes: bool = False,
+        inject_instructions: bool = True,
+        repair: Runnable[Any, Any] | None = None,
+        repair_config: RepairConfig | None = None,
+    ) -> None:
+        self._runnable = runnable
+        self.envelope = envelope or EnvelopeSpec()
+        config = parser_config or ParserConfig()
+        if not config.require_envelope:
+            config = config.model_copy(update={"require_envelope": True})
+        self._parser: StructuredParser[T] = StructuredParser(
+            schema,
+            config=config,
+            envelope=self.envelope,
+            multiple=multiple,
+            multiple_envelopes=multiple_envelopes,
+        )
+        instruction_schema = (
+            inspect_named_schemas(schema) if isinstance(schema, Mapping) else schema
+        )
+        self.instructions = schema_instructions(
+            instruction_schema,
+            envelope=self.envelope.start,
+            multiple=multiple,
+            multiple_envelopes=multiple_envelopes,
+        )
+        self.schema_fingerprint = fingerprint_schema(schema)
+        self.inject_instructions = inject_instructions
+        self.multiple = multiple
+        self.multiple_envelopes = multiple_envelopes
+        self._repairer: Repairer[T] | None = (
+            Repairer(self._parser, repair, repair_config or RepairConfig(), self.instructions)
+            if repair is not None
+            else None
+        )
+
+    def invoke(
+        self,
+        input: Input,  # noqa: A002 - LangChain Runnable API compatibility
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
+    ) -> XStructuredResult[T]:
+        """Invoke the wrapped Runnable and parse its complete response."""
+        raw = self._runnable.invoke(
+            self._prepare_input(input),
+            config=config,
+            **kwargs,
+        )
+        return self._parse_output(raw, config=config)
+
+    async def ainvoke(
+        self,
+        input: Input,  # noqa: A002 - LangChain Runnable API compatibility
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
+    ) -> XStructuredResult[T]:
+        """Asynchronously invoke and parse the wrapped Runnable."""
+        raw = await self._runnable.ainvoke(
+            self._prepare_input(input),
+            config=config,
+            **kwargs,
+        )
+        return await self._aparse_output(raw, config=config)
+
+    def stream(
+        self,
+        input: Input,  # noqa: A002 - LangChain Runnable API compatibility
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
+    ) -> Iterator[StreamEvent[T]]:
+        """Yield ordered protocol events from the wrapped Runnable's stream."""
+        if self.multiple_envelopes:
+            raise ValueError("Streaming multiple named envelopes is not supported")
+        decoder = StreamDecoder(self._parser, self.envelope)
+        raw_chunks: list[Any] = []
+        raw_text_parts: list[str] = []
+        structured: T | None = None
+        structured_found = False
+
+        for raw_chunk in self._runnable.stream(self._prepare_input(input), config=config, **kwargs):
+            raw_chunks.append(raw_chunk)
+            text = output_text(raw_chunk)
+            raw_text_parts.append(text)
+            for event in decoder.feed(text):
+                if event.kind is StreamEventKind.STRUCTURED_END:
+                    structured = event.structured
+                    structured_found = True
+                yield event
+
+        yield from decoder.finalize()
+        if not structured_found:
+            raise RuntimeError("Streaming decoder completed without a structured value")
+        parsed = self._parser.parse("".join(raw_text_parts))
+        result = self._build_result(
+            raw=tuple(raw_chunks),
+            raw_text="".join(raw_text_parts),
+            content=decoder.text,
+            structured=cast(T, structured),
+            json_text=parsed.json_text,
+            recovered=parsed.recovered,
+            schema_name=parsed.schema_name,
+            repaired=parsed.repaired,
+            repair_attempt_count=parsed.repair_attempt_count,
+            stream_completed=True,
+        )
+        yield StreamEvent(
+            kind=StreamEventKind.RESULT,
+            sequence=decoder.next_sequence,
+            result=result,
+        )
+
+    async def astream(
+        self,
+        input: Input,  # noqa: A002 - LangChain Runnable API compatibility
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
+    ) -> AsyncIterator[StreamEvent[T]]:
+        """Asynchronously yield ordered protocol events."""
+        if self.multiple_envelopes:
+            raise ValueError("Streaming multiple named envelopes is not supported")
+        decoder = StreamDecoder(self._parser, self.envelope)
+        raw_chunks: list[Any] = []
+        raw_text_parts: list[str] = []
+        structured: T | None = None
+        structured_found = False
+
+        async for raw_chunk in self._runnable.astream(
+            self._prepare_input(input), config=config, **kwargs
+        ):
+            raw_chunks.append(raw_chunk)
+            text = output_text(raw_chunk)
+            raw_text_parts.append(text)
+            for event in decoder.feed(text):
+                if event.kind is StreamEventKind.STRUCTURED_END:
+                    structured = event.structured
+                    structured_found = True
+                yield event
+
+        for event in decoder.finalize():
+            yield event
+        if not structured_found:
+            raise RuntimeError("Streaming decoder completed without a structured value")
+        parsed = self._parser.parse("".join(raw_text_parts))
+        result = self._build_result(
+            raw=tuple(raw_chunks),
+            raw_text="".join(raw_text_parts),
+            content=decoder.text,
+            structured=cast(T, structured),
+            json_text=parsed.json_text,
+            recovered=parsed.recovered,
+            schema_name=parsed.schema_name,
+            repaired=parsed.repaired,
+            repair_attempt_count=parsed.repair_attempt_count,
+            stream_completed=True,
+        )
+        yield StreamEvent(
+            kind=StreamEventKind.RESULT,
+            sequence=decoder.next_sequence,
+            result=result,
+        )
+
+    def _prepare_input(self, value: Input) -> Input:
+        if not self.inject_instructions:
+            return value
+        if isinstance(value, str):
+            return cast(Input, f"{value}\n\n{self.instructions}")
+        if isinstance(value, PromptValue):
+            messages = [
+                SystemMessage(content=self.instructions),
+                *value.to_messages(),
+            ]
+            return cast(Input, messages)
+        if (
+            isinstance(value, Sequence)
+            and not isinstance(value, (str, bytes))
+            and all(isinstance(item, BaseMessage) for item in value)
+        ):
+            return cast(
+                Input,
+                [SystemMessage(content=self.instructions), *value],
+            )
+        return value
+
+    def _parse_output(
+        self, raw: Any, *, config: RunnableConfig | None = None
+    ) -> XStructuredResult[T]:
+        raw_text = output_text(raw)
+        started = perf_counter()
+        try:
+            parsed = self._parser.parse(raw_text)
+        except ParseError as error:
+            if self._repairer is None:
+                raise
+            parsed = self._repairer.repair(raw_text, error, config=config)
+        parse_duration = perf_counter() - started
+        return self._build_result(
+            raw=raw,
+            raw_text=raw_text,
+            content=_natural_content(raw_text, self.envelope),
+            structured=parsed.value,
+            json_text=parsed.json_text,
+            recovered=parsed.recovered,
+            schema_name=parsed.schema_name,
+            repaired=parsed.repaired,
+            repair_attempt_count=parsed.repair_attempt_count,
+            parse_duration=parse_duration,
+        )
+
+    async def _aparse_output(
+        self, raw: Any, *, config: RunnableConfig | None = None
+    ) -> XStructuredResult[T]:
+        raw_text = output_text(raw)
+        started = perf_counter()
+        try:
+            parsed = self._parser.parse(raw_text)
+        except ParseError as error:
+            if self._repairer is None:
+                raise
+            parsed = await self._repairer.arepair(raw_text, error, config=config)
+        parse_duration = perf_counter() - started
+        return self._build_result(
+            raw=raw,
+            raw_text=raw_text,
+            content=_natural_content(raw_text, self.envelope),
+            structured=parsed.value,
+            json_text=parsed.json_text,
+            recovered=parsed.recovered,
+            schema_name=parsed.schema_name,
+            repaired=parsed.repaired,
+            repair_attempt_count=parsed.repair_attempt_count,
+            parse_duration=parse_duration,
+        )
+
+    def _build_result(
+        self,
+        *,
+        raw: Any,
+        raw_text: str,
+        content: str,
+        structured: T,
+        json_text: str,
+        recovered: bool,
+        schema_name: str | None = None,
+        repaired: bool = False,
+        repair_attempt_count: int = 0,
+        stream_completed: bool = True,
+        parse_duration: float = 0.0,
+    ) -> XStructuredResult[T]:
+        envelope_count = raw_text.count(self.envelope.start)
+        metadata: dict[str, Any] = {
+            "schema_fingerprint": self.schema_fingerprint,
+            "envelope_detected": envelope_count > 0,
+            "envelope_count": envelope_count,
+            "parse_duration": parse_duration,
+            "validation_duration": parse_duration,
+            "repair_attempted": repaired,
+            "repair_attempt_count": repair_attempt_count,
+            "stream_completed": stream_completed,
+        }
+        if isinstance(raw, BaseMessage):
+            metadata.update(
+                {
+                    "message_id": raw.id,
+                    "message_name": raw.name,
+                    "message_type": raw.type,
+                    "response_metadata": dict(raw.response_metadata),
+                }
+            )
+            usage_metadata = getattr(raw, "usage_metadata", None)
+            if usage_metadata is not None:
+                metadata["usage_metadata"] = dict(usage_metadata)
+        return XStructuredResult(
+            content=content,
+            structured=structured,
+            raw=raw,
+            raw_text=raw_text,
+            json_text=json_text,
+            recovered=recovered,
+            metadata=MappingProxyType(metadata),
+            schema_name=schema_name,
+            repaired=repaired,
+            repair_attempt_count=repair_attempt_count,
+        )
+
+
+@overload
+def with_xstructured_output(
+    runnable: Runnable[Input, Any],
+    schema: SchemaTarget | NamedSchemaTargets | NamedSchemas,
+    *,
+    multiple: Literal[True],
+    envelope: EnvelopeSpec | None = None,
+    parser_config: ParserConfig | None = None,
+    inject_instructions: bool = True,
+    repair: Runnable[Any, Any] | None = None,
+    repair_config: RepairConfig | None = None,
+) -> XStructuredRunnable[Input, list[T]]: ...
+
+
+@overload
+def with_xstructured_output(
+    runnable: Runnable[Input, Any],
+    schema: SchemaTarget | NamedSchemaTargets | NamedSchemas,
+    *,
+    multiple: Literal[False] = False,
+    envelope: EnvelopeSpec | None = None,
+    parser_config: ParserConfig | None = None,
+    inject_instructions: bool = True,
+    repair: Runnable[Any, Any] | None = None,
+    repair_config: RepairConfig | None = None,
+) -> XStructuredRunnable[Input, T]: ...
+
+
+def with_xstructured_output(
+    runnable: Runnable[Input, Any],
+    schema: SchemaTarget | NamedSchemaTargets | NamedSchemas,
+    *,
+    envelope: EnvelopeSpec | None = None,
+    parser_config: ParserConfig | None = None,
+    multiple: bool = False,
+    multiple_envelopes: bool = False,
+    inject_instructions: bool = True,
+    repair: Runnable[Any, Any] | None = None,
+    repair_config: RepairConfig | None = None,
+) -> XStructuredRunnable[Input, T | list[T]]:
+    """Wrap a Runnable with instruction injection and xstructured decoding.
+
+    *schema* may be a single schema target, or a named set of schema targets
+    (`Mapping[str, SchemaTarget]` or `NamedSchemas`) to validate the response
+    against one of several allowed shapes -- see `xstructured.schema.named`.
+
+    *repair* is an optional, strictly opt-in `Runnable` invoked, bounded by
+    `repair_config.max_attempts` (default 1), to ask a model to fix a
+    response that failed schema validation. It is never invoked unless
+    explicitly supplied here; see `xstructured.langchain.repair` for the
+    bounded retry contract.
+    """
+    return XStructuredRunnable(
+        runnable,
+        schema,
+        envelope=envelope,
+        parser_config=parser_config,
+        multiple=multiple,
+        multiple_envelopes=multiple_envelopes,
+        inject_instructions=inject_instructions,
+        repair=repair,
+        repair_config=repair_config,
+    )
+
+
+def _natural_content(text: str, envelope: EnvelopeSpec) -> str:
+    named_pattern = re.compile(
+        r"<xstructured name=\"[A-Za-z0-9_.-]+\">.*?</xstructured>", re.DOTALL
+    )
+    if named_pattern.search(text):
+        return named_pattern.sub("", text)
+    start_at = text.find(envelope.start)
+    if start_at < 0:
+        return text
+    from xstructured.envelope import EnvelopeScanner
+
+    scanner = EnvelopeScanner(envelope)
+    scanner.feed(text)
+    if not scanner.complete or scanner.payload is None:
+        return text[:start_at]
+    end_at = start_at + len(envelope.start) + len(scanner.payload)
+    return text[:start_at] + text[end_at + len(envelope.end) :]

xstructured/parser/__init__.py

@@ -0,0 +1,6 @@
+"""JSON parser and conservative recovery strategies."""
+
+from .parser import StructuredParser
+from .recovery import recovery_candidates
+
+__all__ = ["StructuredParser", "recovery_candidates"]

xstructured/parser/json_safety.py

@@ -0,0 +1,59 @@
+"""Bounded strict JSON decoding helpers."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+
+class JsonSafetyError(ValueError):
+    """Raised when JSON exceeds a configured safety constraint."""
+
+
+def loads_strict(text: str, *, max_nesting_depth: int) -> Any:
+    """Decode standard JSON after rejecting excessive container nesting."""
+    _validate_nesting_depth(text, max_nesting_depth)
+    return json.loads(
+        text,
+        parse_constant=_reject_non_standard_constant,
+        object_pairs_hook=_reject_duplicate_keys,
+    )
+
+
+def _validate_nesting_depth(text: str, max_nesting_depth: int) -> None:
+    depth = 0
+    in_string = False
+    escaped = False
+
+    for character in text:
+        if in_string:
+            if escaped:
+                escaped = False
+            elif character == "\\":
+                escaped = True
+            elif character == '"':
+                in_string = False
+            continue
+        if character == '"':
+            in_string = True
+        elif character in "[{":
+            depth += 1
+            if depth > max_nesting_depth:
+                raise JsonSafetyError(
+                    f"JSON nesting exceeds configured limit of {max_nesting_depth}"
+                )
+        elif character in "]}":
+            depth -= 1
+
+
+def _reject_non_standard_constant(value: str) -> None:
+    raise JsonSafetyError(f"Non-standard JSON constant {value!r} is not permitted")
+
+
+def _reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
+    result: dict[str, Any] = {}
+    for key, value in pairs:
+        if key in result:
+            raise JsonSafetyError(f"Duplicate JSON object key {key!r} is not permitted")
+        result[key] = value
+    return result

xstructured/parser/parser.py

@@ -0,0 +1,198 @@
+"""Schema-aware JSON parser."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+from typing import Any, Generic, TypeVar, cast
+
+from xstructured.core import ParseError, ParserConfig, ParseResult, RecoveryError, SchemaError
+from xstructured.envelope import EnvelopeScanner, EnvelopeSpec
+from xstructured.schema import SchemaInfo, SchemaTarget, inspect_schema
+from xstructured.schema.named import NamedSchemas, NamedSchemaTargets, inspect_named_schemas
+
+from .json_safety import JsonSafetyError, loads_strict
+from .recovery import recovery_candidates
+
+T = TypeVar("T")
+
+
+class StructuredParser(Generic[T]):
+    """Parse JSON text and validate it with a Pydantic v2 schema target.
+
+    *schema* is either a single schema target (a `BaseModel` subclass, a
+    `TypeAdapter`, or anything `TypeAdapter` accepts) or a *named* set of
+    schema targets -- a `Mapping[str, SchemaTarget]` or an already-built
+    `NamedSchemas`. With named schemas, the JSON payload must be a single
+    object shaped ``{"schema": "<name>", "payload": <value>}`` and
+    `ParseResult.schema_name` reports which named schema matched. With
+    ``multiple=True``, the JSON value must be an array and each item is
+    validated independently; named schemas may be mixed in that array.
+    """
+
+    def __init__(
+        self,
+        schema: SchemaTarget | NamedSchemaTargets | NamedSchemas,
+        *,
+        config: ParserConfig | None = None,
+        envelope: EnvelopeSpec | None = None,
+        multiple: bool = False,
+        multiple_envelopes: bool = False,
+    ) -> None:
+        if isinstance(schema, NamedSchemas):
+            self.schema: SchemaInfo | None = None
+            self.named_schemas: NamedSchemas | None = schema
+        elif isinstance(schema, Mapping):
+            self.schema = None
+            self.named_schemas = inspect_named_schemas(schema)
+        else:
+            self.schema = inspect_schema(schema)
+            self.named_schemas = None
+        self.config = config or ParserConfig()
+        self.envelope = envelope
+        self.multiple = multiple
+        self.multiple_envelopes = multiple_envelopes
+
+    def parse(self, text: str) -> ParseResult[T]:
+        """Parse one response, optionally extracting its configured envelope."""
+        if not isinstance(text, str):
+            raise TypeError("Parser input must be a string")
+        if len(text) > self.config.max_input_chars:
+            raise ParseError(
+                f"Input exceeds configured limit of {self.config.max_input_chars} characters", text
+            )
+
+        if self.multiple_envelopes:
+            return self._parse_named_envelopes(text)
+        json_text, envelope_found = self._extract_envelope(text)
+        candidates = [json_text]
+        if self.config.recovery.enabled:
+            candidates = list(
+                recovery_candidates(
+                    json_text,
+                    strip_markdown_fences=self.config.recovery.strip_markdown_fences,
+                    strip_surrounding_text=self.config.recovery.strip_surrounding_text,
+                )
+            )[: self.config.recovery.max_candidates]
+
+        failures: list[str] = []
+        last_error: Exception | None = None
+        for index, candidate in enumerate(candidates):
+            try:
+                if len(candidate) > self.config.max_payload_chars:
+                    raise JsonSafetyError(
+                        "JSON payload exceeds configured limit of "
+                        f"{self.config.max_payload_chars} characters"
+                    )
+                decoded: Any = loads_strict(
+                    candidate,
+                    max_nesting_depth=self.config.max_nesting_depth,
+                )
+                value, schema_name = self._validate(decoded)
+            except (JsonSafetyError, ValueError, TypeError, SchemaError) as exc:
+                failures.append(str(exc))
+                last_error = exc
+                continue
+            return ParseResult(
+                value=value,
+                raw=text,
+                json_text=candidate,
+                recovered=index > 0,
+                envelope_found=envelope_found,
+                schema_name=schema_name,
+            )
+
+        if self.config.recovery.enabled:
+            raise RecoveryError(
+                "Unable to parse and validate any recovery candidate",
+                text,
+                last_error,
+                tuple(failures),
+            ) from last_error
+        detail = f": {last_error}" if last_error is not None else ""
+        raise ParseError(
+            f"Unable to parse and validate JSON{detail}", text, last_error
+        ) from last_error
+
+    def _parse_named_envelopes(self, text: str) -> ParseResult[T]:
+        if self.named_schemas is None:
+            raise SchemaError("multiple_envelopes requires named schema targets")
+        pattern = re.compile(
+            r'<xstructured name="([A-Za-z0-9_.-]+)">(.*?)</xstructured>',
+            re.DOTALL,
+        )
+        matches = list(pattern.finditer(text))
+        if not matches:
+            raise ParseError("No named xstructured envelopes were found", text)
+        values: dict[str, Any] = {}
+        for match in matches:
+            name, payload = match.groups()
+            if name in values:
+                raise ParseError(f"Duplicate named envelope: {name!r}", text)
+            if name not in self.named_schemas.schemas:
+                raise ParseError(f"Unknown named envelope: {name!r}", text)
+            if len(payload) > self.config.max_payload_chars:
+                raise ParseError("Named envelope payload exceeds configured limit", text)
+            try:
+                decoded = loads_strict(payload, max_nesting_depth=self.config.max_nesting_depth)
+                values[name] = self.named_schemas.resolve(name).validate_python(decoded)
+            except (JsonSafetyError, ValueError, TypeError, SchemaError) as exc:
+                raise ParseError(f"Invalid named envelope {name!r}: {exc}", text, exc) from exc
+        return ParseResult(
+            value=cast(T, values),
+            raw=text,
+            json_text="".join(match.group(2) for match in matches),
+            envelope_found=True,
+        )
+
+    def _validate(self, decoded: Any) -> tuple[Any, str | None]:
+        if self.multiple:
+            if not isinstance(decoded, list):
+                raise ValueError("Multiple payload mode requires a JSON array")
+            values: list[Any] = []
+            names: list[str] = []
+            for item in decoded:
+                value, name = self._validate_one(item)
+                values.append(value)
+                if name is not None:
+                    names.append(name)
+            schema_name = names[0] if names and all(name == names[0] for name in names) else None
+            return values, schema_name
+        return self._validate_one(decoded)
+
+    def _validate_one(self, decoded: Any) -> tuple[Any, str | None]:
+        if self.schema is not None:
+            return self.schema.validate_python(decoded), None
+        assert self.named_schemas is not None
+        spec = self.named_schemas.spec
+        if not isinstance(decoded, dict):
+            raise ValueError("Named schema payloads must be a JSON object")
+        if spec.schema_key not in decoded:
+            raise ValueError(f"Missing required {spec.schema_key!r} discriminator key")
+        name = decoded[spec.schema_key]
+        if not isinstance(name, str):
+            raise ValueError(f"{spec.schema_key!r} must be a string schema name")
+        if spec.payload_key not in decoded:
+            raise ValueError(f"Missing required {spec.payload_key!r} payload key")
+        schema_info = self.named_schemas.resolve(name)
+        return schema_info.validate_python(decoded[spec.payload_key]), name
+
+    def _extract_envelope(self, text: str) -> tuple[str, bool]:
+        if self.envelope is None:
+            if self.config.require_envelope:
+                raise ParseError(
+                    "An envelope is required but no envelope spec was configured",
+                    text,
+                )
+            return text, False
+        scanner = EnvelopeScanner(
+            self.envelope,
+            max_envelope_chars=self.config.max_envelope_chars,
+            max_payload_chars=self.config.max_payload_chars,
+        )
+        scanner.feed(text)
+        if scanner.complete:
+            return scanner.finalize(), True
+        if self.config.require_envelope:
+            raise ParseError("No complete envelope was found", text)
+        return text, False