diff --git a/src/interfaze/_chat.py b/src/interfaze/_chat.py index de9ce15..046cffc 100644 --- a/src/interfaze/_chat.py +++ b/src/interfaze/_chat.py @@ -1,6 +1,5 @@ from __future__ import annotations -import re from typing import Any, Dict, Iterable, List, Literal, Optional, Union, cast, overload from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream @@ -10,19 +9,9 @@ from ._errors import InterfazeError from ._guard import guard_tag from ._schema import empty_task_schema -from ._stream import AsyncInterfazeStream, InterfazeStream +from ._stream import AsyncInterfazeStream, InterfazeStream, strip_json_fence from ._types import GuardCode, InterfazeChatCompletion, TaskName -_FENCE = re.compile(r"^```(?:json)?\s*|\s*```$", re.IGNORECASE) - - -def strip_json_fence(content: str) -> str: - """Interfaze wraps ``json_object`` content in a ```json fence; unwrap it.""" - t = content.strip() - if not t.startswith("```"): - return content - return _FENCE.sub("", t).strip() - def _is_non_empty_schema(rf: Any) -> bool: schema = (rf or {}).get("json_schema", {}).get("schema", {}) if isinstance(rf, dict) else {} @@ -142,8 +131,8 @@ def stream( response_format: Optional[Dict[str, Any]] = None, **kwargs: Any, ) -> InterfazeStream: - kw, _ = self._kwargs(messages, model, task, guard, response_format, kwargs) - return InterfazeStream(self._client, kw) + kw, strip = self._kwargs(messages, model, task, guard, response_format, kwargs) + return InterfazeStream(self._client, kw, strip) class AsyncCompletions(_CompletionsBase): @@ -207,8 +196,8 @@ def stream( response_format: Optional[Dict[str, Any]] = None, **kwargs: Any, ) -> AsyncInterfazeStream: - kw, _ = self._kwargs(messages, model, task, guard, response_format, kwargs) - return AsyncInterfazeStream(self._client, kw) + kw, strip = self._kwargs(messages, model, task, guard, response_format, kwargs) + return AsyncInterfazeStream(self._client, kw, strip) class Chat: diff --git a/src/interfaze/_stream.py b/src/interfaze/_stream.py index b000184..c3120b3 100644 --- a/src/interfaze/_stream.py +++ b/src/interfaze/_stream.py @@ -30,6 +30,17 @@ def strip_side_channels(content: str) -> Tuple[str, Optional[str], Optional[List return text.strip(), ("\n".join(thinks) if thinks else None), (pre or None) +_FENCE = re.compile(r"^```(?:json)?\s*|\s*```$", re.IGNORECASE) + + +def strip_json_fence(content: str) -> str: + """Interfaze wraps ``json_object`` content in a ```json fence; unwrap it.""" + t = content.strip() + if not t.startswith("```"): + return content + return _FENCE.sub("", t).strip() + + class _State: def __init__(self) -> None: self.content = "" @@ -66,8 +77,10 @@ def accumulate(self, chunk: ChatCompletionChunk) -> None: if tc.function and tc.function.arguments: acc["arguments"] += tc.function.arguments - def build(self) -> InterfazeChatCompletion: + def build(self, strip_fence: bool = False) -> InterfazeChatCompletion: text, reasoning, precontext = strip_side_channels(self.content) + if strip_fence: + text = strip_json_fence(text) tool_calls = [ {"id": t["id"], "type": "function", "function": {"name": t["name"], "arguments": t["arguments"]}} for t in self.tool_calls.values() @@ -95,9 +108,10 @@ def build(self) -> InterfazeChatCompletion: class InterfazeStream: """Sync streaming helper — iterate chunks, then ``get_final_completion()``.""" - def __init__(self, client: OpenAI, kwargs: Dict[str, Any]) -> None: + def __init__(self, client: OpenAI, kwargs: Dict[str, Any], strip_fence: bool = False) -> None: self._client = client self._kwargs = kwargs + self._strip_fence = strip_fence self._state = _State() self._started = False self._done = False @@ -131,15 +145,16 @@ def get_final_completion(self) -> InterfazeChatCompletion: raise InterfazeError( "Call get_final_completion() after fully iterating, or instead of iterating." ) - return self._state.build() + return self._state.build(self._strip_fence) class AsyncInterfazeStream: """Async streaming helper — ``async for`` chunks, then ``await get_final_completion()``.""" - def __init__(self, client: AsyncOpenAI, kwargs: Dict[str, Any]) -> None: + def __init__(self, client: AsyncOpenAI, kwargs: Dict[str, Any], strip_fence: bool = False) -> None: self._client = client self._kwargs = kwargs + self._strip_fence = strip_fence self._state = _State() self._started = False self._done = False @@ -171,4 +186,4 @@ async def get_final_completion(self) -> InterfazeChatCompletion: raise InterfazeError( "Call get_final_completion() after fully iterating, or instead of iterating." ) - return self._state.build() + return self._state.build(self._strip_fence) diff --git a/tests/test_stream.py b/tests/test_stream.py index 09aa11d..234246b 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -1,12 +1,20 @@ from __future__ import annotations import asyncio +import json import respx -from conftest import STREAM_CHUNKS, STREAM_THINK, mock_sse +from conftest import STREAM_CHUNKS, STREAM_THINK, _chunk, mock_sse from interfaze import AsyncInterfaze, Interfaze +FENCED_JSON = [ + _chunk({"content": "```json\n"}), + _chunk({"content": '{"city": "Tokyo"}'}), + _chunk({"content": "\n```"}), + _chunk({}, finish_reason="stop"), +] + @respx.mock def test_stream_iterates_and_accumulates(): @@ -46,3 +54,36 @@ async def go(): n, final = asyncio.run(go()) assert n == len(STREAM_CHUNKS) assert final.precontext and final.precontext[0].name == "ocr" + + +@respx.mock +def test_stream_json_object_strips_fence(): + mock_sse(FENCED_JSON) + stream = Interfaze(api_key="t").chat.completions.stream( + messages=[{"role": "user", "content": "x"}], response_format={"type": "json_object"} + ) + content = stream.get_final_completion().choices[0].message.content or "" + assert not content.lstrip().startswith("```") + assert json.loads(content)["city"] == "Tokyo" + + +@respx.mock +def test_stream_without_json_object_keeps_fence(): + mock_sse(FENCED_JSON) + stream = Interfaze(api_key="t").chat.completions.stream(messages=[{"role": "user", "content": "x"}]) + content = stream.get_final_completion().choices[0].message.content or "" + assert content.lstrip().startswith("```") + + +@respx.mock +def test_async_stream_json_object_strips_fence(): + mock_sse(FENCED_JSON) + + async def go(): + stream = AsyncInterfaze(api_key="t").chat.completions.stream( + messages=[{"role": "user", "content": "x"}], response_format={"type": "json_object"} + ) + return await stream.get_final_completion() + + content = asyncio.run(go()).choices[0].message.content or "" + assert not content.lstrip().startswith("```")