chore: import upstream snapshot with attribution

tests/llm_proxy/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

tests/llm_proxy/test_llm_proxy_cpu.py

@@ -0,0 +1,382 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import multiprocessing
+import random
+from typing import Any, List, cast
+
+import litellm
+import openai
+import opentelemetry.trace as trace_api
+import pytest
+from agentops.sdk.core import BatchSpanProcessor
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import ModelResponse
+from litellm.utils import custom_llm_setup
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+from agentlightning.llm_proxy import LightningSpanExporter, LLMProxy
+from agentlightning.store import LightningStoreServer
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.store.threading import LightningStoreThreaded
+from agentlightning.types import Span
+from agentlightning.utils.server_launcher import PythonServerLauncherArgs
+
+from ..common.network import get_free_port
+from ..common.tracer import clear_tracer_provider
+
+pytestmark = pytest.mark.llmproxy
+
+
+class _FakeSpanContext:
+    def __init__(self, span_id: int):
+        self.span_id = span_id
+
+
+class _FakeParent:
+    def __init__(self, span_id: int):
+        self.span_id = span_id
+
+
+class _FakeReadableSpan:
+    def __init__(self, span_id: int, parent_id: int | None, attrs: dict[str, str]):
+        self._ctx = _FakeSpanContext(span_id)
+        self.parent = None if parent_id is None else _FakeParent(parent_id)
+        self.attributes = attrs
+        self.name = f"span-{span_id}"
+
+    def get_span_context(self):
+        return self._ctx
+
+
+class _FakeStore(InMemoryLightningStore):
+    def __init__(self):
+        super().__init__()
+        self.added: list[tuple[str, str, int, _FakeReadableSpan]] = []
+
+    async def add_otel_span(
+        self, rollout_id: str, attempt_id: str, readable_span: ReadableSpan, sequence_id: int | None = None
+    ) -> Span:
+        assert isinstance(sequence_id, int)
+        assert isinstance(readable_span, _FakeReadableSpan)
+        self.added.append((rollout_id, attempt_id, sequence_id, readable_span))
+        return cast(Span, None)
+
+
+@pytest.mark.asyncio
+async def test_exporter_tree_and_flush_headers_parsing():
+    store = _FakeStore()
+    exporter = LightningSpanExporter(store)
+
+    # Build a root and two children. Headers distributed across spans.
+    root = _FakeReadableSpan(1, None, {"metadata.requester_custom_headers": "{'x-rollout-id': 'r1'}"})
+    child_a = _FakeReadableSpan(2, 1, {"metadata.requester_custom_headers": "{'x-attempt-id': 'a9'}"})
+    child_b = _FakeReadableSpan(3, 1, {"metadata.requester_custom_headers": "{'x-sequence-id': '7'}"})
+
+    # Push to buffer and export
+    res = exporter.export(cast(List[ReadableSpan], [root, child_a, child_b]))
+    assert res.name == "SUCCESS"
+
+    # Give event loop a moment to run exporter coroutine
+    await asyncio.sleep(0.1)
+
+    # Should have flushed all three with merged headers
+    assert len(store.added) == 3
+    for rid, aid, sid, sp in store.added:
+        assert rid == "r1"
+        assert aid == "a9"
+        assert sid == 7
+        assert isinstance(sp, _FakeReadableSpan)
+
+    exporter.shutdown()
+
+
+def test_exporter_helpers():
+    store = _FakeStore()
+    exporter = LightningSpanExporter(store)
+
+    # Tree: 10(root) -> 11(child) -> 12(grandchild); 20(root2)
+    s10 = _FakeReadableSpan(10, None, {})
+    s11 = _FakeReadableSpan(11, 10, {})
+    s12 = _FakeReadableSpan(12, 11, {})
+    s20 = _FakeReadableSpan(20, None, {})
+
+    for _ in range(10):
+        exporter._buffer = cast(List[ReadableSpan], [s10, s11, s12, s20])  # pyright: ignore[reportPrivateUsage]
+        random.shuffle(exporter._buffer)  # pyright: ignore[reportPrivateUsage]
+
+        roots = list(exporter._get_root_span_ids())  # pyright: ignore[reportPrivateUsage]
+        assert set(roots) == {10, 20}
+
+        subtree_ids = set(exporter._get_subtrees(10))  # pyright: ignore[reportPrivateUsage]
+        assert subtree_ids == {10, 11, 12}
+
+        popped = exporter._pop_subtrees(10)  # pyright: ignore[reportPrivateUsage]
+        assert {sp.get_span_context().span_id for sp in popped} == {  # pyright: ignore[reportOptionalMemberAccess]
+            10,
+            11,
+            12,
+        }
+        # Remaining buffer has only s20
+        assert {
+            sp.get_span_context().span_id  # pyright: ignore[reportOptionalMemberAccess]
+            for sp in exporter._buffer  # pyright: ignore[reportPrivateUsage]
+        } == {20}
+
+    exporter.shutdown()
+
+    # TODO: add more complex tests for the exporter helper
+
+
+@pytest.mark.asyncio
+async def test_update_model_list():
+    store = InMemoryLightningStore()
+    proxy = LLMProxy(
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "openai/gpt-4o",
+                },
+            }
+        ],
+        launch_mode="asyncio",
+        port=get_free_port(),
+        store=store,
+    )
+    await proxy.start()
+    assert proxy.is_running()
+    assert proxy.model_list == [
+        {
+            "model_name": "gpt-4o-arbitrary",
+            "litellm_params": {
+                "model": "openai/gpt-4o",
+            },
+        }
+    ]
+    proxy.update_model_list(
+        [
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "openai/gpt-4o-mini",
+                },
+            }
+        ]
+    )
+    assert proxy.model_list == [
+        {
+            "model_name": "gpt-4o-arbitrary",
+            "litellm_params": {
+                "model": "openai/gpt-4o-mini",
+            },
+        }
+    ]
+    assert proxy.is_running()
+    await proxy.stop()
+
+
+@pytest.mark.asyncio
+async def test_restart_resets_litellm_logging_worker() -> None:
+    """LLMProxy.start() should recreate LiteLLM's logging worker on each run."""
+
+    try:
+        from litellm.litellm_core_utils import logging_worker as litellm_logging_worker
+    except ImportError:
+        pytest.skip("LiteLLM logging worker not available")
+
+    store = InMemoryLightningStore()
+    proxy = LLMProxy(
+        model_list=[
+            {
+                "model_name": "dummy-model",
+                # The backend is never invoked; only the proxy lifecycle matters here.
+                "litellm_params": {"model": "gpt-3.5-turbo"},
+            }
+        ],
+        store=store,
+        launcher_args=PythonServerLauncherArgs(
+            port=get_free_port(),
+            launch_mode="asyncio",
+            healthcheck_url="/health",
+            startup_timeout=10.0,
+            process_join_timeout=10.0,
+        ),
+    )
+
+    try:
+        await proxy.start()
+        first_worker = litellm_logging_worker.GLOBAL_LOGGING_WORKER
+        await proxy.stop()
+
+        await proxy.start()
+        second_worker = litellm_logging_worker.GLOBAL_LOGGING_WORKER
+    finally:
+        await proxy.stop()
+
+    assert first_worker is not second_worker, "LiteLLM logging worker should be refreshed after restart"
+
+
+class TestLLM(CustomLLM):
+    def __init__(self, content: str) -> None:
+        super().__init__()
+        self.content = content
+
+    def completion(self, *args: Any, **kwargs: Any) -> ModelResponse:
+        return litellm.completion(  # type: ignore
+            model="gpt-4o",
+            messages=[{"role": "user", "content": "Hello world"}],
+            mock_response=self.content,
+        )
+
+    async def acompletion(self, *args: Any, **kwargs: Any) -> ModelResponse:
+        return litellm.completion(  # type: ignore
+            model="gpt-4o",
+            messages=[{"role": "user", "content": "Hello world"}],
+            mock_response=self.content,
+        )
+
+
+@pytest.mark.asyncio
+async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixture) -> None:
+    clear_tracer_provider()
+
+    restart_times: int = 30
+
+    store = LightningStoreThreaded(InMemoryLightningStore())
+    caplog.set_level(logging.WARNING)
+
+    port = get_free_port()
+    try:
+        llm_proxy = LLMProxy(
+            model_list=[
+                {
+                    "model_name": "gpt-4o-arbitrary",
+                    "litellm_params": {
+                        # NOTE: The model after "/" cannot be an openai model like gpt-4o
+                        # This might be a bug with litellm
+                        "model": "test-llm/any-llm",
+                    },
+                }
+            ],
+            launcher_args=PythonServerLauncherArgs(
+                launch_mode="thread",
+                healthcheck_url="/health",
+                port=port,
+            ),
+            store=store,
+        )
+        for restart_idx in range(restart_times):
+            llm_instance = TestLLM(f"Hi! {restart_idx}")
+            litellm.custom_provider_map = [{"provider": "test-llm", "custom_handler": llm_instance}]
+            custom_llm_setup()
+            await llm_proxy.restart()
+            assert llm_proxy.is_running()
+
+            openai_client = openai.AsyncOpenAI(
+                base_url=llm_proxy.server_launcher.access_endpoint,
+                api_key="token-abc123",
+                timeout=5,
+                max_retries=0,
+            )
+            response = await openai_client.chat.completions.create(
+                model="gpt-4o-arbitrary",
+                messages=[{"role": "user", "content": "Hello world"}],
+                stream=False,
+            )
+            assert response.choices[0].message.content == f"Hi! {restart_idx}"
+
+            error_logs = [record.message for record in caplog.records if record.levelno >= logging.ERROR]
+            assert not error_logs, f"Found error logs: {error_logs}"
+            assert not any("Cannot add callback" in record.message for record in caplog.records)
+
+        await llm_proxy.stop()
+    finally:
+        litellm.custom_provider_map = []
+        custom_llm_setup()
+
+
+async def llm_proxy_span_exporter_loop(otlp_enabled: bool = False):
+    store = LightningStoreThreaded(InMemoryLightningStore())
+
+    if otlp_enabled:
+        store = LightningStoreServer(store, "127.0.0.1", get_free_port())
+        await store.start()
+
+    llm_instance = TestLLM(f"Hi! I'm a test LLM")
+    litellm.custom_provider_map = [{"provider": "test-llm", "custom_handler": llm_instance}]
+    custom_llm_setup()
+    proxy = LLMProxy(
+        launcher_args=PythonServerLauncherArgs(
+            launch_mode="thread",
+            healthcheck_url="/health",
+            port=get_free_port(),
+        ),
+        store=store,
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "test-llm/any-llm",
+                },
+            }
+        ],
+    )
+    await proxy.start()
+
+    rollout = await store.start_rollout(None)
+    resource = proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id)
+
+    client = openai.AsyncOpenAI(
+        base_url=resource.endpoint,
+        api_key="token-abc123",
+        timeout=5,
+        max_retries=0,
+    )
+    response = await client.chat.completions.create(
+        model="gpt-4o-arbitrary",
+        messages=[{"role": "user", "content": "Hello world"}],
+        stream=False,
+    )
+    assert response.choices[0].message.content == "Hi! I'm a test LLM"
+
+    spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+    assert len(spans) > 0, "Should have captured spans"
+    for span in spans:
+        assert span.rollout_id == rollout.rollout_id, f"Span {span.name} has incorrect rollout_id"
+        assert span.attempt_id == rollout.attempt.attempt_id, f"Span {span.name} has incorrect attempt_id"
+        assert span.sequence_id == 1, f"Span {span.name} has incorrect sequence_id"
+
+    tracer_provider = trace_api.get_tracer_provider()
+
+    have_asserted_loop = False
+    for span_processor in tracer_provider._active_span_processor._span_processors:  # type: ignore
+        if isinstance(span_processor, (SimpleSpanProcessor, BatchSpanProcessor)):
+            if isinstance(span_processor.span_exporter, LightningSpanExporter):
+                if otlp_enabled:
+                    assert span_processor.span_exporter._loop is None  # type: ignore
+                else:
+                    assert span_processor.span_exporter._loop is not None  # type: ignore
+                have_asserted_loop = True
+                break
+    assert have_asserted_loop, f"LightningSpanExporter should be used with otlp_enabled={otlp_enabled}"
+
+    await proxy.stop()
+
+    if isinstance(store, LightningStoreServer):
+        await store.stop()
+
+
+def llm_proxy_span_exporter_loop_sync(otlp_enabled: bool = False):
+    asyncio.run(llm_proxy_span_exporter_loop(otlp_enabled))
+
+
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+def test_llm_proxy_span_exporter_loop(otlp_enabled: bool):
+    context = multiprocessing.get_context("spawn")
+    process = context.Process(target=llm_proxy_span_exporter_loop_sync, args=(otlp_enabled,))
+    process.start()
+    process.join(timeout=30.0)
+    assert process.exitcode == 0

