usegradient 1.3.1__py3-none-any.whl

gradient_sdk/integrations/agent.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+import json
+from typing import Any, Literal
+
+from ..tools import Tool, execute_tool_calls
+from ..tracing import get_tracer
+
+
+def generate(
+    *,
+    provider: Literal["openai", "anthropic"],
+    model: str,
+    interface: Literal["responses", "chat", "messages"] | str | None = None,
+    messages: list[dict[str, Any]] | str,
+    tools: list[Tool] | None = None,
+    client: Any | None = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Run a traced provider call and return a normalized response."""
+
+    tracer = get_tracer()
+    with tracer.start_span(f"llm.{provider}.{model}", kind="LLM") as span:
+        span.set_attribute("llm.provider", provider)
+        span.set_attribute("llm.model_name", model)
+        span.set_attribute(
+            "input.value",
+            messages if isinstance(messages, str) else json.dumps(messages, default=str),
+        )
+        if provider == "openai":
+            result = _generate_openai(
+                model=model,
+                interface=interface or "responses",
+                messages=messages,
+                tools=tools,
+                client=client,
+                **kwargs,
+            )
+        else:
+            result = _generate_anthropic(model=model, messages=messages, tools=tools, client=client, **kwargs)
+        usage = result.get("usage", {})
+        span.set_attribute("llm.token_count.prompt", int(usage.get("prompt_tokens", 0)))
+        span.set_attribute("llm.token_count.completion", int(usage.get("completion_tokens", 0)))
+        span.set_attribute("llm.token_count.total", int(usage.get("total_tokens", 0)))
+        span.set_output(result.get("text", ""))
+        return result
+
+
+def run_agent_loop(
+    *,
+    provider: Literal["openai", "anthropic"],
+    model: str,
+    messages: list[dict[str, Any]],
+    tools: list[Tool],
+    client: Any | None = None,
+    max_turns: int = 8,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Run a traced tool-calling loop similar to LangChain AgentExecutor."""
+
+    tracer = get_tracer()
+    with tracer.start_span("agent.loop", kind="AGENT") as span:
+        span.set_attribute("llm.provider", provider)
+        span.set_attribute("llm.model_name", model)
+        span.set_attribute("input.value", json.dumps(messages, default=str))
+        history = list(messages)
+        final_text = ""
+        usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        tool_results: list[dict[str, Any]] = []
+        for _ in range(max_turns):
+            result = generate(
+                provider=provider,
+                model=model,
+                messages=history,
+                tools=tools,
+                client=client,
+                **kwargs,
+            )
+            _merge_usage(usage, result.get("usage", {}))
+            tool_calls = result.get("tool_calls") or []
+            final_text = result.get("text", "")
+            if not tool_calls:
+                break
+            history.append({"role": "assistant", "content": final_text, "tool_calls": tool_calls})
+            executed = execute_tool_calls(tools, tool_calls)
+            tool_results.extend(executed)
+            for item in executed:
+                history.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": item.get("tool_call_id"),
+                        "name": item.get("name"),
+                        "content": json.dumps(item.get("output", item.get("error")), default=str),
+                    }
+                )
+        payload = {"text": final_text, "messages": history, "usage": usage, "tool_results": tool_results}
+        span.set_output(payload)
+        return payload
+
+
+def _merge_usage(target: dict[str, int], incoming: dict[str, int]) -> None:
+    for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
+        target[key] = int(target.get(key, 0)) + int(incoming.get(key, 0))
+
+
+def _generate_openai(
+    *,
+    model: str,
+    interface: str,
+    messages: list[dict[str, Any]] | str,
+    tools: list[Tool] | None,
+    client: Any | None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    if client is None:
+        from openai import OpenAI
+
+        client = OpenAI()
+    if interface == "responses" and hasattr(client, "responses"):
+        payload: dict[str, Any] = {"model": model, "input": messages, **kwargs}
+        if tools:
+            payload["tools"] = [_openai_responses_tool(tool) for tool in tools]
+        response = client.responses.create(**payload)
+        return _parse_openai_responses_response(response)
+
+    payload = {"model": model, "messages": messages, **kwargs}
+    if tools:
+        payload["tools"] = [tool.to_openai() for tool in tools]
+    response = client.chat.completions.create(**payload)
+    text = ""
+    tool_calls: list[dict[str, Any]] = []
+    if response.choices:
+        message = response.choices[0].message
+        text = message.content or ""
+        for call in getattr(message, "tool_calls", None) or []:
+            tool_calls.append(
+                {
+                    "id": call.id,
+                    "name": call.function.name,
+                    "arguments": call.function.arguments,
+                }
+            )
+    usage = getattr(response, "usage", None)
+    prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0) if usage else 0
+    completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0) if usage else 0
+    return {
+        "text": text,
+        "tool_calls": tool_calls,
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        },
+    }
+
+
+def _parse_openai_responses_response(response: Any) -> dict[str, Any]:
+    text = str(_get(response, "output_text", "") or "")
+    tool_calls: list[dict[str, Any]] = []
+    for item in _get(response, "output", []) or []:
+        item_type = _get(item, "type", "")
+        if item_type == "message" and not text:
+            text = _extract_response_message_text(item)
+        if item_type in {"function_call", "tool_call"}:
+            tool_calls.append(
+                {
+                    "id": _get(item, "id", _get(item, "call_id", None)),
+                    "name": _get(item, "name", ""),
+                    "arguments": _get(item, "arguments", _get(item, "input", {})),
+                }
+            )
+    usage = _get(response, "usage", {}) or {}
+    prompt_tokens = int(_get(usage, "input_tokens", _get(usage, "prompt_tokens", 0)) or 0)
+    completion_tokens = int(_get(usage, "output_tokens", _get(usage, "completion_tokens", 0)) or 0)
+    total_tokens = int(_get(usage, "total_tokens", prompt_tokens + completion_tokens) or 0)
+    return {
+        "text": text,
+        "tool_calls": tool_calls,
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+        },
+    }
+
+
+def _openai_responses_tool(tool: Tool) -> dict[str, Any]:
+    raw = tool.to_openai()
+    fn = raw.get("function", {})
+    return {
+        "type": "function",
+        "name": fn.get("name", tool.name),
+        "description": fn.get("description", tool.description),
+        "parameters": fn.get("parameters", tool.parameters),
+    }
+
+
+def _extract_response_message_text(item: Any) -> str:
+    parts: list[str] = []
+    for content in _get(item, "content", []) or []:
+        value = _get(content, "text", None)
+        if value is None:
+            value = _get(content, "output_text", "")
+        if value:
+            parts.append(str(value))
+    return "".join(parts)
+
+
+def _generate_anthropic(
+    *,
+    model: str,
+    messages: list[dict[str, Any]] | str,
+    tools: list[Tool] | None,
+    client: Any | None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    if client is None:
+        from anthropic import Anthropic
+
+        client = Anthropic()
+    if isinstance(messages, str):
+        anthropic_messages = [{"role": "user", "content": messages}]
+    else:
+        anthropic_messages = messages
+    payload: dict[str, Any] = {
+        "model": model,
+        "max_tokens": kwargs.pop("max_tokens", 1024),
+        "messages": anthropic_messages,
+        **kwargs,
+    }
+    if tools:
+        payload["tools"] = [tool.to_anthropic() for tool in tools]
+    response = client.messages.create(**payload)
+    content = _get(response, "content", []) or []
+    text = "".join(str(_get(block, "text", "")) for block in content if _get(block, "type", "") == "text")
+    tool_calls: list[dict[str, Any]] = []
+    for block in content:
+        if _get(block, "type", "") == "tool_use":
+            tool_calls.append(
+                {
+                    "id": _get(block, "id", None),
+                    "name": _get(block, "name", ""),
+                    "input": _get(block, "input", {}),
+                }
+            )
+    usage = _get(response, "usage", {}) or {}
+    prompt_tokens = int(_get(usage, "input_tokens", 0) or 0)
+    completion_tokens = int(_get(usage, "output_tokens", 0) or 0)
+    return {
+        "text": text,
+        "tool_calls": tool_calls,
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        },
+    }
+
+
+def _get(value: Any, key: str, default: Any = None) -> Any:
+    if isinstance(value, dict):
+        return value.get(key, default)
+    return getattr(value, key, default)

