usegradient 1.3.1__py3-none-any.whl

gradient_sdk/session.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import atexit
+from dataclasses import dataclass
+from typing import Any, Mapping
+
+from .client import GradientClient
+from .models import CreateTraceRunResponse, MachineCreateRequest
+from .tracing import SpanExporter, Tracer, set_tracer
+
+_current_session: "GradientSession | None" = None
+
+
+@dataclass
+class GradientSession:
+    project: str
+    client: GradientClient
+    run: CreateTraceRunResponse
+    tracer: Tracer
+    machine: dict[str, Any] | None = None
+    full_proxy_url: str | None = None
+    session_id: str | None = None
+    user_id: str | None = None
+    _shutdown: bool = False
+
+    def flush(self) -> None:
+        self.tracer.flush()
+
+    def shutdown(self) -> None:
+        if self._shutdown:
+            return
+        self._shutdown = True
+        self.flush()
+        self.tracer.shutdown()
+        try:
+            self.client.finish_trace_run(self.run.id)
+        except Exception:
+            pass
+        global _current_session
+        if _current_session is self:
+            _current_session = None
+            set_tracer(None)
+
+    def __enter__(self) -> "GradientSession":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> bool:
+        self.shutdown()
+        return False
+
+
+def current_session() -> GradientSession | None:
+    return _current_session
+
+
+def init(
+    project: str = "default",
+    *,
+    client: GradientClient | None = None,
+    environment: MachineCreateRequest | Mapping[str, Any] | None = None,
+    trace_mode: str = "full",
+    session_id: str | None = None,
+    user_id: str | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    capture_policy: Mapping[str, Any] | None = None,
+    wait_for_environment: bool = False,
+    wait_timeout_seconds: int = 30,
+) -> GradientSession:
+    """Initialize a Gradient project session with automatic tracing."""
+
+    global _current_session
+    if _current_session is not None:
+        _current_session.shutdown()
+
+    resolved_client = client or GradientClient.from_credentials()
+    run_metadata = {"project": project, **dict(metadata or {})}
+    machine: dict[str, Any] | None = None
+    full_proxy_url: str | None = None
+
+    if environment is not None:
+        traced = resolved_client.create_traced_machine(
+            environment,
+            trace_mode=trace_mode,
+            metadata=run_metadata,
+            wait=wait_for_environment,
+            wait_timeout_seconds=wait_timeout_seconds,
+        )
+        run = CreateTraceRunResponse.from_dict(traced["trace_run"])
+        run.ingest_token = traced["ingest_token"]
+        run.ingest_url = traced["ingest_url"]
+        machine = traced["machine"]
+        full_proxy_url = traced["full_proxy_url"]
+    else:
+        run = resolved_client.create_trace_run(
+            trace_mode=trace_mode,
+            metadata=run_metadata,
+            capture_policy=capture_policy,
+        )
+
+    ingest_endpoint = resolved_client.api_url.rstrip("/") + "/v1/traces"
+    exporter = SpanExporter(
+        endpoint=ingest_endpoint,
+        ingest_token=run.ingest_token,
+        resource={
+            "service.name": project,
+            "openinference.project.name": project,
+            "gradient.run.id": run.id,
+        },
+        transport=resolved_client._transport,
+        timeout=resolved_client.timeout,
+    )
+    tracer = Tracer(
+        project=project,
+        session_id=session_id,
+        user_id=user_id,
+        exporter=exporter,
+    )
+    set_tracer(tracer)
+
+    session = GradientSession(
+        project=project,
+        client=resolved_client,
+        run=run,
+        tracer=tracer,
+        machine=machine,
+        full_proxy_url=full_proxy_url,
+        session_id=session_id,
+        user_id=user_id,
+    )
+    _current_session = session
+    atexit.register(session.shutdown)
+    return session