tests/llm_proxy/test_llm_proxy_gpu.py

@@ -0,0 +1,566 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Test the LLMProxy class. Still under development.
+
+General TODOs:
+
+1. Add tests for retries
+2. Add tests for timeout
+3. Add tests for multiple models in model list
+4. Add tests for multi-modal models
+
+There are some specific TODOs for each test function.
+"""
+
+import ast
+import asyncio
+import json
+from typing import Any, Dict, List, Sequence, Type, Union, cast
+
+import anthropic
+import openai
+import pytest
+from litellm.integrations.custom_logger import CustomLogger
+from portpicker import pick_unused_port
+
+from agentlightning import LlmProxyTraceToTriplet
+from agentlightning.llm_proxy import LLMProxy, _reset_litellm_logging_worker  # pyright: ignore[reportPrivateUsage]
+from agentlightning.store import LightningStore, LightningStoreServer, LightningStoreThreaded
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.types import LLM, Span
+
+from ..common.tracer import clear_tracer_provider
+from ..common.vllm import VLLM_VERSION, RemoteOpenAIServer
+
+pytestmark = [pytest.mark.gpu, pytest.mark.llmproxy]
+
+
+@pytest.fixture(scope="module")
+def qwen25_model():
+    with RemoteOpenAIServer(
+        model="Qwen/Qwen2.5-0.5B-Instruct",
+        vllm_serve_args=[
+            "--gpu-memory-utilization",
+            "0.7",
+            "--enable-auto-tool-choice",
+            "--tool-call-parser",
+            "hermes",
+            "--port",
+            str(pick_unused_port()),
+        ],
+    ) as server:
+        yield server
+
+
+def test_qwen25_model_sanity(qwen25_model: RemoteOpenAIServer):
+    client = qwen25_model.get_client()
+    response = client.chat.completions.create(
+        model="Qwen/Qwen2.5-0.5B-Instruct",
+        messages=[{"role": "user", "content": "Hello, world!"}],
+        stream=False,
+    )
+    assert response.choices[0].message.content is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_basic_integration(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    clear_tracer_provider()
+    inmemory_store = InMemoryLightningStore()
+    if otlp_enabled:
+        store = LightningStoreServer(store=inmemory_store, host="127.0.0.1", port=pick_unused_port())
+        await store.start()
+    else:
+        store = LightningStoreThreaded(inmemory_store)
+    proxy = LLMProxy(
+        port=pick_unused_port(),
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "hosted_vllm/" + qwen25_model.model,
+                    "api_base": qwen25_model.url_for("v1"),
+                },
+            }
+        ],
+        store=store,
+        launch_mode="thread" if not otlp_enabled else "mp",
+    )
+
+    rollout = await store.start_rollout(None)
+
+    await proxy.start()
+
+    resource = proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id)
+
+    client = openai.OpenAI(base_url=resource.endpoint, api_key="token-abc123")
+    response = client.chat.completions.create(
+        model="gpt-4o-arbitrary",
+        messages=[{"role": "user", "content": "Repeat after me: Hello, world!"}],
+        stream=False,
+    )
+    assert response.choices[0].message.content is not None
+    assert "hello, world" in response.choices[0].message.content.lower()
+
+    await proxy.stop()
+
+    spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+
+    if isinstance(store, LightningStoreServer):
+        await store.stop()
+
+    # Verify all spans have correct rollout_id, attempt_id, and sequence_id
+    assert len(spans) > 0, "Should have captured spans"
+    for span in spans:
+        assert span.rollout_id == rollout.rollout_id, f"Span {span.name} has incorrect rollout_id"
+        assert span.attempt_id == rollout.attempt.attempt_id, f"Span {span.name} has incorrect attempt_id"
+        assert span.sequence_id == 1, f"Span {span.name} has incorrect sequence_id"
+
+        # Verify start time and end time
+        # TODO: Remove this when this PR is merged: https://github.com/BerriAI/litellm/pull/16558
+        print(f">>> Span: {span.name}")
+        print(f">>> Start time: {span.start_time}")
+        print(f">>> End time: {span.end_time}")
+        print(f">>> Attributes: {span.attributes.keys()}")
+        assert span.start_time is not None, f"Span {span.name} has no start time"
+        assert span.end_time is not None, f"Span {span.name} has no end time"
+
+    # Find the raw_gen_ai_request span and verify token IDs
+    raw_gen_ai_spans = [s for s in spans if s.name == "raw_gen_ai_request"]
+    assert len(raw_gen_ai_spans) == 1, f"Expected 1 raw_gen_ai_request span, found {len(raw_gen_ai_spans)}"
+    raw_span = raw_gen_ai_spans[0]
+
+    # Verify prompt_token_ids is present and non-empty
+    assert (
+        "llm.hosted_vllm.prompt_token_ids" in raw_span.attributes
+    ), "prompt_token_ids not found in raw_gen_ai_request span"
+    prompt_token_ids: list[int] = ast.literal_eval(raw_span.attributes["llm.hosted_vllm.prompt_token_ids"])  # type: ignore
+    assert isinstance(prompt_token_ids, list), "prompt_token_ids should be a list"
+    assert len(prompt_token_ids) > 0, "prompt_token_ids should not be empty"
+    assert all(isinstance(tid, int) for tid in prompt_token_ids), "All prompt token IDs should be integers"
+
+    # Verify response token_ids is present in choices
+    assert "llm.hosted_vllm.choices" in raw_span.attributes, "choices not found in raw_gen_ai_request span"
+    choices: list[dict[str, Any]] = ast.literal_eval(raw_span.attributes["llm.hosted_vllm.choices"])  # type: ignore
+    assert len(choices) > 0, "Should have at least one choice"
+    if VLLM_VERSION >= (0, 10, 2):
+        assert "token_ids" in choices[0], "token_ids not found in choice"
+        response_token_ids: list[int] = choices[0]["token_ids"]
+    else:
+        assert (
+            "llm.hosted_vllm.response_token_ids" in raw_span.attributes
+        ), "response_token_ids not found in raw_gen_ai_request span"
+        response_token_ids_list: list[list[int]] = ast.literal_eval(raw_span.attributes["llm.hosted_vllm.response_token_ids"])  # type: ignore
+        assert isinstance(response_token_ids_list, list), "response_token_ids_list should be a list"
+        assert len(response_token_ids_list) > 0, "response_token_ids_list should not be empty"
+        assert all(
+            isinstance(tid_list, list) for tid_list in response_token_ids_list
+        ), "All response token IDs should be lists"
+        assert all(
+            isinstance(tid, int) for tid_list in response_token_ids_list for tid in tid_list
+        ), "All response token IDs should be integers"
+        response_token_ids = response_token_ids_list[0]
+    assert isinstance(response_token_ids, list), "response token_ids should be a list"
+    assert len(response_token_ids) > 0, "response token_ids should not be empty"
+    assert all(isinstance(tid, int) for tid in response_token_ids), "All response token IDs should be integers"
+
+    # Find the litellm_request span and verify gen_ai prompts/completions
+    litellm_spans = [s for s in spans if s.name == "litellm_request"]
+    assert len(litellm_spans) == 1, f"Expected 1 litellm_request span, found {len(litellm_spans)}"
+    litellm_span = litellm_spans[0]
+
+    # Verify gen_ai.prompt attributes
+    assert "gen_ai.prompt.0.role" in litellm_span.attributes, "gen_ai.prompt.0.role not found"
+    assert litellm_span.attributes["gen_ai.prompt.0.role"] == "user", "Expected user role in prompt"
+    assert "gen_ai.prompt.0.content" in litellm_span.attributes, "gen_ai.prompt.0.content not found"
+    assert litellm_span.attributes["gen_ai.prompt.0.content"] == "Repeat after me: Hello, world!"
+
+    # Verify gen_ai.completion attributes
+    assert "gen_ai.completion.0.role" in litellm_span.attributes, "gen_ai.completion.0.role not found"
+    assert litellm_span.attributes["gen_ai.completion.0.role"] == "assistant", "Expected assistant role in completion"
+    assert "gen_ai.completion.0.content" in litellm_span.attributes, "gen_ai.completion.0.content not found"
+    assert "gen_ai.completion.0.finish_reason" in litellm_span.attributes, "gen_ai.completion.0.finish_reason not found"
+
+
+async def _make_proxy_and_store(
+    qwen25_model: RemoteOpenAIServer,
+    *,
+    retries: int = 0,
+    gunicorn: bool = False,
+    callbacks: List[Union[Type[CustomLogger], str]] | None = None,
+    otlp_enabled: bool = False,
+):
+    clear_tracer_provider()
+    _reset_litellm_logging_worker()  # type: ignore
+    store = InMemoryLightningStore()
+    if otlp_enabled:
+        store = LightningStoreServer(store=store, host="127.0.0.1", port=pick_unused_port())
+        # When the server is forked into subprocess, it automatically becomes a client of the store
+        await store.start()
+    else:
+        # Backward compatibility with legacy thread + non-otlp mode
+        store = LightningStoreThreaded(store)
+    proxy = LLMProxy(
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "hosted_vllm/" + qwen25_model.model,
+                    "api_base": qwen25_model.url_for("v1"),
+                },
+            }
+        ],
+        launch_mode="thread" if not otlp_enabled else "mp",
+        port=pick_unused_port(),
+        num_workers=4 if gunicorn else 1,
+        store=store,
+        num_retries=retries,
+        callbacks=callbacks,
+    )
+    await proxy.start()
+    return proxy, store
+
+
+async def _new_resource(proxy: LLMProxy, store: LightningStore):
+    rollout = await store.start_rollout(None)
+    return proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id), rollout
+
+
+def _get_client_for_resource(resource: LLM):
+    return openai.OpenAI(base_url=resource.endpoint, api_key="token-abc123", timeout=120, max_retries=0)
+
+
+def _get_async_client_for_resource(resource: LLM):
+    return openai.AsyncOpenAI(base_url=resource.endpoint, api_key="token-abc123", timeout=120, max_retries=0)
+
+
+def _find_span(spans: Sequence[Span], name: str):
+    return [s for s in spans if s.name == name]
+
+
+def _attr(s: Span, key: str, default: Any = None):  # type: ignore
+    return s.attributes.get(key, default)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+        client = _get_client_for_resource(resource)
+
+        for i in range(3):
+            r = client.chat.completions.create(
+                model="gpt-4o-arbitrary",
+                messages=[{"role": "user", "content": f"Say ping {i}"}],
+                stream=False,
+            )
+            assert r.choices[0].message.content
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        assert len(spans) > 0
+        # Different requests have different sequence_ids
+        assert {s.sequence_id for s in spans} == {1, 2, 3}
+        # At least 3 requests recorded
+        assert len(_find_span(spans, "raw_gen_ai_request")) == 3
+        # TODO: Check response contents and token ids for the 3 requests respectively
+    finally:
+        await proxy.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["gunicorn", "thread", "uvicorn"])
+async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, mode: str):
+    proxy, store = await _make_proxy_and_store(qwen25_model, gunicorn=mode == "gunicorn", otlp_enabled=mode != "thread")
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+        aclient = _get_async_client_for_resource(resource)
+
+        async def _one(i: int):
+            r = await aclient.chat.completions.create(
+                model="gpt-4o-arbitrary",
+                messages=[{"role": "user", "content": f"Return #{i}"}],
+                stream=False,
+            )
+            return r.choices[0].message.content
+
+        outs = await asyncio.gather(*[_one(i) for i in range(10)])
+        assert len([o for o in outs if o]) == 10
+        await asyncio.sleep(1.0)  # Allow some extra time for the spans to be recorded
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        assert len(_find_span(spans, "raw_gen_ai_request")) == 10
+        assert {s.sequence_id for s in spans} == {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
+        # TODO: Check whether the sequence ids get mixed up or not
+    finally:
+        await proxy.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    # litellm proxy accepts Anthropic schema and forwards to OpenAI backend
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+
+        a = anthropic.Anthropic(base_url=resource.endpoint, api_key="token-abc123", timeout=120)
+        msg = a.messages.create(
+            model="gpt-4o-arbitrary",
+            max_tokens=64,
+            messages=[{"role": "user", "content": "Respond with the word: OK"}],
+        )
+        # Anthropic SDK returns content list
+        txt = "".join([b.text for b in msg.content if b.type == "text"])
+        assert "OK" in txt.upper()
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        assert len(spans) > 0
+    finally:
+        await proxy.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+        client = _get_client_for_resource(resource)
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "echo",
+                    "description": "Echo a string",
+                    "parameters": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]},
+                },
+            }
+        ]
+
+        r1 = client.chat.completions.create(
+            model="gpt-4o-arbitrary",
+            messages=[{"role": "user", "content": "Call the echo tool with text=hello"}],
+            tools=cast(Any, tools),
+            tool_choice="auto",
+            stream=False,
+        )
+        # If the small model does not tool-call, skip gracefully
+        tool_calls = r1.choices[0].message.tool_calls or []
+        if not tool_calls:
+            pytest.skip("model did not emit tool calls in this environment")
+
+        call = tool_calls[0]
+        assert call.type == "function"
+        assert call.function and call.function.name == "echo"
+        args = json.loads(call.function.arguments)
+        assert "text" in args
+
+        r2 = client.chat.completions.create(
+            model="gpt-4o-arbitrary",
+            messages=cast(
+                Any,
+                [
+                    {"role": "user", "content": "Call the echo tool with text=hello"},
+                    {
+                        "role": "assistant",
+                        "tool_calls": [
+                            {
+                                "id": call.id,
+                                "type": "function",
+                                "function": {"name": "echo", "arguments": call.function.arguments},
+                            }
+                        ],
+                    },
+                    {"role": "tool", "tool_call_id": call.id, "name": "echo", "content": args["text"]},
+                ],
+            ),
+            stream=False,
+        )
+        assert args["text"] in (r2.choices[0].message.content or "")
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        assert len(_find_span(spans, "litellm_request")) == 2
+        assert len(_find_span(spans, "raw_gen_ai_request")) == 2
+
+        # TODO: Check response contents and token ids for the 2 requests respectively
+    finally:
+        await proxy.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+        client = _get_client_for_resource(resource)
+
+        stream = client.chat.completions.create(
+            model="gpt-4o-arbitrary",
+            messages=[{"role": "user", "content": "Say the word 'apple'"}],
+            stream=True,
+        )
+        collected: list[str] = []
+        for evt in stream:
+            print(f">>> Event: {evt}")
+            for c in evt.choices:
+                if c.delta and getattr(c.delta, "content", None):
+                    assert isinstance(c.delta.content, str)
+                    collected.append(c.delta.content)
+        # Sometimes the model responds with "hello" instead of "apple"
+        assert "apple" in "".join(collected).lower() or "hello" in "".join(collected).lower()
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        assert len(spans) > 0
+        for span in spans:
+            print(f">>> Span {span.name}: {span.attributes}")
+            if span.name == "raw_gen_ai_request":
+                assert "llm.hosted_vllm.prompt_token_ids" in span.attributes
+                assert "llm.hosted_vllm.choices" in span.attributes
+            if span.name == "litellm_request":
+                assert "gen_ai.completion.0.content" in span.attributes
+    finally:
+        await proxy.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+        adapter = LlmProxyTraceToTriplet()
+        client = anthropic.Anthropic(base_url=resource.endpoint, api_key="token-abc123", timeout=120)
+
+        # non-stream
+        response = client.messages.create(
+            model="gpt-4o-arbitrary",
+            max_tokens=64,
+            messages=[{"role": "user", "content": "Say the word: banana"}],
+        )
+
+        txt = "".join([b.text for b in response.content if b.type == "text"])
+        assert "banana" in txt.lower(), f"Response does not contain 'banana': {txt}"
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        for i, span in enumerate(spans):
+            print(f">>> Span {i}: {span.name}, attributes: {span.attributes}")
+        assert len(spans) > 0
+
+        triplets = adapter.adapt(spans)
+        for i, triplet in enumerate(triplets):
+            print(f">>> Triplet {i}: {triplet}")
+        assert len(triplets) == 1
+        assert triplets[0].prompt["token_ids"]
+        assert triplets[0].response["token_ids"]
+
+        # stream
+        response = client.messages.create(
+            model="gpt-4o-arbitrary",
+            max_tokens=64,
+            messages=[{"role": "user", "content": "Say the word: banana"}],
+            stream=True,
+        )
+        chunk_number: int = 0
+        for chunk in response:
+            print(f">>> Chunk: {chunk}")
+            chunk_number += 1
+        assert chunk_number >= 1
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        for i, span in enumerate(spans):
+            print(f">>> Span {i}: {span.name}, attributes: {span.attributes}")
+            if span.name == "raw_gen_ai_request":
+                assert "llm.hosted_vllm.prompt_token_ids" in span.attributes
+                assert "llm.hosted_vllm.choices" in span.attributes
+            if span.name == "litellm_request":
+                assert "gen_ai.completion.0.content" in span.attributes
+        assert len(spans) > 0
+        triplets = adapter.adapt(spans)
+        for i, triplet in enumerate(triplets):
+            print(f">>> Triplet {i}: {triplet}")
+            assert triplet.prompt["token_ids"]
+            assert triplet.response["token_ids"]
+        assert len(triplets) == 2
+    finally:
+        await proxy.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()
+
+
+class LogprobsCallback(CustomLogger):
+
+    async def async_pre_call_hook(self, data: Dict[str, Any], **kwargs: Any) -> Dict[str, Any]:  # type: ignore
+        return {**data, "logprobs": 1}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_anthropic_logprobs(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(
+        qwen25_model, callbacks=[LogprobsCallback, "return_token_ids", "opentelemetry"], otlp_enabled=otlp_enabled
+    )
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+        client = anthropic.Anthropic(base_url=resource.endpoint, api_key="token-abc123", timeout=120)
+        adapter = LlmProxyTraceToTriplet()
+
+        # test streaming case only
+        response = client.messages.create(
+            model="gpt-4o-arbitrary",
+            max_tokens=64,
+            messages=[{"role": "user", "content": "Say the word: banana"}],
+            stream=True,
+        )
+
+        chunk_number: int = 0
+        for chunk in response:
+            print(f">>> Chunk: {chunk}")
+            chunk_number += 1
+        assert chunk_number >= 1
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        for i, span in enumerate(spans):
+            print(f">>> Span {i}: {span.name}, attributes: {span.attributes}")
+            if span.name == "raw_gen_ai_request":
+                assert "llm.hosted_vllm.prompt_token_ids" in span.attributes
+                assert "llm.hosted_vllm.choices" in span.attributes
+                choices: list[dict[str, Any]] = ast.literal_eval(span.attributes["llm.hosted_vllm.choices"])  # type: ignore
+
+                # Check for token IDs and logprobs in the first choice
+                assert len(choices) > 0
+                if VLLM_VERSION >= (0, 10, 2):
+                    assert "token_ids" in choices[0]
+                    assert choices[0]["token_ids"]
+                assert "logprobs" in choices[0]
+                assert "content" in choices[0]["logprobs"]
+                assert len(choices[0]["logprobs"]["content"]) > 0
+                assert isinstance(choices[0]["logprobs"]["content"][0], dict)
+                assert "token" in choices[0]["logprobs"]["content"][0]
+                assert "logprob" in choices[0]["logprobs"]["content"][0]
+                assert isinstance(choices[0]["logprobs"]["content"][0]["logprob"], float)
+
+        assert len(spans) > 0
+
+        triplets = adapter.adapt(spans)
+        for i, triplet in enumerate(triplets):
+            print(f">>> Triplet {i}: {triplet}")
+            assert triplet.prompt["token_ids"]
+            assert triplet.response["token_ids"]
+            # TODO: Check logprobs
+    finally:
+        await proxy.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()

tests/llm_proxy/test_stream.py

@@ -0,0 +1,714 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+from typing import Any, AsyncGenerator, Dict, Iterator, List, Optional, cast
+
+import pytest
+
+from agentlightning.llm_proxy import StreamConversionMiddleware
+
+
+def merge_openai_streaming(chunks: Iterator[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Merge chunks from OpenAI chat completion streaming into a single message dict.
+
+    Returns a dict with keys:
+      - role: "assistant" (or whatever)
+      - content: full concatenated content string
+      - function_call: optional dict with keys name, arguments (string or JSON parsed)
+    """
+    role: Optional[str] = None
+    content_parts: List[str] = []
+    function_name: Optional[str] = None
+    function_args_str: Optional[str] = None
+
+    for chunk in chunks:
+        choice = chunk.get("choices", [])[0]
+        delta = choice.get("delta", {})
+
+        if "role" in delta and delta["role"] is not None:
+            role = delta["role"]
+
+        if "content" in delta and delta["content"] is not None:
+            content_parts.append(delta["content"])
+
+        # existing format: function_call
+        if "function_call" in delta and delta["function_call"] is not None:
+            fn = delta["function_call"]
+            if function_name is None:
+                function_name = fn.get("name")
+                function_args_str = fn.get("arguments", "")
+            else:
+                function_args_str += fn.get("arguments", "")
+
+        # new format: tool_calls array
+        if "tool_calls" in delta and delta["tool_calls"]:
+            for tc in delta["tool_calls"]:
+                func = tc.get("function", {})
+                # set name if first time
+                if function_name is None and func.get("name"):
+                    function_name = func["name"]
+                # accumulate arguments
+                if func.get("arguments") is not None:
+                    if function_args_str is None:
+                        function_args_str = func["arguments"]
+                    else:
+                        function_args_str += func["arguments"]
+
+    full_content = "".join(content_parts)
+    result: Dict[str, Any] = {"role": role or "assistant", "content": full_content}
+    if function_name is not None:
+        try:
+            function_args = json.loads(function_args_str or "")  # type: ignore
+        except Exception:
+            function_args = function_args_str
+        result["function_call"] = {"name": function_name, "arguments": function_args}
+    return result
+
+
+def merge_anthropic_streaming(chunks: Iterator[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Merge chunks from Anthropic streaming into a single message dict.
+
+    Returns a dict with keys:
+      - role: "assistant"
+      - content_text: full content text (concatenated)
+      - tool_calls: list of dicts { name, input } if any
+    """
+    role: Optional[str] = None
+    content_text_parts: List[str] = []
+    tool_calls: List[Dict[str, Any]] = []
+    current_tool: Optional[Dict[str, Any]] = None
+    current_tool_input_str: Optional[str] = None
+
+    for chunk in chunks:
+        # role
+        if role is None and "role" in chunk:
+            role = chunk["role"]
+
+        # handle content_block style (fine-grained)
+        typ = chunk.get("type")
+        if typ == "content_block_start":
+            block = chunk.get("content_block", {})
+            if block.get("type") == "tool_use":
+                # finish previous tool if exists
+                if current_tool is not None:
+                    try:
+                        input_obj = json.loads(current_tool_input_str or "")
+                    except Exception:
+                        input_obj = current_tool_input_str
+                    current_tool["input"] = input_obj
+                    tool_calls.append(current_tool)
+                current_tool = {"name": block.get("name"), "id": block.get("id"), "input": None}
+                current_tool_input_str = ""
+            continue
+
+        if typ == "content_block_delta":
+            delta = chunk.get("delta", {})
+            dtyp = delta.get("type")
+            if dtyp == "input_json_delta":
+                current_tool_input_str = (current_tool_input_str or "") + delta.get("partial_json", "")
+            elif dtyp == "text_delta":
+                content_text_parts.append(delta.get("text", ""))
+            continue
+
+        if typ == "content_block_stop":
+            if current_tool is not None:
+                try:
+                    input_obj = json.loads(current_tool_input_str or "")
+                except Exception:
+                    input_obj = current_tool_input_str
+                current_tool["input"] = input_obj
+                tool_calls.append(current_tool)
+                current_tool = None
+                current_tool_input_str = None
+            continue
+
+        # handle normal content items
+        content_items = chunk.get("content", [])
+        for item in content_items:
+            t = item.get("type")
+            if t == "text":
+                content_text_parts.append(item.get("text", ""))
+            elif t == "tool_use":
+                tool_id = item.get("id")
+                name = item.get("name")
+                inp = item.get("input", {})
+                if current_tool and current_tool.get("id") == tool_id:
+                    # merge into same tool
+                    try:
+                        existing = json.loads(current_tool_input_str or "{}")
+                    except Exception:
+                        existing: Dict[str, Any] = {}
+                    if isinstance(existing, dict):  # type: ignore
+                        existing.update(inp)
+                        current_tool_input_str = json.dumps(existing)
+                    else:
+                        # fallback: treat as string concatenation
+                        current_tool_input_str += json.dumps(inp)
+                else:
+                    # finish previous tool
+                    if current_tool is not None:
+                        try:
+                            input_obj = json.loads(current_tool_input_str or "")
+                        except Exception:
+                            input_obj = current_tool_input_str
+                        current_tool["input"] = input_obj
+                        tool_calls.append(current_tool)
+                    current_tool = {"name": name, "id": tool_id, "input": None}
+                    current_tool_input_str = json.dumps(inp)
+            # else: ignore
+    # end loop
+
+    # finish any open tool
+    if current_tool is not None:
+        try:
+            input_obj = json.loads(current_tool_input_str or "")
+        except Exception:
+            input_obj = current_tool_input_str
+        current_tool["input"] = input_obj
+        tool_calls.append(current_tool)
+
+    full_text = "".join(content_text_parts)
+    result: Dict[str, Any] = {"role": role or "assistant", "content_text": full_text}
+    if tool_calls:
+        result["tool_calls"] = tool_calls
+    return result
+
+
+def test_openai_text_only_short():
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]},
+                {"choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}]},
+                {"choices": [{"index": 0, "delta": {"content": " world!"}, "finish_reason": None}]},
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["role"] == "assistant"
+    assert merged["content"] == "Hello world!"
+    assert "function_call" not in merged
+
+
+def test_openai_text_and_function_call_arguments_split():
+    # Mixed content + function_call arguments spread over multiple deltas
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}}]},
+                {"choices": [{"index": 0, "delta": {"content": "Starting… "}}]},
+                {
+                    "choices": [
+                        {"index": 0, "delta": {"function_call": {"name": "get_weather", "arguments": '{"city": "'}}}
+                    ]
+                },
+                {"choices": [{"index": 0, "delta": {"function_call": {"arguments": 'Singapore", "unit": "'}}}]},
+                {"choices": [{"index": 0, "delta": {"function_call": {"arguments": 'celsius"}'}}}]},
+                {"choices": [{"index": 0, "delta": {"content": "done."}}]},
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["content"] == "Starting… done."
+    assert merged["function_call"]["name"] == "get_weather"
+    assert merged["function_call"]["arguments"] == {"city": "Singapore", "unit": "celsius"}
+
+
+def test_openai_tool_calls_via_tool_calls_field():
+    # Newer shape: delta.tool_calls with function.name/arguments segments
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}}]},
+                {
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "tool_calls": [
+                                    {"index": 0, "id": "call_1", "type": "function", "function": {"name": "search"}}
+                                ]
+                            },
+                        }
+                    ]
+                },
+                {
+                    "choices": [
+                        {"index": 0, "delta": {"tool_calls": [{"index": 0, "function": {"arguments": '{"q": "'}}]}}
+                    ]
+                },
+                {
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"tool_calls": [{"index": 0, "function": {"arguments": 'python streaming"}'}}]},
+                        }
+                    ]
+                },
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["function_call"]["name"] == "search"
+    assert merged["function_call"]["arguments"] == {"q": "python streaming"}
+
+
+def test_openai_invalid_json_arguments_falls_back_to_string():
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}}]},
+                {
+                    "choices": [{"index": 0, "delta": {"function_call": {"name": "do", "arguments": '{"bad": '}}}]
+                },  # truncated
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["function_call"]["name"] == "do"
+    # Should be raw string because JSON parsing fails
+    assert isinstance(merged["function_call"]["arguments"], str)
+    assert merged["function_call"]["arguments"].startswith('{"bad": ')
+
+
+def test_anthropic_text_only_multiple_blocks():
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"role": "assistant", "content": [{"type": "text", "text": "Hello "}]},
+                {"content": [{"type": "text", "text": "world!"}]},
+                {"type": "message_delta", "delta": {"stop_reason": "end_turn"}},
+            ],
+        )
+    )
+    merged = merge_anthropic_streaming(chunks)
+    assert merged["role"] == "assistant"
+    assert merged["content_text"] == "Hello world!"
+    assert "tool_calls" not in merged
+
+
+def test_anthropic_tool_use_split_inputs_merge():
+    # Tool input is delivered as multiple content fragments that should be merged
+    chunks = iter(
+        [
+            {"role": "assistant", "content": [{"type": "text", "text": "Working… "}]},
+            {"content": [{"type": "tool_use", "id": "toolu_1", "name": "calculate", "input": {"a": 1}}]},
+            {"content": [{"type": "tool_use", "id": "toolu_1", "name": "calculate", "input": {"b": 2}}]},
+            {"content": [{"type": "text", "text": "done."}]},
+            {"type": "message_stop"},
+        ]
+    )
+    merged = merge_anthropic_streaming(chunks)
+    assert merged["content_text"] == "Working… done."
+    assert merged["tool_calls"][0]["name"] == "calculate"
+    assert merged["tool_calls"][0]["input"] == {"a": 1, "b": 2}
+
+
+def test_anthropic_fine_grained_input_json_delta():
+    # Simulate SSE-style events: content_block_start(tool_use) + multiple input_json_delta pieces
+    chunks = iter(
+        [
+            {
+                "type": "content_block_start",
+                "index": 1,
+                "content_block": {"type": "tool_use", "id": "toolu_x", "name": "fetch"},
+            },
+            {
+                "type": "content_block_delta",
+                "index": 1,
+                "delta": {"type": "input_json_delta", "partial_json": '{"url": "'},
+                "active_tool_id": "toolu_x",
+            },
+            {
+                "type": "content_block_delta",
+                "index": 1,
+                "delta": {"type": "input_json_delta", "partial_json": 'https://example.com"}'},
+                "active_tool_id": "toolu_x",
+            },
+            {"type": "content_block_stop", "index": 1},
+            {"type": "message_stop"},
+        ]
+    )
+    merged = merge_anthropic_streaming(chunks)
+    [tool] = merged["tool_calls"]
+    assert tool["id"] == "toolu_x"
+    assert tool["name"] == "fetch"
+    assert tool["input"] == {"url": "https://example.com"}
+
+
+def test_anthropic_text_and_tool_interleaved_with_text_deltas():
+    # Mix text via text_delta and plain text content items
+    chunks = iter(
+        [
+            {"role": "assistant", "content": [{"type": "text", "text": "Start "}]},
+            {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "middle "}},
+            {"content": [{"type": "text", "text": "end."}]},
+            {"type": "message_stop"},
+        ]
+    )
+    merged = merge_anthropic_streaming(chunks)
+    assert merged["content_text"] == "Start middle end."
+
+
+def test_anthropic_partial_json_left_as_string_when_invalid():
+    # Provide malformed JSON parts; merger should keep raw string for tool input
+    chunks = iter(
+        [
+            {
+                "type": "content_block_start",
+                "index": 2,
+                "content_block": {"type": "tool_use", "id": "toolu_bad", "name": "ingest"},
+            },
+            {
+                "type": "content_block_delta",
+                "index": 2,
+                "delta": {"type": "input_json_delta", "partial_json": '{"alpha": 1, '},
+                "active_tool_id": "toolu_bad",
+            },
+            {
+                "type": "content_block_delta",
+                "index": 2,
+                "delta": {"type": "input_json_delta", "partial_json": '"beta": 2'},
+                "active_tool_id": "toolu_bad",
+            },
+            # missing closing brace
+            {"type": "content_block_stop", "index": 2},
+            {"type": "message_stop"},
+        ]
+    )
+    merged = merge_anthropic_streaming(chunks)
+    [tool] = merged["tool_calls"]
+    assert tool["id"] == "toolu_bad"
+    assert isinstance(tool["input"], str)
+    assert tool["input"].startswith('{"alpha": 1, ')
+
+
+@pytest.mark.parametrize("text_len", [1, 50, 500])
+def test_openai_long_text_stream_rounds_up(text_len: int):
+    # Create a synthetic long content split into ~20-40 char pieces as the merger would see
+    text = "x" * text_len
+    # Simulate content arriving in three chunks
+    part1, part2, part3 = text[: text_len // 3], text[text_len // 3 : 2 * text_len // 3], text[2 * text_len // 3 :]
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}}]},
+                {"choices": [{"index": 0, "delta": {"content": part1}}]},
+                {"choices": [{"index": 0, "delta": {"content": part2}}]},
+                {"choices": [{"index": 0, "delta": {"content": part3}}]},
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["content"] == text
+
+
+async def collect_sse(gen: AsyncGenerator[str, Any]) -> List[str]:
+    """Drain an async generator of SSE strings into a list."""
+    out: List[str] = []
+    async for s in gen:
+        assert isinstance(s, str)
+        out.append(s)
+    return out
+
+
+def parse_openai_sse_to_json_events(sse_chunks: List[str]) -> List[Dict[str, Any]]:
+    """From the OpenAI stream (which uses only 'data:' lines), return JSON events.
+    Filters out the literal DONE sentinel.
+    """
+    events: List[Dict[str, Any]] = []
+    for chunk in sse_chunks:
+        # each chunk looks like 'data: {...}\n\n' OR 'data: [DONE]\n\n'
+        for line in chunk.splitlines():
+            line = line.strip()
+            if not line.startswith("data:"):
+                continue
+            payload = line[len("data:") :].strip()
+            if payload == "[DONE]":
+                continue
+            events.append(json.loads(payload))
+    return events
+
+
+def parse_anthropic_sse_to_json_payloads(sse_chunks: List[str]) -> List[Dict[str, Any]]:
+    """Extract the JSON payload from each Anthropic SSE event (ignore pings)."""
+    out: List[Dict[str, Any]] = []
+    for chunk in sse_chunks:
+        # chunks look like 'event: <name>\ndata: {json}\n\n'
+        if "data:" not in chunk:
+            continue
+        data_line = [ln for ln in chunk.splitlines() if ln.startswith("data:")]
+        if not data_line:
+            continue
+        payload = data_line[0][len("data:") :].strip()
+        obj = json.loads(payload)
+        if obj.get("type") == "ping":
+            continue
+        out.append(obj)
+    return out
+
+
+@pytest.fixture
+def mw() -> StreamConversionMiddleware:
+    # BaseHTTPMiddleware requires an ASGI app; we only need the instance for bound methods.
+    class _DummyApp:
+        async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
+            pass
+
+    return StreamConversionMiddleware(_DummyApp())
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "text, finish_reason",
+    [
+        ("Hello world.", "stop"),
+        ("This answer was cut off on purpose.", "length"),
+    ],
+)
+async def test_openai_content_only_stream_roundtrip(mw: StreamConversionMiddleware, text: str, finish_reason: str):
+    response_json = {
+        "id": "chatcmpl-test",
+        "object": "chat.completion",
+        "model": "gpt-4o-mini",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": text},
+                "finish_reason": finish_reason,
+                # include logprobs to ensure it doesn't interfere with streaming
+                "logprobs": None,
+            }
+        ],
+    }
+
+    sse_chunks = await collect_sse(mw.openai_stream_generator(response_json))
+
+    # basic shape checks
+    assert any('"delta": {"role": ' in s for s in sse_chunks)
+    assert any("[DONE]" in s for s in sse_chunks)
+
+    events = parse_openai_sse_to_json_events(sse_chunks)
+    assert events, "Expected JSON events from stream"
+
+    # the last JSON event before [DONE] should contain the finish_reason
+    last = events[-1]
+    assert last["choices"][0]["finish_reason"] == finish_reason
+
+    merged = merge_openai_streaming(iter(events))
+    assert merged["role"] == "assistant"
+    assert merged["content"] == text
+    assert "function_call" not in merged
+
+
+@pytest.mark.asyncio
+async def test_openai_long_text_chunking_and_reassembly(mw: StreamConversionMiddleware):
+    long_text = """
+        This is a deliberately long sentence that should be broken into multiple streaming deltas by the
+        chunking logic so that we can verify reassembly yields the exact same content without loss. """.strip()
+
+    response_json = {
+        "id": "chatcmpl-long",
+        "object": "chat.completion",
+        "model": "gpt-4o-mini",
+        "choices": [{"index": 0, "message": {"role": "assistant", "content": long_text}, "finish_reason": "stop"}],
+    }
+
+    sse_chunks = await collect_sse(mw.openai_stream_generator(response_json))
+    events = parse_openai_sse_to_json_events(sse_chunks)
+
+    # ensure multiple content delta chunks were emitted
+    content_deltas = [ev for ev in events if ev["choices"][0]["delta"].get("content")]
+    assert len(content_deltas) > 1
+
+    merged = merge_openai_streaming(iter(events))
+    assert merged["content"] == long_text
+
+
+@pytest.mark.asyncio
+async def test_openai_tool_call_only_stream_roundtrip(mw: StreamConversionMiddleware):
+    response_json = {
+        "id": "chatcmpl-tool",
+        "object": "chat.completion",
+        "model": "gpt-4o-mini",
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "",
+                    "tool_calls": [
+                        {
+                            "id": "call_1",
+                            "type": "function",
+                            "function": {
+                                "name": "get_weather",
+                                "arguments": json.dumps({"location": "Boston"}),
+                            },
+                        }
+                    ],
+                },
+                "finish_reason": "tool_calls",
+            }
+        ],
+    }
+
+    sse_chunks = await collect_sse(mw.openai_stream_generator(response_json))
+    events = parse_openai_sse_to_json_events(sse_chunks)
+
+    # expect at least one tool_calls delta with name, followed by deltas with arguments
+    assert any(
+        (tc := ev["choices"][0]["delta"].get("tool_calls")) and tc[0].get("function", {}).get("name") == "get_weather"
+        for ev in events
+    )
+    assert any(
+        (tc := ev["choices"][0]["delta"].get("tool_calls")) and "arguments" in tc[0].get("function", {})
+        for ev in events
+    )
+
+    merged = merge_openai_streaming(iter(events))
+    assert merged["function_call"]["name"] == "get_weather"
+    assert merged["function_call"]["arguments"] == {"location": "Boston"}
+
+
+@pytest.mark.asyncio
+async def test_openai_content_and_tool_call_stream_roundtrip(mw: StreamConversionMiddleware):
+    response_json = {
+        "id": "chatcmpl-mixed",
+        "object": "chat.completion",
+        "model": "gpt-4o-mini",
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "I'll call the weather tool now...",
+                    "tool_calls": [
+                        {
+                            "id": "call_1",
+                            "type": "function",
+                            "function": {
+                                "name": "get_weather",
+                                "arguments": json.dumps({"location": "Singapore", "units": "metric"}),
+                            },
+                        }
+                    ],
+                },
+                "finish_reason": "tool_calls",
+            }
+        ],
+    }
+
+    sse_chunks = await collect_sse(mw.openai_stream_generator(response_json))
+    events = parse_openai_sse_to_json_events(sse_chunks)
+
+    merged = merge_openai_streaming(iter(events))
+    assert merged["content"].startswith("I'll call the weather tool")
+    assert merged["function_call"]["name"] == "get_weather"
+    assert merged["function_call"]["arguments"] == {"location": "Singapore", "units": "metric"}
+
+
+@pytest.mark.asyncio
+async def test_anthropic_text_only_stream_roundtrip(mw: StreamConversionMiddleware):
+    original_response = {
+        "id": "msg_123",
+        "model": "claude-3.5-sonnet",
+        "content": [
+            {"type": "text", "text": "Hello there from Claude."},
+        ],
+        "usage": {"input_tokens": 0, "output_tokens": 7},
+        "stop_reason": "end_turn",
+    }
+
+    sse_chunks = await collect_sse(mw.anthropic_stream_generator(original_response))
+
+    # sanity: stream contains lifecycle events
+    assert any("event: message_start" in s for s in sse_chunks)
+    assert any("event: message_stop" in s for s in sse_chunks)
+
+    payloads = parse_anthropic_sse_to_json_payloads(sse_chunks)
+    merged = merge_anthropic_streaming(iter(payloads))
+
+    assert merged["role"] == "assistant"
+    assert merged["content_text"] == "Hello there from Claude."
+    assert "tool_calls" not in merged
+
+
+@pytest.mark.asyncio
+async def test_anthropic_tool_use_only_stream_roundtrip(mw: StreamConversionMiddleware):
+    original_response = {
+        "id": "msg_tool",
+        "model": "claude-3.5-sonnet",
+        "content": [
+            {
+                "type": "tool_use",
+                "id": "toolu_1",
+                "name": "get_weather",
+                "input": {"location": "Boston"},
+            }
+        ],
+        "usage": {"input_tokens": 0, "output_tokens": 0},
+        "stop_reason": "end_turn",
+    }
+
+    sse_chunks = await collect_sse(mw.anthropic_stream_generator(original_response))
+    payloads = parse_anthropic_sse_to_json_payloads(sse_chunks)
+
+    merged = merge_anthropic_streaming(iter(payloads))
+
+    assert merged["tool_calls"][0]["name"] == "get_weather"
+    assert merged["tool_calls"][0]["id"] == "toolu_1"
+    assert merged["tool_calls"][0]["input"] == {"location": "Boston"}
+
+
+@pytest.mark.asyncio
+async def test_anthropic_mixed_text_and_tool_use_roundtrip(mw: StreamConversionMiddleware):
+    # tool input is long to ensure multiple input_json_delta chunks
+    long_input = {
+        "location": "Singapore",
+        "units": "metric",
+        "details": {"hourly": True, "with_forecast": True, "days": 5},
+    }
+    original_response = {
+        "id": "msg_mixed",
+        "model": "claude-3.5-sonnet",
+        "content": [
+            {"type": "text", "text": "I'll check the weather tool for you."},
+            {"type": "tool_use", "id": "toolu_2", "name": "get_weather", "input": long_input},
+        ],
+        "usage": {"input_tokens": 0, "output_tokens": 0},
+        "stop_reason": "end_turn",
+    }
+
+    sse_chunks = await collect_sse(mw.anthropic_stream_generator(original_response))
+    payloads = parse_anthropic_sse_to_json_payloads(sse_chunks)
+
+    # Verify we saw content_block_start/stop and deltas for both text and tool input
+    types = [p.get("type") for p in payloads]
+    assert "content_block_start" in types
+    assert "content_block_delta" in types
+    assert "content_block_stop" in types
+    assert any(p.get("delta", {}).get("type") == "text_delta" for p in payloads)
+    assert any(p.get("delta", {}).get("type") == "input_json_delta" for p in payloads)
+
+    merged = merge_anthropic_streaming(iter(payloads))
+    assert merged["content_text"].startswith("I'll check the weather tool")
+    tool = merged["tool_calls"][0]
+    assert tool["name"] == "get_weather"
+    assert tool["id"] == "toolu_2"
+    assert tool["input"] == long_input