data-designer 0.1.0__py3-none-any.whl

data_designer/engine/models/parsers/postprocessors.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional, Type
+
+import json_repair
+from pydantic import BaseModel, ValidationError
+
+from data_designer.engine.models.parsers.types import (
+    CodeBlock,
+    LLMStructuredResponse,
+    PydanticTypeBlock,
+    StructuredDataBlock,
+    TextBlock,
+)
+
+
+def merge_text_blocks(
+    structured_response: LLMStructuredResponse,
+) -> LLMStructuredResponse:
+    processed_response = structured_response.model_copy()
+    processed_response.parsed = []
+    accumulator = None
+    for block in structured_response.parsed:
+        if isinstance(block, TextBlock):
+            if accumulator is not None:
+                accumulator = TextBlock(text=accumulator.text + block.text)
+            else:
+                accumulator = block
+        else:
+            if accumulator is not None:
+                processed_response.parsed.append(accumulator)
+                accumulator = None
+
+            processed_response.parsed.append(block)
+
+    if accumulator:
+        processed_response.parsed.append(accumulator)
+
+    return processed_response
+
+
+def deserialize_json_code(
+    structured_response: LLMStructuredResponse,
+) -> LLMStructuredResponse:
+    processed_response = structured_response.model_copy()
+    processed_response.parsed = []
+
+    for block in structured_response.parsed:
+        if isinstance(block, CodeBlock) and block.code_lang == "json":
+            deserialized = json_repair.loads(block.code)
+
+            block = StructuredDataBlock(serialized=block.code, obj=deserialized)
+
+            processed_response.parsed.append(block)
+        else:
+            processed_response.parsed.append(block)
+
+    return processed_response
+
+
+class RealizePydanticTypes:
+    types: list[Type[BaseModel]]
+
+    def __init__(self, types: list[Type[BaseModel]]):
+        self.types = types
+
+    def _fit_types(self, obj: dict) -> Optional[BaseModel]:
+        final_obj = None
+
+        for t in self.types:
+            try:
+                final_obj = t.model_validate(obj)
+            except ValidationError:
+                pass
+
+        return final_obj
+
+    def __call__(self, structured_response: LLMStructuredResponse) -> LLMStructuredResponse:
+        processed_response = structured_response.model_copy()
+        processed_response.parsed = []
+
+        for block in structured_response.parsed:
+            if isinstance(block, StructuredDataBlock):
+                new_block = block
+                pydantic_obj = self._fit_types(block.obj)
+                if pydantic_obj:
+                    new_block = PydanticTypeBlock(serialized=block.serialized, obj=pydantic_obj)
+                processed_response.parsed.append(new_block)
+            else:
+                processed_response.parsed.append(block)
+
+        return processed_response

data_designer/engine/models/parsers/tag_parsers.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from lxml.etree import _Element
+
+from data_designer.engine.models.parsers.types import CodeBlock, TextBlock
+
+
+def text_parser(element: _Element) -> TextBlock:
+    return TextBlock(text=element.text if element.text else "")
+
+
+def text_parser_keep_markup(element: _Element) -> TextBlock:
+    body = element.text if element.text else ""
+    return TextBlock(text=f"<{element.tag}>{body}</{element.tag}>")
+
+
+def inline_code_parser(element: _Element) -> TextBlock:
+    return TextBlock(text=f"`{element.text if element.text else ''}`")
+
+
+def code_block_parser(element: _Element) -> CodeBlock:
+    """Parse a <pre><code> element node.
+
+    This parser handles the special case of Markdown->HTML conversion
+    for fenced code blocks. These take on the form:
+
+       ```xx
+       ...
+       ```
+
+        <pre><code class="language-xx">...</code></pre>
+
+    This parser is intended to be attached to the special case of "pre.code"
+    tag hierarchies.
+
+    Syntax Handling
+
+        If the syntax is not specified, e.g. ``<code>...</code>`` or
+        ``<code class="">...</code>``, then the syntax field is returned
+        as None. However, the parser does not _enforce_ the prefix
+        `language-` on the value of the class attribute.
+        If it is not present, then the entire value
+
+    Args:
+        element (lxml.etree._Element): An element of the lxml-parsed
+            element tree.
+
+    Returns:
+        CodeBlock: Datat structured containing both the body of the code
+            as well as the specified synax of the code block.
+
+    """
+    prefix = "language-"
+    language_identifier = element.attrib.get("class", "")
+    language_identifier = language_identifier.removeprefix(prefix)
+    return CodeBlock(
+        code=element.text.strip() if element.text else "",
+        code_lang=language_identifier if language_identifier else None,
+    )

data_designer/engine/models/parsers/types.py

