openhands-sdk 1.8.2__py3-none-any.whl → 1.9.1__py3-none-any.whl

openhands/sdk/critic/impl/api/chat_template.py

@@ -0,0 +1,232 @@
+"""
+Standalone chat template implementation using Jinja2.
+
+This module provides a lightweight implementation of chat template rendering
+that is compatible with HuggingFace transformers but removes the dependency
+on the full transformers library.
+
+The implementation follows the same approach as transformers:
+- Uses Jinja2 for template rendering
+- Loads templates dynamically from tokenizer_config.json
+- Supports caching of compiled templates and fetched configs
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Sequence
+from functools import lru_cache
+from pathlib import Path
+from typing import Any
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+import jinja2
+from jinja2.ext import loopcontrols
+from jinja2.sandbox import ImmutableSandboxedEnvironment
+
+
+# Cache directory for downloaded tokenizer configs
+CACHE_DIR = Path.home() / ".cache" / "chat_templates"
+
+
+def _get_cache_path(tokenizer_name: str) -> Path:
+    """Get the cache path for a tokenizer config."""
+    # Create a safe filename from the tokenizer name
+    safe_name = hashlib.md5(tokenizer_name.encode()).hexdigest()
+    return CACHE_DIR / f"{safe_name}_tokenizer_config.json"
+
+
+def _fetch_tokenizer_config(
+    tokenizer_name: str, use_cache: bool = True
+) -> dict[str, Any]:
+    """
+    Fetch tokenizer_config.json from HuggingFace Hub.
+
+    Args:
+        tokenizer_name: The HuggingFace model/tokenizer name
+            (e.g., "Qwen/Qwen3-4B-Instruct-2507")
+        use_cache: Whether to use cached config if available
+
+    Returns:
+        The parsed tokenizer config dictionary
+    """
+    cache_path = _get_cache_path(tokenizer_name)
+
+    # Try to load from cache
+    if use_cache and cache_path.exists():
+        with open(cache_path, encoding="utf-8") as f:
+            return json.load(f)
+
+    # Fetch from HuggingFace Hub
+    url = f"https://huggingface.co/{tokenizer_name}/raw/main/tokenizer_config.json"
+
+    try:
+        request = Request(url, headers={"User-Agent": "chat_template/1.0"})
+        with urlopen(request, timeout=30) as response:
+            config = json.loads(response.read().decode("utf-8"))
+    except URLError as e:
+        raise RuntimeError(f"Failed to fetch tokenizer config from {url}: {e}")
+
+    # Cache the config
+    if use_cache:
+        CACHE_DIR.mkdir(parents=True, exist_ok=True)
+        with open(cache_path, "w", encoding="utf-8") as f:
+            json.dump(config, f)
+
+    return config
+
+
+@lru_cache(maxsize=16)
+def _compile_jinja_template(chat_template: str) -> jinja2.Template:
+    """
+    Compile a Jinja2 chat template.
+
+    This matches the transformers implementation with custom tojson filter
+    and other utilities.
+    """
+
+    def raise_exception(message: str) -> None:
+        raise jinja2.exceptions.TemplateError(message)
+
+    def tojson(
+        x: Any,
+        ensure_ascii: bool = False,
+        indent: int | None = None,
+        separators: tuple[str, str] | None = None,
+        sort_keys: bool = False,
+    ) -> str:
+        # Match the transformers implementation - no HTML escaping
+        return json.dumps(
+            x,
+            ensure_ascii=ensure_ascii,
+            indent=indent,
+            separators=separators,
+            sort_keys=sort_keys,
+        )
+
+    jinja_env = ImmutableSandboxedEnvironment(
+        trim_blocks=True,
+        lstrip_blocks=True,
+        extensions=[loopcontrols],
+    )
+    jinja_env.filters["tojson"] = tojson
+    jinja_env.globals["raise_exception"] = raise_exception
+
+    return jinja_env.from_string(chat_template)
+
+
+class ChatTemplateRenderer:
+    """
+    A lightweight chat template renderer compatible with HuggingFace transformers.
+
+    This class can dynamically load templates from HuggingFace Hub or use
+    provided templates directly.
+    """
+
+    def __init__(
+        self,
+        tokenizer_name: str | None = None,
+        chat_template: str | None = None,
+        use_cache: bool = True,
+    ):
+        """
+        Initialize the renderer.
+
+        Args:
+            tokenizer_name: HuggingFace tokenizer name to load template from.
+                If provided, will fetch tokenizer_config.json from
+                HuggingFace Hub.
+            chat_template: Direct Jinja2 template string.
+                If provided, tokenizer_name is ignored.
+            use_cache: Whether to cache fetched tokenizer configs.
+        """
+        if chat_template is not None:
+            self._chat_template = chat_template
+        elif tokenizer_name is not None:
+            config = _fetch_tokenizer_config(tokenizer_name, use_cache=use_cache)
+            self._chat_template = config.get("chat_template")
+            if self._chat_template is None:
+                raise ValueError(
+                    f"No chat_template found in tokenizer config for {tokenizer_name}"
+                )
+        else:
+            raise ValueError("Either tokenizer_name or chat_template must be provided")
+
+        self._compiled_template = _compile_jinja_template(self._chat_template)
+
+    @property
+    def chat_template(self) -> str:
+        """The raw Jinja2 chat template string."""
+        assert self._chat_template is not None
+        return self._chat_template
+
+    def apply_chat_template(
+        self,
+        messages: Sequence[dict[str, Any]],
+        tools: Sequence[dict[str, Any]] | None = None,
+        add_generation_prompt: bool = False,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Apply the chat template to format messages.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content' keys.
+            tools: Optional list of tool definitions for function calling.
+            add_generation_prompt: If True, append assistant prompt at the end.
+            **kwargs: Additional template variables.
+
+        Returns:
+            Formatted string ready for tokenization.
+        """
+        return self._compiled_template.render(
+            messages=messages,
+            tools=tools,
+            add_generation_prompt=add_generation_prompt,
+            **kwargs,
+        )
+
+
+# Convenience function for simple use cases
+def apply_chat_template(
+    messages: Sequence[dict[str, Any]],
+    tokenizer_name: str | None = None,
+    chat_template: str | None = None,
+    tools: Sequence[dict[str, Any]] | None = None,
+    add_generation_prompt: bool = False,
+    use_cache: bool = True,
+    **kwargs: Any,
+) -> str:
+    """
+    Apply a chat template to format messages.
+
+    This is a convenience function that creates a renderer and applies the
+    template. For repeated use with the same tokenizer, prefer using
+    ChatTemplateRenderer directly.
+
+    Args:
+        messages: List of message dicts with 'role' and 'content' keys.
+        tokenizer_name: HuggingFace tokenizer name to load template from.
+        chat_template: Direct Jinja2 template string.
+            If provided, tokenizer_name is ignored.
+        tools: Optional list of tool definitions for function calling.
+        add_generation_prompt: If True, append assistant prompt at the end.
+        use_cache: Whether to cache fetched tokenizer configs.
+        **kwargs: Additional template variables.
+
+    Returns:
+        Formatted string ready for tokenization.
+    """
+    renderer = ChatTemplateRenderer(
+        tokenizer_name=tokenizer_name,
+        chat_template=chat_template,
+        use_cache=use_cache,
+    )
+    return renderer.apply_chat_template(
+        messages=messages,
+        tools=tools,
+        add_generation_prompt=add_generation_prompt,
+        **kwargs,
+    )

