From f092bd6257bca176990a63243d4070660949bff4 Mon Sep 17 00:00:00 2001 From: Joe Nudell Date: Thu, 16 Apr 2026 11:48:34 -0400 Subject: [PATCH 1/4] better compatibility with responses api --- bc2/core/common/openai.py | 84 ++++++++++++++------------------------- 1 file changed, 30 insertions(+), 54 deletions(-) diff --git a/bc2/core/common/openai.py b/bc2/core/common/openai.py index eb5a66b..7d48f30 100644 --- a/bc2/core/common/openai.py +++ b/bc2/core/common/openai.py @@ -6,28 +6,15 @@ from abc import abstractmethod from dataclasses import dataclass from functools import cached_property -from typing import Any, Generic, Literal, Sequence, Type, TypeAlias, TypeVar, cast +from typing import Any, Generic, Literal, Sequence, Type, TypeVar, cast from openai import AsyncOpenAI, OpenAI -from openai.types.chat import ( - ChatCompletionAssistantMessageParam as _OpenAIChatCompletionAssistantMessageParam, +from openai.types.responses import ( + ResponseInputImage as _OpenAIResponseInputImage, ) -from openai.types.chat import ( - ChatCompletionContentPartImageParam as _OpenAIChatImageMessagePart, +from openai.types.responses import ( + ResponseInputText as _OpenAIResponseInputText, ) -from openai.types.chat import ( - ChatCompletionContentPartTextParam as _OpenAIChatTextMessagePart, -) -from openai.types.chat import ( - ChatCompletionMessageParam as _OpenAIChatCompletionMessageParam, -) -from openai.types.chat import ( - ChatCompletionSystemMessageParam as _OpenAIChatCompletionSystemMessageParam, -) -from openai.types.chat import ( - ChatCompletionUserMessageParam as _OpenAIChatCompletionUserMessageParam, -) -from openai.types.chat.chat_completion_content_part_image_param import ImageURL from pydantic import BaseModel, Field, PositiveInt, SerializationInfo, model_serializer from .datafile import DataType, load_data_file, load_data_file_from_path @@ -40,10 +27,6 @@ TResult = TypeVar("TResult") -_OpenAIChatMessagePart: TypeAlias = ( - _OpenAIChatTextMessagePart | _OpenAIChatImageMessagePart -) - class FilteredContentError(Exception): """Error we throw when OpenAI content moderation is triggered.""" @@ -114,9 +97,9 @@ class OpenAIChatInputText(BaseModel): type: Literal["text"] = "text" text: str - def as_chat_message_part(self) -> _OpenAIChatTextMessagePart: + def as_chat_message_part(self) -> _OpenAIResponseInputText: """Convert the input to a chat message.""" - return _OpenAIChatTextMessagePart(type=self.type, text=self.text) + return _OpenAIResponseInputText(type="input_text", text=self.text) class OpenAIUrl(BaseModel): @@ -131,14 +114,12 @@ class OpenAIChatInputImageUrl(BaseModel): type: Literal["image_url"] = "image_url" image_url: OpenAIUrl - def as_chat_message_part(self) -> _OpenAIChatImageMessagePart: + def as_chat_message_part(self) -> _OpenAIResponseInputImage: """Convert the input to a chat message.""" - return _OpenAIChatImageMessagePart( - type=self.type, - image_url=ImageURL( - url=self.image_url.url, - detail="high", - ), + return _OpenAIResponseInputImage( + type="input_image", + detail="high", + image_url=self.image_url.url, ) @@ -148,36 +129,24 @@ def as_chat_message_part(self) -> _OpenAIChatImageMessagePart: class OpenAIChatTurn(BaseModel): """A chat turn for an OpenAI model.""" - role: Literal["assistant", "user", "system"] + role: Literal["user", "system"] content: str | list[OpenAIChatInput] - def as_chat_message(self) -> _OpenAIChatCompletionMessageParam: + def as_chat_message(self) -> dict[str, Any]: """Convert the turn to a chat message.""" - match self.role: - case "assistant": - return _OpenAIChatCompletionAssistantMessageParam( - role=self.role, - content=self._format_content_no_images(), - ) - case "user": - return _OpenAIChatCompletionUserMessageParam( - role=self.role, - content=self._format_content(), - ) - case "system": - return _OpenAIChatCompletionSystemMessageParam( - role=self.role, - content=self._format_content_no_images(), - ) + return { + "role": self.role, + "content": self._format_content(), + } - def _format_content(self) -> str | list[_OpenAIChatMessagePart]: + def _format_content(self) -> str | list[OpenAIChatInput]: if isinstance(self.content, str): return self.content return [ c if isinstance(c, str) else c.as_chat_message_part() for c in self.content ] - def _format_content_no_images(self) -> str | list[_OpenAIChatTextMessagePart]: + def _format_content_no_images(self) -> str | list[OpenAIChatInput]: if isinstance(self.content, str): return self.content return [ @@ -513,14 +482,21 @@ def invoke( # Configure max tokens and submit the query. max_tokens = self.token_cap - response = client.responses.parse( + + call_params = dict( **openai_api_settings, max_output_tokens=max_tokens, input=messages, - text_format=response_format, store=False, ) + # Call the API using `parse` or `create` depending on structured output. + if response_format: + call_params["text_format"] = response_format + response = client.responses.parse(**call_params) + else: + response = client.responses.create(**call_params) + # Interpret response stop_reason = response.incomplete_details and response.incomplete_details.reason if stop_reason == "max_output_tokens": @@ -548,7 +524,7 @@ def invoke( max_tokens=max_tokens, content=response.output_text or "", completion_tokens=completion_tokens, - parsed=response.output_parsed, + parsed=getattr(response, "output_parsed", None), ) From ed7208616256b84dc3d518e4db24ab6f1223378f Mon Sep 17 00:00:00 2001 From: Joe Nudell Date: Thu, 16 Apr 2026 14:25:34 -0400 Subject: [PATCH 2/4] fix typing --- bc2/core/common/openai.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/bc2/core/common/openai.py b/bc2/core/common/openai.py index 7d48f30..9b2a93e 100644 --- a/bc2/core/common/openai.py +++ b/bc2/core/common/openai.py @@ -9,6 +9,9 @@ from typing import Any, Generic, Literal, Sequence, Type, TypeVar, cast from openai import AsyncOpenAI, OpenAI +from openai.types.responses import ( + EasyInputMessageParam as _OpenAIEasyInputMessageParam, +) from openai.types.responses import ( ResponseInputImage as _OpenAIResponseInputImage, ) @@ -132,12 +135,12 @@ class OpenAIChatTurn(BaseModel): role: Literal["user", "system"] content: str | list[OpenAIChatInput] - def as_chat_message(self) -> dict[str, Any]: + def as_chat_message(self) -> _OpenAIEasyInputMessageParam: """Convert the turn to a chat message.""" - return { - "role": self.role, - "content": self._format_content(), - } + return _OpenAIEasyInputMessageParam( + role=self.role, + content=self._format_content(), + ) def _format_content(self) -> str | list[OpenAIChatInput]: if isinstance(self.content, str): From 9853eb3d421c1dcb89b95cc39010ff818bc1fe0a Mon Sep 17 00:00:00 2001 From: Joe Nudell Date: Thu, 16 Apr 2026 14:42:52 -0400 Subject: [PATCH 3/4] fix tests --- bc2/core/inspect/test_masked_subjects.py | 97 ++++++++++++------------ bc2/core/redact/test_openai.py | 32 ++++---- 2 files changed, 68 insertions(+), 61 deletions(-) diff --git a/bc2/core/inspect/test_masked_subjects.py b/bc2/core/inspect/test_masked_subjects.py index 0b5a042..01a5f08 100644 --- a/bc2/core/inspect/test_masked_subjects.py +++ b/bc2/core/inspect/test_masked_subjects.py @@ -1,6 +1,8 @@ import json from unittest.mock import patch +from openai.types.responses import ResponseInputText + from ..common.context import Context from ..common.datafile import DataType, load_data_file from ..common.name_map import IdToMaskMap, IdToNameMap @@ -34,20 +36,20 @@ def test_inspect_subject_masks(openai_mock): ) ctx = Context() - openai_mock.return_value.responses.parse.return_value.output_text = json.dumps( + openai_mock.return_value.responses.create.return_value.output_text = json.dumps( { "A": "Subject 1", "B": "Subject 2", "C": "Subject 3", } ) - openai_mock.return_value.responses.parse.return_value.status = "completed" - openai_mock.return_value.responses.parse.return_value.usage = type( + openai_mock.return_value.responses.create.return_value.status = "completed" + openai_mock.return_value.responses.create.return_value.usage = type( "Usage", (), {"output_tokens": 10} )() - openai_mock.return_value.responses.parse.return_value.incomplete_details = None - openai_mock.return_value.responses.parse.return_value.output_parsed = None - openai_mock.return_value.responses.parse.return_value.error = None + openai_mock.return_value.responses.create.return_value.incomplete_details = None + openai_mock.return_value.responses.create.return_value.output_parsed = None + openai_mock.return_value.responses.create.return_value.error = None result = cfg.driver( rt, @@ -68,10 +70,9 @@ def test_inspect_subject_masks(openai_mock): "C": "Subject 3", } ) - openai_mock.return_value.responses.parse.assert_called_once_with( + openai_mock.return_value.responses.create.assert_called_once_with( model="gpt-4o", max_output_tokens=None, - text_format=None, store=False, input=[ { @@ -81,45 +82,47 @@ def test_inspect_subject_masks(openai_mock): { "role": "user", "content": [ - { - "type": "text", - "text": ( - "[COLLECTION#1]\n" - "" - "" - "ALeopold" - "" - "" - "BPollock" - "" - "" - "CAbbott" - "" - "\n\n" - "[COLLECTION#2]\n" - "" - "" - "Leopold" - "Subject 1" - "" - "" - "Pollock" - "Subject 2" - "" - "" - "Abbott" - "Subject 3" - "" - "" - "Poldy" - "Subject 1" - "" - "\n\n" - "[NARRATIVE]\n" - "Leopold is first, then Pollock, then Abbott, " - "then Poldy again." - ), - } + ResponseInputText.model_validate( + { + "type": "input_text", + "text": ( + "[COLLECTION#1]\n" + "" + "" + "ALeopold" + "" + "" + "BPollock" + "" + "" + "CAbbott" + "" + "\n\n" + "[COLLECTION#2]\n" + "" + "" + "Leopold" + "Subject 1" + "" + "" + "Pollock" + "Subject 2" + "" + "" + "Abbott" + "Subject 3" + "" + "" + "Poldy" + "Subject 1" + "" + "\n\n" + "[NARRATIVE]\n" + "Leopold is first, then Pollock, then Abbott, " + "then Poldy again." + ), + } + ) ], }, ], diff --git a/bc2/core/redact/test_openai.py b/bc2/core/redact/test_openai.py index 4d8f763..b182d68 100644 --- a/bc2/core/redact/test_openai.py +++ b/bc2/core/redact/test_openai.py @@ -1,6 +1,7 @@ from unittest.mock import MagicMock, patch import pytest +from openai.types.responses import ResponseInputText from ..common.context import Context from ..common.name_map import NameToMaskMap @@ -48,7 +49,7 @@ def mock_create(*args, **kwargs): response.error = None return response - openai_mock.return_value.responses.parse.side_effect = mock_create + openai_mock.return_value.responses.create.side_effect = mock_create cfg = OpenAIRedactConfig.model_validate( { @@ -85,10 +86,9 @@ def mock_create(*args, **kwargs): "Leopold, Pollock, and Abbott went to the store.", ("[", "]"), ) - openai_mock.return_value.responses.parse.assert_called_once_with( + openai_mock.return_value.responses.create.assert_called_once_with( model="gpt-4o-2024-05-13", max_output_tokens=4_096, - text_format=None, store=False, input=[ { @@ -102,10 +102,12 @@ def mock_create(*args, **kwargs): { "role": "user", "content": [ - { - "type": "text", - "text": "Leopold, Pollock, and Abbott went to the store.", - } + ResponseInputText.model_validate( + { + "type": "input_text", + "text": "Leopold, Pollock, and Abbott went to the store.", + } + ) ], }, ], @@ -137,7 +139,7 @@ def mock_create(*args, **kwargs): response.error = None return response - openai_mock.return_value.responses.parse.side_effect = mock_create + openai_mock.return_value.responses.create.side_effect = mock_create cfg = OpenAIRedactConfig.model_validate( { @@ -176,10 +178,9 @@ def mock_create(*args, **kwargs): "Leopold, Pollock, and Abbott went to the store.", ("{", "}"), ) - openai_mock.return_value.responses.parse.assert_called_once_with( + openai_mock.return_value.responses.create.assert_called_once_with( model="gpt-4o-2024-05-13", max_output_tokens=4_096, - text_format=None, store=False, input=[ { @@ -199,10 +200,12 @@ def mock_create(*args, **kwargs): { "role": "user", "content": [ - { - "type": "text", - "text": "Leopold, Pollock, and Abbott went to the store.", - } + ResponseInputText.model_validate( + { + "type": "input_text", + "text": "Leopold, Pollock, and Abbott went to the store.", + } + ) ], }, ], @@ -239,3 +242,4 @@ def test_redact_raises_on_empty_narrative(openai_mock, narrative): ) openai_mock.return_value.responses.parse.assert_not_called() + openai_mock.return_value.responses.create.assert_not_called() From 05fd17113841e9568a76ceb98edde5b462987e9b Mon Sep 17 00:00:00 2001 From: Joe Nudell Date: Thu, 16 Apr 2026 14:48:37 -0400 Subject: [PATCH 4/4] add tests for ontology extraction --- bc2/core/ontology/test_openai.py | 279 +++++++++++++++++++++++++++++++ 1 file changed, 279 insertions(+) create mode 100644 bc2/core/ontology/test_openai.py diff --git a/bc2/core/ontology/test_openai.py b/bc2/core/ontology/test_openai.py new file mode 100644 index 0000000..763773d --- /dev/null +++ b/bc2/core/ontology/test_openai.py @@ -0,0 +1,279 @@ +import json +from unittest.mock import MagicMock, patch + +import pytest +from azure.ai.documentintelligence.models import AnalyzeResult +from openai import OpenAI +from openai.types.responses import ResponseInputText + +from ..common.context import Context +from ..common.file import MemoryFile +from ..common.ontology import ( + Cited, + Offense, + PoliceReport, + Subject, +) +from .base import EmptyOntologyError +from .openai import OpenAIOntologyConfig + + +def _make_report() -> PoliceReport: + return PoliceReport( + reporting_agency=Cited(ids=[0], content="SFPD"), + case_number=Cited(ids=[0], content="2024-0001"), + location=Cited(ids=[1], content="101 Main St"), + incident_type=Cited(ids=[1], content="Theft"), + subjects=[ + Subject( + seq=Cited(ids=[2], content=1), + type=Cited(ids=[2], content="Victim"), + name=Cited(ids=[2], content="Leopold"), + address=Cited(ids=[2], content="101 Main St"), + phone=Cited(ids=[2], content="555-1234"), + race=Cited(ids=[2], content="Unknown"), + sex=Cited(ids=[2], content="M"), + dob=Cited(ids=[2], content="1990-01-01"), + ) + ], + narratives=[Cited(ids=[3], content="Incident narrative here.")], + offenses=[ + Offense( + crime=Cited(ids=[1], content="Theft"), + statute=None, + code=None, + ) + ], + ) + + +def _analyze_result_dict() -> dict: + return { + "paragraphs": [ + { + "content": "Reporting Agency: SFPD", + "spans": [{"offset": 0, "length": 22}], + "boundingRegions": [ + { + "pageNumber": 1, + "polygon": [0, 0, 4.25, 0, 4.25, 0.5, 0, 0.5], + } + ], + }, + { + "content": "Incident: Theft at 101 Main St", + "spans": [{"offset": 24, "length": 30}], + "boundingRegions": [ + { + "pageNumber": 1, + "polygon": [0, 1.0, 8.5, 1.0, 8.5, 1.5, 0, 1.5], + } + ], + }, + { + "content": "Narrative paragraph on page 2.", + "spans": [{"offset": 56, "length": 30}], + "boundingRegions": [ + { + "pageNumber": 2, + "polygon": [0, 0, 8.5, 0, 8.5, 11.0, 0, 11.0], + } + ], + }, + ], + "pages": [ + {"pageNumber": 1, "width": 8.5, "height": 11.0}, + {"pageNumber": 2, "width": 8.5, "height": 11.0}, + ], + } + + +def _make_config() -> OpenAIOntologyConfig: + return OpenAIOntologyConfig.model_validate( + { + "engine": "ontology:openai", + "client": { + "api_key": "abc123", + "base_url": "http://openai.local", + }, + "generator": { + "method": "chat", + "model": "gpt-4o-2024-05-13", + "system": { + "engine": "string", + "prompt": "Extract the ontology.", + }, + }, + }, + ) + + +def _mock_parse_response(parsed: PoliceReport | None) -> MagicMock: + response = MagicMock() + response.output_text = parsed.model_dump_json() if parsed is not None else "" + response.output_parsed = parsed + response.status = "completed" + response.usage = type("Usage", (), {"output_tokens": 10})() + response.incomplete_details = None + response.error = None + return response + + +def _install_openai_mock( + openai_mock: MagicMock, parsed: PoliceReport | None +) -> MagicMock: + """Install a spec'd OpenAI client as the return value of the patched class. + + Using `spec=OpenAI` ensures that `hasattr(client, 'mime_type')` is False so + the preprocessor mixin does not confuse the client with a preprocessor + method while iterating driver attributes. + """ + client = MagicMock(spec=OpenAI) + client.responses.parse.return_value = _mock_parse_response(parsed) + openai_mock.return_value = client + return client + + +@patch("bc2.core.common.openai.OpenAI") +def test_extract_formats_xml_and_builds_chunks(openai_mock): + report = _make_report() + client = _install_openai_mock(openai_mock, report) + + cfg = _make_config() + + analyze_result = AnalyzeResult(_analyze_result_dict()) + result = cfg.driver.extract(analyze_result) + + assert result.report == report + assert len(result.chunks) == 3 + assert result.chunks[0].content == "Reporting Agency: SFPD" + assert result.chunks[0].spans[0].offset == 0 + assert result.chunks[0].spans[0].length == 22 + assert result.chunks[0].regions[0].page == 0 + # Polygon normalized by page width (8.5) and height (11.0). + assert result.chunks[0].regions[0].points == [ + (0.0, 0.0), + (0.5, 0.0), + (0.5, 0.5 / 11.0), + (0.0, 0.5 / 11.0), + ] + # Second paragraph spans the full page width. + assert result.chunks[1].regions[0].points == [ + (0.0, 1.0 / 11.0), + (1.0, 1.0 / 11.0), + (1.0, 1.5 / 11.0), + (0.0, 1.5 / 11.0), + ] + # Third paragraph is on page index 1. + assert result.chunks[2].regions[0].page == 1 + + client.responses.parse.assert_called_once() + call_kwargs = client.responses.parse.call_args.kwargs + assert call_kwargs["model"] == "gpt-4o-2024-05-13" + assert call_kwargs["store"] is False + assert call_kwargs["text_format"] is PoliceReport + assert call_kwargs["input"] == [ + {"role": "system", "content": "Extract the ontology."}, + { + "role": "user", + "content": [ + ResponseInputText.model_validate( + { + "type": "input_text", + "text": ( + "XML DOCUMENT\n===\n" + "" + '

Reporting Agency: SFPD

' + '

Incident: Theft at 101 Main St

' + '

Narrative paragraph on page 2.

' + "" + ), + } + ) + ], + }, + ] + + +@patch("bc2.core.common.openai.OpenAI") +def test_driver_call_persists_ontology_in_context(openai_mock): + report = _make_report() + _install_openai_mock(openai_mock, report) + + cfg = _make_config() + + file = MemoryFile( + content=json.dumps(_analyze_result_dict()).encode("utf-8"), + mime_type="application/x-analyze-result", + ) + context = Context() + + output = cfg.driver(file, context) + + assert output.mime_type == "application/x-ontology" + assert context.ontology is not None + assert context.ontology.report == report + assert len(context.ontology.chunks) == 3 + + serialized = json.loads(output.content().decode("utf-8")) + assert serialized["report"]["case_number"]["content"] == "2024-0001" + assert len(serialized["chunks"]) == 3 + + +@patch("bc2.core.common.openai.OpenAI") +def test_extract_raises_when_parsed_is_none(openai_mock): + _install_openai_mock(openai_mock, None) + + cfg = _make_config() + + analyze_result = AnalyzeResult(_analyze_result_dict()) + + with pytest.raises(ValueError, match="no structured output"): + cfg.driver.extract(analyze_result) + + +@patch("bc2.core.common.openai.OpenAI") +def test_driver_raises_empty_ontology_when_no_paragraphs(openai_mock): + report = _make_report() + _install_openai_mock(openai_mock, report) + + cfg = _make_config() + + file = MemoryFile( + content=json.dumps({"paragraphs": [], "pages": []}).encode("utf-8"), + mime_type="application/x-analyze-result", + ) + + with pytest.raises(EmptyOntologyError): + cfg.driver(file, Context()) + + +@patch("bc2.core.common.openai.OpenAI") +def test_extract_raises_when_bounding_page_missing(openai_mock): + report = _make_report() + _install_openai_mock(openai_mock, report) + + cfg = _make_config() + + from azure.ai.documentintelligence.models import AnalyzeResult + + analyze_result = AnalyzeResult( + { + "paragraphs": [ + { + "content": "Stray paragraph", + "spans": [{"offset": 0, "length": 15}], + "boundingRegions": [ + { + "pageNumber": 99, + "polygon": [0, 0, 1, 0, 1, 1, 0, 1], + } + ], + } + ], + "pages": [{"pageNumber": 1, "width": 8.5, "height": 11.0}], + } + ) + + with pytest.raises(ValueError, match="Page 99 not found"): + cfg.driver.extract(analyze_result)