@@ -0,0 +1,82 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Optional, Protocol, Type, runtime_checkable
+
+from lxml.etree import _Element
+from pydantic import BaseModel, Field
+from typing_extensions import Self
+
+
+class LLMStructuredResponse(BaseModel):
+    """Output format for the LLM Response Parser."""
+
+    response: str = Field(description="Raw Markdown/Markup response received from the LLM and input to the parser.")
+    markup: str = Field(description="Markup/HTML resulting from running Markdown parsing on response.")
+    parsed: list[BaseModel] = Field(
+        default_factory=list,
+        description="Structured content parsed from markup. Elements of this list are in document-order.",
+    )
+
+    def head(self, n: int) -> Self:
+        """Retain only the first n elements of the parsed response."""
+        out = self.model_copy()
+        out.parsed = out.parsed[:n]
+        return out
+
+    def tail(self, n: int) -> Self:
+        """Retain only the last n elements of the parsed response."""
+        out = self.model_copy()
+        out.parsed = out.parsed[-n:]
+        return out
+
+    def filter(self, block_types: list[Type[BaseModel]]) -> Self:
+        out = self.model_copy()
+        out.parsed = [b for b in out.parsed if isinstance(b, tuple(block_types))]
+        return out
+
+
+@runtime_checkable
+class TagParser(Protocol):
+    """Protocol for tag parsing implementations.
+
+    All TagParsers are objects which can take as input an `lxml`
+    element, do some computation, and return some kind of structured
+    output, represented as a subclass of Pydantic `BaseModel`.
+    This protocol implementation can cover both classes as well
+    as curried fuctions as parsers (e.g. `partial`).
+    """
+
+    def __call__(self, element: _Element) -> BaseModel: ...
+
+
+@runtime_checkable
+class PostProcessor(Protocol):
+    """Protocol for parsed output postprocessing implementations.
+
+    Implementations of this protocol are used to transform the results of
+    the LLM response parser while retaining the same output structure.
+    This is done so that PostProcessor implementations can be chained
+    together.
+    """
+
+    def __call__(self, structured_response: LLMStructuredResponse) -> LLMStructuredResponse: ...
+
+
+class TextBlock(BaseModel):
+    text: str
+
+
+class CodeBlock(BaseModel):
+    code: str
+    code_lang: Optional[str] = None
+
+
+class StructuredDataBlock(BaseModel):
+    serialized: str
+    obj: Any
+
+
+class PydanticTypeBlock(BaseModel):
+    serialized: str
+    obj: BaseModel

data_designer/engine/models/recipes/base.py

@@ -0,0 +1,79 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+from collections.abc import Callable
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class ResponseRecipe(abc.ABC, Generic[T]):
+    """Base class for defining response recipes.
+
+    Response recipes contain all necessary information for
+    getting an LLM to perform a particular common task,
+    like outputting code in a desired format or following
+    structured output.
+    """
+
+    @abc.abstractmethod
+    def _build_parser_fn(self) -> Callable[[str], T]:
+        """Build the recipe's output parser function."""
+        ...
+
+    @property
+    @abc.abstractmethod
+    def example_template(self) -> str: ...
+
+    @abc.abstractmethod
+    def serialize_output(self, output: T) -> str:
+        """Serialize an instance of the parser output."""
+        ...
+
+    @abc.abstractmethod
+    def deserialize_output(self, serialized_output: str) -> T:
+        """Deserialize a serialized instance of the parser output."""
+        ...
+
+    def __init__(self):
+        self._parse_fn = self._build_parser_fn()
+
+    @property
+    def task_instructions(self) -> str | None:
+        """Specifies task instructions.
+
+        These instructions lay out the particular task information the
+        LLM requires in order to carry out the function of the recipe.
+        """
+        return None
+
+    def parse(self, response: str) -> T:
+        """Apply the recipe's parser to a raw model output."""
+        return self._parse_fn(response)
+
+    def generate_response_example(self, example: T) -> str:
+        """Create a serialized response example that the parser would admit."""
+        return self.example_template.format(example=example)
+
+    def apply_recipe_to_user_prompt(self, user_prompt: str) -> str:
+        """Appends recipe specific task instructions if applicable.
+
+        Args:
+            user_prompt (str): User prompt to be appended with recipe specific task instructions if applicable.
+
+        Returns:
+            str: Final user prompt
+        """
+        return f"{user_prompt}\n\n{self.task_instructions}" if self.task_instructions else user_prompt
+
+    def apply_recipe_to_system_prompt(self, system_prompt: str | None) -> str:
+        """Appends recipe specific task instructions if applicable.
+
+        Args:
+            system_prompt (str): System prompt to be appended with recipe specific task instructions if applicable.
+
+        Returns:
+            str: Final system prompt
+        """
+        return system_prompt

