mantisdk 0.1.0__py3-none-any.whl

mantisdk/types/resources.py

@@ -0,0 +1,204 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+"""Typed representations of tunable resources shared between Mantisdk components."""
+
+import inspect
+import logging
+from typing import (
+    Annotated,
+    Any,
+    Dict,
+    Literal,
+    Optional,
+    Union,
+)
+
+from pydantic import BaseModel, Field
+
+from .core import AttemptedRollout
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "Resource",
+    "LLM",
+    "ProxyLLM",
+    "PromptTemplate",
+    "ResourceUnion",
+    "NamedResources",
+    "ResourcesUpdate",
+]
+
+
+class Resource(BaseModel):
+    """Base class for tunable resources distributed to executors."""
+
+    resource_type: Any
+    """Alias of the resource type."""
+
+
+class LLM(Resource):
+    """Resource that identifies an LLM endpoint and its configuration."""
+
+    resource_type: Literal["llm"] = "llm"
+    endpoint: str
+    """The URL of the LLM API endpoint."""
+    model: str
+    """The identifier for the model to be used (e.g., 'gpt-4o')."""
+    api_key: Optional[str] = None
+    """Optional secret used to authenticate requests."""
+    sampling_parameters: Dict[str, Any] = Field(default_factory=dict)
+    """A dictionary of hyperparameters for model inference, such as temperature, top_p, etc."""
+
+    def get_base_url(self, *args: Any, **kwargs: Any) -> str:
+        """Return the base URL consumed by OpenAI-compatible clients.
+
+        Users are encouraged to use `get_base_url(rollout_id, attempt_id)` to get
+        the LLM endpoint instead of accessing `.endpoint` directly.
+        """
+        return self.endpoint
+
+
+class ProxyLLM(LLM):
+    """LLM resource that rewrites endpoints through [`LLMProxy`][mantisdk.LLMProxy].
+
+    The proxy injects rollout- and attempt-specific routing information into the
+    endpoint so that downstream services can attribute requests correctly.
+    """
+
+    resource_type: Literal["proxy_llm"] = "proxy_llm"  # type: ignore
+    _initialized: bool = False
+
+    def model_post_init(self, __context: Any) -> None:
+        """Mark initialization as complete after Pydantic finishes setup."""
+        super().model_post_init(__context)
+        object.__setattr__(self, "_initialized", True)
+
+    def __getattribute__(self, name: str) -> Any:
+        """Emit a warning when `endpoint` is accessed directly after initialization."""
+        # Check if we're accessing endpoint after initialization and not from base_url
+        if name == "endpoint":
+            try:
+                initialized = object.__getattribute__(self, "_initialized")
+            except AttributeError:
+                initialized = False
+
+            if initialized:
+                # Check the call stack to see if we're being called from base_url
+                frame = inspect.currentframe()
+                if frame and frame.f_back:
+                    caller_name = frame.f_back.f_code.co_name
+                    if caller_name != "get_base_url":
+                        logger.warning(
+                            "Accessing 'endpoint' directly on ProxyLLM is discouraged. "
+                            "Use 'get_base_url(rollout_id, attempt_id)' instead to get the properly formatted endpoint."
+                        )
+        return super().__getattribute__(name)
+
+    def with_attempted_rollout(self, rollout: AttemptedRollout) -> LLM:
+        """Bake rollout metadata into a concrete [`LLM`][mantisdk.LLM] instance."""
+        return LLM(
+            endpoint=self.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+            model=self.model,
+            sampling_parameters=self.sampling_parameters,
+            api_key=self.api_key,
+        )
+
+    def get_base_url(self, rollout_id: Optional[str], attempt_id: Optional[str]) -> str:
+        """Return the routed endpoint for a specific rollout/attempt pair.
+
+        Args:
+            rollout_id: Identifier of the rollout making the request.
+            attempt_id: Identifier of the attempt within that rollout.
+
+        Returns:
+            Fully qualified endpoint including rollout metadata.
+
+        Raises:
+            ValueError: If exactly one of ``rollout_id`` or ``attempt_id`` is provided.
+        """
+        if rollout_id is None and attempt_id is None:
+            return self.endpoint
+
+        if not (isinstance(rollout_id, str) and isinstance(attempt_id, str)):
+            raise ValueError("rollout_id and attempt_id must be strings or all be empty")
+
+        prefix = self.endpoint
+        if prefix.endswith("/"):
+            prefix = prefix[:-1]
+        if prefix.endswith("/v1"):
+            prefix = prefix[:-3]
+            has_v1 = True
+        else:
+            has_v1 = False
+        # Now the prefix should look like "http://localhost:11434"
+
+        # Append the rollout and attempt id to the prefix
+        prefix = prefix + f"/rollout/{rollout_id}/attempt/{attempt_id}"
+        if has_v1:
+            prefix = prefix + "/v1"
+        return prefix
+
+
+class PromptTemplate(Resource):
+    """Resource describing a reusable prompt template."""
+
+    resource_type: Literal["prompt_template"] = "prompt_template"
+    template: str
+    """The template string. The format depends on the engine."""
+    engine: Literal["jinja", "f-string", "poml"]
+    """The templating engine to use for rendering the prompt."""
+
+    def format(self, **kwargs: Any) -> str:
+        """Format the prompt using keyword arguments.
+
+        !!! warning
+            Only the `f-string` engine is supported for now.
+        """
+        if self.engine == "f-string":
+            return self.template.format(**kwargs)
+        else:
+            raise NotImplementedError(
+                "Formatting prompt templates for non-f-string engines with format() helper is not supported yet."
+            )
+
+
+# Use discriminated union for proper deserialization
+# TODO: migrate to use a registry
+ResourceUnion = Annotated[Union[LLM, ProxyLLM, PromptTemplate], Field(discriminator="resource_type")]
+NamedResources = Dict[str, ResourceUnion]
+"""Mapping from resource names to their configured instances.
+
+Examples:
+    ```python
+    resources: NamedResources = {
+        "main_llm": LLM(
+            endpoint="http://localhost:8080",
+            model="llama3",
+            sampling_parameters={"temperature": 0.7, "max_tokens": 100},
+        ),
+        "system_prompt": PromptTemplate(
+            template="You are a helpful assistant.",
+            engine="f-string",
+        ),
+    }
+    ```
+"""
+
+
+class ResourcesUpdate(BaseModel):
+    """Update payload broadcast to clients when resources change."""
+
+    resources_id: str
+    """Identifier used to version the resources."""
+    create_time: float
+    """Timestamp of the creation time of the resources."""
+    update_time: float
+    """Timestamp of the last update time of the resources."""
+    version: int
+    """Version of the resources."""
+    resources: NamedResources
+    """Mapping of resource names to their definitions."""