openhands/sdk/critic/impl/api/client.py

@@ -0,0 +1,313 @@
+import copy
+from collections.abc import Sequence
+from typing import Any, cast
+
+import httpx
+from litellm import ChatCompletionToolParam
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    PrivateAttr,
+    SecretStr,
+    field_validator,
+)
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
+
+from .chat_template import ChatTemplateRenderer
+
+
+# ============================================================
+# Typed API response models
+# ============================================================
+
+
+class UsageTokens(BaseModel):
+    prompt_tokens: int | None = None
+    total_tokens: int | None = None
+    completion_tokens: int | None = None
+    prompt_tokens_details: dict | None = None
+    model_config = ConfigDict(extra="allow")
+
+
+class ClassificationItem(BaseModel):
+    """One per-label or flat classification result."""
+
+    index: int | None = None
+    label: str | None = None
+    probs: list[float]
+    num_classes: int | None = None
+    model_config = ConfigDict(extra="allow")
+
+
+class ClassificationResponse(BaseModel):
+    id: str | None = None
+    object: str | None = None
+    created: int | None = None
+    model: str | None = None
+    data: list[ClassificationItem] = Field(default_factory=list)
+    usage: UsageTokens | None = None
+    model_config = ConfigDict(extra="allow")
+
+
+class LabelProbMap(BaseModel):
+    """Normalized probability map label -> value, with optional ordering."""
+
+    probs: dict[str, float]  # {"label": probability}
+    order: list[str] | None = None  # if you requested a specific order
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================
+# CriticClient
+# ============================================================
+
+
+class CriticClient(BaseModel):
+    """
+    Core inference client for the Critic classification service.
+
+    Owns:
+      - Configuration (server URL, API key, model, tokenizer, etc.)
+      - Label space (for predictions only)
+      - Message normalization and chat template formatting
+      - Inference via vLLM /classify endpoint
+
+    Does NOT handle:
+      - Dataset loading
+      - Ground truth extraction
+      - Evaluation / metrics
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="ignore")
+
+    # --- connection / model config ---
+    server_url: str = Field(
+        default="https://all-hands-ai--critic-qwen3-4b-serve.modal.run",
+        description="Base URL of the vLLM classification service",
+    )
+    api_key: str | SecretStr = Field(
+        ..., description="API key for authenticating with the vLLM service"
+    )
+    model_name: str = Field(
+        default="critic-qwen3-4b", description="Name of the model to use"
+    )
+    tokenizer_name: str = Field(
+        default="Qwen/Qwen3-4B-Instruct-2507",
+        description="HuggingFace tokenizer name for loading chat template",
+    )
+    pass_tools_definitions: bool = Field(
+        default=True, description="Whether to pass tool definitions to the model"
+    )
+    timeout_seconds: float = Field(
+        default=300.0, description="Timeout for requests to the model"
+    )
+    has_success_label: bool = Field(
+        default=True, description="Whether the model predicts success label at index 0"
+    )
+
+    # --- runtime fields ---
+    _client: httpx.Client = PrivateAttr(default_factory=httpx.Client)
+    _template_renderer: ChatTemplateRenderer | None = PrivateAttr(default=None)
+
+    # --- label space ---
+    sentiment_labels: tuple[str, ...] = (
+        "sentiment_positive",
+        "sentiment_neutral",
+        "sentiment_negative",
+    )
+    agent_issue_labels: tuple[str, ...] = (
+        "misunderstood_intention",
+        "did_not_follow_instruction",
+        "insufficient_analysis",
+        "insufficient_clarification",
+        "improper_tool_use_or_setup",
+        "loop_behavior",
+        "insufficient_testing",
+        "insufficient_debugging",
+        "incomplete_implementation",
+        "file_management_errors",
+        "scope_creep",
+        "risky_actions_or_permission",
+        "other_agent_issue",
+    )
+    infra_labels: tuple[str, ...] = (
+        "infrastructure_external_issue",
+        "infrastructure_agent_caused_issue",
+    )
+    user_followup_labels: tuple[str, ...] = (
+        "clarification_or_restatement",
+        "correction",
+        "direction_change",
+        "vcs_update_requests",
+        "progress_or_scope_concern",
+        "frustration_or_complaint",
+        "removal_or_reversion_request",
+        "other_user_issue",
+    )
+    sentiment_map: dict[str, str] = {
+        "Positive": "sentiment_positive",
+        "Neutral": "sentiment_neutral",
+        "Negative": "sentiment_negative",
+    }
+
+    # ---------------------
+    # Validation
+    # ---------------------
+    @field_validator("api_key", mode="before")
+    @classmethod
+    def _validate_and_convert_api_key(cls, v: str | SecretStr) -> SecretStr:
+        """Convert str to SecretStr and validate non-empty."""
+        if isinstance(v, SecretStr):
+            secret_value = v.get_secret_value()
+        else:
+            secret_value = v
+
+        if not secret_value or not secret_value.strip():
+            raise ValueError("api_key must be non-empty")
+
+        return SecretStr(secret_value) if isinstance(v, str) else v
+
+    # ---------------------
+    # Label helpers
+    # ---------------------
+    @property
+    def all_labels(self) -> tuple[str, ...]:
+        base_labels = (
+            self.sentiment_labels
+            + self.agent_issue_labels
+            + self.infra_labels
+            + self.user_followup_labels
+        )
+        if self.has_success_label:
+            return ("success",) + base_labels
+        return base_labels
+
+    # ---------------------
+    # Tokenizer / formatting
+    # ---------------------
+    def _get_template_renderer(self) -> ChatTemplateRenderer:
+        """Lazily initialize the chat template renderer."""
+        if self._template_renderer is None:
+            self._template_renderer = ChatTemplateRenderer(
+                tokenizer_name=self.tokenizer_name
+            )
+        return self._template_renderer
+
+    @staticmethod
+    def normalize_messages(messages: Sequence[dict]) -> Sequence[dict]:
+        """Ensure messages all have string content and flatten text blocks."""
+        out: list[dict] = []
+        for msg in messages or []:
+            content = msg.get("content", "") or ""
+            if isinstance(content, list):
+                text_parts = [
+                    block.get("text", "")
+                    for block in content
+                    if isinstance(block, dict) and block.get("type") == "text"
+                ]
+                content = "\n".join(text_parts)
+            if not isinstance(content, str):
+                content = str(content)
+            out.append({"role": msg.get("role", ""), "content": content})
+        return out
+
+    def apply_chat_template(
+        self,
+        messages: Sequence[dict],
+        tools: Sequence[ChatCompletionToolParam] | None = None,
+    ) -> str:
+        renderer = self._get_template_renderer()
+        msgs = self.normalize_messages(copy.deepcopy(messages))
+        # Cast tools to Sequence[dict[str, Any]] for type compatibility
+        # ChatCompletionToolParam is a TypedDict which is structurally compatible
+        tools_dicts: Sequence[dict[str, Any]] | None = (
+            cast(Sequence[dict[str, Any]], tools) if tools is not None else None
+        )
+        if self.pass_tools_definitions and tools_dicts:
+            return renderer.apply_chat_template(
+                msgs, tools=tools_dicts, add_generation_prompt=False
+            )
+        return renderer.apply_chat_template(msgs, add_generation_prompt=False)
+
+    # ---------------------
+    # Inference
+    # ---------------------
+    def classify_trace(
+        self,
+        messages: Sequence[dict],
+        tools: Sequence[ChatCompletionToolParam] | None = None,
+    ) -> ClassificationResponse:
+        """POST /classify and parse response into ClassificationResponse."""
+        formatted = self.apply_chat_template(messages, tools)
+
+        def should_retry(exc: BaseException) -> bool:
+            # Retry only on 500 Internal Server Error
+            if isinstance(exc, httpx.HTTPStatusError):
+                return exc.response.status_code == 500
+            return False
+
+        @retry(
+            retry=retry_if_exception(should_retry),
+            stop=stop_after_attempt(3),  # up to 3 tries
+            wait=wait_exponential(
+                multiplier=1, min=1, max=8
+            ),  # exponential backoff: 1s, 2s, 4s, 8s
+            reraise=True,  # re-raise the last exception if all retries fail
+        )
+        def _post_with_retry():
+            api_key_value = (
+                self.api_key.get_secret_value()
+                if isinstance(self.api_key, SecretStr)
+                else self.api_key
+            )
+            resp = self._client.post(
+                f"{self.server_url}/classify",
+                headers={
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {api_key_value}",
+                },
+                json={"model": self.model_name, "input": formatted},
+                timeout=self.timeout_seconds,
+            )
+            resp.raise_for_status()
+            return resp
+
+        resp = _post_with_retry()
+        return ClassificationResponse.model_validate(resp.json())
+
+    # ---------------------
+    # Post-processing helpers
+    # ---------------------
+    def extract_prob_map(self, response: ClassificationResponse) -> LabelProbMap:
+        """
+        Server format (flat-only, strict):
+          response.data == [ ClassificationItem(probs=[p0, p1, ..., pN-1],
+                            num_classes=N) ]
+        We align probs directly to self.all_labels (same length, same order).
+        """
+        if not response.data:
+            raise ValueError("empty response.data from server")
+
+        item = response.data[0]
+        if not item.probs:
+            raise ValueError("server returned empty 'probs'")
+        if item.num_classes is not None and item.num_classes != len(item.probs):
+            raise ValueError(
+                f"num_classes ({item.num_classes}) does not match "
+                f"len(probs) ({len(item.probs)})"
+            )
+
+        probs = [float(x) for x in item.probs]
+        if len(probs) != len(self.all_labels):
+            raise ValueError(
+                f"len(probs) ({len(probs)}) != len(all_labels) "
+                f"({len(self.all_labels)}). "
+                "Ensure server label space matches client label space."
+            )
+
+        mapping = {lbl: probs[i] for i, lbl in enumerate(self.all_labels)}
+        return LabelProbMap(probs=mapping, order=list(self.all_labels))
+
+    def predict_labels(self, probs: list[float], threshold: float = 0.5) -> list[int]:
+        return [1 if p > threshold else 0 for p in probs]

openhands/sdk/critic/impl/api/critic.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from openhands.sdk.critic.base import CriticBase, CriticResult
+from openhands.sdk.critic.impl.api.client import CriticClient
+from openhands.sdk.critic.impl.api.taxonomy import categorize_features
+
+
+if TYPE_CHECKING:
+    from openhands.sdk.event import LLMConvertibleEvent, SystemPromptEvent
+
+
+class APIBasedCritic(CriticBase, CriticClient):
+    def evaluate(
+        self,
+        events: Sequence[LLMConvertibleEvent],
+        git_patch: str | None = None,  # noqa: ARG002
+    ) -> CriticResult:
+        # Local imports to avoid circular dependencies during module load
+        from openhands.sdk.context.view import View
+        from openhands.sdk.event import LLMConvertibleEvent, SystemPromptEvent
+
+        system_prompt_event: SystemPromptEvent | None = None
+        tools = []
+        for event in events:
+            if isinstance(event, SystemPromptEvent):
+                system_prompt_event = event
+                tools = event.tools
+                break
+        if system_prompt_event is None:
+            raise ValueError(
+                "SystemPromptEvent is required for APIBasedCritic evaluation"
+            )
+        if not tools:
+            raise ValueError(
+                "APIBasedCritic requires tools to be defined in SystemPromptEvent. "
+                "Ensure your agent configuration includes tool definitions."
+            )
+            raise ValueError("Tools are required for APIBasedCritic evaluation")
+
+        # This will only retain events that are kept by the condenser
+        view = View.from_events(events)
+        llm_convertible_events = view.events
+
+        # Convert events to messages
+        messages = LLMConvertibleEvent.events_to_messages(llm_convertible_events)
+
+        # Serialize messages to dicts for API
+        for message in messages:
+            message.cache_enabled = False
+            message.vision_enabled = False  # Critic does not support vision currently
+            message.function_calling_enabled = True
+            message.force_string_serializer = False
+            message.send_reasoning_content = False
+        formatted_messages = [message.to_chat_dict() for message in messages]
+
+        # Convert ToolDefinition objects to ChatCompletionToolParam format
+        tools_for_api = [tool.to_openai_tool() for tool in tools]
+        response = self.classify_trace(formatted_messages, tools_for_api)
+        prob_map = self.extract_prob_map(response)
+
+        explanation = []
+
+        if "success" not in prob_map.probs:
+            raise ValueError("APIBasedCritic requires 'success' label in the response.")
+
+        score = prob_map.probs["success"]
+        explanation.append(f"Success: {score:.2f}")
+
+        # Add top labels to explanation
+        sorted_probs = sorted(prob_map.probs.items(), key=lambda x: x[1], reverse=True)
+        explanation.append(json.dumps(dict(sorted_probs)))
+
+        # Collect event IDs for reproducibility
+        event_ids = [event.id for event in llm_convertible_events]
+
+        # Categorize features for visualization
+        categorized = categorize_features(prob_map.probs)
+
+        return CriticResult(
+            score=score,
+            message="; ".join(explanation),
+            metadata={
+                "event_ids": event_ids,
+                "categorized_features": categorized,
+            },
+        )