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

langfun/core/structured/schema/python_test.py

@@ -0,0 +1,410 @@
+# Copyright 2025 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.
+import dataclasses
+import inspect
+import typing
+import unittest
+
+import langfun.core as lf
+from langfun.core.llms import fake
+from langfun.core.structured.schema import base
+from langfun.core.structured.schema import python
+import pyglove as pg
+
+
+class Activity(pg.Object):
+  description: str
+
+
+class Itinerary(pg.Object):
+  """A travel itinerary for a day."""
+
+  day: pg.typing.Int[1, None]
+  type: pg.typing.Enum['daytime', 'nighttime']
+  activities: list[Activity]
+  hotel: pg.typing.Annotated[
+      pg.typing.Str['.*Hotel'] | None,
+      'Hotel to stay if applicable.'
+  ]
+
+
+class PlaceOfInterest(pg.Object):
+  """The name of a place of interest.
+
+  A place of interest is a place that people ususally visit during their
+  travels.
+  """
+
+  name: str
+
+
+Itinerary.__serialization_key__ = 'Itinerary'
+
+
+class Node(pg.Object):
+  children: list['Node']
+
+
+class SchemaReprTest(unittest.TestCase):
+
+  def test_class_definition(self):
+    self.assertEqual(
+        python.class_definition(Activity, allowed_dependencies=set()),
+        'class Activity:\n  description: str\n',
+    )
+    self.assertEqual(
+        python.class_definition(Itinerary),
+        inspect.cleandoc("""
+            class Itinerary(Object):
+              \"\"\"A travel itinerary for a day.\"\"\"
+              day: int(min=1)
+              type: Literal['daytime', 'nighttime']
+              activities: list[Activity]
+              # Hotel to stay if applicable.
+              hotel: str(regex='.*Hotel') | None
+            """) + '\n',
+    )
+    self.assertEqual(
+        python.class_definition(
+            PlaceOfInterest, allowed_dependencies=set()
+        ),
+        inspect.cleandoc("""
+            class PlaceOfInterest:
+              \"\"\"The name of a place of interest.
+
+              A place of interest is a place that people ususally visit during their
+              travels.
+              \"\"\"
+              name: str
+            """) + '\n',
+    )
+
+    class A(pg.Object):
+      pass
+
+    self.assertEqual(
+        python.class_definition(A, allowed_dependencies=set()),
+        'class A:\n  pass\n',
+    )
+    self.assertEqual(
+        python.class_definition(A),
+        'class A(Object):\n  pass\n',
+    )
+
+    class C(pg.Object):
+      x: str
+      __kwargs__: typing.Any
+
+    self.assertEqual(
+        python.class_definition(C), 'class C(Object):\n  x: str\n'
+    )
+
+    class D(pg.Object):
+      x: str
+      @python.include_method_in_prompt
+      def __call__(self, y: int) -> int:
+        return len(self.x) + y
+
+    self.assertEqual(
+        python.class_definition(D),
+        inspect.cleandoc(
+            """
+            class D(Object):
+              x: str
+
+              def __call__(self, y: int) -> int:
+                return len(self.x) + y
+            """) + '\n'
+    )
+
+    class E(pg.Object):
+      x: str
+      y: typing.Annotated[int, 'y', dict(exclude_from_prompt=True)]
+
+    self.assertEqual(
+        python.class_definition(E),
+        inspect.cleandoc(
+            """
+            class E(Object):
+              x: str
+            """) + '\n'
+    )
+
+  def test_repr(self):
+    class Foo(pg.Object):
+      x: int
+
+    @dataclasses.dataclass(frozen=True)
+    class Bar:
+      """Class Bar."""
+      y: str
+
+    @dataclasses.dataclass(frozen=True)
+    class Baz(Bar):  # pylint: disable=unused-variable
+      pass
+
+    class A(pg.Object):
+      foo: Foo
+
+      @python.include_method_in_prompt
+      def foo_value(self) -> int:
+        return self.foo.x
+
+      def baz_value(self) -> str:
+        return 'baz'
+
+    class B(A):
+      bar: Bar
+      foo2: Foo
+
+      @python.include_method_in_prompt
+      def bar_value(self) -> str:
+        return self.bar.y
+
+    schema = base.Schema([B])
+    self.assertEqual(
+        python.PythonPromptingProtocol().class_definitions(schema),
+        inspect.cleandoc('''
+            class Foo:
+              x: int
+
+            class Bar:
+              """Class Bar."""
+              y: str
+
+            class Baz(Bar):
+              """Baz(y: str)"""
+              y: str
+
+            class B:
+              foo: Foo
+              bar: Bar
+              foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
+            ''') + '\n',
+    )
+
+    self.assertEqual(
+        python.PythonPromptingProtocol().result_definition(schema), 'list[B]'
+    )
+
+    self.assertEqual(
+        base.schema_repr(schema),
+        inspect.cleandoc('''
+            list[B]
+
+            ```python
+            class Foo:
+              x: int
+
+            class Bar:
+              """Class Bar."""
+              y: str
+
+            class Baz(Bar):
+              """Baz(y: str)"""
+              y: str
+
+            class B:
+              foo: Foo
+              bar: Bar
+              foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
+            ```
+            '''),
+    )
+    self.assertEqual(
+        python.PythonPromptingProtocol().schema_repr(
+            schema,
+            include_result_definition=False,
+            markdown=False,
+        ),
+        inspect.cleandoc('''
+            class Foo:
+              x: int
+
+            class Bar:
+              """Class Bar."""
+              y: str
+
+            class Baz(Bar):
+              """Baz(y: str)"""
+              y: str
+
+            class B:
+              foo: Foo
+              bar: Bar
+              foo2: Foo
+
+              def bar_value(self) -> str:
+                return self.bar.y
+
+              def foo_value(self) -> int:
+                return self.foo.x
+            '''),
+    )
+
+
+class ValuePythonReprTest(unittest.TestCase):
+
+  def test_repr(self):
+    class Foo(pg.Object):
+      x: int
+
+    class A(pg.Object):
+      foo: list[Foo]
+      y: str | None
+
+    self.assertEqual(
+        base.value_repr(1, base.Schema(int)),
+        '```python\n1\n```'
+    )
+    self.assertEqual(
+        base.value_repr(lf.Template('hi, {{a}}', a='foo')),
+        'hi, foo'
+    )
+    self.assertEqual(
+        base.value_repr(
+            A([Foo(1), Foo(2)], 'bar'), base.Schema(A), markdown=False,
+        ),
+        "A(foo=[Foo(x=1), Foo(x=2)], y='bar')",
+    )
+    self.assertEqual(
+        base.value_repr(
+            A([Foo(1), Foo(2)], 'bar'),
+            base.Schema(A),
+            markdown=True,
+            compact=False,
+            assign_to_var='output',
+        ),
+        inspect.cleandoc("""
+            ```python
+            output = A(
+              foo=[
+                Foo(
+                  x=1
+                ),
+                Foo(
+                  x=2
+                )
+              ],
+              y='bar'
+            )
+            ```
+            """),
+    )
+    self.assertEqual(
+        base.value_repr(A),
+        inspect.cleandoc("""
+            ```python
+            class Foo(Object):
+              x: int
+
+            class A(Object):
+              foo: list[Foo]
+              y: str | None
+            ```
+            """),
+    )
+    self.assertEqual(python.source_form(int), 'int')
+
+  def test_parse(self):
+    class Foo(pg.Object):
+      x: int
+
+    class A(pg.Object):
+      foo: list[Foo]
+      y: str | None
+
+    self.assertEqual(
+        base.parse_value(
+            "A(foo=[Foo(x=1), Foo(x=2)], y='bar')", base.Schema(A)
+        ),
+        A([Foo(1), Foo(2)], y='bar'),
+    )
+
+  def test_parse_with_correction(self):
+    class Foo(pg.Object):
+      x: int
+
+    class A(pg.Object):
+      foo: list[Foo]
+      y: str | None
+
+    self.assertEqual(
+        base.parse_value(
+            "A(foo=[Foo(x=1), Foo(x=2)], y='bar'",
+            base.Schema(A),
+            autofix=1,
+            autofix_lm=fake.StaticResponse(
+                inspect.cleandoc(
+                    """
+                    CorrectedCode(
+                        corrected_code='A(foo=[Foo(x=1), Foo(x=2)], y=\\\'bar\\\')',
+                    )
+                    """
+                )
+            ),
+        ),
+        A([Foo(1), Foo(2)], y='bar'),
+    )
+
+  def test_parse_class_def(self):
+    self.assertTrue(
+        inspect.isclass(
+            base.parse_value(
+                """
+                class A:
+                  x: Dict[str, Any]
+                  y: Optional[Sequence[str]]
+                  z: Union[int, List[int], Tuple[int]]
+                """,
+                permission=pg.coding.CodePermission.ALL,
+            )
+        )
+    )
+
+
+class StructureFromPythonTest(unittest.TestCase):
+
+  def test_parse_class_def(self):
+
+    class B:
+      pass
+
+    schema = base.Schema([B])
+    v = python.structure_from_python(
+        """
+        class C(B):
+          pass
+        """,
+        global_vars=dict(B=B),
+        permission=pg.coding.CodePermission.ALL,
+    )
+    self.assertEqual(v.__module__, 'builtins')
+    self.assertEqual(schema.class_dependencies(), [B])
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/structured/schema_generation.py

