fix(inference): return requested model id in streamed completion chunks (#6320)

# What does this PR do?

Re-opens #6314 (previously closed). @mattf confirmed this as a real bug
and asked to continue it — reopening with the same fix.

Fixes an inconsistency in `InferenceRouter.openai_completion`: streamed
`/v1/completions` chunks report the provider-internal model id instead
of the model id the client requested.

The router rewrites `params.model` to the provider's
`provider_resource_id` before calling the provider. For non-streaming
requests it then restores the requested fully qualified id
(`response.model = request_model_id`), and the chat-completions
streaming path does the same per chunk in
`stream_tokens_and_compute_metrics_openai_chat` (`chunk.model =
fully_qualified_model_id`). But the `stream=True` branch of
`openai_completion` returned the provider stream untouched, so every
streamed chunk leaked the provider-internal id — e.g. a client
requesting `openai/gpt-4o` gets `model="openai/gpt-4o"` when not
streaming but `model="gpt-4o"` on every streamed chunk. Clients that
correlate responses by the model they requested (or that multiplex
providers behind aliases) see different model ids depending solely on
whether they streamed.

The fix wraps the provider stream in a small async generator that
restores the requested model id on each chunk before yielding, mirroring
the existing chat-streaming behavior — including skipping `None` chunks,
exactly as `stream_tokens_and_compute_metrics_openai_chat` does.
Provider exceptions and stream termination propagate through unchanged.
The return type annotation is also corrected to `OpenAICompletion |
AsyncIterator[OpenAICompletion]`, matching both the chat sibling and the
`Inference` API definition in `src/ogx_api/inference/api.py`.

## Test Plan

Added unit tests in `tests/unit/core/routers/test_inference_router.py`:

- `test_openai_completion_streaming_rewrites_model_id` — the core case
(requested fully qualified id differs from the provider's internal id,
multi-chunk); fails on `main`, passes with this fix
- `test_openai_completion_streaming_empty_stream` — a provider stream
yielding zero chunks produces an empty stream without errors
- `test_openai_completion_streaming_model_id_already_correct` — chunks
already carrying the requested id pass through unchanged
- `test_openai_completion_streaming_skips_none_chunks` — `None` chunks
are skipped, matching the chat streaming path; fails on `main`
- `test_openai_completion_streaming_propagates_provider_errors` — a
mid-stream provider error propagates to the caller after earlier chunks
are delivered
- `test_openai_completion_non_streaming_rewrites_model_id` — regression
guard for the already-correct non-streaming path

```
$ uv run pytest tests/unit/core/routers/ -q
87 passed, 1 warning in 1.21s
```

`uv run pre-commit run --files ...` passes on both touched files (ruff,
ruff format, mypy, and all repo-local hooks).

---------

Signed-off-by: Sohum Trivedi <trivsohum@gmail.com>

Author: sohumt123

src/ogx/core/routers/inference.py

@@ -189,7 +189,7 @@ async def rerank(
     async def openai_completion(
         self,
         params: Annotated[OpenAICompletionRequestWithExtraBody, Body(...)],
-    ) -> OpenAICompletion:
+    ) -> OpenAICompletion | AsyncIterator[OpenAICompletion]:
         logger.debug(
             "InferenceRouter.openai_completion: model=, stream=, prompt",
             model=params.model,
@@ -204,12 +204,29 @@ async def openai_completion(
         params.model = provider_resource_id
 
         if params.stream:
-            return await provider.openai_completion(params)
+            response_stream = await provider.openai_completion(params)
+            # Providers respond with their internal model id, so rewrite each
+            # chunk to carry the fully qualified model id the client requested
+            # (mirrors the non-streaming path below and the chat-streaming path).
+            return self._rewrite_completion_stream_model_id(response_stream, request_model_id)
 
         response = await provider.openai_completion(params)
         response.model = request_model_id
         return response
 
+    async def _rewrite_completion_stream_model_id(
+        self,
+        response: AsyncIterator[OpenAICompletion],
+        fully_qualified_model_id: str,
+    ) -> AsyncIterator[OpenAICompletion]:
+        """Yield streamed completion chunks with the requested model id restored."""
+        async for chunk in response:
+            # Skip None chunks, mirroring stream_tokens_and_compute_metrics_openai_chat
+            if chunk is None:
+                continue
+            chunk.model = fully_qualified_model_id
+            yield chunk
+
     async def openai_chat_completion(
         self,
         params: Annotated[OpenAIChatCompletionRequestWithExtraBody, Body(...)],

tests/unit/core/routers/test_inference_router.py

@@ -23,11 +23,14 @@
 from ogx.core.routers.inference import InferenceRouter
 from ogx_api import (
     ModelType,
+    OpenAICompletion,
+    OpenAICompletionRequestWithExtraBody,
     RerankData,
     RerankResponse,
     RoutingTable,
 )
 from ogx_api.inference import RerankRequest
+from ogx_api.inference.models import OpenAICompletionChoice
 
 
 @pytest.fixture
@@ -49,6 +52,192 @@ def mock_routing_table():
     return routing_table, mock_provider
 
 
+@pytest.fixture
+def mock_llm_routing_table():
+    """Create a mock routing table with an LLM model registered under a fully qualified id"""
+    routing_table = MagicMock(spec=RoutingTable)
+
+    mock_model = MagicMock()
+    mock_model.identifier = "test_provider/test-llm-model"
+    mock_model.model_type = ModelType.llm
+    mock_model.provider_resource_id = "test-llm-model"
+
+    mock_provider = MagicMock()
+    mock_provider.__provider_id__ = "test_provider"
+
+    routing_table.get_object_by_identifier = AsyncMock(return_value=mock_model)
+    routing_table.get_provider_impl = AsyncMock(return_value=mock_provider)
+
+    return routing_table, mock_provider
+
+
+def _make_completion_chunk(text: str, model: str) -> OpenAICompletion:
+    return OpenAICompletion(
+        id="cmpl-test",
+        choices=[OpenAICompletionChoice(finish_reason="stop", text=text, index=0)],
+        created=0,
+        model=model,
+        object="text_completion",
+    )
+
+
+async def test_openai_completion_streaming_rewrites_model_id(mock_llm_routing_table):
+    """
+    Test that streamed /v1/completions chunks report the fully qualified model id
+    that the client requested, not the provider-internal resource id.
+
+    This mirrors the non-streaming path in openai_completion (which sets
+    response.model = request_model_id) and the chat streaming path
+    (stream_tokens_and_compute_metrics_openai_chat, which rewrites chunk.model).
+    """
+    routing_table, mock_provider = mock_llm_routing_table
+    router = InferenceRouter(routing_table=routing_table)
+
+    async def provider_stream():
+        # Providers respond with their internal model id
+        yield _make_completion_chunk("Hello", model="test-llm-model")
+        yield _make_completion_chunk(" world", model="test-llm-model")
+
+    mock_provider.openai_completion = AsyncMock(return_value=provider_stream())
+
+    params = OpenAICompletionRequestWithExtraBody(
+        model="test_provider/test-llm-model",
+        prompt="Say hello",
+        stream=True,
+    )
+
+    stream = await router.openai_completion(params)
+    chunks = [chunk async for chunk in stream]
+
+    assert len(chunks) == 2
+    assert [chunk.model for chunk in chunks] == ["test_provider/test-llm-model", "test_provider/test-llm-model"], (
+        "Streamed completion chunks should carry the requested model id, not the provider resource id"
+    )
+    assert [choice.text for chunk in chunks for choice in chunk.choices] == ["Hello", " world"]
+
+    # The provider itself should still be called with its own resource id
+    called_params = mock_provider.openai_completion.call_args.args[0]
+    assert called_params.model == "test-llm-model"
+
+
+async def test_openai_completion_streaming_empty_stream(mock_llm_routing_table):
+    """A provider stream that yields no chunks produces an empty stream without errors."""
+    routing_table, mock_provider = mock_llm_routing_table
+    router = InferenceRouter(routing_table=routing_table)
+
+    async def provider_stream():
+        return
+        yield  # unreachable; makes this function an async generator
+
+    mock_provider.openai_completion = AsyncMock(return_value=provider_stream())
+
+    params = OpenAICompletionRequestWithExtraBody(
+        model="test_provider/test-llm-model",
+        prompt="Say hello",
+        stream=True,
+    )
+
+    stream = await router.openai_completion(params)
+    chunks = [chunk async for chunk in stream]
+
+    assert chunks == []
+
+
+async def test_openai_completion_streaming_model_id_already_correct(mock_llm_routing_table):
+    """Chunks that already carry the fully qualified model id are passed through unchanged."""
+    routing_table, mock_provider = mock_llm_routing_table
+    router = InferenceRouter(routing_table=routing_table)
+
+    async def provider_stream():
+        yield _make_completion_chunk("Hello", model="test_provider/test-llm-model")
+
+    mock_provider.openai_completion = AsyncMock(return_value=provider_stream())
+
+    params = OpenAICompletionRequestWithExtraBody(
+        model="test_provider/test-llm-model",
+        prompt="Say hello",
+        stream=True,
+    )
+
+    stream = await router.openai_completion(params)
+    chunks = [chunk async for chunk in stream]
+
+    assert len(chunks) == 1
+    assert chunks[0].model == "test_provider/test-llm-model"
+    assert chunks[0].choices[0].text == "Hello"
+
+
+async def test_openai_completion_streaming_skips_none_chunks(mock_llm_routing_table):
+    """None chunks from a provider are skipped, mirroring the chat streaming path."""
+    routing_table, mock_provider = mock_llm_routing_table
+    router = InferenceRouter(routing_table=routing_table)
+
+    async def provider_stream():
+        yield _make_completion_chunk("Hello", model="test-llm-model")
+        yield None
+        yield _make_completion_chunk(" world", model="test-llm-model")
+
+    mock_provider.openai_completion = AsyncMock(return_value=provider_stream())
+
+    params = OpenAICompletionRequestWithExtraBody(
+        model="test_provider/test-llm-model",
+        prompt="Say hello",
+        stream=True,
+    )
+
+    stream = await router.openai_completion(params)
+    chunks = [chunk async for chunk in stream]
+
+    assert [chunk.model for chunk in chunks] == ["test_provider/test-llm-model", "test_provider/test-llm-model"]
+    assert [choice.text for chunk in chunks for choice in chunk.choices] == ["Hello", " world"]
+
+
+async def test_openai_completion_streaming_propagates_provider_errors(mock_llm_routing_table):
+    """Errors raised by the provider mid-stream propagate to the caller after earlier chunks are delivered."""
+    routing_table, mock_provider = mock_llm_routing_table
+    router = InferenceRouter(routing_table=routing_table)
+
+    async def provider_stream():
+        yield _make_completion_chunk("Hello", model="test-llm-model")
+        raise RuntimeError("provider stream failed")
+
+    mock_provider.openai_completion = AsyncMock(return_value=provider_stream())
+
+    params = OpenAICompletionRequestWithExtraBody(
+        model="test_provider/test-llm-model",
+        prompt="Say hello",
+        stream=True,
+    )
+
+    stream = await router.openai_completion(params)
+    chunks = []
+    with pytest.raises(RuntimeError, match="provider stream failed"):
+        async for chunk in stream:
+            chunks.append(chunk)
+
+    assert len(chunks) == 1
+    assert chunks[0].model == "test_provider/test-llm-model"
+
+
+async def test_openai_completion_non_streaming_rewrites_model_id(mock_llm_routing_table):
+    """Non-streaming /v1/completions responses report the requested model id (regression guard)."""
+    routing_table, mock_provider = mock_llm_routing_table
+    router = InferenceRouter(routing_table=routing_table)
+
+    mock_provider.openai_completion = AsyncMock(
+        return_value=_make_completion_chunk("Hello world", model="test-llm-model")
+    )
+
+    params = OpenAICompletionRequestWithExtraBody(
+        model="test_provider/test-llm-model",
+        prompt="Say hello",
+    )
+
+    response = await router.openai_completion(params)
+
+    assert response.model == "test_provider/test-llm-model"
+
+
 async def test_rerank_calls_provider_correctly(mock_routing_table):
     """
     Test that InferenceRouter.rerank() calls the provider's rerank method with the correct RerankRequest.