usegradient 1.3.1__py3-none-any.whl

gradient_sdk/context.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import secrets
+import time
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import Any, Iterator, Mapping
+
+
+@dataclass(slots=True)
+class RunContext:
+    project: str = "gradient"
+    run_id: str = field(default_factory=lambda: "local-" + secrets.token_hex(8))
+    session_id: str | None = None
+    user_id: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+    tags: list[str] = field(default_factory=list)
+    secrets: dict[str, str] = field(default_factory=dict, repr=False)
+    trace: Any | None = None
+    timeout_seconds: float | None = None
+    replay_config: dict[str, Any] = field(default_factory=dict)
+    environment: dict[str, Any] = field(default_factory=dict)
+    cancelled: bool = False
+    started_at: float = field(default_factory=time.time)
+
+    def child(self, **updates: Any) -> "RunContext":
+        values = {
+            "project": self.project,
+            "run_id": self.run_id,
+            "session_id": self.session_id,
+            "user_id": self.user_id,
+            "metadata": dict(self.metadata),
+            "tags": list(self.tags),
+            "secrets": dict(self.secrets),
+            "trace": self.trace,
+            "timeout_seconds": self.timeout_seconds,
+            "replay_config": dict(self.replay_config),
+            "environment": dict(self.environment),
+            "cancelled": self.cancelled,
+            "started_at": self.started_at,
+        }
+        values.update(updates)
+        return RunContext(**values)
+
+    def check_cancelled(self) -> None:
+        if self.cancelled:
+            raise RuntimeError("Gradient run was cancelled")
+
+    def to_event(self) -> dict[str, Any]:
+        return {
+            "project": self.project,
+            "run_id": self.run_id,
+            "session_id": self.session_id,
+            "user_id": self.user_id,
+            "metadata": dict(self.metadata),
+            "tags": list(self.tags),
+            "replay_config": dict(self.replay_config),
+            "environment": dict(self.environment),
+        }
+
+    @classmethod
+    def from_value(
+        cls,
+        value: "RunContext | Mapping[str, Any] | None" = None,
+        **defaults: Any,
+    ) -> "RunContext":
+        if isinstance(value, RunContext):
+            return value.child(**defaults) if defaults else value
+        payload = dict(value or {})
+        payload.update({key: item for key, item in defaults.items() if item is not None})
+        return RunContext(**payload)
+
+
+_current_context: ContextVar[RunContext | None] = ContextVar(
+    "gradient_run_context",
+    default=None,
+)
+
+
+def current_context() -> RunContext | None:
+    return _current_context.get()
+
+
+@contextmanager
+def use_context(context: RunContext) -> Iterator[RunContext]:
+    token = _current_context.set(context)
+    try:
+        yield context
+    finally:
+        _current_context.reset(token)