data_designer/engine/models/recipes/response_recipes.py

@@ -0,0 +1,291 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+import json
+
+from pydantic import BaseModel
+
+from data_designer.config.utils.code_lang import CodeLang
+from data_designer.engine.models.parsers.errors import ParserException
+from data_designer.engine.models.parsers.parser import LLMResponseParser
+from data_designer.engine.models.parsers.postprocessors import (
+    StructuredDataBlock,
+    deserialize_json_code,
+    merge_text_blocks,
+)
+from data_designer.engine.models.parsers.types import CodeBlock
+from data_designer.engine.models.recipes.base import (
+    ResponseRecipe,
+)
+from data_designer.engine.processing.gsonschema.validators import JSONSchemaValidationError, validate
+
+
+class TextResponseRecipe(ResponseRecipe[str]):
+    """Default text-parser.
+
+    This parser is meant to cover the "pass-through" case of natural language LLM responses.
+    """
+
+    @property
+    def example_template(self) -> str:
+        return "{example}"
+
+    def serialize_output(self, output: str) -> str:
+        return output
+
+    def deserialize_output(self, serialized_output: str) -> str:
+        return serialized_output
+
+    def _build_parser_fn(self) -> Callable[[str], str]:
+        parser = LLMResponseParser(
+            postprocessors=[
+                merge_text_blocks,
+            ]
+        )
+
+        return lambda x: parser.parse(x).response
+
+
+class StructuredResponseRecipe(ResponseRecipe[dict]):
+    """Recipe for structured responses.
+
+    This recipe is intended to cover the generic case of
+    prompting-based requests for structured data outputs,
+    and the structure in question is determined by a
+    provided JSON Schema.
+
+    The LLM's response us validated against the provided
+    JSON Schema, however the object returned is python
+    dictionary obtained from deserializing the LLM's
+    JSON response.
+    """
+
+    json_schema: dict
+    pruning: bool
+    no_extra_properties: bool
+
+    def __init__(
+        self,
+        json_schema: dict,
+        pruning: bool = True,
+        no_extra_properties: bool = True,
+        **kwargs,
+    ):
+        """Initialize StructuredResponseRecipe.
+
+        Args:
+            json_schema (dict): A target JSON schema that the LLM
+                should adhere to when making its response.
+            pruning (bool): If `True`, then any extra fields in the returned
+                JSON object will be removed. Otherwise, they are retained,
+                which could raise validation errors. Default=True
+            no_extra_properties (bool) If `True`, then validation will fail
+                if extra properties are encountered in the returned JSON response.
+                Default=True.
+        """
+        super().__init__(**kwargs)
+        self.json_schema = json_schema
+        self.pruning = pruning
+        self.no_extra_properties = no_extra_properties
+
+    @property
+    def task_instructions(self) -> str:
+        return (
+            "* Your response must be in JSON format.\n"
+            "* Your JSON response must be returned within a Markdown, ```json code fence.\n"
+            "* The JSON format is given as a JSON Schema description within <response_schema> markup tags.\n\n"
+            f"<response_schema>\n{self.schema}\n</response_schema>"
+        )
+
+    @property
+    def example_template(self) -> str:
+        return "```json\n{example}\n```"
+
+    def generate_response_example(self, example: dict) -> str:
+        return self.example_template.format(example=json.dumps(example))
+
+    @property
+    def schema(self) -> str:
+        return json.dumps(self.json_schema)
+
+    def serialize_output(self, output: dict) -> str:
+        return json.dumps(output, ensure_ascii=False)
+
+    def deserialize_output(self, serialized_output: str) -> dict:
+        return json.loads(serialized_output)
+
+    @property
+    def _validate_args(self):
+        return {
+            "schema": self.json_schema,
+            "pruning": self.pruning,
+            "no_extra_properties": self.no_extra_properties,
+        }
+
+    def _build_parser_fn(self) -> Callable[[str], dict]:
+        parser = LLMResponseParser(
+            postprocessors=[
+                merge_text_blocks,
+                deserialize_json_code,
+            ]
+        )
+
+        def parse_fn(response: str) -> dict:
+            try:
+                obj = parser.parse(response).filter([StructuredDataBlock]).parsed.pop().obj
+                return validate(obj, **self._validate_args)
+            except IndexError:
+                raise ParserException(
+                    "No parsable JSON structure within ```json markdown fence.",
+                    source=response,
+                ) from None
+            except JSONSchemaValidationError as exc:
+                raise ParserException(
+                    "Response doesn't match requested <response_schema>\n" + str(exc),
+                    source=response,
+                ) from None
+
+        return parse_fn
+
+
+class PydanticResponseRecipe(ResponseRecipe[BaseModel]):
+    """Recipe for Pydantic responses.
+
+    This recipe covers the case that we have a Pydantic
+    data type (BaseModel) already specified in the runtime
+    making LLM calls, and we want to obtain an object of
+    that same data type as the output from the parser.
+
+    This recipe operates in a very similar fashion to
+    `StructuredResponseRecipe` except that it is initialized
+    from a Pydantic `BaseModel` and does the extra step of
+    validating against that `BaseModel` using
+    `BaseModel.model_validate` for its return.
+    """
+
+    data_type: type[BaseModel]
+
+    def __init__(self, data_type: type[BaseModel], **kwargs):
+        """Initialize a PydanticResponseRecipe.
+
+        Args:
+            data_type (type(BaseModel)): The target Pydantic BaseModel
+                subclass that the LLM should adhere to in its response,
+                and defines the output type of the parser.
+        """
+        super().__init__(**kwargs)
+        self.data_type = data_type
+
+    @property
+    def schema(self) -> str:
+        return json.dumps(self.data_type.model_json_schema())
+
+    @property
+    def task_instructions(self) -> str:
+        return (
+            "* Your response must be in JSON format.\n"
+            "* Your JSON response must be returned within a Markdown, ```json code fence.\n"
+            "* The JSON format is given as a JSON Schema description within <response_schema> markup tags.\n\n"
+            f"<response_schema>\n{self.schema}\n</response_schema>"
+        )
+
+    @property
+    def example_template(self) -> str:
+        return "```json\n{example}\n```"
+
+    def generate_response_example(self, example: BaseModel) -> str:
+        return self.example_template.format(example=example.model_dump_json())
+
+    def serialize_output(self, output: BaseModel) -> str:
+        return output.model_dump_json()
+
+    def deserialize_output(self, serialized_output: str) -> BaseModel:
+        return self.data_type.model_validate_json(serialized_output)
+
+    def _build_parser_fn(self) -> Callable[[str], BaseModel]:
+        parser = LLMResponseParser(
+            postprocessors=[
+                merge_text_blocks,
+                deserialize_json_code,
+            ]
+        )
+
+        def parse_fn(response: str) -> BaseModel:
+            try:
+                obj = parser.parse(response).filter([StructuredDataBlock]).parsed.pop().obj
+                return self.data_type.model_validate(obj)
+            except IndexError:
+                raise ParserException(
+                    "No parsable JSON structure within ```json markdown fence.",
+                    source=response,
+                ) from None
+            except Exception as exc:
+                raise ParserException(
+                    "Response doesn't match requested <response_schema>\n" + str(exc),
+                    source=response,
+                ) from None
+
+        return parse_fn
+
+
+class CodeResponseRecipe(ResponseRecipe[str]):
+    """Obtain a code snippet from an LLM."""
+
+    def __init__(self, syntax: str | CodeLang, **kwargs):
+        """Initialize a CodeResponseRecipe.
+
+        Args:
+            syntax (str | CodeLang): The code syntax that the
+                LLM should adhere to, e.g. `"python"`, `"sql"`, etc.
+        """
+        super().__init__(**kwargs)
+        self.syntax = CodeLang.parse_lang(syntax)
+
+    @property
+    def task_instructions(self) -> str:
+        return (
+            f"* Your response must be code written in {self.syntax}.\n"
+            "* You will follow accepted and common syntax and best-practices.\n"
+            f"* Your response will be given in markdown code fences specifying the correct language.\n"
+            "* Only respond with a SINGLE code block."
+        )
+
+    @property
+    def example_template(self) -> str:
+        return f"```{self.syntax}\n{{example}}\n```\n"
+
+    def serialize_output(self, output: str) -> str:
+        return output
+
+    def deserialize_output(self, serialized_output: str) -> str:
+        return serialized_output
+
+    def _build_parser_fn(self) -> Callable[[str], str]:
+        parser = LLMResponseParser(
+            postprocessors=[
+                merge_text_blocks,
+            ]
+        )
+
+        def parse_fn(response: str) -> str:
+            try:
+                code_block = parser.parse(response).filter([CodeBlock]).parsed.pop()
+                # For the type checker -- should always pass
+                assert isinstance(code_block, CodeBlock)
+            except IndexError:
+                raise ParserException(
+                    "No parsable code response.",
+                    source=response,
+                ) from None
+
+            # Only report this as a parser error if there was a mismatch.
+            if code_block.code_lang and code_block.code_lang != self.syntax:
+                raise ParserException(
+                    f"Responded with code not matching the requested syntax ({self.syntax}).",
+                    source=response,
+                )
+
+            return code_block.code.strip()
+
+        return parse_fn