harbor-langsmith 0.1.0__tar.gz

harbor_langsmith-0.1.0/PKG-INFO

@@ -0,0 +1,42 @@
+Metadata-Version: 2.4
+Name: harbor-langsmith
+Version: 0.1.0
+Summary: LangSmith plugin for Harbor jobs.
+Author: Alex Shaw
+Author-email: Alex Shaw <alexgshaw64@gmail.com>
+License-Expression: Apache-2.0
+Requires-Dist: harbor>=0.13.0
+Requires-Dist: requests>=2.32.4
+Requires-Python: >=3.12
+Project-URL: Repository, https://github.com/harbor-framework/harbor
+Project-URL: Issues, https://github.com/harbor-framework/harbor/issues
+Description-Content-Type: text/markdown
+
+# harbor-langsmith
+
+LangSmith plugin for Harbor jobs.
+
+```bash
+pip install "harbor[langsmith]"
+export LANGSMITH_API_KEY=...
+harbor run ... --plugin langsmith
+```
+
+You can also pass the full import path:
+
+```bash
+harbor run ... --plugin harbor_langsmith:LangSmithPlugin
+```
+
+Optional environment variables:
+
+- `HARBOR_LANGSMITH_DATASET`
+- `HARBOR_LANGSMITH_EXPERIMENT`
+- `LANGSMITH_ENDPOINT`
+- `LANGSMITH_WORKSPACE_ID`
+- `HARBOR_LANGSMITH_SYNC_DATASET=false`
+- `HARBOR_LANGSMITH_FAIL_FAST=true`
+
+Plugin kwargs (CLI `--pk` or job config `kwargs:`) mirror the constructor options:
+`dataset_name`, `experiment_name`, `endpoint`, `api_key`, `workspace_id`,
+`sync_dataset`, and `fail_fast`.

harbor_langsmith-0.1.0/README.md

@@ -0,0 +1,28 @@
+# harbor-langsmith
+
+LangSmith plugin for Harbor jobs.
+
+```bash
+pip install "harbor[langsmith]"
+export LANGSMITH_API_KEY=...
+harbor run ... --plugin langsmith
+```
+
+You can also pass the full import path:
+
+```bash
+harbor run ... --plugin harbor_langsmith:LangSmithPlugin
+```
+
+Optional environment variables:
+
+- `HARBOR_LANGSMITH_DATASET`
+- `HARBOR_LANGSMITH_EXPERIMENT`
+- `LANGSMITH_ENDPOINT`
+- `LANGSMITH_WORKSPACE_ID`
+- `HARBOR_LANGSMITH_SYNC_DATASET=false`
+- `HARBOR_LANGSMITH_FAIL_FAST=true`
+
+Plugin kwargs (CLI `--pk` or job config `kwargs:`) mirror the constructor options:
+`dataset_name`, `experiment_name`, `endpoint`, `api_key`, `workspace_id`,
+`sync_dataset`, and `fail_fast`.

harbor_langsmith-0.1.0/pyproject.toml