gradient_sdk/tools.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import inspect
+import json
+from dataclasses import dataclass, field
+from typing import Any, Callable, get_args, get_origin, get_type_hints
+
+from .context import current_context
+from .tracing import get_tracer
+
+
+@dataclass
+class Tool:
+    name: str
+    description: str
+    fn: Callable[..., Any]
+    parameters: dict[str, Any] = field(default_factory=dict)
+
+    def invoke(self, **kwargs: Any) -> Any:
+        tracer = get_tracer()
+        with tracer.start_span(f"tool.{self.name}", kind="TOOL") as span:
+            span.set_attribute("tool.name", self.name)
+            span.set_attribute("input.value", json.dumps(kwargs, default=str))
+            context = current_context()
+            if context is not None:
+                span.set_attribute("gradient.run.id", context.run_id)
+                span.set_attribute("gradient.project", context.project)
+                if context.session_id:
+                    span.set_attribute("session.id", context.session_id)
+                if context.user_id:
+                    span.set_attribute("user.id", context.user_id)
+            try:
+                result = self.fn(**kwargs)
+                span.set_output(result)
+                return result
+            except Exception as exc:
+                span.record_exception(exc)
+                raise
+
+    def to_openai(self) -> dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters,
+            },
+        }
+
+    def to_anthropic(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "description": self.description,
+            "input_schema": self.parameters,
+        }
+
+
+def _json_type(annotation: Any) -> str:
+    origin = get_origin(annotation)
+    if origin is list:
+        return "array"
+    if origin is dict:
+        return "object"
+    if origin is tuple:
+        return "array"
+    if annotation in (str, inspect.Parameter.empty):
+        return "string"
+    if annotation is int:
+        return "integer"
+    if annotation is float:
+        return "number"
+    if annotation is bool:
+        return "boolean"
+    if origin is type(None) or annotation is type(None):
+        return "null"
+    args = get_args(annotation)
+    if args:
+        non_none = [arg for arg in args if arg is not type(None)]
+        if non_none:
+            return _json_type(non_none[0])
+    return "string"
+
+
+def infer_parameters(fn: Callable[..., Any]) -> dict[str, Any]:
+    signature = inspect.signature(fn)
+    try:
+        hints = get_type_hints(fn)
+    except Exception:
+        hints = {}
+    properties: dict[str, Any] = {}
+    required: list[str] = []
+    for name, param in signature.parameters.items():
+        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+            continue
+        annotation = hints.get(name, param.annotation)
+        properties[name] = {"type": _json_type(annotation)}
+        if param.default is inspect.Parameter.empty:
+            required.append(name)
+    schema: dict[str, Any] = {"type": "object", "properties": properties}
+    if required:
+        schema["required"] = required
+    return schema
+
+
+def tool_from_function(
+    fn: Callable[..., Any],
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    parameters: dict[str, Any] | None = None,
+) -> Tool:
+    resolved_name = name or getattr(fn, "__name__", "tool")
+    resolved_description = description or inspect.getdoc(fn) or f"Tool {resolved_name}"
+    return Tool(
+        name=resolved_name,
+        description=resolved_description,
+        fn=fn,
+        parameters=parameters or infer_parameters(fn),
+    )
+
+
+def execute_tool_calls(tools: list[Tool], calls: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    by_name = {tool.name: tool for tool in tools}
+    results: list[dict[str, Any]] = []
+    for call in calls:
+        name = call.get("name") or call.get("function", {}).get("name")
+        if not name:
+            continue
+        tool = by_name.get(name)
+        if tool is None:
+            results.append({"tool_call_id": call.get("id"), "name": name, "error": f"unknown tool: {name}"})
+            continue
+        raw_args = call.get("arguments") or call.get("function", {}).get("arguments") or call.get("input") or {}
+        if isinstance(raw_args, str):
+            try:
+                raw_args = json.loads(raw_args)
+            except json.JSONDecodeError:
+                raw_args = {}
+        try:
+            output = tool.invoke(**raw_args)
+            results.append(
+                {
+                    "tool_call_id": call.get("id"),
+                    "name": name,
+                    "output": output,
+                }
+            )
+        except Exception as exc:
+            results.append(
+                {
+                    "tool_call_id": call.get("id"),
+                    "name": name,
+                    "error": str(exc),
+                }
+            )
+    return results

gradient_sdk/tracing.py

@@ -0,0 +1,361 @@
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import secrets
+import threading
+import time
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import Any, Callable
+from urllib.request import Request
+
+Transport = Callable[[Request, float], Any]
+
+_current_span: ContextVar["Span | None"] = ContextVar("gradient_current_span", default=None)
+_global_tracer: "Tracer | None" = None
+
+
+def _now_ns() -> int:
+    return time.time_ns()
+
+
+def _otlp_value(value: Any) -> dict[str, Any]:
+    if isinstance(value, bool):
+        return {"boolValue": value}
+    if isinstance(value, int):
+        return {"intValue": str(value)}
+    if isinstance(value, float):
+        return {"doubleValue": value}
+    if isinstance(value, (list, tuple)):
+        return {"arrayValue": {"values": [_otlp_value(item) for item in value]}}
+    return {"stringValue": str(value)}
+
+
+def _otlp_attributes(attributes: dict[str, Any]) -> list[dict[str, Any]]:
+    return [{"key": key, "value": _otlp_value(value)} for key, value in attributes.items()]
+
+
+def _serialize_span(span: "Span", resource: dict[str, Any], scope: dict[str, str]) -> dict[str, Any]:
+    return {
+        "traceId": span.trace_id,
+        "spanId": span.span_id,
+        "parentSpanId": span.parent_span_id or "",
+        "name": span.name,
+        "kind": "SPAN_KIND_INTERNAL",
+        "startTimeUnixNano": str(span.start_time_ns),
+        "endTimeUnixNano": str(span.end_time_ns or _now_ns()),
+        "attributes": _otlp_attributes(span.attributes),
+        "events": [
+            {
+                "timeUnixNano": str(event.get("time_unix_nano", _now_ns())),
+                "name": event.get("name", "exception"),
+                "attributes": _otlp_attributes(event.get("attributes", {})),
+            }
+            for event in span.events
+        ],
+        "status": {
+            "code": "STATUS_CODE_ERROR" if span.status_code == "ERROR" else "STATUS_CODE_OK",
+            "message": span.status_message or "",
+        },
+    }
+
+
+@dataclass
+class Span:
+    name: str
+    trace_id: str
+    span_id: str
+    parent_span_id: str | None = None
+    kind: str = "CHAIN"
+    start_time_ns: int = field(default_factory=_now_ns)
+    end_time_ns: int | None = None
+    attributes: dict[str, Any] = field(default_factory=dict)
+    events: list[dict[str, Any]] = field(default_factory=list)
+    status_code: str = "OK"
+    status_message: str = ""
+    _ended: bool = False
+    _token: Any = None
+    _tracer: "Tracer | None" = None
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        self.attributes[key] = value
+
+    def set_attributes(self, values: dict[str, Any]) -> None:
+        self.attributes.update(values)
+
+    def set_output(self, value: Any, *, mime_type: str = "application/json") -> None:
+        if isinstance(value, str):
+            self.set_attribute("output.value", value)
+            self.set_attribute("output.mime_type", "text/plain")
+            return
+        self.set_attribute("output.value", json.dumps(value, default=str))
+        self.set_attribute("output.mime_type", mime_type)
+
+    def record_exception(self, exc: BaseException) -> None:
+        self.status_code = "ERROR"
+        self.status_message = str(exc)
+        self.events.append(
+            {
+                "name": "exception",
+                "time_unix_nano": _now_ns(),
+                "attributes": {
+                    "exception.type": type(exc).__name__,
+                    "exception.message": str(exc),
+                },
+            }
+        )
+
+    def end(self) -> None:
+        if self._ended:
+            return
+        self.end_time_ns = _now_ns()
+        self._ended = True
+        if self._token is not None:
+            _current_span.reset(self._token)
+            self._token = None
+        if self._tracer is not None:
+            self._tracer._on_span_end(self)
+
+    def __enter__(self) -> "Span":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> bool:
+        if exc is not None and exc_type is not None:
+            self.record_exception(exc)  # type: ignore[arg-type]
+        self.end()
+        return False
+
+
+class NoOpSpan:
+    def set_attribute(self, key: str, value: Any) -> None:
+        return None
+
+    def set_attributes(self, values: dict[str, Any]) -> None:
+        return None
+
+    def set_output(self, value: Any, *, mime_type: str = "application/json") -> None:
+        return None
+
+    def record_exception(self, exc: BaseException) -> None:
+        return None
+
+    def end(self) -> None:
+        return None
+
+    def __enter__(self) -> "NoOpSpan":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> bool:
+        return False
+
+
+class SpanExporter:
+    def __init__(
+        self,
+        *,
+        endpoint: str,
+        ingest_token: str,
+        resource: dict[str, Any],
+        scope_name: str = "gradient-sdk",
+        scope_version: str = "0.3.0",
+        transport: Transport | None = None,
+        timeout: float = 10.0,
+        batch_size: int = 32,
+        flush_interval_seconds: float = 1.0,
+    ) -> None:
+        self.endpoint = endpoint.rstrip("/")
+        self.ingest_token = ingest_token
+        self.resource = resource
+        self.scope = {"name": scope_name, "version": scope_version}
+        self.transport = transport
+        self.timeout = timeout
+        self.batch_size = batch_size
+        self.flush_interval_seconds = flush_interval_seconds
+        self._queue: list[Span] = []
+        self._lock = threading.Lock()
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._worker, daemon=True)
+        self._thread.start()
+        atexit.register(self.shutdown)
+
+    def enqueue(self, span: Span) -> None:
+        with self._lock:
+            self._queue.append(span)
+            if len(self._queue) >= self.batch_size:
+                self._flush_locked()
+
+    def flush(self) -> None:
+        with self._lock:
+            self._flush_locked()
+
+    def shutdown(self) -> None:
+        self._stop.set()
+        self.flush()
+        if self._thread.is_alive():
+            self._thread.join(timeout=2.0)
+
+    def _worker(self) -> None:
+        while not self._stop.wait(self.flush_interval_seconds):
+            self.flush()
+
+    def _flush_locked(self) -> None:
+        if not self._queue:
+            return
+        spans = self._queue[:]
+        self._queue.clear()
+        payload = {
+            "resourceSpans": [
+                {
+                    "resource": {"attributes": _otlp_attributes(self.resource)},
+                    "scopeSpans": [
+                        {
+                            "scope": self.scope,
+                            "spans": [_serialize_span(span, self.resource, self.scope) for span in spans],
+                        }
+                    ],
+                }
+            ]
+        }
+        req = Request(
+            self.endpoint,
+            data=json.dumps(payload).encode("utf-8"),
+            method="POST",
+            headers={
+                "Authorization": f"Bearer {self.ingest_token}",
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+            },
+        )
+        try:
+            if self.transport is not None:
+                with self.transport(req, self.timeout) as res:
+                    _ = res.read()
+            else:
+                from urllib.request import urlopen
+
+                with urlopen(req, timeout=self.timeout) as res:
+                    _ = res.read()
+        except Exception:
+            return
+
+
+class Tracer:
+    def __init__(
+        self,
+        *,
+        project: str,
+        session_id: str | None = None,
+        user_id: str | None = None,
+        exporter: SpanExporter | None = None,
+    ) -> None:
+        self.project = project
+        self.session_id = session_id
+        self.user_id = user_id
+        self.exporter = exporter
+        self.resource = {
+            "service.name": project,
+            "openinference.project.name": project,
+        }
+
+    def start_span(
+        self,
+        name: str,
+        *,
+        kind: str = "CHAIN",
+        attributes: dict[str, Any] | None = None,
+    ) -> Span | NoOpSpan:
+        if self.exporter is None:
+            return NoOpSpan()
+
+        parent = _current_span.get()
+        trace_id = parent.trace_id if parent is not None else secrets.token_hex(16)
+        parent_span_id = parent.span_id if parent is not None else None
+        span = Span(
+            name=name,
+            trace_id=trace_id,
+            span_id=secrets.token_hex(8),
+            parent_span_id=parent_span_id,
+            kind=kind,
+            _tracer=self,
+        )
+        span.set_attribute("openinference.span.kind", kind)
+        if self.session_id:
+            span.set_attribute("session.id", self.session_id)
+        if self.user_id:
+            span.set_attribute("user.id", self.user_id)
+        if attributes:
+            span.set_attributes(attributes)
+        span._token = _current_span.set(span)
+        return span
+
+    def _on_span_end(self, span: Span) -> None:
+        if self.exporter is not None:
+            self.exporter.enqueue(span)
+
+    def flush(self) -> None:
+        if self.exporter is not None:
+            self.exporter.flush()
+
+    def shutdown(self) -> None:
+        if self.exporter is not None:
+            self.exporter.shutdown()
+
+
+class NoOpTracer:
+    def start_span(
+        self,
+        name: str,
+        *,
+        kind: str = "CHAIN",
+        attributes: dict[str, Any] | None = None,
+    ) -> NoOpSpan:
+        return NoOpSpan()
+
+    def flush(self) -> None:
+        return None
+
+    def shutdown(self) -> None:
+        return None
+
+
+_NO_OP_TRACER = NoOpTracer()
+
+
+def get_tracer() -> Tracer | NoOpTracer:
+    global _global_tracer
+    if _global_tracer is None:
+        token = os.environ.get("GRADIENT_TRACE_TOKEN", "")
+        endpoint = os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "")
+        if token and endpoint:
+            project = (
+                os.environ.get("GRADIENT_PROJECT")
+                or os.environ.get("OTEL_SERVICE_NAME")
+                or "gradient"
+            )
+            run_id = os.environ.get("GRADIENT_TRACE_RUN_ID", "")
+            resource = {
+                "service.name": project,
+                "openinference.project.name": project,
+            }
+            if run_id:
+                resource["gradient.run.id"] = run_id
+            _global_tracer = Tracer(
+                project=project,
+                exporter=SpanExporter(
+                    endpoint=endpoint,
+                    ingest_token=token,
+                    resource=resource,
+                ),
+            )
+    return _global_tracer or _NO_OP_TRACER
+
+
+def set_tracer(tracer: Tracer | None) -> None:
+    global _global_tracer
+    _global_tracer = tracer
+
+
+def current_span() -> Span | NoOpSpan | None:
+    return _current_span.get()