@@ -90,16 +90,35 @@ def generate_class(
     skip_lm: bool = False,
     **kwargs,
 ) -> Type[Any] | lf.Message:
-  """Generate a class with specified name based on the prompt.
-
-  Example:
-    ```
-    trip_cls = lf.classgen(
-        'Trip',
-        'A trip plan to visit {{ city }}, city='San Francisco',
-        lm=lf.llms.GeminiPro()
-    )
-    ```
+  """Generates a Python class dynamically from a prompt using an LLM.
+
+  `lf.structured.generate_class` takes a class name and a natural language
+  description (prompt) and uses a language model to generate a Python class
+  (inheriting from `pg.Object`) that matches the description.
+  This is useful for creating structured data types on-the-fly based on
+  dynamic requirements.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+  import pyglove as pg
+
+  trip_plan_cls = lf.structured.generate_class(
+      'TripPlan',
+      'A trip plan to visit San Francisco, including a list of destinations,'
+      'start date, end date, and total budget.',
+      lm=lf.llms.Gemini25Flash())
+
+  # This might generate a class like:
+  # class TripPlan(pg.Object):
+  #   destinations: list[str]
+  #   start_date: str
+  #   end_date: str
+  #   total_budget: float
+
+  print(lf.Schema.from_value(trip_plan_cls).schema_str('python'))
+  ```
 
   Args:
     name: Class name to be generated.
@@ -108,17 +127,17 @@ def generate_class(
     lm: The language model to use. If not specified, the language model from
       `lf.context` context manager will be used.
     examples: An optional list of fewshot examples for helping class generation.
-      If None, a default single shot example will be used. Use
-      `lf.structured.classgen_example` to generate example.
+      If None, a default single-shot example will be used. Use
+      `lf.structured.classgen_example` to generate examples.
     returns_message: If True, returns `lf.Message` as the output, instead of
       returning the structured `message.result`.
     skip_lm: If True, returns the rendered prompt as a UserMessage object.
-      otherwise return the LLM response based on the rendered prompt.
+      otherwise returns the LLM response based on the rendered prompt.
     **kwargs: Template variables passed to `prompt` and keyword arguments passed
       to `lf.structured.GenerateClass`.
 
   Returns:
-    Generated class.
+    The generated Python class, or `lf.Message` if `returns_message` is True.
 
   Raises:
     CodeError: if generation failed.

langfun/core/structured/scoring.py

@@ -31,42 +31,54 @@ def score(
     *,
     lm: lf.LanguageModel | None = None,
     examples: list[mapping.MappingExample] | None = None,
-    protocol: schema_lib.SchemaProtocol = 'python',
+    protocol: str = 'python',
     return_scoring_results: bool = False,
     **kwargs,
 ) -> list[float] | list[lf.LMScoringResult]:
-  """Scores the outputs based on the prompt.
-
-  Examples:
-    ```
-    # Example 1: Scoring text output based on the user prompt.
-    scores = lf.score('{{x}} + {{y}} =', ['1', '2', '3'], lm=lm, x=1, y=2)
-    assert len(scores) == 3
-
-    # Example 2: Scoring int output based on the formulated OOP prompt.
-    scores = lf.score('1 + 1 =', [1, 2, 3], lm=lm)
-    assert len(scores) == 3
-
-    class Answer(pg.Object):
-      result: int
-
-    # Example 3: Scoring object output based on the formulated OOP prompt.
-    scores = lf.score('1 + 1 =', [Answer(1), Answer(2), Answer(3)], lm=lm)
-    assert len(scores) == 3
-
-    # Example 4: Scoring object field value based on the formulated OOP prompt
-    # and the generated tokens before the first `pg.oneof`.
-    scores = lf.score('1 + 1 =', [Answer(pg.oneof([1, 2, 3]))], lm=lm)
-    assert len(scores) == 3
-
-    # Example 5: Scoring multiple prompt/completion pairs.
-    scores = lf.score(
-        ['1 + 1=', '2 + 3='],
-        ['2', '4'],
-        lm=lm
-    )
-    assert len(scores) == 2
-    ```
+  """Scores completions based on a prompt using a language model.
+
+  `lf.score` computes the likelihood of each completion being generated given
+  a prompt, according to the specified language model. It can score text
+  completions or structured objects. If `schema` is provided, Langfun
+  formats the prompt and completions appropriately before scoring.
+
+  **Example 1: Score text completions**
+  ```python
+  import langfun as lf
+  scores = lf.score(
+      '1 + 1 =',
+      ['2', '3', '4'],
+      lm=lf.llms.Gemini25Flash())
+  print([f'{s:.3f}' for s in scores])
+  # Output: ['-0.001', '-2.345', '-3.456']
+  ```
+
+  **Example 2: Score structured completions**
+  ```python
+  import langfun as lf
+  import pyglove as pg
+
+  class Answer(pg.Object):
+    result: int
+
+  scores = lf.score(
+      '1 + 1 =',
+      [Answer(result=2), Answer(result=3), Answer(result=4)],
+      lm=lf.llms.Gemini25Flash())
+  print([f'{s:.3f}' for s in scores])
+  # Output: ['-0.001', '-2.345', '-3.456']
+  ```
+
+  **Example 3: Score multiple prompt/completion pairs**
+  ```python
+  import langfun as lf
+  scores = lf.score(
+      ['1 + 1 =', '2 + 2 ='],
+      ['2', '4'],
+      lm=lf.llms.Gemini25Flash())
+  print([f'{s:.3f}' for s in scores])
+  # Output: ['-0.001', '-0.002']
+  ```
 
   Args:
     prompt: The prompt(s) based on which each completion will be scored.
@@ -74,8 +86,7 @@ def score(
     schema: The schema as the output type. If None, it will be inferred from
       the completions.
     lm: The language model used for scoring.
-    examples: Fewshot exemplars used together with the prompt in getting the
-      completions.
+    examples: Few-shot examples used to construct the prompt for scoring.
     protocol: The protocol for formulating the prompt based on objects.
     return_scoring_results: If True, returns a list of `lf.LMScoringResult`,
       otherwise returns a list of floats as the scores of each completion.
@@ -190,7 +201,7 @@ async def ascore(
     *,
     lm: lf.LanguageModel | None = None,
     examples: list[mapping.MappingExample] | None = None,
-    protocol: schema_lib.SchemaProtocol = 'python',
+    protocol: str = 'python',
     return_scoring_results: bool = False,
     **kwargs,
 ) -> list[float] | list[lf.LMScoringResult]:

langfun/core/structured/tokenization.py

@@ -23,30 +23,45 @@ import pyglove as pg
 
 
 def tokenize(
-    prompt: Union[str, pg.Symbolic] | list[str | pg.Symbolic],
+    prompt: Union[str, pg.Symbolic, list[str | pg.Symbolic]],
     schema: Union[
         schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
     ] = None,
     *,
     lm: lf.LanguageModel | None = None,
     examples: list[mapping.MappingExample] | None = None,
-    protocol: schema_lib.SchemaProtocol = 'python',
+    protocol: str = 'python',
     **kwargs,
 ) -> list[tuple[str | bytes, int]]:
-  """Tokenize the prompt for `lf.query`.
+  """Renders a prompt and tokenizes it using a language model.
+
+  `lf.tokenize` first renders a prompt based on the provided `prompt`,
+  `schema`, and `examples`, similar to `lf.query`, and then uses the
+  specified language model (`lm`) to tokenize the resulting message.
+  This is useful for understanding how a prompt is seen by the model or
+  for estimating token counts before sending requests.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+  tokens = lf.tokenize('Hello world!', lm=lf.llms.Gpt4())
+  print(tokens)
+  # Output might look like: [('Hello', 15339), (' world', 1917), ('!', 0)]
+  ```
 
   Args:
