usegradient 1.3.1__py3-none-any.whl

gradient_sdk/__init__.py

@@ -0,0 +1,122 @@
+from .client import GradientClient, default_machine_template
+from .agent import Agent, AgentResult
+from .context import RunContext, current_context
+from .credentials import (
+    DEFAULT_API_URL,
+    DEFAULT_PROXY_URL,
+    DEFAULT_REGISTRY_URL,
+    Credentials,
+    load_credentials,
+    save_credentials,
+)
+from .decorators import agent, chain, embedding, llm, retriever, task, tool, trace, trace_context, workflow
+from .errors import GradientAPIError, GradientAuthError, GradientError
+from .graph import END, START, StateGraph
+from .integrations import generate, run_agent_loop, wrap_anthropic, wrap_openai
+from .models import (
+    AppSecret,
+    CreateTraceRunResponse,
+    CreateMachineResponse,
+    ExecMachineResponse,
+    Guest,
+    MachineCreateRequest,
+    MachineListItem,
+    MachineSpec,
+    PortEntry,
+    ReplayTraceRunResponse,
+    Restart,
+    ServiceConfig,
+    TraceEvent,
+    TraceRun,
+    TraceSpan,
+    WhoamiResponse,
+)
+from .model_catalog import Model, ModelSpec, supported_models
+from .runtime import (
+    BuildHandle,
+    DeploymentHandle,
+    Environment,
+    Gradient,
+    ProgressEvent,
+    RemoteDeployment,
+    RemoteRunResult,
+    RunHandle,
+    Trace,
+    TracePolicy,
+)
+from .session import GradientSession, current_session, init
+from .tools import Tool, execute_tool_calls, infer_parameters, tool_from_function
+from .tracing import current_span, get_tracer
+
+__all__ = [
+    "AppSecret",
+    "Agent",
+    "AgentResult",
+    "BuildHandle",
+    "CreateMachineResponse",
+    "CreateTraceRunResponse",
+    "Credentials",
+    "DEFAULT_API_URL",
+    "DEFAULT_PROXY_URL",
+    "DEFAULT_REGISTRY_URL",
+    "DeploymentHandle",
+    "END",
+    "GradientAPIError",
+    "GradientAuthError",
+    "Gradient",
+    "GradientClient",
+    "GradientError",
+    "GradientSession",
+    "Environment",
+    "ExecMachineResponse",
+    "Guest",
+    "MachineCreateRequest",
+    "MachineListItem",
+    "MachineSpec",
+    "Model",
+    "ModelSpec",
+    "PortEntry",
+    "ProgressEvent",
+    "ReplayTraceRunResponse",
+    "RemoteDeployment",
+    "RemoteRunResult",
+    "Restart",
+    "START",
+    "ServiceConfig",
+    "StateGraph",
+    "supported_models",
+    "Tool",
+    "Trace",
+    "TraceEvent",
+    "TracePolicy",
+    "TraceRun",
+    "TraceSpan",
+    "WhoamiResponse",
+    "agent",
+    "chain",
+    "current_context",
+    "current_session",
+    "current_span",
+    "default_machine_template",
+    "embedding",
+    "execute_tool_calls",
+    "generate",
+    "get_tracer",
+    "infer_parameters",
+    "init",
+    "llm",
+    "load_credentials",
+    "retriever",
+    "run_agent_loop",
+    "RunHandle",
+    "RunContext",
+    "save_credentials",
+    "task",
+    "tool",
+    "tool_from_function",
+    "trace",
+    "trace_context",
+    "workflow",
+    "wrap_anthropic",
+    "wrap_openai",
+]

gradient_sdk/agent.py