gradient_sdk/integrations/anthropic.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from ..tracing import get_tracer
+
+
+def wrap_anthropic(client: Any) -> Any:
+    """Wrap an Anthropic client so messages.create calls are traced automatically."""
+
+    class MessagesProxy:
+        def __init__(self, messages: Any) -> None:
+            self._messages = messages
+
+        def create(self, *args: Any, **kwargs: Any) -> Any:
+            model = kwargs.get("model", "unknown")
+            prompt = kwargs.get("messages", args[0] if args else "")
+            response = self._messages.create(*args, **kwargs)
+            text = "".join(
+                getattr(block, "text", "")
+                for block in getattr(response, "content", [])
+                if getattr(block, "type", "") == "text"
+            )
+            usage = getattr(response, "usage", None)
+            prompt_tokens = int(getattr(usage, "input_tokens", 0) or 0) if usage else 0
+            completion_tokens = int(getattr(usage, "output_tokens", 0) or 0) if usage else 0
+            tracer = get_tracer()
+            with tracer.start_span(f"anthropic.messages.{model}", kind="LLM") as span:
+                span.set_attribute("llm.provider", "anthropic")
+                span.set_attribute("llm.model_name", model)
+                span.set_attribute(
+                    "input.value",
+                    prompt if isinstance(prompt, str) else json.dumps(prompt, default=str),
+                )
+                span.set_attribute("llm.token_count.prompt", prompt_tokens)
+                span.set_attribute("llm.token_count.completion", completion_tokens)
+                span.set_attribute("llm.token_count.total", prompt_tokens + completion_tokens)
+                span.set_output(text)
+            return response
+
+    class AnthropicProxy:
+        def __init__(self, inner: Any) -> None:
+            self._inner = inner
+            if hasattr(inner, "messages"):
+                self.messages = MessagesProxy(inner.messages)
+
+        def __getattr__(self, name: str) -> Any:
+            return getattr(self._inner, name)
+
+    return AnthropicProxy(client)