-    prompt: The prompt(s) based on which each completion will be scored.
-    schema: The schema as the output type. If None, it will be inferred from
-      the completions.
-    lm: The language model used for scoring.
-    examples: Fewshot exemplars used together with the prompt in getting the
-      completions.
+    prompt: The prompt to render and tokenize. Can be a string, `pg.Symbolic`,
+      or `lf.Template`.
+    schema: The schema for formatting the prompt, if `prompt` is structured or
+      if schema-based formatting is needed.
+    lm: The language model to use for tokenization.
+    examples: Few-shot examples to include in the rendered prompt.
     protocol: The protocol for formulating the prompt based on objects.
     **kwargs: Keyword arguments that are referred by the prompt.
 
   Returns:
-    A list of (text, token_id) tuples.
+    A list of (token_str, token_id) tuples representing the tokenized prompt.
   """
   input_message = querying.query_prompt(
       prompt,
@@ -72,7 +87,7 @@ async def atokenize(
     *,
     lm: lf.LanguageModel | None = None,
     examples: list[mapping.MappingExample] | None = None,
-    protocol: schema_lib.SchemaProtocol = 'python',
+    protocol: str = 'python',
     **kwargs,
 ) -> list[tuple[str | bytes, int]]:
   """Async version of `lf.tokenize`."""

langfun/core/subscription.py

@@ -35,7 +35,7 @@ EventType = TypeVar('EventType')
 
 
 class EventHandler(Generic[EventType], metaclass=abc.ABCMeta):
-  """Interface for event subscriber."""
+  """Interface for event handler."""
 
   @classmethod
   @functools.cache
@@ -51,7 +51,7 @@ class EventHandler(Generic[EventType], metaclass=abc.ABCMeta):
 
   @classmethod
   def accepts(cls, event: Event[Any]) -> bool:
-    """Returns True if current event handler class can accepts an event."""
+    """Returns True if current event handler class can accept an event."""
     return isinstance(event, cls.event_type())
 
   @abc.abstractmethod