@@ -0,0 +1,368 @@
+from __future__ import annotations
+
+import inspect
+import json
+from dataclasses import dataclass, field
+from typing import Any, Callable, Iterable, Iterator, Mapping
+
+from .context import RunContext, use_context
+from .integrations.agent import generate
+from .model_catalog import Model, resolve_model
+from .tools import Tool, execute_tool_calls, tool_from_function
+from .tracing import get_tracer
+
+
+@dataclass(slots=True)
+class AgentResult:
+    output: Any
+    text: str
+    messages: list[dict[str, Any]]
+    tool_results: list[dict[str, Any]]
+    usage: dict[str, int]
+    events: list[dict[str, Any]]
+    context: RunContext
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def __str__(self) -> str:
+        return self.text
+
+
+class Agent:
+    __gradient_agent__ = True
+
+    def __init__(
+        self,
+        model: str | Callable[..., Any] | Any,
+        *,
+        tools: Iterable[Tool | Callable[..., Any]] | None = None,
+        instructions: str | None = None,
+        memory: Any | None = None,
+        middleware: Iterable[Callable[..., Any]] | None = None,
+        trace: bool = True,
+        gradient: Any | None = None,
+        name: str | None = None,
+        max_turns: int = 8,
+        retries: int = 0,
+        fallback_models: Iterable[str | Callable[..., Any] | Any] | None = None,
+        structured_output: Callable[[Any], Any] | None = None,
+        entrypoint: str | None = None,
+        client: Any | None = None,
+    ) -> None:
+        self.model = model
+        self.tools = [_coerce_tool(item) for item in tools or []]
+        self.instructions = instructions
+        self.memory = memory
+        self.middleware = list(middleware or [])
+        self.trace = trace
+        self.gradient = gradient
+        self.name = name or _model_name(model)
+        self.max_turns = max_turns
+        self.retries = retries
+        self.fallback_models = list(fallback_models or [])
+        self.structured_output = structured_output
+        self.entrypoint = entrypoint
+        self.client = client
+
+    def __call__(self, input: Any = None) -> Any:
+        return self.run(input).output
+
+    def bind(self, gradient: Any) -> "Agent":
+        self.gradient = gradient
+        return self
+
+    def run(
+        self,
+        input: Any = None,
+        *,
+        context: RunContext | Mapping[str, Any] | None = None,
+        **metadata: Any,
+    ) -> AgentResult:
+        ctx = RunContext.from_value(
+            context,
+            metadata={**dict(getattr(context, "metadata", {}) or {}), **metadata}
+            if context is not None
+            else metadata,
+        )
+        events: list[dict[str, Any]] = []
+        with use_context(ctx):
+            with _maybe_agent_span(self.name, self.trace, input) as span:
+                result = self._run_loop(input, ctx, events)
+                if span is not None:
+                    span.set_output(result.output)
+                self._save_memory(result)
+                return result
+
+    def stream(
+        self,
+        input: Any = None,
+        *,
+        context: RunContext | Mapping[str, Any] | None = None,
+        **metadata: Any,
+    ) -> Iterator[dict[str, Any]]:
+        result = self.run(input, context=context, **metadata)
+        yield from result.events
+
+    def deploy(self, name: str | None = None, **kwargs: Any) -> Any:
+        if self.gradient is None:
+            raise RuntimeError("Agent.deploy requires an Agent(..., gradient=gradient) or agent.bind(gradient)")
+        return self.gradient.deploy(self, name=name, **kwargs)
+
+    def evaluate(
+        self,
+        dataset: Iterable[Mapping[str, Any]] | Any,
+        evaluators: Mapping[str, Callable[[Mapping[str, Any]], Any]] | None = None,
+        *,
+        input_key: str = "input",
+        expected_key: str = "expected_output",
+    ) -> dict[str, Any]:
+        rows = list(dataset)
+        results: list[dict[str, Any]] = []
+        for index, example in enumerate(rows):
+            payload = example.get(input_key) if isinstance(example, Mapping) else example
+            expected = example.get(expected_key) if isinstance(example, Mapping) else None
+            started = self.run(payload)
+            scores = {
+                name: evaluator(
+                    {
+                        "input": payload,
+                        "output": started.output,
+                        "text": started.text,
+                        "expected_output": expected,
+                        "example": example,
+                    }
+                )
+                for name, evaluator in (evaluators or {}).items()
+            }
+            results.append(
+                {
+                    "index": index,
+                    "input": payload,
+                    "output": started.output,
+                    "text": started.text,
+                    "expected_output": expected,
+                    "scores": scores,
+                    "usage": started.usage,
+                }
+            )
+        return {"count": len(results), "results": results}
+
+    def _run_loop(
+        self,
+        input: Any,
+        context: RunContext,
+        events: list[dict[str, Any]],
+    ) -> AgentResult:
+        messages = self._initial_messages(input)
+        usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        tool_results: list[dict[str, Any]] = []
+        final_text = ""
+        output: Any = ""
+
+        events.append(_event("context", {"context": context.to_event()}))
+        events.append(_event("start", {"input": input, "agent": self.name}))
+
+        for turn in range(max(1, self.max_turns)):
+            context.check_cancelled()
+            response = self._call_model(messages, context)
+            _merge_usage(usage, response.get("usage", {}))
+            final_text = str(response.get("text") or response.get("output") or "")
+            output = response.get("output", final_text)
+            events.append(
+                _event(
+                    "message",
+                    {"role": "assistant", "content": final_text, "turn": turn + 1},
+                )
+            )
+            tool_calls = list(response.get("tool_calls") or [])
+            if not tool_calls:
+                break
+            messages.append({"role": "assistant", "content": final_text, "tool_calls": tool_calls})
+            events.append(_event("tool_calls", {"calls": tool_calls, "turn": turn + 1}))
+            executed = execute_tool_calls(self.tools, tool_calls)
+            tool_results.extend(executed)
+            events.append(_event("tool_results", {"results": executed, "turn": turn + 1}))
+            for item in executed:
+                messages.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),
+                    }
+                )
+
+        if self.structured_output is not None:
+            output = self.structured_output(output)
+        events.append(_event("final", {"output": output, "text": final_text, "usage": usage}))
+        return AgentResult(
+            output=output,
+            text=final_text,
+            messages=messages,
+            tool_results=tool_results,
+            usage=usage,
+            events=events,
+            context=context,
+        )
+
+    def _initial_messages(self, input: Any) -> list[dict[str, Any]]:
+        messages: list[dict[str, Any]] = []
+        if self.instructions:
+            messages.append({"role": "system", "content": self.instructions})
+        loaded_memory = _load_memory(self.memory)
+        messages.extend(loaded_memory)
+        if isinstance(input, Mapping) and isinstance(input.get("messages"), list):
+            messages.extend(dict(item) for item in input["messages"])
+        elif isinstance(input, list) and all(isinstance(item, Mapping) for item in input):
+            messages.extend(dict(item) for item in input)
+        else:
+            messages.append({"role": "user", "content": input if isinstance(input, str) else json.dumps(input, default=str)})
+        return messages
+
+    def _call_model(self, messages: list[dict[str, Any]], context: RunContext) -> dict[str, Any]:
+        def invoke() -> dict[str, Any]:
+            return self._call_model_without_middleware(messages, context)
+
+        call_next = invoke
+        for middleware in reversed(self.middleware):
+            previous = call_next
+            call_next = lambda middleware=middleware, previous=previous: middleware(context, previous)
+        return call_next()
+
+    def _call_model_without_middleware(
+        self,
+        messages: list[dict[str, Any]],
+        context: RunContext,
+    ) -> dict[str, Any]:
+        models = [self.model, *self.fallback_models]
+        last_error: BaseException | None = None
+        for model in models:
+            for _ in range(self.retries + 1):
+                try:
+                    return _normalize_model_result(
+                        _invoke_model(
+                            model,
+                            messages=messages,
+                            tools=self.tools,
+                            context=context,
+                            client=self.client,
+                        )
+                    )
+                except Exception as exc:
+                    last_error = exc
+        if last_error is not None:
+            raise last_error
+        return {"text": "", "tool_calls": [], "usage": {}}
+
+    def _save_memory(self, result: AgentResult) -> None:
+        if self.memory is None:
+            return
+        if hasattr(self.memory, "save"):
+            self.memory.save(result.messages)
+        elif isinstance(self.memory, list):
+            self.memory[:] = result.messages
+
+
+def _coerce_tool(value: Tool | Callable[..., Any]) -> Tool:
+    if isinstance(value, Tool):
+        return value
+    tool_obj = getattr(value, "gradient_tool", None)
+    if isinstance(tool_obj, Tool):
+        return tool_obj
+    return tool_from_function(value)
+
+
+def _load_memory(memory: Any | None) -> list[dict[str, Any]]:
+    if memory is None:
+        return []
+    if hasattr(memory, "load"):
+        loaded = memory.load()
+        return list(loaded or [])
+    if isinstance(memory, list):
+        return [dict(item) for item in memory if isinstance(item, Mapping)]
+    return []
+
+
+def _invoke_model(
+    model: Model | str | Callable[..., Any] | Any,
+    *,
+    messages: list[dict[str, Any]],
+    tools: list[Tool],
+    context: RunContext,
+    client: Any | None = None,
+) -> Any:
+    spec = resolve_model(model)
+    if spec is not None:
+        return generate(
+            provider=spec.provider,  # type: ignore[arg-type]
+            model=spec.model,
+            interface=spec.interface,
+            messages=messages,
+            tools=tools,
+            client=client,
+        )
+    if callable(model):
+        kwargs = {"messages": messages, "tools": tools, "context": context}
+        try:
+            signature = inspect.signature(model)
+            accepted = {key: value for key, value in kwargs.items() if key in signature.parameters}
+            if accepted:
+                return model(**accepted)
+        except (TypeError, ValueError):
+            pass
+        return model(messages)
+    if hasattr(model, "generate"):
+        return model.generate(messages=messages, tools=tools, context=context)
+    if hasattr(model, "invoke"):
+        return model.invoke(messages)
+    raise TypeError("Agent model must be a provider string, callable, or object with generate/invoke")
+
+
+def _normalize_model_result(result: Any) -> dict[str, Any]:
+    if isinstance(result, str):
+        return {"text": result, "output": result, "tool_calls": [], "usage": {}}
+    if isinstance(result, Mapping):
+        text = result.get("text", result.get("content", result.get("output", "")))
+        return {
+            "text": str(text or ""),
+            "output": result.get("output", text),
+            "tool_calls": list(result.get("tool_calls") or result.get("toolCalls") or []),
+            "usage": dict(result.get("usage") or {}),
+        }
+    return {"text": str(result), "output": result, "tool_calls": [], "usage": {}}
+
+
+def _merge_usage(target: dict[str, int], incoming: Mapping[str, Any]) -> None:
+    for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
+        target[key] = int(target.get(key, 0)) + int(incoming.get(key, 0) or 0)
+
+
+def _event(event_type: str, data: Mapping[str, Any]) -> dict[str, Any]:
+    return {"type": event_type, "data": dict(data)}
+
+
+def _model_name(model: Any) -> str:
+    if isinstance(model, str):
+        return model
+    return getattr(model, "__name__", model.__class__.__name__)
+
+
+class _maybe_agent_span:
+    def __init__(self, name: str, enabled: bool, input: Any) -> None:
+        self.name = name
+        self.enabled = enabled
+        self.input = input
+        self._span: Any = None
+
+    def __enter__(self) -> Any:
+        if not self.enabled:
+            return None
+        tracer = get_tracer()
+        self._span = tracer.start_span(f"agent.{self.name}", kind="AGENT")
+        entered = self._span.__enter__()
+        entered.set_attribute("input.value", json.dumps(self.input, default=str))
+        return entered
+
+    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
+        if self._span is not None:
+            return bool(self._span.__exit__(exc_type, exc, tb))
+        return False