Merge branch 'main' into feat/httpx-connection-pool-config

Author: elinacse

src/ogx/providers/inline/responses/builtin/responses/streaming.py

@@ -113,8 +113,10 @@
     ChatCompletionResult,
 )
 from .utils import (
+    StreamingCitationCleaner,
     convert_chat_choice_to_response_message,
     convert_mcp_tool_choice,
+    extract_citations_from_text,
     is_function_tool_call,
     run_guardrails,
     should_summarize_reasoning,
@@ -1103,6 +1105,10 @@ async def _process_streaming_chunks(
         refusal_text_accumulated = []
         pending_guardrail_events: list[OpenAIResponseObjectStream] = []
         chars_since_last_check = 0
+        # Cleans citation markers out of delta text as it streams, so a client
+        # reconstructing output from deltas sees the same text as content_part.done /
+        # output_item.done (which clean the fully accumulated text once streaming ends).
+        citation_cleaner = StreamingCitationCleaner(self.citation_files)
 
         async for raw_chunk in completion_result:
             # Providers returning OpenAIChatCompletionChunkWithReasoning wrap
@@ -1165,21 +1171,27 @@ async def _process_streaming_chunks(
                             ),
                             sequence_number=self.sequence_number,
                         )
-                    self.sequence_number += 1
+                    # Withhold citation markers (and any text that might still turn into
+                    # one) from the delta stream, so it stays consistent with the cleaned
+                    # text in content_part.done / output_item.done below. A marker split
+                    # across chunk boundaries can cause this to yield nothing for a chunk.
+                    cleaned_delta = citation_cleaner.feed(chunk_choice.delta.content)
+                    if cleaned_delta:
+                        self.sequence_number += 1
 
-                    text_delta_event = OpenAIResponseObjectStreamResponseOutputTextDelta(
-                        content_index=content_index,
-                        delta=chunk_choice.delta.content,
-                        item_id=message_item_id,
-                        logprobs=chunk_logprobs if chunk_logprobs is not None else [],
-                        output_index=message_output_index,
-                        sequence_number=self.sequence_number,
-                    )
-                    # Buffer text delta events for guardrail check
-                    if self.enable_guardrails:
-                        pending_guardrail_events.append(text_delta_event)
-                    else:
-                        yield text_delta_event
+                        text_delta_event = OpenAIResponseObjectStreamResponseOutputTextDelta(
+                            content_index=content_index,
+                            delta=cleaned_delta,
+                            item_id=message_item_id,
+                            logprobs=chunk_logprobs if chunk_logprobs is not None else [],
+                            output_index=message_output_index,
+                            sequence_number=self.sequence_number,
+                        )
+                        # Buffer text delta events for guardrail check
+                        if self.enable_guardrails:
+                            pending_guardrail_events.append(text_delta_event)
+                        else:
+                            yield text_delta_event
 
                 # Collect content for final response
                 content_delta = chunk_choice.delta.content or ""
@@ -1368,15 +1380,34 @@ async def _process_streaming_chunks(
 
         # Emit content_part.done event if text content was streamed (before content gets cleared)
         if content_part_emitted:
+            # Flush any text the citation cleaner was still withholding (e.g. a marker-like
+            # sequence that never completed) so delta-reconstructed text catches up with
+            # the cleaned text below before the round closes out. Guardrail buffering
+            # doesn't apply here: the moderation check above already ran over the full raw
+            # accumulated text, which includes whatever this flush contains.
+            flushed_delta = citation_cleaner.flush()
+            if flushed_delta:
+                self.sequence_number += 1
+                yield OpenAIResponseObjectStreamResponseOutputTextDelta(
+                    content_index=content_index,
+                    delta=flushed_delta,
+                    item_id=message_item_id,
+                    logprobs=[],
+                    output_index=message_output_index,
+                    sequence_number=self.sequence_number,
+                )
+
             final_text = "".join(chat_response_content)
+            part_annotations, part_clean_text = extract_citations_from_text(final_text, self.citation_files)
             self.sequence_number += 1
             yield OpenAIResponseObjectStreamResponseContentPartDone(
                 content_index=content_index,
                 response_id=self.response_id,
                 item_id=message_item_id,
                 output_index=message_output_index,
                 part=OpenAIResponseContentPartOutputText(
-                    text=final_text,
+                    text=part_clean_text,
+                    annotations=list(part_annotations),
                     logprobs=[],
                 ),
                 sequence_number=self.sequence_number,
@@ -1411,10 +1442,11 @@ async def _process_streaming_chunks(
             content_parts = []
             if content_part_emitted:
                 final_text = "".join(chat_response_content)
+                final_annotations, final_clean_text = extract_citations_from_text(final_text, self.citation_files)
                 content_parts.append(
                     OpenAIResponseOutputMessageContentOutputText(
-                        text=final_text,
-                        annotations=[],
+                        text=final_clean_text,
+                        annotations=list(final_annotations),
                         logprobs=chat_response_logprobs if chat_response_logprobs else [],
                     )
                 )

src/ogx/providers/inline/responses/builtin/responses/tool_executor.py

@@ -221,8 +221,11 @@ async def search_single_store(vector_store_id):
         )
 
         # handling missing attributes for old versions
+        # Iterate in descending score order so that when a model doesn't cite inline and
+        # extract_citations_from_text falls back to a single file, it picks the most relevant
+        # one: dict insertion order determines which file_id lands first.
         citation_files = {}
-        for result in search_results:
+        for result in sorted(search_results, key=lambda r: r.score, reverse=True):
             file_id = result.file_id
             if not file_id and result.attributes:
                 file_id = result.attributes.get("document_id")
@@ -233,7 +236,8 @@ async def search_single_store(vector_store_id):
             if not filename:
                 filename = "unknown"
 
-            citation_files[file_id] = filename
+            if file_id not in citation_files:
+                citation_files[file_id] = filename
 
         # Cast to proper InterleavedContent type (list invariance)
         return ToolInvocationResult(

src/ogx/providers/inline/responses/builtin/responses/utils.py

@@ -116,7 +116,7 @@ async def convert_chat_choice_to_response_message(
     """Convert an OpenAI Chat Completion choice into an OpenAI Response output message."""
     output_content = choice.message.content or ""
 
-    annotations, clean_text = _extract_citations_from_text(output_content, citation_files or {})
+    annotations, clean_text = extract_citations_from_text(output_content, citation_files or {})
     logprobs = choice.logprobs.content if choice.logprobs and choice.logprobs.content else []
 
     return OpenAIResponseMessage(
@@ -480,19 +480,41 @@ async def get_message_type_by_role(role: str) -> type[OpenAIMessageParam] | None
     return role_to_type.get(role)  # type: ignore[return-value]  # Pydantic models use ModelMetaclass
 
 
+CITATION_MARKER_REGEX = re.compile(
+    r"<\|(?P<file_id_pipe>file-[A-Za-z0-9_-]+)\|>"
+    r"|\[(?P<file_id_bracket>file-[A-Za-z0-9_-]+)\]"
+    r"|\((?P<file_id_paren>file-[A-Za-z0-9_-]+)\)"
+)
+
+# Matches an in-progress citation marker at the very end of a string, including a possible
+# single space right before it (since a known marker's preceding space gets dropped by
+# _extract_citations_from_text, and that only happens correctly if the space and the
+# marker end up cleaned together — see StreamingCitationCleaner). Also matches a bare
+# trailing space on its own, since it might turn out to precede a marker in the next
+# chunk. Used to withhold text from streamed deltas until either a marker completes (and
+# gets cleaned) or a later chunk proves it wasn't a marker after all (and gets flushed
+# through as literal text).
+_PENDING_CITATION_MARKER_TAIL_REGEX = re.compile(
+    r"(?: ?<(?:\|(?:file-[A-Za-z0-9_-]*)?)?| ?\[(?:file-[A-Za-z0-9_-]*)?| ?\((?:file-[A-Za-z0-9_-]*)?| )$"
+)
+
+
 def _extract_citations_from_text(
     text: str, citation_files: dict[str, str]
 ) -> tuple[list[OpenAIResponseAnnotationFileCitation], str]:
     """Extract citation markers from text and create annotations
 
     Args:
-        text: The text containing citation markers like [file-Cn3MSNn72ENTiiq11Qda4A]
+        text: The text containing citation markers like <|file-Cn3MSNn72ENTiiq11Qda4A|>.
+            The primary marker format is `<|file-id|>`, but `[file-id]` and `(file-id)`
+            are also accepted since weaker models often approximate the instructed
+            format rather than reproduce it exactly.
         citation_files: Dictionary mapping file_id to filename
 
     Returns:
         Tuple of (annotations_list, clean_text_without_markers)
     """
-    file_id_regex = re.compile(r"<\|(?P<file_id>file-[A-Za-z0-9_-]+)\|>")
+    file_id_regex = CITATION_MARKER_REGEX
 
     annotations = []
     parts = []
@@ -503,22 +525,32 @@ def _extract_citations_from_text(
         # segment before the marker
         prefix = text[last_end : m.start()]
 
-        # drop one space if it exists (since marker is at sentence end)
-        if prefix.endswith(" "):
+        fid = m.group("file_id_pipe") or m.group("file_id_bracket") or m.group("file_id_paren")
+        is_known = fid in citation_files
+
+        # drop one space if it exists (since marker is at sentence end); only do this when
+        # the marker itself is about to be removed below, otherwise we'd merge the prefix
+        # and the marker text together with no space between them
+        if is_known and prefix.endswith(" "):
             prefix = prefix[:-1]
 
         parts.append(prefix)
         total_len += len(prefix)
 
-        fid = m.group(1)
-        if fid in citation_files:
+        if is_known:
             annotations.append(
                 OpenAIResponseAnnotationFileCitation(
                     file_id=fid,
                     filename=citation_files[fid],
                     index=total_len,  # index points to punctuation
                 )
             )
+        else:
+            # Unrecognized marker (e.g. a stale/mismatched file id): preserve it verbatim
+            # rather than silently deleting user-visible text we can't actually attribute.
+            marker_text = m.group(0)
+            parts.append(marker_text)
+            total_len += len(marker_text)
 
         last_end = m.end()
 
@@ -527,6 +559,78 @@ def _extract_citations_from_text(
     return annotations, cleaned_text
 
 
+def extract_citations_from_text(
+    text: str, citation_files: dict[str, str]
+) -> tuple[list[OpenAIResponseAnnotationFileCitation], str]:
+    """Extract citation markers from text, with a fallback for models that don't cite inline.
+
+    Delegates to `_extract_citations_from_text` for marker-based extraction. Some models
+    (particularly small/local ones served e.g. via Ollama) don't reliably reproduce the
+    inline citation marker even when instructed to. If file_search actually retrieved
+    documents for this response, attribute the answer to the single most relevant one
+    rather than silently returning no annotations just because the model didn't echo the
+    marker. Attributing every retrieved file would imply the whole answer draws equally
+    on all of them, which usually isn't true and isn't what OpenAI's API does.
+
+    Args:
+        text: The text possibly containing citation markers.
+        citation_files: Dictionary mapping file_id to filename for files retrieved this turn,
+            ordered by descending relevance score (see tool_executor.py).
+
+    Returns:
+        Tuple of (annotations_list, clean_text_without_markers)
+    """
+    annotations, clean_text = _extract_citations_from_text(text, citation_files)
+    if not annotations and citation_files:
+        file_id, filename = next(iter(citation_files.items()))
+        annotations = [OpenAIResponseAnnotationFileCitation(file_id=file_id, filename=filename, index=len(clean_text))]
+    return annotations, clean_text
+
+
+class StreamingCitationCleaner:
+    """Incrementally strips citation markers from streamed text deltas.
+
+    content_part.done / output_item.done events clean the fully accumulated text via
+    `extract_citations_from_text`. Without this, delta events would carry the raw,
+    unprocessed text (markers and all), so a client that reconstructs output purely from
+    deltas would end up disagreeing with the final payload. Feeding every delta through
+    this cleaner keeps the two consistent.
+
+    Markers can be split across chunk boundaries (e.g. one chunk ends in "<|file-abc" and
+    the next starts with "123|>"), so a marker-looking sequence at the end of the buffered
+    text is withheld until it either completes (and gets cleaned) or a later feed()/flush()
+    call proves it wasn't a marker after all (and gets emitted as literal text).
+
+    Note: this only cleans complete markers, so if a space that would normally be dropped
+    before a marker (see `_extract_citations_from_text`) lands in a different feed() call
+    than the marker itself, that single space is not retroactively removed. This is a
+    minor cosmetic difference from the final text, not a correctness issue.
+    """
+
+    def __init__(self, citation_files: dict[str, str]):
+        self._citation_files = citation_files
+        self._buffer = ""
+
+    def feed(self, delta: str) -> str:
+        """Feed newly arrived raw text; return the portion now safe to emit to the client."""
+        self._buffer += delta
+        pending_match = _PENDING_CITATION_MARKER_TAIL_REGEX.search(self._buffer)
+        safe_upto = pending_match.start() if pending_match else len(self._buffer)
+        safe_text, self._buffer = self._buffer[:safe_upto], self._buffer[safe_upto:]
+        if not safe_text:
+            return ""
+        _, cleaned = _extract_citations_from_text(safe_text, self._citation_files)
+        return cleaned
+
+    def flush(self) -> str:
+        """Flush any remaining buffered text once no more input is coming this round."""
+        if not self._buffer:
+            return ""
+        _, cleaned = _extract_citations_from_text(self._buffer, self._citation_files)
+        self._buffer = ""
+        return cleaned
+
+
 def is_function_tool_call(
     tool_call: OpenAIChatCompletionToolCall,
     tools: list[OpenAIResponseInputTool],

tests/unit/providers/responses/builtin/fixtures/file_search_final_answer_completion.yaml

@@ -0,0 +1,9 @@
+id: chat-completion-457
+choices:
+  - message:
+      content: 'Global warming is caused by greenhouse gases <|file-abc123|>.'
+      role: assistant
+    finish_reason: stop
+    index: 0
+created: 1234567891
+model: ollama/llama3.2:3b

tests/unit/providers/responses/builtin/fixtures/file_search_tool_call_completion.yaml

@@ -0,0 +1,14 @@
+id: chat-completion-456
+choices:
+  - message:
+      tool_calls:
+        - id: tool_call_456
+          type: function
+          function:
+            name: file_search
+            arguments: '{"query":"What is global warming?"}'
+      role: assistant
+    finish_reason: stop
+    index: 0
+created: 1234567890
+model: ollama/llama3.2:3b