datarobot-genai 0.1.74__tar.gz → 0.1.75__tar.gz

{datarobot_genai-0.1.74 → datarobot_genai-0.1.75}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.74
+Version: 0.1.75
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.

{datarobot_genai-0.1.74 → datarobot_genai-0.1.75}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.1.74"
+version = "0.1.75"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"

{datarobot_genai-0.1.74 → datarobot_genai-0.1.75}/src/datarobot_genai/core/chat/responses.py

@@ -14,13 +14,18 @@
 
 """OpenAI-compatible response helpers for chat interactions."""
 
+import asyncio
+import queue
 import time
 import traceback as tb
 import uuid
 from asyncio import AbstractEventLoop
 from collections.abc import AsyncGenerator
+from collections.abc import AsyncIterator
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+from typing import TypeVar
 
 from openai.types import CompletionUsage
 from openai.types.chat import ChatCompletion
@@ -168,3 +173,97 @@ def to_custom_model_streaming_response(
             choices=[choice],
             usage=None,
         )
+
+
+def streaming_iterator_to_custom_model_streaming_response(
+    streaming_response_iterator: Iterator[tuple[str, MultiTurnSample | None, dict[str, int]]],
+    model: str | object | None,
+) -> Iterator[CustomModelStreamingResponse]:
+    """Convert the OpenAI ChatCompletionChunk response to CustomModelStreamingResponse."""
+    completion_id = str(uuid.uuid4())
+    created = int(time.time())
+
+    last_pipeline_interactions = None
+    last_usage_metrics = None
+
+    while True:
+        try:
+            (
+                response_text,
+                pipeline_interactions,
+                usage_metrics,
+            ) = next(streaming_response_iterator)
+            last_pipeline_interactions = pipeline_interactions
+            last_usage_metrics = usage_metrics
+
+            if response_text:
+                choice = ChunkChoice(
+                    index=0,
+                    delta=ChoiceDelta(role="assistant", content=response_text),
+                    finish_reason=None,
+                )
+                yield CustomModelStreamingResponse(
+                    id=completion_id,
+                    object="chat.completion.chunk",
+                    created=created,
+                    model=model,
+                    choices=[choice],
+                    usage=CompletionUsage(**usage_metrics) if usage_metrics else None,
+                )
+        except StopIteration:
+            break
+    # Yield final chunk indicating end of stream
+    choice = ChunkChoice(
+        index=0,
+        delta=ChoiceDelta(role="assistant"),
+        finish_reason="stop",
+    )
+    yield CustomModelStreamingResponse(
+        id=completion_id,
+        object="chat.completion.chunk",
+        created=created,
+        model=model,
+        choices=[choice],
+        usage=CompletionUsage(**last_usage_metrics) if last_usage_metrics else None,
+        pipeline_interactions=last_pipeline_interactions.model_dump_json()
+        if last_pipeline_interactions
+        else None,
+    )
+
+
+T = TypeVar("T")
+
+
+def async_gen_to_sync_thread(
+    async_iterator: AsyncIterator[T],
+    thread_pool_executor: ThreadPoolExecutor,
+    event_loop: asyncio.AbstractEventLoop,
+) -> Iterator[T]:
+    """Run an async iterator in a separate thread and provide a sync iterator."""
+    # A thread-safe queue for communication
+    sync_queue: queue.Queue[Any] = queue.Queue()
+    # A sentinel object to signal the end of the async generator
+    SENTINEL = object()  # noqa: N806
+
+    async def run_async_to_queue() -> None:
+        """Run in the separate thread's event loop."""
+        try:
+            async for item in async_iterator:
+                sync_queue.put(item)
+        except Exception as e:
+            # Put the exception on the queue to be re-raised in the main thread
+            sync_queue.put(e)
+        finally:
+            # Signal the end of iteration
+            sync_queue.put(SENTINEL)
+
+    thread_pool_executor.submit(event_loop.run_until_complete, run_async_to_queue()).result()
+
+    # The main thread consumes items synchronously
+    while True:
+        item = sync_queue.get()
+        if item is SENTINEL:
+            break
+        if isinstance(item, Exception):
+            raise item
+        yield item