fix tests

Author: Marta

pyproject.toml

@@ -5,7 +5,7 @@ description = "Logprobs for OpenAI Structured Outputs"
 authors = [{ name = "Sarus Technologies", email = "nicolas.grislain@gmail.com" }]
 readme = "README.md"
 keywords = ['python']
-requires-python = ">=3.9,<4.0"
+requires-python = ">=3.10,<4.0"
 classifiers = [
     "Intended Audience :: Developers",
     "Programming Language :: Python",
@@ -18,9 +18,9 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
-    "openai>=1.58.1",
-    "pydantic>=2.10.4",
-    "lark>=1.2.2",
+    "openai~=1.58.1",
+    "pydantic~=2.10.4",
+    "lark~=1.2.2",
 ]
 
 [project.urls]

structured_logprobs/helpers.py

@@ -0,0 +1,176 @@
+from typing import Any
+
+from lark import Lark, Token, Transformer_NonRecursive, Tree, v_args
+from lark.tree import Meta
+from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob
+from pydantic import BaseModel
+
+
+class HasProb(BaseModel):
+    value: Any
+    start: int
+    end: int
+    logprob: float
+
+
+# Define a grammar for JSON
+json_grammar = r"""
+    start: value
+
+    ?value: object              #'?' is a Lark convention indicating that the rule can return the value directly instead of creating a separate parse tree node.
+          | array
+          | string
+          | SIGNED_NUMBER -> number    #'-> number' specifies an alias for the rule
+          | "true"
+          | "false"
+          | "null"
+
+    array  : "[" [value ("," value)*] "]"
+    object : "{" [pair ("," pair)*] "}"
+    pair   : key ":" value
+    key    : ESCAPED_STRING
+
+    string : ESCAPED_STRING
+
+    %import common.ESCAPED_STRING
+    %import common.SIGNED_NUMBER
+    %import common.WS
+    %ignore WS
+"""
+
+
+# Transformer that processes the tree and substitutes each atomic value with the cumulative log-probability of its tokens
+@v_args(meta=True)
+class Extractor(Transformer_NonRecursive):
+    def __init__(self, tokens: list[ChatCompletionTokenLogprob], token_indices: list[int]):
+        super().__init__()
+        self.tokens = tokens
+        self.token_indices = token_indices
+
+    def _compute_logprob_sum(self, start: int, end: int) -> float:
+        token_start = self.token_indices[start]
+        token_end = self.token_indices[end]
+        sum_logporb = sum(self.tokens[i].logprob for i in range(token_start, token_end))
+        return sum_logporb
+
+    def number(self, meta: Meta, children: list[Token]) -> float:
+        logprob_sum = self._compute_logprob_sum(meta.start_pos, meta.end_pos)
+        return logprob_sum
+
+    def string(self, meta: Meta, children: list[Token]) -> float:
+        logprob_sum = self._compute_logprob_sum(meta.start_pos, meta.end_pos)
+        return logprob_sum
+
+    def true(self, meta: Meta, children: list[Token]) -> float:
+        logprob_sum = self._compute_logprob_sum(meta.start_pos, meta.end_pos)
+        return logprob_sum
+
+    def false(self, meta: Meta, children: list[Token]) -> float:
+        logprob_sum = self._compute_logprob_sum(meta.start_pos, meta.end_pos)
+        return logprob_sum
+
+    def null(self, meta: Meta, children: list[Token]) -> None:
+        return None
+
+    def array(self, meta: Meta, children: list[Any]) -> list[float]:
+        return children
+
+    def object(self, meta: Meta, children: list[tuple[str, Any]]) -> dict[str, Any]:
+        result = {}
+        for key, value in children:
+            result[key] = value
+        return result
+
+    def pair(self, meta: Meta, children: list[Any]) -> tuple[str, Any]:
+        value = children[1]
+        key = children[0]
+        if isinstance(value, Tree) and not value.children:  # ['b', Tree(Token('RULE', 'value'), [])]
+            value = None
+        return key, value
+
+    def key(self, meta: Meta, children: list[Token]) -> str:
+        return children[0][1:-1]
+
+    def start(self, meta: Meta, children: list[dict[str, Any]]) -> dict[str, Any]:
+        return children[0]
+
+
+json_parser = Lark(json_grammar, parser="lalr", propagate_positions=True, maybe_placeholders=False)
+
+
+def extract_json_data(
+    json_string: str, tokens: list[ChatCompletionTokenLogprob], token_indices: list[int]
+) -> dict[str, Any]:
+    json_parser = Lark(json_grammar, parser="lalr", propagate_positions=True, maybe_placeholders=False)
+    tree = json_parser.parse(json_string)
+    extractor = Extractor(tokens, token_indices)
+    return extractor.transform(tree)
+
+
+# Transformer that embeds log-probabilities for atomic values as in-line fields in dictionaries
+@v_args(meta=True)
+class ExtractorInline(Transformer_NonRecursive):
+    def __init__(self, tokens: list[ChatCompletionTokenLogprob], token_indices: list[int]):
+        super().__init__()
+        self.tokens = tokens
+        self.token_indices = token_indices
+
+    def _compute_logprob_sum(self, start: int, end: int) -> float:
+        token_start = self.token_indices[start]
+        token_end = self.token_indices[end]
+        sum_logporb = sum(self.tokens[i].logprob for i in range(token_start, token_end))
+        return sum_logporb
+
+    def number(self, meta: Meta, children: list[Token]) -> HasProb:
+        logprob_sum = self._compute_logprob_sum(meta.start_pos, meta.end_pos)
+        return HasProb(value=float(children[0]), start=meta.start_pos, end=meta.end_pos, logprob=logprob_sum)
+
+    def string(self, meta: Meta, children: list[Token]) -> HasProb:
+        logprob_sum = self._compute_logprob_sum(meta.start_pos, meta.end_pos)
+        return HasProb(value=children[0][1:-1], start=meta.start_pos, end=meta.end_pos, logprob=logprob_sum)
+
+    def true(self, meta: Meta, children: list[Token]) -> HasProb:
+        logprob_sum = self._compute_logprob_sum(meta.start_pos, meta.end_pos)
+        return HasProb(value=True, start=meta.start_pos, end=meta.end_pos, logprob=logprob_sum)
+
+    def false(self, meta: Meta, children: list[Token]) -> HasProb:
+        logprob_sum = self._compute_logprob_sum(meta.start_pos, meta.end_pos)
+        return HasProb(value=False, start=meta.start_pos, end=meta.end_pos, logprob=logprob_sum)
+
+    def null(self, meta: Meta, children: list[Token]) -> None:
+        return None
+
+    def array(self, meta: Meta, children: list[dict[str, Any] | Any]) -> list[dict[str, Any] | Any]:
+        return [child.value if isinstance(child, HasProb) else child for child in children]
+
+    def object(self, meta: Meta, children: list[tuple[str, Any]]) -> dict[str, Any]:
+        result = {}
+        for key, value in children:
+            if isinstance(value, HasProb):
+                result[key] = value.value
+                result[f"{key}_logprob"] = value.logprob
+            else:
+                result[key] = value
+        return result
+
+    def pair(self, meta: Meta, children: list[str | Any]) -> tuple[str, Any]:
+        value = children[1]
+        key = children[0]
+        if isinstance(value, Tree) and not value.children:  # ['b', Tree(Token('RULE', 'value'), [])]
+            value = None
+        return key, value
+
+    def key(self, meta: Meta, children: list[Token]) -> str:
+        return children[0][1:-1]
+
+    def start(self, meta: Meta, children: list[dict[str, Any]]) -> dict[str, Any]:
+        return children[0]
+
+
+def extract_json_data_inline(
+    json_string: str, tokens: list[ChatCompletionTokenLogprob], token_indices: list[int]
+) -> dict[str, Any]:
+    json_parser = Lark(json_grammar, parser="lalr", propagate_positions=True, maybe_placeholders=False)
+    tree = json_parser.parse(json_string)
+    extractor = ExtractorInline(tokens, token_indices)
+    return extractor.transform(tree)