@@ -0,0 +1,29 @@
+[project]
+name = "harbor-langsmith"
+version = "0.1.0"
+description = "LangSmith plugin for Harbor jobs."
+readme = "README.md"
+license = "Apache-2.0"
+authors = [{ name = "Alex Shaw", email = "alexgshaw64@gmail.com" }]
+requires-python = ">=3.12"
+dependencies = [
+    "harbor>=0.13.0",
+    "requests>=2.32.4",
+]
+
+[project.entry-points."harbor.plugins"]
+langsmith = "harbor_langsmith:LangSmithPlugin"
+
+[project.urls]
+Repository = "https://github.com/harbor-framework/harbor"
+Issues = "https://github.com/harbor-framework/harbor/issues"
+
+[tool.uv.sources]
+harbor = { workspace = true }
+
+[build-system]
+requires = ["uv_build>=0.10.8,<0.11.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "harbor_langsmith"

harbor_langsmith-0.1.0/src/harbor_langsmith/__init__.py

@@ -0,0 +1,3 @@
+from harbor_langsmith.plugin import LangSmithPlugin
+
+__all__ = ["LangSmithPlugin"]

harbor_langsmith-0.1.0/src/harbor_langsmith/plugin.py

@@ -0,0 +1,496 @@
+import asyncio
+import os
+import tomllib
+from datetime import datetime, timezone
+from typing import Any
+from uuid import NAMESPACE_URL, uuid5
+
+import requests
+
+from harbor.job import Job
+from harbor.models.job.plugin import BaseJobPlugin
+from harbor.models.job.result import JobResult
+from harbor.trial.hooks import TrialEvent, TrialHookEvent
+
+
+class LangSmithPlugin(BaseJobPlugin):
+    def __init__(
+        self,
+        *,
+        dataset_name: str | None = None,
+        experiment_name: str | None = None,
+        endpoint: str | None = None,
+        api_key: str | None = None,
+        workspace_id: str | None = None,
+        sync_dataset: bool | None = None,
+        fail_fast: bool | None = None,
+    ):
+        super().__init__()
+        self.dataset_name = dataset_name or os.getenv("HARBOR_LANGSMITH_DATASET")
+        self.experiment_name = experiment_name or os.getenv(
+            "HARBOR_LANGSMITH_EXPERIMENT"
+        )
+        self.endpoint = endpoint or os.getenv("LANGSMITH_ENDPOINT")
+        self.api_key = api_key or os.getenv("LANGSMITH_API_KEY")
+        self.workspace_id = workspace_id or os.getenv("LANGSMITH_WORKSPACE_ID")
+        self.sync_dataset = (
+            self._env_bool("HARBOR_LANGSMITH_SYNC_DATASET", default=True)
+            if sync_dataset is None
+            else sync_dataset
+        )
+        self.fail_fast = (
+            self._env_bool("HARBOR_LANGSMITH_FAIL_FAST", default=False)
+            if fail_fast is None
+            else fail_fast
+        )
+        self.request_timeout = float(
+            os.getenv("HARBOR_LANGSMITH_REQUEST_TIMEOUT", "120")
+        )
+        self._session = requests.Session()
+        self._base_url = ""
+        self._dataset_id: str | None = None
+        self._experiment_id: str | None = None
+        self._example_ids: dict[str, str] = {}
+        self._run_ids: dict[str, str] = {}
+        self._phase_run_ids: dict[tuple[str, TrialEvent], str] = {}
+
+    async def on_job_start(self, job: Job) -> None:
+        await asyncio.to_thread(self._setup, job)
+        job.on_trial_started(self._handle_event)
+        job.on_environment_started(self._handle_event)
+        job.on_agent_started(self._handle_event)
+        job.on_verification_started(self._handle_event)
+        job.on_trial_ended(self._handle_event)
+        job.on_trial_cancelled(self._handle_event)
+
+    async def on_job_end(self, job_result: JobResult) -> None:
+        if self._experiment_id is None:
+            return
+        try:
+            await asyncio.to_thread(
+                self._request,
+                "PATCH",
+                f"/sessions/{self._experiment_id}",
+                json={
+                    "end_time": self._format_time(
+                        job_result.finished_at or datetime.now(timezone.utc)
+                    )
+                },
+                ok_statuses={200, 202, 204},
+            )
+        except Exception:
+            if self.fail_fast:
+                raise
+
+    def _setup(self, job: Any) -> None:
+        if not self.api_key:
+            raise RuntimeError("LANGSMITH_API_KEY is required for LangSmithPlugin")
+
+        endpoint = (self.endpoint or "https://api.smith.langchain.com").rstrip("/")
+        self._base_url = (
+            endpoint if endpoint.endswith("/api/v1") else f"{endpoint}/api/v1"
+        )
+
+        self._session.headers.update({"x-api-key": self.api_key})
+        if self.workspace_id:
+            self._session.headers.update({"LANGSMITH-WORKSPACE-ID": self.workspace_id})
+
+        if self.sync_dataset:
+            self._dataset_id = self._get_or_create_dataset(job)
+            self._example_ids = self._get_or_create_examples(job)
+
+        experiment_id = self._stable_uuid(job.id, "experiment")
+        payload: dict[str, Any] = {
+            "id": experiment_id,
+            "name": self.experiment_name or job.config.job_name,
+            "start_time": self._format_time(datetime.now(timezone.utc)),
+            "extra": {
+                "metadata": {
+                    "harbor_job_id": str(job.id),
+                    "harbor_job_name": job.config.job_name,
+                    "harbor_job_dir": str(job.job_dir),
+                }
+            },
+        }
+        if self._dataset_id is not None:
+            payload["reference_dataset_id"] = self._dataset_id
+
+        self._request("POST", "/sessions", json=payload, ok_statuses={200, 201, 409})
+        self._experiment_id = experiment_id
+
+    async def _handle_event(self, event: TrialHookEvent) -> None:
+        try:
+            await asyncio.to_thread(self._handle_event_sync, event)
+        except Exception:
+            if self.fail_fast:
+                raise
+
+    def _handle_event_sync(self, event: TrialHookEvent) -> None:
+        if self._experiment_id is None:
+            return
+        if event.event == TrialEvent.START:
+            self._create_root_run(event)
+            return
+        if event.event in {
+            TrialEvent.ENVIRONMENT_START,
+            TrialEvent.AGENT_START,
+            TrialEvent.VERIFICATION_START,
+        }:
+            self._create_phase_run(event)
+            return
+        if event.event in {TrialEvent.END, TrialEvent.CANCEL}:
+            self._finish_trial(event)
+
+    def _get_or_create_dataset(self, job: Any) -> str | None:
+        dataset_name = self.dataset_name or self._default_dataset_name(job)
+        if dataset_name is None:
+            return None
+
+        existing = self._find_dataset(dataset_name)
+        if existing is not None:
+            return existing
+
+        payload = {
+            "name": dataset_name,
+            "description": f"Harbor dataset synced from job {job.config.job_name}",
+            "metadata": {"source": "harbor"},
+        }
+        response = self._request(
+            "POST", "/datasets", json=payload, ok_statuses={200, 201, 409}
+        )
+        if response.status_code == 409:
+            return self._find_dataset(dataset_name)
+        return self._extract_id(response.json())
+
+    def _get_or_create_examples(self, job: Any) -> dict[str, str]:
+        if self._dataset_id is None:
+            return {}
+
+        example_ids: dict[str, str] = {}
+        for task_config in job._task_configs:
+            task_id = task_config.get_task_id()
+            full_task_name = task_id.get_name()
+            task_name = full_task_name.split("/")[-1]
+            example_id = self._stable_uuid(self._dataset_id, "example", task_name)
+            instruction = self._read_instruction(task_config)
+            payload = {
+                "id": example_id,
+                "dataset_id": self._dataset_id,
+                "inputs": {
+                    "task_name": task_name,
+                    "instruction": instruction,
+                    "task_id": task_id.model_dump(mode="json"),
+                },
+                "outputs": {},
+                "metadata": {
+                    "source": "harbor",
+                    "task_config": task_config.model_dump(mode="json"),
+                },
+            }
+            response = self._request(
+                "POST", "/examples", json=payload, ok_statuses={200, 201, 409}
+            )
+            if response.status_code == 409:
+                self._request(
+                    "PATCH",
+                    f"/examples/{example_id}",
+                    json={
+                        "inputs": payload["inputs"],
+                        "outputs": payload["outputs"],
+                        "metadata": payload["metadata"],
+                    },
+                    ok_statuses={200, 202, 204},
+                )
+            example_ids[task_name] = example_id
+            example_ids[full_task_name] = example_id
+            configured_task_name = self._read_configured_task_name(task_config)
+            if configured_task_name is not None:
+                example_ids[configured_task_name] = example_id
+                example_ids[configured_task_name.split("/")[-1]] = example_id
+        return example_ids
+
+    def _create_root_run(self, event: TrialHookEvent) -> None:
+        run_id = self._stable_uuid(
+            event.config.job_id, "trial", event.config.trial_name
+        )
+        self._run_ids[event.config.trial_name] = run_id
+        reference_example_id = self._example_ids.get(
+            event.task_name
+        ) or self._example_ids.get(event.task_name.split("/")[-1])
+        payload = {
+            "id": run_id,
+            "name": event.config.trial_name,
+            "run_type": "chain",
+            "inputs": {
+                "task_name": event.task_name,
+                "instruction": self._read_instruction(event.config.task),
+                "task_id": event.config.task.get_task_id().model_dump(mode="json"),
+                "trial_name": event.config.trial_name,
+                "agent": event.config.agent.name or event.config.agent.import_path,
+                "model": event.config.agent.model_name,
+            },
+            "start_time": self._format_time(event.timestamp),
+            "session_id": self._experiment_id,
+            "reference_example_id": reference_example_id,
+            "extra": {
+                "metadata": self._trial_metadata(event),
+                "tags": ["harbor", "harbor-trial"],
+            },
+        }
+        if payload["reference_example_id"] is None:
+            payload.pop("reference_example_id")
+        self._request("POST", "/runs", json=payload, ok_statuses={200, 201, 409})
+
+    def _create_phase_run(self, event: TrialHookEvent) -> None:
+        parent_run_id = self._run_ids.get(event.config.trial_name)
+        if parent_run_id is None:
+            self._create_root_run(event)
+            parent_run_id = self._run_ids[event.config.trial_name]
+
+        run_id = self._stable_uuid(
+            event.config.job_id,
+            "trial",
+            event.config.trial_name,
+            "phase",
+            event.event.value,
+        )
+        self._phase_run_ids[(event.config.trial_name, event.event)] = run_id
+        payload = {
+            "id": run_id,
+            "name": event.event.value,
+            "run_type": "chain",
+            "inputs": {"phase": event.event.value},
+            "start_time": self._format_time(event.timestamp),
+            "session_id": self._experiment_id,
+            "parent_run_id": parent_run_id,
+            "extra": {
+                "metadata": self._trial_metadata(event),
+                "tags": ["harbor", "harbor-phase", event.event.value],
+            },
+        }
+        self._request("POST", "/runs", json=payload, ok_statuses={200, 201, 409})
+
+    def _finish_trial(self, event: TrialHookEvent) -> None:
+        run_id = self._run_ids.get(event.config.trial_name)
+        if run_id is None:
+            self._create_root_run(event)
+            run_id = self._run_ids[event.config.trial_name]
+
+        result = event.result
+        payload: dict[str, Any] = {
+            "outputs": self._trial_outputs(result),
+            "end_time": self._format_time(
+                result.finished_at if result and result.finished_at else event.timestamp
+            ),
+        }
+        if result and result.exception_info is not None:
+            payload["error"] = (
+                f"{result.exception_info.exception_type}: "
+                f"{result.exception_info.exception_message}"
+            )
+        self._request(
+            "PATCH", f"/runs/{run_id}", json=payload, ok_statuses={200, 202, 204}
+        )
+
+        if result is not None:
+            self._finish_phase_runs(result)
+            self._create_feedback(run_id, result)
+
+    def _finish_phase_runs(self, result: Any) -> None:
+        phases = {
+            TrialEvent.ENVIRONMENT_START: result.environment_setup,
+            TrialEvent.AGENT_START: result.agent_execution,
+            TrialEvent.VERIFICATION_START: result.verifier,
+        }
+        for event, timing in phases.items():
+            if timing is None:
+                continue
+            run_id = self._phase_run_ids.get((result.trial_name, event))
+            if run_id is None or timing.finished_at is None:
+                continue
+            self._request(
+                "PATCH",
+                f"/runs/{run_id}",
+                json={
+                    "outputs": {"phase": event.value},
+                    "end_time": self._format_time(timing.finished_at),
+                },
+                ok_statuses={200, 202, 204},
+            )
+
+    def _create_feedback(self, run_id: str, result: Any) -> None:
+        if result.verifier_result is not None:
+            for key, score in result.verifier_result.rewards.items():
+                self._request(
+                    "POST",
+                    "/feedback",
+                    json={
+                        "id": self._stable_uuid(run_id, "feedback", key),
+                        "run_id": run_id,
+                        "key": key,
+                        "score": score,
+                        "feedback_source_type": "api",
+                    },
+                    ok_statuses={200, 201, 409},
+                )
+        if result.exception_info is not None:
+            self._request(
+                "POST",
+                "/feedback",
+                json={
+                    "id": self._stable_uuid(run_id, "feedback", "harbor_error"),
+                    "run_id": run_id,
+                    "key": "harbor_error",
+                    "score": 1,
+                    "value": result.exception_info.exception_type,
+                    "comment": result.exception_info.exception_message,
+                    "feedback_source_type": "api",
+                },
+                ok_statuses={200, 201, 409},
+            )
+
+    def _trial_outputs(self, result: Any | None) -> dict[str, Any]:
+        if result is None:
+            return {}
+        n_input, n_cache, n_output, cost = result.compute_token_cost_totals()
+        agent_metadata = (
+            result.agent_result.metadata
+            if result.agent_result is not None
+            and result.agent_result.metadata is not None
+            else None
+        )
+        return {
+            "task_name": result.task_name,
+            "trial_name": result.trial_name,
+            "agent_output": (
+                agent_metadata.get("answer_written") if agent_metadata else None
+            ),
+            "agent_metadata": agent_metadata,
+            "rewards": (
+                result.verifier_result.rewards
+                if result.verifier_result is not None
+                else None
+            ),
+            "exception": (
+                result.exception_info.model_dump(mode="json")
+                if result.exception_info is not None
+                else None
+            ),
+            "tokens": {
+                "input": n_input,
+                "cache": n_cache,
+                "output": n_output,
+            },
+            "cost_usd": cost,
+        }
+
+    def _trial_metadata(self, event: TrialHookEvent) -> dict[str, Any]:
+        return {
+            "harbor_trial_id": event.trial_id,
+            "harbor_trial_name": event.config.trial_name,
+            "harbor_task_name": event.task_name,
+            "harbor_job_id": str(event.config.job_id),
+            "harbor_agent": event.config.agent.name or event.config.agent.import_path,
+            "harbor_model": event.config.agent.model_name,
+            "harbor_trial_config": event.config.model_dump(mode="json"),
+            "source": "harbor",
+        }
+
+    @staticmethod
+    def _read_instruction(task_config: Any) -> str | None:
+        try:
+            instruction_path = task_config.get_local_path() / "instruction.md"
+            if instruction_path.exists():
+                return instruction_path.read_text()
+        except Exception:
+            return None
+        return None
+
+    @staticmethod
+    def _read_configured_task_name(task_config: Any) -> str | None:
+        try:
+            config_path = task_config.get_local_path() / "task.toml"
+            if not config_path.exists():
+                return None
+            data = tomllib.loads(config_path.read_text())
+            task = data.get("task")
+            if isinstance(task, dict):
+                name = task.get("name")
+                return name if isinstance(name, str) else None
+        except Exception:
+            return None
+        return None
+
+    def _default_dataset_name(self, job: Any) -> str | None:
+        if len(job.config.datasets) == 1:
+            dataset = job.config.datasets[0]
+            name = dataset.name or (dataset.path.name if dataset.path else None)
+            if name is not None and dataset.version:
+                return f"{name}@{dataset.version}"
+            return name
+        if len(job.config.tasks) == 1:
+            task_id = job.config.tasks[0].get_task_id()
+            return task_id.get_name().split("/")[-1]
+        if len(job.config.tasks) > 1:
+            return "harbor-adhoc-tasks"
+        return None
+
+    def _find_dataset(self, dataset_name: str) -> str | None:
+        for params in ({"name": dataset_name}, {"dataset_name": dataset_name}):
+            response = self._request(
+                "GET", "/datasets", params=params, ok_statuses={200, 404}
+            )
+            if response.status_code == 404:
+                continue
+            datasets = response.json()
+            if isinstance(datasets, dict):
+                datasets = datasets.get("datasets") or datasets.get("items") or []
+            if not isinstance(datasets, list):
+                continue
+            for dataset in datasets:
+                if dataset.get("name") == dataset_name:
+                    return self._extract_id(dataset)
+        return None
+
+    def _request(
+        self,
+        method: str,
+        path: str,
+        *,
+        ok_statuses: set[int],
+        **kwargs: Any,
+    ) -> requests.Response:
+        response = self._session.request(
+            method,
+            f"{self._base_url}{path}",
+            timeout=self.request_timeout,
+            **kwargs,
+        )
+        if response.status_code not in ok_statuses:
+            response.raise_for_status()
+        return response
+
+    @staticmethod
+    def _extract_id(payload: Any) -> str | None:
+        if isinstance(payload, dict):
+            value = payload.get("id") or payload.get("dataset_id")
+            return str(value) if value is not None else None
+        return None
+
+    @staticmethod
+    def _stable_uuid(*parts: Any) -> str:
+        normalized = ":".join(str(part) for part in parts if part is not None)
+        return str(uuid5(NAMESPACE_URL, f"harbor-langsmith:{normalized}"))
+
+    @staticmethod
+    def _format_time(value: datetime) -> str:
+        if value.tzinfo is None:
+            value = value.replace(tzinfo=timezone.utc)
+        return value.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
+
+    @staticmethod
+    def _env_bool(name: str, *, default: bool) -> bool:
+        value = os.getenv(name)
+        if value is None:
+            return default
+        return value.lower() in {"1", "true", "yes", "on"}