gradient_sdk/integrations/openai.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from ..tracing import get_tracer
+
+
+def _record_llm_span(
+    *,
+    name: str,
+    provider: str,
+    model: str,
+    input_value: Any,
+    output_value: Any,
+    prompt_tokens: int = 0,
+    completion_tokens: int = 0,
+) -> None:
+    tracer = get_tracer()
+    with tracer.start_span(name, kind="LLM") as span:
+        span.set_attribute("llm.provider", provider)
+        span.set_attribute("llm.model_name", model)
+        span.set_attribute("input.value", input_value if isinstance(input_value, str) else json.dumps(input_value, default=str))
+        span.set_attribute("llm.token_count.prompt", prompt_tokens)
+        span.set_attribute("llm.token_count.completion", completion_tokens)
+        span.set_attribute("llm.token_count.total", prompt_tokens + completion_tokens)
+        span.set_output(output_value)
+
+
+def wrap_openai(client: Any) -> Any:
+    """Wrap an OpenAI client so chat/responses calls are traced automatically."""
+
+    class ChatCompletionsProxy:
+        def __init__(self, completions: Any) -> None:
+            self._completions = completions
+
+        def create(self, *args: Any, **kwargs: Any) -> Any:
+            model = kwargs.get("model", "unknown")
+            messages = kwargs.get("messages", kwargs.get("input", args[0] if args else ""))
+            response = self._completions.create(*args, **kwargs)
+            text, prompt_tokens, completion_tokens = _extract_openai_chat_response(response)
+            _record_llm_span(
+                name=f"openai.chat.{model}",
+                provider="openai",
+                model=model,
+                input_value=messages,
+                output_value=text,
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+            )
+            return response
+
+    class ResponsesProxy:
+        def __init__(self, responses: Any) -> None:
+            self._responses = responses
+
+        def create(self, *args: Any, **kwargs: Any) -> Any:
+            model = kwargs.get("model", "unknown")
+            prompt = kwargs.get("input", args[0] if args else "")
+            response = self._responses.create(*args, **kwargs)
+            text = getattr(response, "output_text", str(response))
+            usage = getattr(response, "usage", None)
+            prompt_tokens = int(getattr(usage, "input_tokens", 0) or 0) if usage else 0
+            completion_tokens = int(getattr(usage, "output_tokens", 0) or 0) if usage else 0
+            _record_llm_span(
+                name=f"openai.responses.{model}",
+                provider="openai",
+                model=model,
+                input_value=prompt,
+                output_value=text,
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+            )
+            return response
+
+    class ChatProxy:
+        def __init__(self, chat: Any) -> None:
+            self._chat = chat
+            self.completions = ChatCompletionsProxy(getattr(chat, "completions", chat))
+
+    class OpenAIProxy:
+        def __init__(self, inner: Any) -> None:
+            self._inner = inner
+            if hasattr(inner, "chat"):
+                self.chat = ChatProxy(inner.chat)
+            if hasattr(inner, "responses"):
+                self.responses = ResponsesProxy(inner.responses)
+
+        def __getattr__(self, name: str) -> Any:
+            return getattr(self._inner, name)
+
+    return OpenAIProxy(client)
+
+
+def _extract_openai_chat_response(response: Any) -> tuple[str, int, int]:
+    text = ""
+    choices = getattr(response, "choices", None) or []
+    if choices:
+        message = getattr(choices[0], "message", None)
+        if message is not None:
+            text = getattr(message, "content", "") or ""
+    usage = getattr(response, "usage", None)
+    prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0) if usage else 0
+    completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0) if usage else 0
+    return text, prompt_tokens, completion_tokens