gradient_sdk/credentials.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Mapping
+import json
+import os
+import tomllib
+
+from .errors import GradientAuthError
+
+DEFAULT_API_URL = "https://console.usegradient.dev"
+DEFAULT_REGISTRY_URL = "https://registry.usegradient.dev"
+DEFAULT_PROXY_URL = "https://proxy.usegradient.dev"
+
+
+@dataclass(slots=True)
+class Credentials:
+    api_key: str = ""
+    api_url: str = DEFAULT_API_URL
+    registry_url: str = DEFAULT_REGISTRY_URL
+    proxy_url: str = DEFAULT_PROXY_URL
+    slug: str = ""
+
+
+def _credentials_path(home: str | Path | None = None) -> Path:
+    root = Path(home).expanduser() if home else Path.home()
+    return root / ".gradient" / "credentials"
+
+
+def _read_credentials_file(path: Path) -> Credentials:
+    if not path.exists():
+        return Credentials()
+    with path.open("rb") as f:
+        raw = tomllib.load(f)
+    return Credentials(
+        api_key=str(raw.get("api_key", "")),
+        api_url=str(raw.get("api_url", "") or DEFAULT_API_URL),
+        registry_url=str(raw.get("registry_url", "") or DEFAULT_REGISTRY_URL),
+        proxy_url=str(raw.get("proxy_url", "") or DEFAULT_PROXY_URL),
+        slug=str(raw.get("slug", "")),
+    )
+
+
+def load_credentials(
+    *,
+    path: str | Path | None = None,
+    env: Mapping[str, str] | None = None,
+    home: str | Path | None = None,
+    require_api_key: bool = False,
+) -> Credentials:
+    """Load Gradient credentials from disk and environment variables."""
+
+    values = _read_credentials_file(Path(path).expanduser() if path else _credentials_path(home))
+    environ = os.environ if env is None else env
+
+    if environ.get("GRADIENT_API_KEY"):
+        values.api_key = environ["GRADIENT_API_KEY"]
+    if environ.get("GRADIENT_API_URL"):
+        values.api_url = environ["GRADIENT_API_URL"]
+    if environ.get("GRADIENT_REGISTRY_URL"):
+        values.registry_url = environ["GRADIENT_REGISTRY_URL"]
+    if environ.get("GRADIENT_PROXY_URL"):
+        values.proxy_url = environ["GRADIENT_PROXY_URL"]
+
+    values.api_url = (values.api_url or DEFAULT_API_URL).rstrip("/")
+    values.registry_url = (values.registry_url or DEFAULT_REGISTRY_URL).rstrip("/")
+    values.proxy_url = (values.proxy_url or DEFAULT_PROXY_URL).rstrip("/")
+
+    if require_api_key and not values.api_key:
+        raise GradientAuthError(
+            "not logged in. Set GRADIENT_API_KEY or run `gradient login --api-key grad_live_...`."
+        )
+    return values
+
+
+def save_credentials(credentials: Credentials, *, path: str | Path | None = None) -> Path:
+    """Write credentials in the same TOML format used by the Gradient CLI."""
+
+    target = Path(path).expanduser() if path else _credentials_path()
+    target.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+    lines = [
+        f"api_key = {json.dumps(credentials.api_key)}",
+        f"api_url = {json.dumps(credentials.api_url.rstrip('/'))}",
+        f"registry_url = {json.dumps(credentials.registry_url)}",
+        f"proxy_url = {json.dumps(credentials.proxy_url)}",
+        f"slug = {json.dumps(credentials.slug)}",
+        "",
+    ]
+    target.write_text("\n".join(lines), encoding="utf-8")
+    target.chmod(0o600)
+    return target

