datarobot-genai 0.2.0__py3-none-any.whl

datarobot_genai/core/chat/responses.py

@@ -0,0 +1,297 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""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 ag_ui.core import BaseEvent
+from ag_ui.core import Event
+from ag_ui.core import TextMessageChunkEvent
+from ag_ui.core import TextMessageContentEvent
+from openai.types import CompletionUsage
+from openai.types.chat import ChatCompletion
+from openai.types.chat import ChatCompletionChunk
+from openai.types.chat import ChatCompletionMessage
+from openai.types.chat.chat_completion import Choice
+from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
+from openai.types.chat.chat_completion_chunk import ChoiceDelta
+from ragas import MultiTurnSample
+
+from datarobot_genai.core.agents import default_usage_metrics
+
+
+class CustomModelChatResponse(ChatCompletion):
+    pipeline_interactions: str | None = None
+
+
+class CustomModelStreamingResponse(ChatCompletionChunk):
+    pipeline_interactions: str | None = None
+    event: Event | None = None
+
+
+def to_custom_model_chat_response(
+    response_text: str,
+    pipeline_interactions: MultiTurnSample | None,
+    usage_metrics: dict[str, int],
+    model: str | object | None,
+) -> CustomModelChatResponse:
+    """Convert the OpenAI ChatCompletion response to CustomModelChatResponse."""
+    choice = Choice(
+        index=0,
+        message=ChatCompletionMessage(role="assistant", content=response_text),
+        finish_reason="stop",
+    )
+
+    if model is None:
+        model = "unspecified-model"
+    else:
+        model = str(model)
+
+    required_usage_metrics: dict[str, int] = {
+        "completion_tokens": 0,
+        "prompt_tokens": 0,
+        "total_tokens": 0,
+    }
+
+    return CustomModelChatResponse(
+        id=str(uuid.uuid4()),
+        object="chat.completion",
+        choices=[choice],
+        created=int(time.time()),
+        model=model,
+        usage=CompletionUsage.model_validate(required_usage_metrics | usage_metrics),
+        pipeline_interactions=pipeline_interactions.model_dump_json()
+        if pipeline_interactions
+        else None,
+    )
+
+
+def to_custom_model_streaming_response(
+    thread_pool_executor: ThreadPoolExecutor,
+    event_loop: AbstractEventLoop,
+    streaming_response_generator: AsyncGenerator[
+        tuple[str | Event, MultiTurnSample | None, dict[str, int]], None
+    ],
+    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
+
+    if model is None:
+        model = "unspecified-model"
+    else:
+        model = str(model)
+
+    required_usage_metrics = default_usage_metrics()
+    try:
+        agent_response = aiter(streaming_response_generator)
+        while True:
+            try:
+                (
+                    response_text_or_event,
+                    pipeline_interactions,
+                    usage_metrics,
+                ) = thread_pool_executor.submit(
+                    event_loop.run_until_complete, anext(agent_response)
+                ).result()
+                last_pipeline_interactions = pipeline_interactions
+                last_usage_metrics = usage_metrics
+
+                if isinstance(response_text_or_event, str) and response_text_or_event:
+                    choice = ChunkChoice(
+                        index=0,
+                        delta=ChoiceDelta(role="assistant", content=response_text_or_event),
+                        finish_reason=None,
+                    )
+                    yield CustomModelStreamingResponse(
+                        id=completion_id,
+                        object="chat.completion.chunk",
+                        created=created,
+                        model=model,
+                        choices=[choice],
+                        usage=CompletionUsage.model_validate(required_usage_metrics | usage_metrics)
+                        if usage_metrics
+                        else None,
+                    )
+                elif isinstance(response_text_or_event, BaseEvent):
+                    content = ""
+                    if isinstance(
+                        response_text_or_event, (TextMessageContentEvent, TextMessageChunkEvent)
+                    ):
+                        content = response_text_or_event.delta or content
+                    choice = ChunkChoice(
+                        index=0,
+                        delta=ChoiceDelta(role="assistant", content=content),
+                        finish_reason=None,
+                    )
+
+                    yield CustomModelStreamingResponse(
+                        id=completion_id,
+                        object="chat.completion.chunk",
+                        created=created,
+                        model=model,
+                        choices=[choice],
+                        usage=CompletionUsage.model_validate(required_usage_metrics | usage_metrics)
+                        if usage_metrics
+                        else None,
+                        event=response_text_or_event,
+                    )
+            except StopAsyncIteration:
+                break
+        event_loop.run_until_complete(streaming_response_generator.aclose())
+        # 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.model_validate(required_usage_metrics | last_usage_metrics)
+            if last_usage_metrics
+            else None,
+            pipeline_interactions=last_pipeline_interactions.model_dump_json()
+            if last_pipeline_interactions
+            else None,
+        )
+    except Exception as e:
+        tb.print_exc()
+        created = int(time.time())
+        choice = ChunkChoice(
+            index=0,
+            delta=ChoiceDelta(role="assistant", content=str(e), refusal="error"),
+            finish_reason="stop",
+        )
+        yield CustomModelStreamingResponse(
+            id=completion_id,
+            object="chat.completion.chunk",
+            created=created,
+            model=model,
+            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

datarobot_genai/core/cli/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .agent_environment import AgentEnvironment
+from .agent_kernel import AgentKernel
+
+__all__ = ["AgentEnvironment", "AgentKernel"]

datarobot_genai/core/cli/agent_environment.py

@@ -0,0 +1,47 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from datarobot_genai.core.cli.agent_kernel import AgentKernel
+
+
+class AgentEnvironment:
+    def __init__(
+        self,
+        api_token: str | None = None,
+        base_url: str | None = None,
+    ):
+        self.api_token = os.environ.get("DATAROBOT_API_TOKEN") or api_token
+        if not self.api_token:
+            raise ValueError(
+                "Missing DataRobot API token. Set the DATAROBOT_API_TOKEN "
+                "environment variable or provide it explicitly."
+            )
+        self.base_url = (
+            os.environ.get("DATAROBOT_ENDPOINT") or base_url or "https://app.datarobot.com"
+        )
+        if not self.base_url:
+            raise ValueError(
+                "Missing DataRobot endpoint. Set the DATAROBOT_ENDPOINT environment "
+                "variable or provide it explicitly."
+            )
+        self.base_url = self.base_url.replace("/api/v2", "")
+
+    @property
+    def interface(self) -> AgentKernel:
+        return AgentKernel(
+            api_token=str(self.api_token),
+            base_url=str(self.base_url),
+        )

datarobot_genai/core/cli/agent_kernel.py

@@ -0,0 +1,211 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+import time
+from typing import Any
+from typing import cast
+
+import requests
+from openai import OpenAI
+from openai import Stream
+from openai.types.chat import ChatCompletion
+from openai.types.chat import ChatCompletionChunk
+from openai.types.chat import ChatCompletionSystemMessageParam
+from openai.types.chat import ChatCompletionUserMessageParam
+from openai.types.chat.completion_create_params import CompletionCreateParamsNonStreaming
+from openai.types.chat.completion_create_params import CompletionCreateParamsStreaming
+
+
+class AgentKernel:
+    def __init__(
+        self,
+        api_token: str,
+        base_url: str,
+    ):
+        self.base_url = base_url
+        self.api_token = api_token
+
+    @property
+    def headers(self) -> dict[str, str]:
+        return {
+            "Authorization": f"Token {self.api_token}",
+        }
+
+    def load_completion_json(self, completion_json: str) -> CompletionCreateParamsNonStreaming:
+        """Load the completion JSON from a file or return an empty prompt."""
+        if not os.path.exists(completion_json):
+            raise FileNotFoundError(f"Completion JSON file not found: {completion_json}")
+
+        with open(completion_json) as f:
+            completion_data = json.load(f)
+
+        completion_create_params = CompletionCreateParamsNonStreaming(
+            **completion_data,  # type: ignore[typeddict-item]
+        )
+        return cast(CompletionCreateParamsNonStreaming, completion_create_params)
+
+    def construct_prompt(
+        self, user_prompt: str, verbose: bool, stream: bool = False
+    ) -> CompletionCreateParamsNonStreaming | CompletionCreateParamsStreaming:
+        extra_body = {
+            "api_key": self.api_token,
+            "api_base": self.base_url,
+            "verbose": verbose,
+        }
+        if stream:
+            return CompletionCreateParamsStreaming(
+                model="datarobot-deployed-llm",
+                messages=[
+                    ChatCompletionSystemMessageParam(
+                        content="You are a helpful assistant",
+                        role="system",
+                    ),
+                    ChatCompletionUserMessageParam(
+                        content=user_prompt,
+                        role="user",
+                    ),
+                ],
+                n=1,
+                temperature=1,
+                stream=True,
+                extra_body=extra_body,  # type: ignore[typeddict-unknown-key]
+            )
+        else:
+            return CompletionCreateParamsNonStreaming(
+                model="datarobot-deployed-llm",
+                messages=[
+                    ChatCompletionSystemMessageParam(
+                        content="You are a helpful assistant",
+                        role="system",
+                    ),
+                    ChatCompletionUserMessageParam(
+                        content=user_prompt,
+                        role="user",
+                    ),
+                ],
+                n=1,
+                temperature=1,
+                extra_body=extra_body,  # type: ignore[typeddict-unknown-key]
+            )
+
+    def local(
+        self,
+        user_prompt: str,
+        completion_json: str = "",
+        stream: bool = False,
+        config: Any | None = None,
+    ) -> ChatCompletion | Stream[ChatCompletionChunk]:
+        chat_api_url = config.agent_endpoint if config else self.base_url
+        print(chat_api_url)
+
+        return self._do_chat_completion(chat_api_url, user_prompt, completion_json, stream=stream)
+
+    def custom_model(self, custom_model_id: str, user_prompt: str, timeout: float = 300) -> str:
+        chat_api_url = (
+            f"{self.base_url}/api/v2/genai/agents/fromCustomModel/{custom_model_id}/chat/"
+        )
+        print(chat_api_url)
+
+        headers = {
+            "Authorization": f"Bearer {os.environ['DATAROBOT_API_TOKEN']}",
+            "Content-Type": "application/json",
+        }
+        data = {"messages": [{"role": "user", "content": user_prompt}]}
+
+        print(f'Querying custom model with prompt: "{data}"')
+        print(
+            "Please wait... This may take 1-2 minutes the first time "
+            "you run this as a codespace is provisioned "
+            "for the custom model to execute."
+        )
+        response = requests.post(
+            chat_api_url,
+            headers=headers,
+            json=data,
+        )
+
+        if not response.ok or not response.headers.get("Location"):
+            raise Exception(response.text)
+        # Wait for the agent to complete
+        status_location = response.headers["Location"]
+        while response.ok:
+            time.sleep(1)
+            response = requests.get(
+                status_location, headers=headers, allow_redirects=False, timeout=timeout
+            )
+            if response.status_code == 303:
+                agent_response = requests.get(response.headers["Location"], headers=headers).json()
+                # Show the agent response
+                break
+            else:
+                status_response = response.json()
+                if status_response["status"] in ["ERROR", "ABORTED"]:
+                    raise Exception(status_response)
+        else:
+            raise Exception(response.content)
+
+        if "errorMessage" in agent_response and agent_response["errorMessage"]:
+            return (
+                f"Error: "
+                f"{agent_response.get('errorMessage', 'No error message available')}"
+                f"Error details:"
+                f"{agent_response.get('errorDetails', 'No details available')}"
+            )
+        elif "choices" in agent_response:
+            return str(agent_response["choices"][0]["message"]["content"])
+        else:
+            return str(agent_response)
+
+    def deployment(
+        self,
+        deployment_id: str,
+        user_prompt: str,
+        completion_json: str = "",
+        stream: bool = False,
+    ) -> ChatCompletion | Stream[ChatCompletionChunk]:
+        chat_api_url = f"{self.base_url}/api/v2/deployments/{deployment_id}/"
+        print(chat_api_url)
+
+        return self._do_chat_completion(chat_api_url, user_prompt, completion_json, stream=stream)
+
+    def _do_chat_completion(
+        self,
+        url: str,
+        user_prompt: str,
+        completion_json: str = "",
+        stream: bool = False,
+    ) -> ChatCompletion | Stream[ChatCompletionChunk]:
+        if len(user_prompt) > 0:
+            completion_create_params = self.construct_prompt(
+                user_prompt, stream=stream, verbose=True
+            )
+        else:
+            completion_create_params = self.load_completion_json(completion_json)
+
+        openai_client = OpenAI(
+            base_url=url,
+            api_key=self.api_token,
+            _strict_response_validation=False,
+        )
+
+        print(f'Querying deployment with prompt: "{completion_create_params}"')
+        print(
+            "Please wait for the agent to complete the response. "
+            "This may take a few seconds to minutes "
+            "depending on the complexity of the agent workflow."
+        )
+
+        completion = openai_client.chat.completions.create(**completion_create_params)
+        return completion