@agenthippo-ai/agenthippo 0.1.0

package/out/scripts/agentide_litellm_hooks_v1.py

@@ -0,0 +1,604 @@
+#---------------------------------------------------------------------------------------------
+# Copyright (c) AgentHippo.ai. All rights reserved.
+#---------------------------------------------------------------------------------------------
+"""
+AgentIDE LiteLLM proxy hooks - Clean structured assistant output with separate input/output
+Mimics langfuse_otel style: gen_ai.input.messages (history) + gen_ai.output.messages (new reply only)
+
+Usage in litellm config.yaml:
+  callbacks: ["agentide_litellm_hooks_v1.clean_separated_otel_output"]
+"""
+
+from __future__ import annotations
+from typing import Any
+import json
+
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.opentelemetry import OpenTelemetry
+from opentelemetry.trace import Status, StatusCode
+
+_OTEL_TEXT_MAX_LEN = 4000           # per text block
+_OTEL_PREVIEW_LEN = 180             # short preview in event
+
+
+def _read_field(obj: Any, key: str, default: Any = None) -> Any:
+    if isinstance(obj, dict):
+        return obj.get(key, default)
+    return getattr(obj, key, default)
+
+
+def _content_blocks_to_text(content: Any) -> str:
+    if isinstance(content, str):
+        return content
+    if not isinstance(content, (list, tuple)):
+        return ""
+    parts: list[str] = []
+    for block in content:
+        if isinstance(block, str):
+            if block.strip():
+                parts.append(block)
+            continue
+        if not isinstance(block, dict):
+            continue
+        text = block.get("text") or block.get("output_text") or block.get("content")
+        if isinstance(text, str) and text.strip():
+            parts.append(text)
+    return "\n".join(parts)
+
+
+def _extract_assistant_text(response_obj: Any) -> str:
+    if response_obj is None:
+        return ""
+    # OpenAI-style
+    choices = _read_field(response_obj, "choices")
+    if isinstance(choices, list) and choices:
+        msg = _read_field(choices[0], "message")
+        content = _read_field(msg, "content")
+        text = _content_blocks_to_text(content)
+        if text.strip():
+            return text
+    # Direct output fields
+    for key in ("output_text", "content", "output", "text"):
+        val = _read_field(response_obj, key)
+        if val:
+            text = _content_blocks_to_text(val)
+            if text.strip():
+                return text
+    return ""
+
+
+def _extract_text_from_success_kwargs(kwargs: dict) -> str:
+    text = _extract_assistant_text(kwargs.get("response_obj"))
+    if text:
+        return text
+    for k in ("complete_response", "model_response", "response", "original_response"):
+        candidate = kwargs.get(k)
+        if candidate:
+            text = _extract_assistant_text(candidate)
+            if text:
+                return text
+    return ""
+
+
+def _make_output_messages_structured(text: str) -> list:
+    if not text.strip():
+        return []
+    trimmed = text[:_OTEL_TEXT_MAX_LEN]
+    # Align with LiteLLM's GEN_AI_OUTPUT_MESSAGES shape: parts + inner content (not content + text).
+    return [{"role": "assistant", "parts": [{"type": "text", "content": trimmed}]}]
+
+
+def _model_name(model: Any) -> str:
+    return model.lower() if isinstance(model, str) else ""
+
+
+def _is_reasoning_unsupported_model(model: Any) -> bool:
+    """True when model is known to reject reasoning parameters."""
+    return "gpt-4.1-nano" in _model_name(model)
+
+
+def _is_anthropic_route(model: Any) -> bool:
+    """True when the request will be forwarded to Anthropic."""
+    name = _model_name(model)
+    return "claude" in name or "anthropic" in name
+
+
+def _is_web_search_unsupported_model(model: Any) -> bool:
+    """True when the upstream model rejects Codex web-search tools."""
+    return "gpt-4.1-nano" in _model_name(model)
+
+
+# OpenAI Responses API fields that Anthropic (and some gateways) reject if forwarded verbatim.
+_OPENAI_RESPONSES_ONLY_KEYS = ("client_metadata",)
+
+# Codex built-in tools not supported on all OpenAI models.
+_UNSUPPORTED_TOOL_IDENTIFIERS = frozenset({"web_search_preview", "web_search"})
+
+
+def _strip_key_recursive(data: Any, key: str, _visited: set[int] | None = None) -> None:
+    """Remove a key in nested dict/list payloads (bounded; avoids cyclic graphs)."""
+    if _visited is None:
+        _visited = set()
+    if isinstance(data, dict):
+        obj_id = id(data)
+        if obj_id in _visited:
+            return
+        _visited.add(obj_id)
+        data.pop(key, None)
+        for value in list(data.values()):
+            _strip_key_recursive(value, key, _visited)
+        return
+    if isinstance(data, list):
+        obj_id = id(data)
+        if obj_id in _visited:
+            return
+        _visited.add(obj_id)
+        for item in data:
+            _strip_key_recursive(item, key, _visited)
+
+
+def _resolve_route_model(payload: dict, model: Any) -> str:
+    name = _model_name(model)
+    if name:
+        return name
+    for candidate in (
+        payload.get("model"),
+        payload.get("model_group"),
+        payload.get("model_name"),
+        payload.get("deployment"),
+    ):
+        name = _model_name(candidate)
+        if name:
+            return name
+    metadata = payload.get("metadata")
+    if isinstance(metadata, dict):
+        for candidate in (metadata.get("model_group"), metadata.get("model")):
+            name = _model_name(candidate)
+            if name:
+                return name
+    return ""
+
+
+def _tool_is_unsupported(tool: Any) -> bool:
+    if not isinstance(tool, dict):
+        return False
+    tool_type = tool.get("type")
+    if isinstance(tool_type, str) and tool_type in _UNSUPPORTED_TOOL_IDENTIFIERS:
+        return True
+    name = tool.get("name")
+    if isinstance(name, str) and name in _UNSUPPORTED_TOOL_IDENTIFIERS:
+        return True
+    fn = tool.get("function")
+    if isinstance(fn, dict):
+        fn_name = fn.get("name")
+        if isinstance(fn_name, str) and fn_name in _UNSUPPORTED_TOOL_IDENTIFIERS:
+            return True
+    return False
+
+
+def _strip_unsupported_tools(data: dict) -> None:
+    tools = data.get("tools")
+    if not isinstance(tools, list):
+        return
+    filtered = [tool for tool in tools if not _tool_is_unsupported(tool)]
+    if len(filtered) == len(tools):
+        return
+    if filtered:
+        data["tools"] = filtered
+    else:
+        data.pop("tools", None)
+
+
+def _apply_codex_proxy_sanitizers(data: dict, model: Any) -> None:
+    route_model = _resolve_route_model(data, model)
+    if _is_anthropic_route(route_model):
+        for key in _OPENAI_RESPONSES_ONLY_KEYS:
+            _strip_key_recursive(data, key)
+    if _is_web_search_unsupported_model(route_model):
+        _strip_unsupported_tools(data)
+
+
+def _strip_unsupported_reasoning_params(data: Any, _visited: set[int] | None = None) -> None:
+    """
+    Mutates `data` in-place by removing unsupported reasoning controls.
+
+    OpenAI's API error reports `reasoning.effort`, but the payload is sent as:
+      { "reasoning": { "effort": ... } }
+    """
+    if _visited is None:
+        _visited = set()
+    if isinstance(data, dict):
+        obj_id = id(data)
+        if obj_id in _visited:
+            return
+        _visited.add(obj_id)
+        data.pop("reasoning", None)
+        data.pop("reasoning_effort", None)
+        for v in list(data.values()):
+            _strip_unsupported_reasoning_params(v, _visited)
+        return
+    if isinstance(data, list):
+        obj_id = id(data)
+        if obj_id in _visited:
+            return
+        _visited.add(obj_id)
+        for item in data:
+            _strip_unsupported_reasoning_params(item, _visited)
+
+
+def _extract_text_content(message: dict[str, Any]) -> str:
+    """Extract text content from OpenAI-style message content."""
+    content = message.get("content")
+    if isinstance(content, str):
+        return content
+    if not isinstance(content, list):
+        return ""
+    parts: list[str] = []
+    for block in content:
+        if isinstance(block, dict) and block.get("type") == "text":
+            text = block.get("text")
+            if isinstance(text, str) and text.strip():
+                parts.append(text)
+    return "\n".join(parts)
+
+
+def _set_text_content(message: dict[str, Any], text: str) -> None:
+    """Set text content while preserving existing content shape when possible."""
+    content = message.get("content")
+    if isinstance(content, list):
+        # Keep non-text blocks and append merged text as final text block.
+        preserved = [b for b in content if not (isinstance(b, dict) and b.get("type") == "text")]
+        if text.strip():
+            preserved.append({"type": "text", "text": text})
+        message["content"] = preserved
+        return
+    message["content"] = text
+
+
+def _is_text_only_assistant(message: Any) -> bool:
+    """True when assistant message has text and no tool calls."""
+    if not isinstance(message, dict) or message.get("role") != "assistant":
+        return False
+    tool_calls = message.get("tool_calls")
+    if isinstance(tool_calls, list) and len(tool_calls) > 0:
+        return False
+    return bool(_extract_text_content(message).strip())
+
+
+def _tool_result_ids_from_message(message: Any) -> set[str]:
+    """Collect tool result ids from OpenAI or Anthropic-style tool result messages."""
+    if not isinstance(message, dict):
+        return set()
+
+    role = message.get("role")
+    if role == "tool":
+        tool_call_id = message.get("tool_call_id")
+        return {tool_call_id} if isinstance(tool_call_id, str) and tool_call_id else set()
+
+    content = message.get("content")
+    if not isinstance(content, list):
+        return set()
+    ids: set[str] = set()
+    for block in content:
+        if isinstance(block, dict) and block.get("type") == "tool_result":
+            tool_use_id = block.get("tool_use_id")
+            if isinstance(tool_use_id, str) and tool_use_id:
+                ids.add(tool_use_id)
+    return ids
+
+
+def sanitize_messages(messages: Any) -> Any:
+    """
+    Lightweight tool-order sanitizer for Codex/OpenAI-style tool calls.
+
+    Ensures assistant tool calls are followed immediately by tool results.
+    If text-only assistant messages appear between tool call and tool result,
+    merge text into the tool-call assistant message.
+    If no immediate tool result exists, inject placeholder tool results.
+    """
+    if not isinstance(messages, list) or not messages:
+        return messages
+
+    result: list[Any] = []
+    i = 0
+    changed = False
+
+    while i < len(messages):
+        msg = messages[i]
+        if not isinstance(msg, dict):
+            result.append(msg)
+            i += 1
+            continue
+
+        role = msg.get("role")
+        tool_calls = msg.get("tool_calls")
+        if role != "assistant" or not isinstance(tool_calls, list) or len(tool_calls) == 0:
+            result.append(msg)
+            i += 1
+            continue
+
+        # Assistant tool-call message.
+        tool_ids: list[str] = []
+        for tc in tool_calls:
+            if isinstance(tc, dict):
+                tid = tc.get("id")
+                if isinstance(tid, str) and tid:
+                    tool_ids.append(tid)
+
+        merged = dict(msg)
+        merged_text = _extract_text_content(merged)
+        j = i + 1
+
+        # Merge consecutive text-only assistant messages.
+        while j < len(messages) and _is_text_only_assistant(messages[j]):
+            extra = _extract_text_content(messages[j])
+            if extra.strip():
+                merged_text = f"{merged_text}\n{extra}" if merged_text.strip() else extra
+            j += 1
+            changed = True
+
+        if merged_text != _extract_text_content(merged):
+            _set_text_content(merged, merged_text)
+
+        result.append(merged)
+
+        # Immediate next message must carry tool results for all tool ids.
+        next_ids = _tool_result_ids_from_message(messages[j]) if j < len(messages) else set()
+        missing_ids = [tid for tid in tool_ids if tid not in next_ids]
+        if missing_ids:
+            changed = True
+            for tid in missing_ids:
+                result.append({
+                    "role": "tool",
+                    "tool_call_id": tid,
+                    "content": "[Placeholder: tool result was not provided]",
+                })
+
+        i = j
+
+    return result if changed else messages
+
+
+class ReasoningParamsSanitizer(CustomLogger):
+    """LiteLLM proxy hook: strip reasoning controls for gpt-4.1-nano."""
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: Any,
+        cache: Any,
+        data: dict,
+        call_type: Any,
+    ) -> dict:
+        if not isinstance(data, dict):
+            return data
+        model = data.get("model")
+        if not _is_reasoning_unsupported_model(model):
+            return data
+        _strip_unsupported_reasoning_params(data)
+        return data
+
+    async def async_pre_call_deployment_hook(
+        self,
+        request_kwargs: dict,
+        call_type: Any,
+    ) -> dict:
+        """Extra safety: strip during deployment-layer payload building."""
+        if not isinstance(request_kwargs, dict):
+            return request_kwargs
+        model = request_kwargs.get("model")
+        if not _is_reasoning_unsupported_model(model):
+            return request_kwargs
+        _strip_unsupported_reasoning_params(request_kwargs)
+        return request_kwargs
+
+    async def async_pre_request_hook(
+        self,
+        model: str,
+        messages: list[Any],
+        request_kwargs: dict,
+    ) -> dict:
+        """Extra safety: strip right before provider call."""
+        if not _is_reasoning_unsupported_model(model):
+            return request_kwargs
+        if isinstance(request_kwargs, dict):
+            _strip_unsupported_reasoning_params(request_kwargs)
+        return request_kwargs
+
+
+class CodexProxySanitizer(CustomLogger):
+    """
+    LiteLLM proxy hook: strip Codex/OpenAI Responses fields/tools that break proxy routes.
+
+    - Anthropic: drop client_metadata (OpenAI-only Responses telemetry).
+    - gpt-4.1-nano: drop web_search_preview and related built-in search tools.
+    """
+
+    def _sanitize(self, payload: dict, model: Any) -> dict:
+        if isinstance(payload, dict):
+            _apply_codex_proxy_sanitizers(payload, model)
+        return payload
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: Any,
+        cache: Any,
+        data: dict,
+        call_type: Any,
+    ) -> dict:
+        if not isinstance(data, dict):
+            return data
+        return self._sanitize(data, data.get("model"))
+
+    async def async_pre_call_deployment_hook(
+        self,
+        request_kwargs: dict,
+        call_type: Any,
+    ) -> dict:
+        if not isinstance(request_kwargs, dict):
+            return request_kwargs
+        return self._sanitize(request_kwargs, request_kwargs.get("model"))
+
+    async def async_pre_request_hook(
+        self,
+        model: str,
+        messages: list[Any],
+        request_kwargs: dict,
+    ) -> dict:
+        if isinstance(request_kwargs, dict):
+            self._sanitize(request_kwargs, model)
+        return request_kwargs
+
+
+class ToolOrderSanitizer(CustomLogger):
+    """LiteLLM proxy hook: sanitize tool call/result ordering for Claude routes."""
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: Any,
+        cache: Any,
+        data: dict,
+        call_type: Any,
+    ) -> dict:
+        if not isinstance(data, dict):
+            return data
+        messages = data.get("messages")
+        if isinstance(messages, list) and messages:
+            data["messages"] = sanitize_messages(messages)
+        return data
+
+
+def strip_empty_thinking_blocks(messages: list[Any]) -> list[Any]:
+    """
+    Strip empty thinking/reasoning blocks from OpenAI-format messages.
+
+    Anthropic's API rejects empty `thinking` content blocks and empty
+    `reasoning_content` strings. LiteLLM proxies these fields unchanged,
+    so we sanitize outbound messages here at the proxy layer.
+
+    Handles three shapes:
+      1. Top-level `reasoning_content: ""` on assistant messages
+      2. `{ type: "thinking", thinking: "" }` inside content arrays
+      3. Removes assistant messages that become entirely empty after stripping
+    """
+    result: list[Any] = []
+
+    for msg in messages:
+        if not isinstance(msg, dict):
+            result.append(msg)
+            continue
+
+        cleaned = msg
+
+        # 1. Strip empty top-level reasoning_content on assistant messages
+        if cleaned.get("role") == "assistant" and "reasoning_content" in cleaned:
+            rc = cleaned["reasoning_content"]
+            if not rc or (isinstance(rc, str) and not rc.strip()):
+                cleaned = {k: v for k, v in cleaned.items() if k != "reasoning_content"}
+
+        # 2. Filter empty thinking blocks from content arrays
+        content = cleaned.get("content")
+        if isinstance(content, list):
+            original_len = len(content)
+            filtered = [
+                block for block in content
+                if not (
+                    isinstance(block, dict)
+                    and block.get("type") == "thinking"
+                    and (not block.get("thinking") or (isinstance(block.get("thinking"), str) and not block["thinking"].strip()))
+                )
+            ]
+            if len(filtered) != original_len:
+                # Preserve turn structure: if assistant becomes empty, keep a placeholder
+                if not filtered and cleaned.get("role") == "assistant":
+                    filtered = [{"type": "text", "text": ""}]
+                cleaned = {**cleaned, "content": filtered}
+
+        # 3. Skip entirely empty assistant messages (no content, no tool_calls)
+        if cleaned.get("role") == "assistant" and not cleaned.get("tool_calls"):
+            c = cleaned.get("content")
+            if (isinstance(c, str) and not c.strip()) or (isinstance(c, list) and len(c) == 0):
+                continue
+
+        result.append(cleaned)
+
+    return result
+
+
+class EmptyThinkingBlockSanitizer(CustomLogger):
+    """LiteLLM proxy hook: strip empty thinking/reasoning blocks before Anthropic dispatch."""
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: Any,
+        cache: Any,
+        data: dict,
+        call_type: Any,
+    ) -> dict:
+        if not isinstance(data, dict):
+            return data
+        messages = data.get("messages")
+        if isinstance(messages, list) and messages:
+            data["messages"] = strip_empty_thinking_blocks(messages)
+        return data
+
+
+class CleanSeparatedOtelOutput(OpenTelemetry):
+    """
+    LiteLLM OTEL with clean separation:
+    - gen_ai.input.messages  → full history (as set by LiteLLM)
+    - gen_ai.output.messages → only the new assistant reply (stringified structured array)
+    - Short preview event for quick visibility
+    """
+
+    def _start_primary_span(
+        self,
+        kwargs: dict,
+        response_obj: Any,
+        start_time: Any,
+        end_time: Any,
+        context: Any,
+    ):
+        otel_tracer = self.get_tracer_to_use_for_request(kwargs)
+
+        span = otel_tracer.start_span(
+            name=self._get_span_name(kwargs),
+            start_time=self._to_ns(start_time),
+            context=context,
+        )
+        span.set_status(Status(StatusCode.OK))
+        self.set_attributes(span, kwargs, response_obj)
+
+        # -- Extract only the new assistant output -------------------------------
+        text = _extract_text_from_success_kwargs(kwargs) or \
+               _extract_assistant_text(response_obj)
+
+        if text.strip():
+            # Structured output (separate from input history)
+            structured_msgs = _make_output_messages_structured(text)
+            span.set_attribute(
+                "gen_ai.output.messages",
+                json.dumps(structured_msgs, ensure_ascii=False)
+            )
+
+            # Short preview event - easy to spot, low noise
+            preview = text[:_OTEL_PREVIEW_LEN] + ("..." if len(text) > _OTEL_PREVIEW_LEN else "")
+            span.add_event(
+                name="gen_ai.output.generated",
+                attributes={
+                    "gen_ai.output.preview": preview,
+                    "gen_ai.output.length": len(text),
+                    "gen_ai.output.truncated": len(text) > _OTEL_TEXT_MAX_LEN,
+                    "gen_ai.response.model": _read_field(kwargs, "model", "unknown")
+                }
+            )
+
+        span.end(end_time=self._to_ns(end_time))
+        return span
+
+
+# Register
+structured_genai_otel = CleanSeparatedOtelOutput()
+reasoning_params_sanitizer = ReasoningParamsSanitizer()
+codex_proxy_sanitizer = CodexProxySanitizer()
+tool_order_sanitizer = ToolOrderSanitizer()
+empty_thinking_sanitizer = EmptyThinkingBlockSanitizer()