gradient_sdk/decorators.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import functools
+import inspect
+import json
+from contextlib import contextmanager
+from typing import Any, Callable, TypeVar
+
+from .tools import tool_from_function
+from .tracing import get_tracer
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _serialize_call_args(fn: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> str:
+    signature = inspect.signature(fn)
+    bound = signature.bind_partial(*args, **kwargs)
+    bound.apply_defaults()
+    payload = {key: value for key, value in bound.arguments.items() if key != "self"}
+    return json.dumps(payload, default=str)
+
+
+def _wrap_callable(
+    fn: F,
+    *,
+    name: str | None,
+    kind: str,
+    extra_attributes: dict[str, Any] | None = None,
+    register_tool: bool = False,
+    description: str | None = None,
+    parameters: dict[str, Any] | None = None,
+) -> F:
+    span_name = name or fn.__name__
+    tool_obj = tool_from_function(fn, name=span_name, description=description, parameters=parameters) if register_tool else None
+
+    if inspect.iscoroutinefunction(fn):
+
+        @functools.wraps(fn)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            tracer = get_tracer()
+            with tracer.start_span(span_name, kind=kind, attributes=extra_attributes or {}) as span:
+                span.set_attribute("input.value", _serialize_call_args(fn, args, kwargs))
+                try:
+                    result = await fn(*args, **kwargs)
+                    span.set_output(result)
+                    return result
+                except Exception as exc:
+                    span.record_exception(exc)
+                    raise
+
+        if tool_obj is not None:
+            async_wrapper.gradient_tool = tool_obj  # type: ignore[attr-defined]
+        return async_wrapper  # type: ignore[return-value]
+
+    @functools.wraps(fn)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        tracer = get_tracer()
+        with tracer.start_span(span_name, kind=kind, attributes=extra_attributes or {}) as span:
+            span.set_attribute("input.value", _serialize_call_args(fn, args, kwargs))
+            try:
+                result = fn(*args, **kwargs)
+                span.set_output(result)
+                return result
+            except Exception as exc:
+                span.record_exception(exc)
+                raise
+
+    if tool_obj is not None:
+        sync_wrapper.gradient_tool = tool_obj  # type: ignore[attr-defined]
+    return sync_wrapper  # type: ignore[return-value]
+
+
+def tool(
+    fn: F | None = None,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    parameters: dict[str, Any] | None = None,
+) -> Callable[[F], F] | F:
+    def decorator(inner: F) -> F:
+        return _wrap_callable(
+            inner,
+            name=name or inner.__name__,
+            kind="TOOL",
+            register_tool=True,
+            description=description,
+            parameters=parameters,
+        )
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+def llm(
+    fn: F | None = None,
+    *,
+    name: str | None = None,
+    provider: str | None = None,
+    model: str | None = None,
+) -> Callable[[F], F] | F:
+    attrs: dict[str, Any] = {}
+    if provider:
+        attrs["llm.provider"] = provider
+    if model:
+        attrs["llm.model_name"] = model
+
+    def decorator(inner: F) -> F:
+        return _wrap_callable(inner, name=name or inner.__name__, kind="LLM", extra_attributes=attrs)
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+def chain(fn: F | None = None, *, name: str | None = None) -> Callable[[F], F] | F:
+    def decorator(inner: F) -> F:
+        return _wrap_callable(inner, name=name or inner.__name__, kind="CHAIN")
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+def task(fn: F | None = None, *, name: str | None = None) -> Callable[[F], F] | F:
+    def decorator(inner: F) -> F:
+        return _wrap_callable(inner, name=name or inner.__name__, kind="CHAIN")
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+workflow = chain
+
+
+def agent(fn: F | None = None, *, name: str | None = None) -> Callable[[F], F] | F:
+    def decorator(inner: F) -> F:
+        return _wrap_callable(inner, name=name or inner.__name__, kind="AGENT")
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+agent.step = task  # type: ignore[attr-defined]
+
+
+def retriever(fn: F | None = None, *, name: str | None = None) -> Callable[[F], F] | F:
+    def decorator(inner: F) -> F:
+        return _wrap_callable(inner, name=name or inner.__name__, kind="RETRIEVER")
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+def embedding(fn: F | None = None, *, name: str | None = None) -> Callable[[F], F] | F:
+    def decorator(inner: F) -> F:
+        return _wrap_callable(inner, name=name or inner.__name__, kind="EMBEDDING")
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+class _TraceContext:
+    def __init__(self, name: str, kind: str = "CHAIN", attributes: dict[str, Any] | None = None) -> None:
+        self.name = name
+        self.kind = kind
+        self.attributes = attributes or {}
+        self._span = None
+
+    def __enter__(self):
+        tracer = get_tracer()
+        self._span = tracer.start_span(self.name, kind=self.kind, attributes=self.attributes)
+        return self._span.__enter__()
+
+    def __exit__(self, exc_type, exc, tb):
+        if self._span is not None:
+            return self._span.__exit__(exc_type, exc, tb)
+        return False
+
+
+def trace(name: str | None = None, *, kind: str = "CHAIN", attributes: dict[str, Any] | None = None):
+    if callable(name):
+        fn = name
+        return _wrap_callable(fn, name=fn.__name__, kind=kind, extra_attributes=attributes)
+
+    resolved_name = name or "trace"
+
+    def decorator(fn: F) -> F:
+        return _wrap_callable(fn, name=resolved_name, kind=kind, extra_attributes=attributes)
+
+    return decorator if name is not None else _TraceContext(resolved_name, kind=kind, attributes=attributes)
+
+
+@contextmanager
+def trace_context(name: str, *, kind: str = "CHAIN", attributes: dict[str, Any] | None = None):
+    tracer = get_tracer()
+    with tracer.start_span(name, kind=kind, attributes=attributes or {}) as span:
+        yield span

gradient_sdk/errors.py

@@ -0,0 +1,30 @@
+class GradientError(Exception):
+    """Base exception for Gradient SDK errors."""
+
+
+class GradientAuthError(GradientError):
+    """Raised when the API key is missing, invalid, or revoked."""
+
+
+class GradientAPIError(GradientError):
+    """Raised for non-2xx Gradient API responses."""
+
+    def __init__(
+        self,
+        status: int,
+        method: str,
+        path: str,
+        message: str,
+        *,
+        fly_request_id: str | None = None,
+        response_body: str | None = None,
+    ) -> None:
+        self.status = status
+        self.method = method
+        self.path = path
+        self.message = message
+        self.fly_request_id = fly_request_id
+        self.response_body = response_body
+        suffix = f" (fly-request-id={fly_request_id})" if fly_request_id else ""
+        super().__init__(f"Gradient API {status} {method} {path}: {message}{suffix}")
+

gradient_sdk/graph.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Callable
+
+from .tracing import get_tracer
+
+START = "__start__"
+END = "__end__"
+
+
+class CompiledGraph:
+    def __init__(
+        self,
+        *,
+        name: str,
+        nodes: dict[str, Callable[..., Any]],
+        edges: dict[str, list[str]],
+        conditional_edges: dict[str, tuple[Callable[..., Any], dict[str, str] | None]],
+        entry_point: str,
+        recursion_limit: int,
+    ) -> None:
+        self.name = name
+        self.nodes = nodes
+        self.edges = edges
+        self.conditional_edges = conditional_edges
+        self.entry_point = entry_point
+        self.recursion_limit = recursion_limit
+
+    def invoke(self, state: dict[str, Any], *, config: dict[str, Any] | None = None) -> dict[str, Any]:
+        tracer = get_tracer()
+        with tracer.start_span(f"workflow.{self.name}", kind="CHAIN") as root_span:
+            root_span.set_attribute("input.value", str(state))
+            current = dict(state)
+            current_node = self.entry_point
+            steps = 0
+            while current_node not in (END, None):
+                if steps >= self.recursion_limit:
+                    raise RuntimeError(f"workflow recursion limit exceeded ({self.recursion_limit})")
+                steps += 1
+                if current_node not in self.nodes:
+                    raise KeyError(f"unknown workflow node: {current_node}")
+                node_fn = self.nodes[current_node]
+                with tracer.start_span(f"workflow.node.{current_node}", kind="CHAIN") as node_span:
+                    node_span.set_attribute("workflow.node", current_node)
+                    update = self._call_node(node_fn, current, config or {})
+                    if update:
+                        current.update(update)
+                    node_span.set_output(update or {})
+                current_node = self._resolve_next(current_node, current)
+            root_span.set_output(current)
+            return current
+
+    async def ainvoke(self, state: dict[str, Any], *, config: dict[str, Any] | None = None) -> dict[str, Any]:
+        tracer = get_tracer()
+        with tracer.start_span(f"workflow.{self.name}", kind="CHAIN") as root_span:
+            root_span.set_attribute("input.value", str(state))
+            current = dict(state)
+            current_node = self.entry_point
+            steps = 0
+            while current_node not in (END, None):
+                if steps >= self.recursion_limit:
+                    raise RuntimeError(f"workflow recursion limit exceeded ({self.recursion_limit})")
+                steps += 1
+                if current_node not in self.nodes:
+                    raise KeyError(f"unknown workflow node: {current_node}")
+                node_fn = self.nodes[current_node]
+                with tracer.start_span(f"workflow.node.{current_node}", kind="CHAIN") as node_span:
+                    node_span.set_attribute("workflow.node", current_node)
+                    update = await self._acall_node(node_fn, current, config or {})
+                    if update:
+                        current.update(update)
+                    node_span.set_output(update or {})
+                current_node = self._resolve_next(current_node, current)
+            root_span.set_output(current)
+            return current
+
+    def stream(self, state: dict[str, Any], *, config: dict[str, Any] | None = None) -> Iterator[tuple[str, dict[str, Any]]]:
+        current = dict(state)
+        current_node = self.entry_point
+        steps = 0
+        while current_node not in (END, None):
+            if steps >= self.recursion_limit:
+                raise RuntimeError(f"workflow recursion limit exceeded ({self.recursion_limit})")
+            steps += 1
+            if current_node not in self.nodes:
+                raise KeyError(f"unknown workflow node: {current_node}")
+            update = self._call_node(self.nodes[current_node], current, config or {})
+            if update:
+                current.update(update)
+            yield current_node, update or {}
+            current_node = self._resolve_next(current_node, current)
+
+    async def astream(
+        self,
+        state: dict[str, Any],
+        *,
+        config: dict[str, Any] | None = None,
+    ) -> AsyncIterator[tuple[str, dict[str, Any]]]:
+        current = dict(state)
+        current_node = self.entry_point
+        steps = 0
+        while current_node not in (END, None):
+            if steps >= self.recursion_limit:
+                raise RuntimeError(f"workflow recursion limit exceeded ({self.recursion_limit})")
+            steps += 1
+            if current_node not in self.nodes:
+                raise KeyError(f"unknown workflow node: {current_node}")
+            update = await self._acall_node(self.nodes[current_node], current, config or {})
+            if update:
+                current.update(update)
+            yield current_node, update or {}
+            current_node = self._resolve_next(current_node, current)
+
+    def _resolve_next(self, node: str, state: dict[str, Any]) -> str | None:
+        if node in self.conditional_edges:
+            router, mapping = self.conditional_edges[node]
+            route = router(state)
+            if mapping is not None:
+                return mapping.get(route, route)
+            return route
+        next_nodes = self.edges.get(node, [])
+        if not next_nodes:
+            return END
+        return next_nodes[0]
+
+    def _call_node(self, fn: Callable[..., Any], state: dict[str, Any], config: dict[str, Any]) -> dict[str, Any] | None:
+        result = fn(state, config=config) if _accepts_config(fn) else fn(state)
+        if inspect.isawaitable(result):
+            raise TypeError("async node functions require CompiledGraph.ainvoke()")
+        return result
+
+    async def _acall_node(self, fn: Callable[..., Any], state: dict[str, Any], config: dict[str, Any]) -> dict[str, Any] | None:
+        if inspect.iscoroutinefunction(fn):
+            return await fn(state, config=config) if _accepts_config(fn) else await fn(state)
+        return fn(state, config=config) if _accepts_config(fn) else fn(state)
+
+
+def _accepts_config(fn: Callable[..., Any]) -> bool:
+    try:
+        return "config" in inspect.signature(fn).parameters
+    except (TypeError, ValueError):
+        return False
+
+
+class StateGraph:
+    def __init__(self, state_schema: Any = dict, *, name: str = "workflow") -> None:
+        self.state_schema = state_schema
+        self.name = name
+        self._nodes: dict[str, Callable[..., Any]] = {}
+        self._edges: dict[str, list[str]] = {}
+        self._conditional_edges: dict[str, tuple[Callable[..., Any], dict[str, str] | None]] = {}
+        self._entry_point: str | None = None
+
+    def add_node(self, name: str, fn: Callable[..., Any]) -> "StateGraph":
+        self._nodes[name] = fn
+        return self
+
+    def add_edge(self, source: str, target: str) -> "StateGraph":
+        if source == START:
+            self._entry_point = target
+            return self
+        self._edges.setdefault(source, []).append(target)
+        return self
+
+    def add_conditional_edges(
+        self,
+        source: str,
+        router: Callable[..., Any],
+        mapping: dict[str, str] | None = None,
+    ) -> "StateGraph":
+        self._conditional_edges[source] = (router, mapping)
+        return self
+
+    def set_entry_point(self, name: str) -> "StateGraph":
+        self._entry_point = name
+        return self
+
+    def compile(self, *, recursion_limit: int = 25) -> CompiledGraph:
+        if not self._entry_point:
+            raise ValueError("workflow entry point is not set")
+        return CompiledGraph(
+            name=self.name,
+            nodes=dict(self._nodes),
+            edges={key: list(value) for key, value in self._edges.items()},
+            conditional_edges=dict(self._conditional_edges),
+            entry_point=self._entry_point,
+            recursion_limit=recursion_limit,
+        )

gradient_sdk/integrations/__init__.py

@@ -0,0 +1,7 @@
+"""Provider integrations for automatic LLM tracing."""
+
+from .agent import generate, run_agent_loop
+from .anthropic import wrap_anthropic
+from .openai import wrap_openai
+
+__all__ = ["generate", "run_agent_loop", "wrap_anthropic", "wrap_openai"]