structured_logprobs/main.py

@@ -5,7 +5,7 @@
 from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob
 from pydantic import BaseModel
 
-from helpers import extract_json_data, extract_json_data_inline
+from structured_logprobs.helpers import extract_json_data, extract_json_data_inline
 
 MISSING_LOGPROBS_MESSAGE = "The 'logprobs' field is missing"
 
@@ -48,16 +48,11 @@ def map_characters_to_token_indices(extracted_data_token: list[ChatCompletionTok
         [0, 1, 1, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 4]
     """
 
-    json_output = "".join(token_data.token for token_data in extracted_data_token)
-
-    token_indices = [-1] * len(json_output)
-    current_char_pos = 0
+    token_indices = []
 
     for token_idx, token_data in enumerate(extracted_data_token):
         token_text = token_data.token
-        for _ in range(len(token_text)):
-            token_indices[current_char_pos] = token_idx
-            current_char_pos += 1
+        token_indices.extend([token_idx] * len(token_text))
 
     return token_indices
 

tests/conftest.py

@@ -21,7 +21,7 @@ class CalendarEvent(BaseModel):
 def chat_completion(pytestconfig) -> ChatCompletion:
     client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
     base_path = Path(pytestconfig.rootdir)  # Base directory where pytest was run
-    schema_path = base_path / "tests" / "questions_json_schema.json"
+    schema_path = base_path / "tests" / "resources" / "questions_json_schema.json"
     with open(schema_path) as f:
         schema_content = json.load(f)
 
@@ -74,7 +74,7 @@ class CalendarEvent(BaseModel):
 @pytest.fixture
 def simple_parsed_completion(pytestconfig) -> ParsedChatCompletion[CalendarEvent] | None:
     base_path = Path(pytestconfig.rootdir)  # Base directory where pytest was run
-    with open(base_path / "tests" / "simple_parsed_completion.json") as f:
+    with open(base_path / "tests" / "resources" / "simple_parsed_completion.json") as f:
         return ParsedChatCompletion[CalendarEvent].model_validate_json(f.read())
     return None
 

tests/resources/questions_json_schema.json

@@ -0,0 +1,28 @@
+{
+    "type": "json_schema",
+    "json_schema": {
+        "name": "answears",
+        "description": "Response to questions in JSON format",
+        "schema": {
+            "type": "object",
+            "properties": {
+                "capital_of_France": { "type": "string" },
+                "the_two_nicest_colors": {
+                    "type": "array",
+                    "items": {
+                        "type": "string",
+                        "enum": ["red", "blue", "green", "yellow", "purple"]
+                    }
+                },
+                "die_shows": { "type": "integer" }
+            },
+            "required": [
+                "capital_of_France",
+                "the_two_nicest_colors",
+                "die_shows"
+            ],
+            "additionalProperties": false
+        },
+        "strict": true
+    }
+}

tests/resources/simple_json_schema.json

@@ -0,0 +1,21 @@
+{
+    "type": "json_schema",
+    "json_schema": {
+        "name": "event_extraction",
+        "description": "Extract details about an event, including participants, event name, and date.",
+        "schema": {
+            "type": "object",
+            "properties": {
+                "name": { "type": "string" },
+                "date": { "type": "string" },
+                "participants": {
+                    "type": "array",
+                    "items": { "type": "string" }
+                }
+            },
+            "required": ["name", "date", "participants"],
+            "additionalProperties": false
+        },
+        "strict": true
+    }
+}