Added feature to utilize Python syntax for field comments in schema instead of verbose Annotated[Doc()]-based comment style. Currently implemented for examples/math

Author: rohankalbag

python/examples/math/demo.py

@@ -4,7 +4,7 @@
 import sys
 from typing import cast
 from dotenv import dotenv_values
-import schema as math
+import schema_with_comments as math
 from typechat import Failure, create_language_model, process_requests
 from program import TypeChatProgramTranslator, TypeChatProgramValidator, evaluate_json_program
 

python/examples/math/program.py

@@ -121,6 +121,7 @@ class TypeChatProgramTranslator(TypeChatJsonTranslator[JsonProgram]):
     _api_declaration_str: str
 
     def __init__(self, model: TypeChatLanguageModel, validator: TypeChatProgramValidator, api_type: type):
+        api_type = self._convert_pythonic_comments_to_annotated_docs(api_type)
         super().__init__(model=model, validator=validator, target_type=api_type, _raise_on_schema_errors = False)
         # TODO: the conversion result here has errors!
         conversion_result = python_type_to_typescript_schema(api_type)

python/examples/math/pythonic_comment_handling.py

@@ -0,0 +1,46 @@
+import re
+import inspect
+from schema_with_comments import MathAPI
+
+
+def _convert_pythonic_comments_to_annotated_docs(schema_class, debug=True):
+
+    schema_path = inspect.getfile(schema_class)
+
+    with open(schema_path, 'r') as file:
+        schema_class_source = file.read()
+
+    if debug:
+        print("File contents before modification:")
+        print("--"*50)
+        print(schema_class_source)
+        print("--"*50)
+
+    pattern = r"(\w+\s*:\s*.*?)(?=\s*#\s*(.+?)(?:\n|\Z))"
+    commented_fields = re.findall(pattern, schema_class_source)
+    annotated_fields = []
+
+    for field, comment in commented_fields:
+        field_separator = field.split(":")
+        field_name = field_separator[0].strip()
+        field_type = field_separator[1].strip()
+
+        annotated_fields.append(
+            f"{field_name}: Annotated[{field_type}, Doc(\"{comment}\")]")
+
+    for field, annotation in zip(commented_fields, annotated_fields):
+        schema_class_source = schema_class_source.replace(field[0], annotation)
+
+    if debug:
+        print("File contents after modification:")
+        print("--"*50)
+        print(schema_class_source)
+        print("--"*50)
+
+    namespace = {}
+    exec(schema_class_source, namespace)
+    return namespace[schema_class.__name__]
+
+
+if __name__ == "__main__":
+    print(_convert_pythonic_comments_to_annotated_docs(MathAPI))

python/examples/math/schema.py

@@ -1,6 +1,5 @@
 from typing_extensions import TypedDict, Annotated, Callable, Doc
 
-
 class MathAPI(TypedDict):
     """
     This is API for a simple calculator
@@ -12,4 +11,4 @@ class MathAPI(TypedDict):
     div: Annotated[Callable[[float, float], float], Doc("Divide two numbers")]
     neg: Annotated[Callable[[float], float], Doc("Negate a number")]
     id: Annotated[Callable[[float], float], Doc("Identity function")]
-    unknown: Annotated[Callable[[str], float], Doc("Unknown request")]
+    unknown: Annotated[Callable[[str], float], Doc("Unknown request")]

python/examples/math/schema_with_comments.py

@@ -0,0 +1,16 @@
+from typing_extensions import TypedDict, Annotated, Callable, Doc
+
+class MathAPI(TypedDict):
+    """
+    This is API for a simple calculator
+    """
+
+    # this is a comment
+
+    add: Callable[[float, float], float] # Add two numbers
+    sub: Callable[[float, float], float] # Subtract two numbers
+    mul: Callable[[float, float], float] # Multiply two numbers
+    div: Callable[[float, float], float] # Divide two numbers
+    neg: Callable[[float], float] # Negate a number
+    id: Callable[[float], float] # Identity function
+    unknown: Callable[[str], float] # Unknown request

python/src/typechat/_internal/translator.py

@@ -1,6 +1,8 @@
 from typing_extensions import Generic, TypeVar
 
 import pydantic_core
+import re
+import inspect
 
 from typechat._internal.model import PromptSection, TypeChatLanguageModel
 from typechat._internal.result import Failure, Result, Success
@@ -123,3 +125,41 @@ def _create_repair_prompt(self, validation_error: str) -> str:
 The following is a revised JSON object:
 """
         return prompt
+    
+    def _convert_pythonic_comments_to_annotated_docs(schema_class, debug=False):
+
+        schema_path = inspect.getfile(schema_class)
+
+        with open(schema_path, 'r') as file:
+            schema_class_source = file.read()
+
+        if debug:
+            print("File contents before modification:")
+            print("--"*50)
+            print(schema_class_source)
+            print("--"*50)
+
+        pattern = r"(\w+\s*:\s*.*?)(?=\s*#\s*(.+?)(?:\n|\Z))"
+        commented_fields = re.findall(pattern, schema_class_source)
+        annotated_fields = []
+
+        for field, comment in commented_fields:
+            field_separator = field.split(":")
+            field_name = field_separator[0].strip()
+            field_type = field_separator[1].strip()
+
+            annotated_fields.append(
+                f"{field_name}: Annotated[{field_type}, Doc(\"{comment}\")]")
+
+        for field, annotation in zip(commented_fields, annotated_fields):
+            schema_class_source = schema_class_source.replace(field[0], annotation)
+
+        if debug:
+            print("File contents after modification:")
+            print("--"*50)
+            print(schema_class_source)
+            print("--"*50)
+
+        namespace = {}
+        exec(schema_class_source, namespace)
+        return namespace[schema_class.__name__]