langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__py3-none-any.whl

langfun/core/structured/schema.py

@@ -1,987 +0,0 @@
-# Copyright 2023 The Langfun Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Schema for structured data."""
-
-import abc
-import inspect
-import io
-import re
-import sys
-import textwrap
-import typing
-from typing import Any, Literal, Sequence, Type, Union
-import langfun.core as lf
-from langfun.core.coding.python import correction
-import pyglove as pg
-
-
-def include_method_in_prompt(method):
-  """Decorator to include a method in the class definition of the prompt."""
-  setattr(method, '__show_in_prompt__', True)
-  return method
-
-
-def should_include_method_in_prompt(method):
-  """Returns true if the method should be shown in the prompt."""
-  return getattr(method, '__show_in_prompt__', False)
-
-
-def parse_value_spec(value) -> pg.typing.ValueSpec:
-  """Parses a PyGlove ValueSpec equivalence into a ValueSpec."""
-  if isinstance(value, pg.typing.ValueSpec):
-    return value
-
-  if isinstance(value, dict) and len(value) == 1 and 'result' in value:
-    value = value['result']
-
-  def _parse_node(v) -> pg.typing.ValueSpec:
-    if isinstance(v, dict):
-      return pg.typing.Dict([(k, _parse_node(cv)) for k, cv in v.items()])
-    elif isinstance(v, list):
-      if len(v) != 1:
-        raise ValueError(
-            'Annotation with list must be a list of a single element. '
-            f'Encountered: {v}'
-        )
-      return pg.typing.List(_parse_node(v[0]))
-    else:
-      spec = pg.typing.ValueSpec.from_annotation(v, auto_typing=True)
-      if isinstance(
-          spec,
-          (
-              pg.typing.Any,
-              pg.typing.Callable,
-              pg.typing.Tuple,
-              pg.typing.Type,
-              pg.typing.Union,
-          ),
-      ):
-        raise ValueError(f'Unsupported schema specification: {v}')
-      return spec
-
-  return _parse_node(value)
-
-
-SchemaProtocol = Literal['json', 'python']
-
-
-class SchemaError(Exception):   # pylint: disable=g-bad-exception-name
-  """Schema error."""
-
-  def __init__(self,
-               schema: 'Schema',
-               value: Any,
-               protocol: SchemaProtocol,
-               cause: Exception):
-    self.schema = schema
-    self.value = value
-    self.protocol = protocol
-    self.cause = cause
-
-  def __str__(self):
-    r = io.StringIO()
-    r.write(
-        pg.colored(
-            f'{self.cause.__class__.__name__}: {self.cause}', 'magenta'
-        )
-    )
-
-    r.write('\n')
-    r.write(pg.colored('Schema:', 'red'))
-    r.write('\n\n')
-    r.write(textwrap.indent(
-        pg.colored(
-            schema_repr(self.protocol).repr(self.schema), 'magenta'
-        ),
-        ' ' * 2
-    ))
-    r.write('\n\n')
-    r.write(pg.colored('Generated value:', 'red'))
-    r.write('\n\n')
-    r.write(textwrap.indent(
-        pg.colored(value_repr(self.protocol).repr(self.value), 'magenta'),
-        ' ' * 2
-    ))
-    return r.getvalue()
-
-
-class Schema(
-    lf.NaturalLanguageFormattable,
-    pg.Object,
-    pg.views.HtmlTreeView.Extension
-):
-  """Base class for structured data schema."""
-
-  spec: pg.typing.Annotated[
-      pg.typing.Object(pg.typing.ValueSpec, transform=parse_value_spec),
-      (
-          'A PyGlove ValueSpec object representing the spec for the value '
-          'to be parsed.'
-      ),
-  ]
-
-  def schema_str(self, protocol: SchemaProtocol = 'json', **kwargs) -> str:
-    """Returns the representation of the schema."""
-    return schema_repr(protocol).repr(self, **kwargs)
-
-  def value_str(
-      self, value: Any, protocol: SchemaProtocol = 'json', **kwargs
-  ) -> str:
-    """Returns the representation of a structured value."""
-    return value_repr(protocol).repr(value, self, **kwargs)
-
-  def parse(
-      self, text: str, protocol: SchemaProtocol = 'json', **kwargs
-  ) -> Any:
-    """Parse a LM generated text into a structured value."""
-    value = value_repr(protocol).parse(text, self, **kwargs)
-
-    # TODO(daiyip): support autofix for schema error.
-    try:
-      return self.spec.apply(value)
-    except Exception as e:
-      raise SchemaError(self, value, protocol, e)  # pylint: disable=raise-missing-from
-
-  def natural_language_format(self) -> str:
-    return self.schema_str()
-
-  def schema_dict(self) -> dict[str, Any]:
-    """Returns the dict representation of the schema."""
-
-    def _node(vs: pg.typing.ValueSpec) -> Any:
-      if isinstance(vs, pg.typing.PrimitiveType):
-        return vs
-      elif isinstance(vs, pg.typing.Dict):
-        assert vs.schema is not None
-        return {str(k): _node(f.value) for k, f in vs.schema.fields.items()}
-      elif isinstance(vs, pg.typing.List):
-        return [_node(vs.element.value)]
-      elif isinstance(vs, pg.typing.Object):
-        if issubclass(vs.cls, pg.Object):
-          d = {pg.JSONConvertible.TYPE_NAME_KEY: vs.cls.__serialization_key__}
-          d.update(
-              {
-                  str(k): _node(f.value)
-                  for k, f in vs.cls.__schema__.fields.items()
-              }
-          )
-          return d
-      raise TypeError(
-          'Unsupported value spec being used as the schema for '
-          f'structured data: {vs}.')
-
-    return {'result': _node(self.spec)}
-
-  def class_dependencies(
-      self,
-      include_base_classes: bool = True,
-      include_subclasses: bool = True,
-      include_generated_subclasses: bool = False) -> list[Type[Any]]:
-    """Returns a list of class dependencies for current schema."""
-    return class_dependencies(
-        self.spec,
-        include_base_classes,
-        include_subclasses,
-        include_generated_subclasses
-    )
-
-  @classmethod
-  def from_value(cls, value) -> 'Schema':
-    """Creates a schema from an equivalent representation."""
-    if isinstance(value, Schema):
-      return value
-    return cls(parse_value_spec(value))
-
-  def _html_tree_view_content(
-      self,
-      *,
-      view: pg.views.HtmlTreeView,
-      **kwargs,
-  ):
-    return pg.Html.element(
-        'div',
-        [pg.Html.escape(self.schema_str(protocol='python'))],
-        css_classes=['lf-schema-definition']
-    ).add_style(
-        """
-        .lf-schema-definition {
-            color: blue;
-            margin: 5px;
-            white-space: pre-wrap;
-        }
-        """
-    )
-
-
-SchemaType = Union[Schema, Type[Any], list[Type[Any]], dict[str, Any]]
-
-
-def _top_level_object_specs_from_value(value: pg.Symbolic) -> list[Type[Any]]:
-  """Returns a list of top level value specs from a symbolic value."""
-  top_level_object_specs = []
-
-  def _collect_top_level_object_specs(k, v, p):
-    del k, p
-    if isinstance(v, pg.Object):
-      top_level_object_specs.append(pg.typing.Object(v.__class__))
-      return pg.TraverseAction.CONTINUE
-    return pg.TraverseAction.ENTER
-
-  pg.traverse(value, _collect_top_level_object_specs)
-  return top_level_object_specs
-
-
-def class_dependencies(
-    value_or_spec: Union[
-        pg.Symbolic,
-        Schema,
-        pg.typing.ValueSpec,
-        Type[pg.Object],
-        tuple[Union[pg.typing.ValueSpec, Type[pg.Object]], ...],
-    ],
-    include_base_classes: bool = True,
-    include_subclasses: bool = True,
-    include_generated_subclasses: bool = False,
-) -> list[Type[Any]]:
-  """Returns a list of class dependencies from a value or specs."""
-  if isinstance(value_or_spec, Schema):
-    value_or_spec = value_or_spec.spec
-
-  if inspect.isclass(value_or_spec) or isinstance(
-      value_or_spec, pg.typing.ValueSpec
-  ):
-    value_or_spec = (value_or_spec,)
-
-  if isinstance(value_or_spec, tuple):
-    value_specs = []
-    for v in value_or_spec:
-      if isinstance(v, pg.typing.ValueSpec):
-        value_specs.append(v)
-      elif inspect.isclass(v):
-        value_specs.append(pg.typing.Object(v))
-      else:
-        raise TypeError(f'Unsupported spec type: {v!r}')
-  else:
-    value_specs = _top_level_object_specs_from_value(value_or_spec)
-
-  seen = set()
-  dependencies = []
-
-  def _add_dependency(cls_or_classes):
-    if isinstance(cls_or_classes, type):
-      cls_or_classes = [cls_or_classes]
-    for cls in cls_or_classes:
-      if cls not in dependencies:
-        dependencies.append(cls)
-
-  def _fill_dependencies(vs: pg.typing.ValueSpec, include_subclasses: bool):
-    if isinstance(vs, pg.typing.Object):
-      if vs.cls not in seen:
-        seen.add(vs.cls)
-
-        if include_base_classes:
-          # Add base classes as dependencies.
-          for base_cls in vs.cls.__bases__:
-            # We only keep track of user-defined symbolic classes.
-            if base_cls is not object and base_cls is not pg.Object:
-              _fill_dependencies(
-                  pg.typing.Object(base_cls), include_subclasses=False
-              )
-
-        # Add members as dependencies.
-        for field in pg.schema(vs.cls).values():
-          _fill_dependencies(field.value, include_subclasses)
-      _add_dependency(vs.cls)
-
-      # Check subclasses if available.
-      if include_subclasses:
-        for cls in vs.cls.__subclasses__():
-          # NOTE(daiyip): To prevent LLM-generated "hallucinated" classes from
-          # polluting the generation space, classes dynamically created by
-          # 'eval' (which have __module__ == 'builtins') are excluded from
-          # dependencies by default.
-          if ((include_generated_subclasses or cls.__module__ != 'builtins')
-              and cls not in dependencies):
-            _fill_dependencies(pg.typing.Object(cls), include_subclasses=True)
-
-    if isinstance(vs, pg.typing.List):
-      _fill_dependencies(vs.element.value, include_subclasses)
-    elif isinstance(vs, pg.typing.Tuple):
-      for elem in vs.elements:
-        _fill_dependencies(elem.value, include_subclasses)
-    elif isinstance(vs, pg.typing.Dict) and vs.schema:
-      for v in vs.schema.values():
-        _fill_dependencies(v.value, include_subclasses)
-    elif isinstance(vs, pg.typing.Union):
-      for v in vs.candidates:
-        _fill_dependencies(v, include_subclasses)
-
-  for value_spec in value_specs:
-    _fill_dependencies(value_spec, include_subclasses)
-  return dependencies
-
-
-def schema_spec(noneable: bool = False) -> pg.typing.ValueSpec:  # pylint: disable=unused-argument
-  if typing.TYPE_CHECKING:
-    return Any
-  return pg.typing.Object(
-      Schema, transform=Schema.from_value, is_noneable=noneable
-  )  # pylint: disable=unreachable-code
-
-
-#
-# Schema representations.
-#
-
-
-class SchemaRepr(metaclass=abc.ABCMeta):
-  """Base class for schema representation."""
-
-  @abc.abstractmethod
-  def repr(self, schema: Schema) -> str:
-    """Returns the representation of the schema."""
-
-
-class SchemaPythonRepr(SchemaRepr):
-  """Python-representation for a schema."""
-
-  def repr(
-      self,
-      schema: Schema,
-      *,
-      include_result_definition: bool = True,
-      markdown: bool = True,
-      **kwargs,
-  ) -> str:
-    ret = ''
-    if include_result_definition:
-      ret += self.result_definition(schema)
-    class_definition_str = self.class_definitions(
-        schema, markdown=markdown, **kwargs
-    )
-    if class_definition_str:
-      ret += f'\n\n{class_definition_str}'
-    return ret.strip()
-
-  def class_definitions(
-      self,
-      schema: Schema,
-      additional_dependencies: list[Type[Any]] | None = None,
-      **kwargs
-  ) -> str | None:
-    """Returns a string containing of class definitions from a schema."""
-    deps = schema.class_dependencies(
-        include_base_classes=False, include_subclasses=True
-    )
-    allowed_dependencies = set(deps)
-    if additional_dependencies:
-      allowed_dependencies.update(additional_dependencies)
-    return class_definitions(
-        deps, allowed_dependencies=allowed_dependencies, **kwargs)
-
-  def result_definition(self, schema: Schema) -> str:
-    return annotation(schema.spec)
-
-
-def source_form(value, compact: bool = True, markdown: bool = False) -> str:
-  """Returns the source code form of an object."""
-  return ValuePythonRepr().repr(value, compact=compact, markdown=markdown)
-
-
-def class_definitions(
-    classes: Sequence[Type[Any]],
-    *,
-    allowed_dependencies: set[Type[Any]] | None = None,
-    strict: bool = False,
-    markdown: bool = False,
-) -> str | None:
-  """Returns a str for class definitions."""
-  if not classes:
-    return None
-  def_str = io.StringIO()
-  for i, cls in enumerate(classes):
-    if i > 0:
-      def_str.write('\n')
-    def_str.write(
-        class_definition(
-            cls,
-            strict=strict,
-            allowed_dependencies=allowed_dependencies,
-        )
-    )
-  ret = def_str.getvalue()
-  if markdown and ret:
-    ret = f'```python\n{ret}```'
-  return ret
-
-
-def class_definition(
-    cls,
-    strict: bool = False,
-    allowed_dependencies: set[Type[Any]] | None = None,
-) -> str:
-  """Returns the Python class definition."""
-  out = io.StringIO()
-  schema = pg.schema(cls)
-  eligible_bases = []
-  for base_cls in cls.__bases__:
-    if base_cls is not object:
-      if allowed_dependencies is None or base_cls in allowed_dependencies:
-        eligible_bases.append(base_cls.__name__)
-
-  if eligible_bases:
-    base_cls_str = ', '.join(eligible_bases)
-    out.write(f'class {cls.__name__}({base_cls_str}):\n')
-  else:
-    out.write(f'class {cls.__name__}:\n')
-
-  if cls.__doc__:
-    doc_lines = cls.__doc__.strip().split('\n')
-    if len(doc_lines) == 1:
-      out.write(f'  """{cls.__doc__}"""\n')
-    else:
-      out.write('  """')
-
-      # Since Python 3.13, the indentation of docstring lines is removed.
-      # Therefore, we add two spaces to each non-empty line to keep the
-      # indentation consistent with the class definition.
-      if sys.version_info >= (3, 13):
-        for i in range(1, len(doc_lines)):
-          if doc_lines[i]:
-            doc_lines[i] = ' ' * 2 + doc_lines[i]
-
-      for line in doc_lines:
-        out.write(line)
-        out.write('\n')
-      out.write('  """\n')
-
-  empty_class = True
-  if schema.fields:
-    for key, field in schema.items():
-      if not isinstance(key, pg.typing.ConstStrKey):
-        pg.logging.warning(
-            'Variable-length keyword arguments is not supported in '
-            f'structured parsing or query. Encountered: {cls}, Schema: {schema}'
-        )
-        continue
-
-      # Skip fields that are marked as excluded from the prompt sent to LLM
-      # for OOP.
-      if field.metadata.get('exclude_from_prompt', False):
-        continue
-
-      # Write field doc string as comments before the field definition.
-      if field.description:
-        for line in field.description.split('\n'):
-          if line:
-            out.write('  # ')
-            out.write(line)
-            out.write('\n')
-
-      annotation_str = annotation(
-          field.value, strict=strict, allowed_dependencies=allowed_dependencies
-      )
-      out.write(f'  {field.key}: {annotation_str}')
-      out.write('\n')
-      empty_class = False
-
-  for method in _iter_newly_defined_methods(cls, allowed_dependencies):
-    source = inspect.getsource(method)
-    # Remove decorators from the method definition.
-    source = re.sub(r'\s*@.*\.include_method_in_prompt.*\n', '', source)
-    out.write('\n')
-    out.write(
-        textwrap.indent(
-            inspect.cleandoc('\n' + source), ' ' * 2)
-    )
-    out.write('\n')
-    empty_class = False
-
-  if empty_class:
-    out.write('  pass\n')
-  return out.getvalue()
-
-
-def _iter_newly_defined_methods(
-    cls, allowed_dependencies: set[Type[Any]] | None):
-  names = {attr_name: True for attr_name in dir(cls)}
-  for base in cls.__bases__:
-    if allowed_dependencies is None or base in allowed_dependencies:
-      for name in dir(base):
-        names.pop(name, None)
-  for name in names.keys():
-    attr = getattr(cls, name)
-    if callable(attr) and should_include_method_in_prompt(attr):
-      yield attr
-
-
-def annotation(
-    vs: pg.typing.ValueSpec,
-    annotate_optional: bool = True,
-    strict: bool = False,
-    allowed_dependencies: set[Type[Any]] | None = None,
-) -> str:
-  """Returns the annotation string for a value spec."""
-  child_annotation_kwargs = dict(
-      strict=strict, allowed_dependencies=allowed_dependencies
-  )
-  if isinstance(vs, pg.typing.Any):
-    return 'Any'
-  elif isinstance(vs, pg.typing.Enum):
-    candidate_str = ', '.join([repr(v) for v in vs.values])
-    return f'Literal[{candidate_str}]'
-  elif isinstance(vs, pg.typing.Union):
-    candidate_str = ', '.join(
-        [
-            annotation(c, annotate_optional=False, **child_annotation_kwargs)
-            for c in vs.candidates
-        ]
-    )
-    if vs.is_noneable:
-      candidate_str += ', None'
-    return f'Union[{candidate_str}]'
-
-  if isinstance(vs, pg.typing.Bool):
-    x = 'bool'
-  elif isinstance(vs, pg.typing.Str):
-    if vs.regex is None:
-      x = 'str'
-    else:
-      if strict:
-        x = f"pg.typing.Str(regex='{vs.regex.pattern}')"
-      else:
-        x = f"str(regex='{vs.regex.pattern}')"
-  elif isinstance(vs, pg.typing.Number):
-    constraints = []
-    min_label = 'min_value' if strict else 'min'
-    max_label = 'max_value' if strict else 'max'
-    if vs.min_value is not None:
-      constraints.append(f'{min_label}={vs.min_value}')
-    if vs.max_value is not None:
-      constraints.append(f'{max_label}={vs.max_value}')
-    x = 'int' if isinstance(vs, pg.typing.Int) else 'float'
-    if constraints:
-      if strict:
-        x = (
-            'pg.typing.Int'
-            if isinstance(vs, pg.typing.Int)
-            else 'pg.typing.Float'
-        )
-      x += '(' + ', '.join(constraints) + ')'
-  elif isinstance(vs, pg.typing.Object):
-    if allowed_dependencies is None or vs.cls in allowed_dependencies:
-      x = vs.cls.__name__
-    else:
-      x = 'Any'
-  elif isinstance(vs, pg.typing.List):
-    item_str = annotation(vs.element.value, **child_annotation_kwargs)
-    x = f'list[{item_str}]'
-  elif isinstance(vs, pg.typing.Tuple):
-    elem_str = ', '.join(
-        [annotation(el.value, **child_annotation_kwargs) for el in vs.elements]
-    )
-    x = f'tuple[{elem_str}]'
-  elif isinstance(vs, pg.typing.Dict):
-    kv_pairs = None
-    if vs.schema is not None:
-      kv_pairs = [
-          (k, annotation(f.value, **child_annotation_kwargs))
-          for k, f in vs.schema.items()
-          if isinstance(k, pg.typing.ConstStrKey)
-      ]
-
-    if kv_pairs:
-      kv_str = ', '.join(f"'{k}': {v}" for k, v in kv_pairs)
-      x = '{' + kv_str + '}'
-      if strict:
-        x = f'pg.typing.Dict({x})'
-    elif vs.schema and vs.schema.dynamic_field:
-      v = annotation(vs.schema.dynamic_field.value, **child_annotation_kwargs)
-      x = f'dict[str, {v}]'
-    else:
-      x = 'dict[str, Any]'
-
-  else:
-    raise TypeError(f'Unsupported value spec being used as schema: {vs}.')
-
-  if annotate_optional and vs.is_noneable:
-    x += ' | None'
-  return x
-
-
-class SchemaJsonRepr(SchemaRepr):
-  """JSON-representation for a schema."""
-
-  def repr(self, schema: Schema, **kwargs) -> str:
-    del kwargs
-    out = io.StringIO()
-    def _visit(node: Any) -> None:
-      if isinstance(node, str):
-        out.write(f'"{node}"')
-      elif isinstance(node, list):
-        assert len(node) == 1, node
-        out.write('[')
-        _visit(node[0])
-        out.write(']')
-      elif isinstance(node, dict):
-        out.write('{')
-        for i, (k, v) in enumerate(node.items()):
-          if i != 0:
-            out.write(', ')
-          out.write(f'"{k}": ')
-          _visit(v)
-        out.write('}')
-      elif isinstance(node, pg.typing.Enum):
-        out.write(' | '.join(
-            f'"{v}"' if isinstance(v, str) else repr(v)
-            for v in node.values))
-      elif isinstance(node, pg.typing.PrimitiveType):
-        x = node.value_type.__name__
-        if isinstance(node, pg.typing.Number):
-          params = []
-          if node.min_value is not None:
-            params.append(f'min={node.min_value}')
-          if node.max_value is not None:
-            params.append(f'max={node.max_value}')
-          if params:
-            x += f'({", ".join(params)})'
-        elif isinstance(node, pg.typing.Str):
-          if node.regex is not None:
-            x += f'(regex={node.regex.pattern})'
-        if node.is_noneable:
-          x = x + ' | None'
-        out.write(x)
-      else:
-        raise ValueError(
-            f'Unsupported value spec being used as schema: {node}.')
-    _visit(schema.schema_dict())
-    return out.getvalue()
-
-
-#
-# Value representations.
-#
-
-
-class ValueRepr(metaclass=abc.ABCMeta):
-  """Base class for value representation."""
-
-  @abc.abstractmethod
-  def repr(self, value: Any, schema: Schema | None = None, **kwargs) -> str:
-    """Returns the representation of a structured value."""
-
-  @abc.abstractmethod
-  def parse(self, text: str, schema: Schema | None = None, **kwargs) -> Any:
-    """Parse a LM generated text into a structured value."""
-
-
-class ValuePythonRepr(ValueRepr):
-  """Python-representation for value."""
-
-  def repr(self,
-           value: Any,
-           schema: Schema | None = None,
-           *,
-           compact: bool = True,
-           verbose: bool = False,
-           markdown: bool = True,
-           assign_to_var: str | None = None,
-           **kwargs) -> str:
-    del schema
-    if inspect.isclass(value):
-      cls_schema = Schema.from_value(value)
-      if isinstance(cls_schema.spec, pg.typing.Object):
-        object_code = SchemaPythonRepr().class_definitions(
-            cls_schema,
-            markdown=markdown,
-            # We add `pg.Object` as additional dependencies to the class
-            # definition so exemplars for class generation could show
-            # pg.Object as their bases.
-            additional_dependencies=[pg.Object]
-        )
-        assert object_code is not None
-        return object_code
-      else:
-        object_code = SchemaPythonRepr().result_definition(cls_schema)
-    elif isinstance(value, lf.Template):
-      return str(value)
-    else:
-      object_code = pg.format(
-          value, compact=compact, verbose=verbose, python_format=True
-      )
-      if assign_to_var is not None:
-        object_code = f'{assign_to_var} = {object_code}'
-    if markdown:
-      return f'```python\n{ object_code }\n```'
-    return object_code
-
-  def parse(
-      self,
-      text: str,
-      schema: Schema | None = None,
-      *,
-      additional_context: dict[str, Type[Any]] | None = None,
-      permission: pg.coding.CodePermission = (
-          pg.coding.CodePermission.ASSIGN | pg.coding.CodePermission.CALL
-      ),
-      autofix=0,
-      autofix_lm: lf.LanguageModel = lf.contextual(),
-      **kwargs,
-  ) -> Any:
-    """Parse a Python string into a structured object."""
-    del kwargs
-    global_vars = additional_context or {}
-    if schema is not None:
-      dependencies = schema.class_dependencies()
-      global_vars.update({d.__name__: d for d in dependencies})
-    return structure_from_python(
-        text,
-        global_vars=global_vars,
-        autofix=autofix,
-        autofix_lm=autofix_lm,
-        permission=permission,
-    )
-
-
-def structure_from_python(
-    code: str,
-    *,
-    global_vars: dict[str, Any] | None = None,
-    permission: pg.coding.CodePermission = (
-        pg.coding.CodePermission.ASSIGN | pg.coding.CodePermission.CALL
-    ),
-    autofix=0,
-    autofix_lm: lf.LanguageModel = lf.contextual(),
-) -> Any:
-  """Evaluates structure from Python code with access to symbols."""
-  global_vars = global_vars or {}
-  global_vars.update({
-      'pg': pg,
-      'Object': pg.Object,
-      'Any': typing.Any,
-      'List': typing.List,
-      'Tuple': typing.Tuple,
-      'Dict': typing.Dict,
-      'Sequence': typing.Sequence,
-      'Optional': typing.Optional,
-      'Union': typing.Union,
-      # Special value markers.
-      'UNKNOWN': UNKNOWN,
-  })
-  # We are creating objects here, so we execute the code without a sandbox.
-  return correction.run_with_correction(
-      code,
-      global_vars=global_vars,
-      sandbox=False,
-      max_attempts=autofix,
-      lm=autofix_lm,
-      permission=permission,
-  )
-
-
-class JsonError(Exception):
-  """Json parsing error."""
-
-  def __init__(self, json: str, cause: Exception):
-    self.json = json
-    self.cause = cause
-
-  def __str__(self) -> str:
-    r = io.StringIO()
-    r.write(
-        pg.colored(
-            f'{self.cause.__class__.__name__}: {self.cause}', 'magenta'
-        )
-    )
-
-    r.write('\n\n')
-    r.write(pg.colored('JSON text:', 'red'))
-    r.write('\n\n')
-    r.write(textwrap.indent(pg.colored(self.json, 'magenta'), ' ' * 2))
-    return r.getvalue()
-
-
-class ValueJsonRepr(ValueRepr):
-  """JSON-representation for value."""
-
-  def repr(self, value: Any, schema: Schema | None = None, **kwargs) -> str:
-    del schema
-    return pg.to_json_str(dict(result=value))
-
-  def parse(self, text: str, schema: Schema | None = None, **kwargs) -> Any:
-    """Parse a JSON string into a structured object."""
-    del schema
-    try:
-      text = cleanup_json(text)
-      v = pg.from_json_str(text, **kwargs)
-    except Exception as e:
-      raise JsonError(text, e)  # pylint: disable=raise-missing-from
-
-    if not isinstance(v, dict) or 'result' not in v:
-      raise JsonError(text, ValueError(
-          'The root node of the JSON must be a dict with key `result`. '
-          f'Encountered: {v}'
-      ))
-    return v['result']
-
-
-def cleanup_json(json_str: str) -> str:
-  """Clean up the LM responded JSON string."""
-  # Treatments:
-  # 1. Extract the JSON string with a top-level dict from the response.
-  #    This prevents the leading and trailing texts in the response to
-  #    be counted as part of the JSON.
-  # 2. Escape new lines in JSON values.
-
-  curly_brackets = 0
-  under_json = False
-  under_str = False
-  str_begin = -1
-
-  cleaned = io.StringIO()
-  for i, c in enumerate(json_str):
-    if c == '{' and not under_str:
-      cleaned.write(c)
-      curly_brackets += 1
-      under_json = True
-      continue
-    elif not under_json:
-      continue
-
-    if c == '}' and not under_str:
-      cleaned.write(c)
-      curly_brackets -= 1
-      if curly_brackets == 0:
-        break
-    elif c == '"' and json_str[i - 1] != '\\':
-      under_str = not under_str
-      if under_str:
-        str_begin = i
-      else:
-        assert str_begin > 0
-        str_value = json_str[str_begin : i + 1].replace('\n', '\\n')
-        cleaned.write(str_value)
-        str_begin = -1
-    elif not under_str:
-      cleaned.write(c)
-
-  if not under_json:
-    raise ValueError(f'No JSON dict in the output: {json_str}')
-
-  if curly_brackets > 0:
-    raise ValueError(
-        f'Malformated JSON: missing {curly_brackets} closing curly braces.'
-    )
-
-  return cleaned.getvalue()
-
-
-def schema_repr(protocol: SchemaProtocol) -> SchemaRepr:
-  """Gets a SchemaRepr object from protocol."""
-  if protocol == 'json':
-    return SchemaJsonRepr()
-  elif protocol == 'python':
-    return SchemaPythonRepr()
-  raise ValueError(f'Unsupported protocol: {protocol}.')
-
-
-def value_repr(protocol: SchemaProtocol) -> ValueRepr:
-  if protocol == 'json':
-    return ValueJsonRepr()
-  elif protocol == 'python':
-    return ValuePythonRepr()
-  raise ValueError(f'Unsupported protocol: {protocol}.')
-
-
-#
-# Special value markers.
-#
-
-
-class Missing(pg.Object, pg.typing.CustomTyping):
-  """Value marker for a missing field.
-
-  This class differs from pg.MISSING_VALUE in two aspects:
-  * When a field is assigned with lf.Missing(), it's considered non-partial.
-  * lf.Missing() could format the value spec as Python annotations that are
-    consistent with `lf.structured.Schema.schema_repr()`.
-  """
-
-  def _on_bound(self):
-    super()._on_bound()
-    self._value_spec = None
-
-  @property
-  def value_spec(self) -> pg.ValueSpec | None:
-    """Returns the value spec that applies to the current missing value."""
-    return self._value_spec
-
-  def custom_apply(
-      self, path: pg.KeyPath, value_spec: pg.ValueSpec, *args, **kwargs
-  ) -> tuple[bool, Any]:
-    self._value_spec = value_spec
-    return (False, self)
-
-  def format(self, *args, **kwargs) -> str:
-    if self._value_spec is None:
-      return 'MISSING'
-    return f'MISSING({annotation(self._value_spec)})'
-
-  @classmethod
-  def find_missing(cls, value: Any) -> dict[str, 'Missing']:
-    """Lists all missing values contained in the value."""
-    missing = {}
-
-    def _visit(k, v, p):
-      del p
-      if isinstance(v, Missing):
-        missing[k] = v
-      return pg.TraverseAction.ENTER
-
-    pg.traverse(value, _visit)
-    return missing
-
-
-MISSING = Missing()
-
-
-def mark_missing(value: Any) -> Any:
-  """Replaces pg.MISSING within the value with lf.structured.Missing objects."""
-  if isinstance(value, list):
-    value = pg.List(value)
-  elif isinstance(value, dict):
-    value = pg.Dict(value)
-  if isinstance(value, pg.Symbolic):
-
-    def _mark_missing(k, v, p):
-      del k, p
-      if pg.MISSING_VALUE == v:
-        v = Missing()
-      return v
-
-    return value.rebind(_mark_missing, raise_on_no_change=False)
-  return value
-
-
-class Unknown(pg.Object, pg.typing.CustomTyping):
-  """Value marker for a field that LMs could not provide."""
-
-  def custom_apply(self, *args, **kwargs) -> tuple[bool, Any]:
-    return (False, self)
-
-  def format(self, *args, **kwargs) -> str:
-    return 'UNKNOWN'
-
-
-UNKNOWN = Unknown()