datagol-lens 0.2.0__tar.gz

datagol_lens-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 DataGOL
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

datagol_lens-0.2.0/PKG-INFO

@@ -0,0 +1,107 @@
+Metadata-Version: 2.4
+Name: datagol-lens
+Version: 0.2.0
+Summary: Python SDK for DataGOL Lens: tracing, cost and budgets for LLM applications
+Author: DataGOL
+License-Expression: MIT
+Project-URL: Homepage, https://datagol.ai
+Keywords: llm,observability,tracing,langchain,datagol
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: System :: Monitoring
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: httpx>=0.27.0
+Provides-Extra: langchain
+Requires-Dist: langchain-core>=0.3; extra == "langchain"
+Dynamic: license-file
+
+# datagol-lens
+
+Python SDK for [DataGOL Lens](https://datagol.ai): traces, token usage, cost and
+budgets for LLM applications. Wrap your Anthropic or OpenAI client, or attach a
+LangChain callback, and every call shows up in Lens as a trace with its spans.
+
+```bash
+pip install datagol-lens
+pip install "datagol-lens[langchain]"   # for the LangChain / LangGraph handler
+```
+
+## Configure
+
+| Variable | Description |
+|---|---|
+| `DATAGOL_LENS_API_URL` | Your Lens backend URL (default `http://localhost:8000`) |
+| `DATAGOL_LENS_API_KEY` | A project API key (`dgl_...`), created in Lens for the project you trace into |
+
+Both can also be passed to `DatagolLensClient(api_url=..., api_key=...)`.
+
+## Trace one call per request
+
+```python
+import anthropic
+from datagol_lens import DatagolLensClient
+
+sdk = DatagolLensClient()
+client = sdk.wrap(anthropic.Anthropic(), project_id="<project-id>", name="support-bot")
+
+resp = client.messages.create(
+    model="claude-haiku-4-5-20251001",
+    messages=[{"role": "user", "content": "What is Python?"}],
+    max_tokens=200,
+)
+```
+
+The same works for `openai.OpenAI()`.
+
+## Trace a multi-step run
+
+```python
+from datagol_lens import DatagolLensClient, TraceContext
+
+sdk = DatagolLensClient()
+
+with TraceContext(sdk, "<project-id>", name="research-agent",
+                  session_id="sess-abc", user_id="user-123") as ctx:
+    llm = ctx.wrap(anthropic.Anthropic())      # each call becomes a span
+    resp = llm.messages.create(model="claude-haiku-4-5-20251001", max_tokens=300,
+                               messages=[{"role": "user", "content": "Summarize recent AI news"}])
+    ctx.add_span(name="vector_search", span_type="retriever",
+                 input={"query": "AI news"}, output={"chunks": 2})
+    ctx.set_output({"summary": resp.content[0].text})
+# the trace is sent when the block exits
+```
+
+## LangChain and LangGraph
+
+```python
+from datagol_lens import DatagolLensClient, TraceContext
+from datagol_lens.langchain import DatagolLensCallbackHandler
+
+with TraceContext(DatagolLensClient(), "<project-id>", name="langchain-agent") as ctx:
+    handler = DatagolLensCallbackHandler(ctx)
+    answer = chain.invoke("What is Python?", config={"callbacks": [handler]})
+```
+
+Every chat-model call is recorded as an `llm` span with its token usage, and every
+tool call as a `tool` span. Attach the same handler to a LangGraph agent's `config`.
+
+## Budgets and attribution
+
+Pass `company_id`, `agent_key` and `user_id` to `TraceContext` and Lens attributes
+the trace's cost to that company, agent and user. Before a run, check the budgets:
+
+```python
+from datagol_lens import BudgetExceeded
+
+result = sdk.check_budget("<project-id>", company_id="acme", user_id="u-1", agent_key="support-bot")
+if not result["allowed"]:
+    raise BudgetExceeded("over budget", result)
+```
+
+Lens prices each span from its model and token counts, so the SDK only reports usage.
+
+## License
+
+MIT

datagol_lens-0.2.0/README.md

@@ -0,0 +1,88 @@
+# datagol-lens
+
+Python SDK for [DataGOL Lens](https://datagol.ai): traces, token usage, cost and
+budgets for LLM applications. Wrap your Anthropic or OpenAI client, or attach a
+LangChain callback, and every call shows up in Lens as a trace with its spans.
+
+```bash
+pip install datagol-lens
+pip install "datagol-lens[langchain]"   # for the LangChain / LangGraph handler
+```
+
+## Configure
+
+| Variable | Description |
+|---|---|
+| `DATAGOL_LENS_API_URL` | Your Lens backend URL (default `http://localhost:8000`) |
+| `DATAGOL_LENS_API_KEY` | A project API key (`dgl_...`), created in Lens for the project you trace into |
+
+Both can also be passed to `DatagolLensClient(api_url=..., api_key=...)`.
+
+## Trace one call per request
+
+```python
+import anthropic
+from datagol_lens import DatagolLensClient
+
+sdk = DatagolLensClient()
+client = sdk.wrap(anthropic.Anthropic(), project_id="<project-id>", name="support-bot")
+
+resp = client.messages.create(
+    model="claude-haiku-4-5-20251001",
+    messages=[{"role": "user", "content": "What is Python?"}],
+    max_tokens=200,
+)
+```
+
+The same works for `openai.OpenAI()`.
+
+## Trace a multi-step run
+
+```python
+from datagol_lens import DatagolLensClient, TraceContext
+
+sdk = DatagolLensClient()
+
+with TraceContext(sdk, "<project-id>", name="research-agent",
+                  session_id="sess-abc", user_id="user-123") as ctx:
+    llm = ctx.wrap(anthropic.Anthropic())      # each call becomes a span
+    resp = llm.messages.create(model="claude-haiku-4-5-20251001", max_tokens=300,
+                               messages=[{"role": "user", "content": "Summarize recent AI news"}])
+    ctx.add_span(name="vector_search", span_type="retriever",
+                 input={"query": "AI news"}, output={"chunks": 2})
+    ctx.set_output({"summary": resp.content[0].text})
+# the trace is sent when the block exits
+```
+
+## LangChain and LangGraph
+
+```python
+from datagol_lens import DatagolLensClient, TraceContext
+from datagol_lens.langchain import DatagolLensCallbackHandler
+
+with TraceContext(DatagolLensClient(), "<project-id>", name="langchain-agent") as ctx:
+    handler = DatagolLensCallbackHandler(ctx)
+    answer = chain.invoke("What is Python?", config={"callbacks": [handler]})
+```
+
+Every chat-model call is recorded as an `llm` span with its token usage, and every
+tool call as a `tool` span. Attach the same handler to a LangGraph agent's `config`.
+
+## Budgets and attribution
+
+Pass `company_id`, `agent_key` and `user_id` to `TraceContext` and Lens attributes
+the trace's cost to that company, agent and user. Before a run, check the budgets:
+
+```python
+from datagol_lens import BudgetExceeded
+
+result = sdk.check_budget("<project-id>", company_id="acme", user_id="u-1", agent_key="support-bot")
+if not result["allowed"]:
+    raise BudgetExceeded("over budget", result)
+```
+
+Lens prices each span from its model and token counts, so the SDK only reports usage.
+
+## License
+
+MIT

datagol_lens-0.2.0/datagol_lens/__init__.py

@@ -0,0 +1,17 @@
+from datagol_lens.client import DatagolLensClient, LensRequestError, BudgetExceeded
+from datagol_lens.gateway_client import GatewayClient
+from datagol_lens.tracer import TraceContext, estimate_cost
+from datagol_lens.decorators import trace, trace_llm_call
+from datagol_lens.wrappers import TracedAnthropicClient, TracedOpenAIClient, TracedLiteLLMCompletion
+from datagol_lens.redaction import Redactor, RedactionRule, built_in_rule
+from datagol_lens.guardrails import Guardrails, GuardrailRule, GuardrailViolation, built_in_jailbreak_rule
+from datagol_lens.langchain import DatagolLensCallbackHandler
+
+__all__ = [
+    "DatagolLensClient", "LensRequestError", "BudgetExceeded", "GatewayClient", "TraceContext", "estimate_cost",
+    "trace", "trace_llm_call",
+    "TracedAnthropicClient", "TracedOpenAIClient", "TracedLiteLLMCompletion",
+    "Redactor", "RedactionRule", "built_in_rule",
+    "Guardrails", "GuardrailRule", "GuardrailViolation", "built_in_jailbreak_rule",
+    "DatagolLensCallbackHandler",
+]

datagol_lens-0.2.0/datagol_lens/client.py

@@ -0,0 +1,285 @@
+import atexit
+import json
+import os
+import queue
+import threading
+import time
+import httpx
+from typing import Any, Dict, List, Optional, Union
+from datagol_lens.redaction import Redactor
+from datagol_lens.guardrails import Guardrails
+
+# Provider -> env var that Anthropic/OpenAI/litellm SDKs read their API key from by default.
+# Finished traces waiting for the background sender; beyond this they are dropped
+# rather than holding memory while Lens is unreachable.
+_MAX_QUEUE = 1000
+
+_PROVIDER_ENV_VARS = {
+    "anthropic": "ANTHROPIC_API_KEY",
+    "openai": "OPENAI_API_KEY",
+}
+
+
+class LensRequestError(Exception):
+    """A Lens API call failed: network error, timeout, non-2xx status, or bad response body.
+
+    ``status_code`` is set when the server answered with a non-2xx status.
+    """
+
+    def __init__(self, message: str, status_code: Optional[int] = None):
+        super().__init__(message)
+        self.status_code = status_code
+
+
+class BudgetExceeded(Exception):
+    """Raised by callers when a budget check blocks a request.
+
+    ``result`` is the full ``check_budget`` response and ``budget`` its ``blocked_by``
+    BudgetStatus (or None, e.g. when the check itself failed and the caller failed closed).
+    """
+
+    def __init__(self, message: str, result: Optional[Dict[str, Any]] = None):
+        super().__init__(message)
+        self.result = result
+        self.budget = (result or {}).get("blocked_by")
+
+
+class DatagolLensClient:
+    """HTTP client for the Datagol Lens API."""
+
+    def __init__(
+        self,
+        api_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        virtual_key: Optional[str] = None,
+        auto_configure_provider_env: bool = True,
+        redactor: Optional[Redactor] = None,
+        guardrails: Optional[Guardrails] = None,
+        background: bool = True,
+    ):
+        """``background`` (default) sends finished traces from a daemon thread so the
+        traced code never waits on Lens; call :meth:`flush` before a short-lived
+        process exits (an exit hook also tries). Pass False to send inline, for
+        example in a serverless function that is frozen right after returning."""
+        self.background = background
+        self._queue: "queue.Queue[tuple[str, bytes]]" = queue.Queue(maxsize=_MAX_QUEUE)
+        self._sender: Optional[threading.Thread] = None
+        self._sender_lock = threading.Lock()
+        self.dropped_traces = 0
+        self.failed_traces = 0
+        self.api_url = (api_url or os.getenv("DATAGOL_LENS_API_URL", "http://localhost:8000")).rstrip("/")
+        self.api_key = api_key or os.getenv("DATAGOL_LENS_API_KEY", "")
+        self._client = httpx.Client(
+            base_url=self.api_url,
+            headers={"Authorization": f"Bearer {self.api_key}"} if self.api_key else {},
+            timeout=10.0,
+        )
+
+        # Set so every TraceContext created from this client — including via wrap()/@trace,
+        # which don't take a redactor/virtual_key_id argument themselves — picks it up automatically.
+        self.redactor = redactor
+        self.guardrails = guardrails
+        self.virtual_key_id: Optional[str] = None
+        self.resolved_provider: Optional[str] = None
+        self.resolved_allowed_models: Optional[List[str]] = None
+        self._prices: Dict[str, Optional[Dict[str, Any]]] = {}
+
+        virtual_key = virtual_key or os.getenv("DATAGOL_LENS_VIRTUAL_KEY", "")
+        if virtual_key:
+            self.activate_virtual_key(virtual_key, auto_configure_env=auto_configure_provider_env)
+
+    def post_trace(self, project_id: str, payload: dict[str, Any]) -> dict[str, Any]:
+        resp = self._client.post(f"/api/v1/projects/{project_id}/traces", json=payload)
+        resp.raise_for_status()
+        return resp.json()
+
+    def enqueue_trace(self, project_id: str, payload: dict[str, Any]) -> None:
+        """Send a finished trace without blocking the caller (or inline when
+        ``background=False``). Never raises. A full queue drops the trace and
+        counts it in ``dropped_traces``."""
+        # Serialised here, in the caller's thread, so later changes to the caller's
+        # dicts can't alter what gets sent.
+        try:
+            body = json.dumps(payload).encode()
+        except (TypeError, ValueError):
+            self.failed_traces += 1
+            return
+        if not self.background:
+            self._send(project_id, body)
+            return
+        self._ensure_sender()
+        try:
+            self._queue.put_nowait((project_id, body))
+        except queue.Full:
+            self.dropped_traces += 1
+
+    def flush(self, timeout: float = 5.0) -> bool:
+        """Wait for queued traces to be sent. True if the queue drained in time."""
+        deadline = time.monotonic() + timeout
+        while self._queue.unfinished_tasks:
+            if time.monotonic() > deadline:
+                return False
+            time.sleep(0.01)
+        return True
+
+    def _send(self, project_id: str, body: bytes) -> None:
+        # The trace carries its own id, so a retry after a transport failure can't
+        # double-write; one short retry covers a dropped connection or a 5xx blip.
+        for attempt in range(2):
+            try:
+                resp = self._client.post(f"/api/v1/projects/{project_id}/traces", content=body,
+                                         headers={"Content-Type": "application/json"})
+                resp.raise_for_status()
+                return
+            except Exception:
+                if attempt == 0:
+                    time.sleep(0.2)
+        self.failed_traces += 1
+
+    def _ensure_sender(self) -> None:
+        if self._sender is not None and self._sender.is_alive():
+            return
+        with self._sender_lock:
+            if self._sender is None or not self._sender.is_alive():
+                if self._sender is None:
+                    atexit.register(self.flush)
+                self._sender = threading.Thread(target=self._drain, name="datagol-lens-sender", daemon=True)
+                self._sender.start()
+
+    def _drain(self) -> None:
+        while True:
+            project_id, body = self._queue.get()
+            try:
+                self._send(project_id, body)
+            finally:
+                self._queue.task_done()
+
+    def resolve_price(self, model: str) -> Optional[Dict[str, Any]]:
+        """Per-token USD rates for ``model`` from the server's price table
+        (``GET /api/v1/pricing/{model}``), looked up once per model per client.
+        None when the server is unreachable or doesn't know the model, in which
+        case the tracer falls back to its builtin table."""
+        if model in self._prices:
+            return self._prices[model]
+        price = None
+        try:
+            resp = self._client.get(f"/api/v1/pricing/{model}")
+            if resp.status_code == 200:
+                price = resp.json()
+        except Exception:
+            price = None
+        self._prices[model] = price
+        return price
+
+    def post_feedback(self, project_id: str, payload: dict[str, Any]) -> dict[str, Any]:
+        resp = self._client.post(f"/api/v1/projects/{project_id}/feedback", json=payload)
+        resp.raise_for_status()
+        return resp.json()
+
+    def _request_json(self, method: str, path: str, timeout: float, **kwargs) -> Dict[str, Any]:
+        """Send a request and return its JSON body, raising :class:`LensRequestError` on any
+        network error, timeout, non-2xx status, or non-JSON body."""
+        try:
+            resp = self._client.request(method, path, timeout=timeout, **kwargs)
+            resp.raise_for_status()
+            return resp.json()
+        except httpx.HTTPStatusError as exc:
+            raise LensRequestError(
+                f"{method} {path} returned {exc.response.status_code}", exc.response.status_code
+            ) from exc
+        except (httpx.HTTPError, ValueError) as exc:
+            raise LensRequestError(f"{method} {path} failed: {exc}") from exc
+
+    def check_budget(self, project_id: str, company_id: str, user_id: Optional[str] = None,
+                     agent_key: Optional[str] = None, timeout: float = 2.0) -> Dict[str, Any]:
+        """Check whether a company/agent/user is within its Lens budgets.
+
+        Returns ``{"allowed", "decision", "blocked_by", "warnings"}``. Raises
+        :class:`LensRequestError` on network errors, timeouts, and non-2xx responses so the
+        caller can fail closed.
+        """
+        return self._request_json(
+            "POST", f"/api/v1/projects/{project_id}/budget/check", timeout,
+            json={"company_id": company_id, "user_id": user_id, "agent_key": agent_key},
+        )
+
+    def get_enforced_companies(self, project_id: str, timeout: float = 2.0) -> Dict[str, Any]:
+        """Return ``{"company_ids": [...], "version": str}`` — the DataGOL company ids with
+        budget enforcement on. Raises :class:`LensRequestError` on failure."""
+        return self._request_json(
+            "GET", f"/api/v1/projects/{project_id}/budget/enforced-companies", timeout,
+        )
+
+    def resolve_virtual_key(self, token: str, model: Optional[str] = None) -> dict[str, Any]:
+        """Resolve a virtual key ("vk_...") to its real provider key + metadata.
+
+        The real key is returned once per call and never logged by the SDK — callers
+        should generally use :meth:`activate_virtual_key` instead of calling this directly.
+        """
+        resp = self._client.post("/api/v1/virtual-keys/resolve", json={"token": token, "model": model})
+        resp.raise_for_status()
+        return resp.json()
+
+    def activate_virtual_key(self, token: str, auto_configure_env: bool = True) -> None:
+        """Resolve a virtual key and (by default) inject the real provider key into the
+        environment variable that provider's own SDK reads, so ``anthropic.Anthropic()``,
+        ``openai.OpenAI()``, or ``litellm.completion(...)`` work without ever touching
+        the real secret in application code.
+        """
+        result = self.resolve_virtual_key(token)
+        self.virtual_key_id = result["id"]
+        self.resolved_provider = result["provider"]
+        self.resolved_allowed_models = result.get("allowed_models")
+        if auto_configure_env:
+            env_var = _PROVIDER_ENV_VARS.get(self.resolved_provider)
+            if env_var:
+                os.environ[env_var] = result["api_key"]
+
+    def gateway(self, virtual_key: Optional[str] = None) -> "GatewayClient":
+        """Return a :class:`~datagol_lens.gateway_client.GatewayClient` (the
+        LiteLLM/OpenAI-compatible ``/v1/...`` gateway) sharing this client's ``api_url``.
+        """
+        from datagol_lens.gateway_client import GatewayClient
+        return GatewayClient(api_url=self.api_url, virtual_key=virtual_key)
+
+    def get_redaction_rules(self, project_id: str) -> List[Dict[str, Any]]:
+        """Fetch active PII/sensitive-data redaction rules configured for a project."""
+        resp = self._client.get(f"/api/v1/projects/{project_id}/redaction-rules")
+        resp.raise_for_status()
+        return resp.json()
+
+    def get_guardrails(self, project_id: str) -> List[Dict[str, Any]]:
+        """Fetch active prompt guardrail rules configured for a project."""
+        resp = self._client.get(f"/api/v1/projects/{project_id}/guardrails")
+        resp.raise_for_status()
+        return resp.json()
+
+    def wrap(self, llm_client, project_id: str,
+             name: Optional[str] = None,
+             tags: Optional[Union[List[str], Dict[str, str]]] = None,
+             metadata: Optional[Dict[str, Any]] = None,
+             guardrails: Optional[Guardrails] = None):
+        """Wrap an LLM client so each call automatically creates a trace.
+
+        If guardrails are configured (via ``guardrails=`` here or on the client), every
+        prompt is validated against them before it reaches the LLM; a violation of a
+        ``block``-action rule raises :class:`GuardrailViolation` and the call never happens.
+
+        Usage::
+
+            client = sdk.wrap(anthropic.Anthropic(), project_id="proj-123")
+            client.messages.create(model="claude-haiku-4-5-20251001", messages=[...], max_tokens=100)
+        """
+        from datagol_lens.wrappers import wrap_for_trace
+        return wrap_for_trace(llm_client, self, project_id, name=name, tags=tags, metadata=metadata,
+                               guardrails=guardrails)
+
+    def close(self):
+        self.flush()
+        self._client.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        self.close()

datagol_lens-0.2.0/datagol_lens/decorators.py

@@ -0,0 +1,62 @@
+import functools
+import time
+from typing import Any, Callable, Dict, List, Optional, Union
+from datagol_lens.client import DatagolLensClient
+from datagol_lens.tracer import TraceContext, estimate_cost
+
+
+def trace(
+    project_id: str,
+    name: Optional[str] = None,
+    tags: Optional[Union[List[str], Dict[str, str]]] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    client: Optional[DatagolLensClient] = None,
+):
+    """Decorator: wraps a function in a TraceContext and ships the trace on completion.
+
+    Usage::
+
+        @trace(project_id="proj-123", tags=["chat"])
+        def answer_question(question: str) -> str:
+            ...
+    """
+    def decorator(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            _client = client or DatagolLensClient()
+            trace_name = name or fn.__qualname__
+            ctx = TraceContext(
+                client=_client,
+                project_id=project_id,
+                name=trace_name,
+                input={"args": list(args), "kwargs": kwargs},
+                tags=tags or [],
+                metadata=metadata or {},
+            )
+            with ctx:
+                result = fn(*args, **kwargs)
+                ctx.set_output({"result": result})
+            return result
+        return wrapper
+    return decorator
+
+
+def trace_llm_call(
+    ctx: TraceContext,
+    model: str,
+    messages: list[dict],
+    response_text: str,
+    prompt_tokens: int,
+    completion_tokens: int,
+    span_name: str = "llm_call",
+):
+    """Helper to record a single LLM call as a span inside an existing TraceContext."""
+    ctx.add_span(
+        name=span_name,
+        span_type="llm",
+        input={"messages": messages},
+        output={"content": response_text},
+        model=model,
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+    )