gradient_sdk/model_catalog.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+
+class Model(str, Enum):
+    GPT_5_6 = "openai:gpt-5.6"
+    GPT_5_6_SOL = "openai:gpt-5.6-sol"
+    GPT_5_6_TERRA = "openai:gpt-5.6-terra"
+    GPT_5_6_LUNA = "openai:gpt-5.6-luna"
+    GPT_5_5 = "openai:gpt-5.5"
+    GPT_5_5_PRO = "openai:gpt-5.5-pro"
+    GPT_5_4 = "openai:gpt-5.4"
+    GPT_5_4_PRO = "openai:gpt-5.4-pro"
+    GPT_5_4_MINI = "openai:gpt-5.4-mini"
+    GPT_5_4_NANO = "openai:gpt-5.4-nano"
+    GPT_5_3_CODEX = "openai:gpt-5.3-codex"
+    GPT_5_2 = "openai:gpt-5.2"
+    GPT_5_2_PRO = "openai:gpt-5.2-pro"
+    GPT_5_1 = "openai:gpt-5.1"
+    GPT_5 = "openai:gpt-5"
+    GPT_5_MINI = "openai:gpt-5-mini"
+    GPT_5_NANO = "openai:gpt-5-nano"
+    GPT_5_PRO = "openai:gpt-5-pro"
+    GPT_4_1 = "openai:gpt-4.1"
+    GPT_4_1_MINI = "openai:gpt-4.1-mini"
+    GPT_4O_MINI = "openai:gpt-4o-mini"
+
+    CLAUDE_FABLE_5 = "anthropic:claude-fable-5"
+    CLAUDE_OPUS_4_8 = "anthropic:claude-opus-4-8"
+    CLAUDE_SONNET_5 = "anthropic:claude-sonnet-5"
+    CLAUDE_HAIKU_4_5 = "anthropic:claude-haiku-4-5"
+    CLAUDE_HAIKU_4_5_20251001 = "anthropic:claude-haiku-4-5-20251001"
+
+
+@dataclass(frozen=True, slots=True)
+class ModelSpec:
+    provider: str
+    model: str
+    interface: str
+
+
+MODEL_SPECS: dict[Model, ModelSpec] = {
+    item: ModelSpec(
+        provider=item.value.split(":", 1)[0],
+        model=item.value.split(":", 1)[1],
+        interface="messages" if item.value.startswith("anthropic:") else "responses",
+    )
+    for item in Model
+}
+
+MODEL_ALIASES: dict[str, Model] = {}
+for item in Model:
+    MODEL_ALIASES[item.value] = item
+    MODEL_ALIASES[item.value.split(":", 1)[1]] = item
+
+
+def resolve_model(value: Model | str | Any) -> ModelSpec | None:
+    if isinstance(value, Model):
+        return MODEL_SPECS[value]
+    if not isinstance(value, str):
+        return None
+    model = MODEL_ALIASES.get(value)
+    if model is None:
+        supported = ", ".join(item.value for item in Model)
+        raise ValueError(
+            f"unsupported model {value!r}; use one of Model.* or these ids: {supported}"
+        )
+    return MODEL_SPECS[model]
+
+
+def supported_models() -> list[str]:
+    return [item.value for item in Model]