docent-python 0.1.5a0__py3-none-any.whl → 0.1.6a0__py3-none-any.whl

docent/data_models/agent_run.py

@@ -15,6 +15,7 @@ from pydantic import (
 from docent.data_models._tiktoken_util import get_token_count, group_messages_into_ranges
 from docent.data_models.transcript import (
     Transcript,
+    TranscriptGroup,
     TranscriptWithoutMetadataValidator,
     fake_model_dump,
 )
@@ -36,6 +37,7 @@ class AgentRun(BaseModel):
         name: Optional human-readable name for the agent run.
         description: Optional description of the agent run.
         transcripts: Dict mapping transcript IDs to Transcript objects.
+        transcript_groups: Dict mapping transcript group IDs to TranscriptGroup objects.
         metadata: Additional structured metadata about the agent run as a JSON-serializable dictionary.
     """
 
@@ -44,6 +46,7 @@ class AgentRun(BaseModel):
     description: str | None = None
 
     transcripts: dict[str, Transcript]
+    transcript_groups: dict[str, TranscriptGroup] = Field(default_factory=dict)
     metadata: dict[str, Any] = Field(default_factory=dict)
 
     @field_serializer("metadata")

docent/data_models/transcript.py

@@ -1,4 +1,5 @@
 import sys
+from datetime import datetime
 from typing import Any
 from uuid import uuid4
 
@@ -73,6 +74,8 @@ class TranscriptGroup(BaseModel):
         id: Unique identifier for the transcript group, auto-generated by default.
         name: Optional human-readable name for the transcript group.
         description: Optional description of the transcript group.
+        collection_id: ID of the collection this transcript group belongs to.
+        agent_run_id: ID of the agent run this transcript group belongs to.
         parent_transcript_group_id: Optional ID of the parent transcript group.
         metadata: Additional structured metadata about the transcript group.
     """
@@ -80,7 +83,10 @@ class TranscriptGroup(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid4()))
     name: str | None = None
     description: str | None = None
+    collection_id: str
+    agent_run_id: str
     parent_transcript_group_id: str | None = None
+    created_at: datetime | None = None
     metadata: dict[str, Any] = Field(default_factory=dict)
 
     @field_serializer("metadata")
@@ -129,6 +135,7 @@ class Transcript(BaseModel):
     name: str | None = None
     description: str | None = None
     transcript_group_id: str | None = None
+    created_at: datetime | None = None
 
     messages: list[ChatMessage]
     metadata: dict[str, Any] = Field(default_factory=dict)

docent/loaders/load_inspect.py

@@ -1,4 +1,7 @@
-from typing import Any
+import json
+from pathlib import Path
+from typing import Any, BinaryIO, Generator, Tuple
+from zipfile import ZipFile
 
 from inspect_ai.log import EvalLog
 from inspect_ai.scorer import CORRECT, INCORRECT, NOANSWER, PARTIAL, Score
@@ -7,9 +10,9 @@ from docent.data_models import AgentRun, Transcript
 from docent.data_models.chat import parse_chat_message
 
 
-def _normalize_inspect_score(score: Score) -> Any:
+def _normalize_inspect_score(score: Score | dict[str, Any]) -> Any:
     """
-    Normalize an inspect score to a float. This implements the same logic as inspect_ai.scorer._metric.value_to_float, but fails more conspicuously.
+    Normalize an inspect score to a float. Logic mirrors inspect_ai.scorer._metric.value_to_float.
 
     Args:
         score: The inspect score to normalize.
@@ -18,7 +21,7 @@ def _normalize_inspect_score(score: Score) -> Any:
         The normalized score as a float, or None if the score is not a valid value.
     """
 
-    def _leaf_normalize(value: int | float | bool | str | None) -> float | str | None:
+    def _leaf_normalize(value: Any) -> Any:
         if value is None:
             return None
         if isinstance(value, int | float | bool):
@@ -38,12 +41,17 @@ def _normalize_inspect_score(score: Score) -> Any:
             return float(value)
         return value
 
-    if isinstance(score.value, int | float | bool | str):
-        return _leaf_normalize(score.value)
-    if isinstance(score.value, list):
-        return [_leaf_normalize(v) for v in score.value]
-    assert isinstance(score.value, dict), "Inspect score must be leaf value, list, or dict"
-    return {k: _leaf_normalize(v) for k, v in score.value.items()}
+    if isinstance(score, dict):
+        value = score["value"]
+    else:
+        value = score.value
+
+    if isinstance(value, int | float | bool | str):
+        return _leaf_normalize(value)
+    if isinstance(value, list):
+        return [_leaf_normalize(v) for v in value]  # type: ignore
+    assert isinstance(value, dict), "Inspect score must be leaf value, list, or dict"
+    return {k: _leaf_normalize(v) for k, v in value.items()}  # type: ignore
 
 
 def load_inspect_log(log: EvalLog) -> list[AgentRun]:
@@ -86,3 +94,117 @@ def load_inspect_log(log: EvalLog) -> list[AgentRun]:
         )
 
     return agent_runs
+
+
+def _read_sample_as_run(data: dict[str, Any], header_metadata: dict[str, Any] = {}) -> AgentRun:
+    if "scores" in data:
+        normalized_scores = {k: _normalize_inspect_score(v) for k, v in data["scores"].items()}
+    else:
+        normalized_scores = {}
+
+    if "metadata" in data:
+        sample_metadata = data["metadata"]
+    else:
+        sample_metadata = {}
+
+    run_metadata: dict[str, Any] = {
+        "sample_id": data.get("id"),
+        "epoch": data.get("epoch"),
+        "target": data.get("target"),
+        # Scores could have answers, explanations, and other metadata besides the values we extract
+        "scoring_metadata": data.get("scores"),
+        "scores": normalized_scores,
+        # If a key exists in header and sample, sample takes precedence
+        **header_metadata,
+        **sample_metadata,
+    }
+
+    run = AgentRun(
+        transcripts={
+            "main": Transcript(
+                messages=[parse_chat_message(m) for m in data["messages"]], metadata={}
+            ),
+        },
+        metadata=run_metadata,
+    )
+    return run
+
+
+def _run_metadata_from_header(header: dict[str, Any]) -> dict[str, Any]:
+    """
+    Inspect logs often have a lot of metadata.
+    This function tries to get the most important stuff without adding clutter.
+    """
+    m: dict[str, Any] = {}
+    if e := header.get("eval"):
+        m["task"] = e["task"]
+        m["model"] = e["model"]
+    return m
+
+
+def get_total_samples(file_path: Path, format: str = "json") -> int:
+    """Return the total number of samples in the provided file."""
+    with open(file_path, "rb") as f:
+        if format == "json":
+            data = json.load(f)
+            return len(data.get("samples", []))
+        elif format == "eval":
+            z = ZipFile(f, mode="r")
+            try:
+                return sum(
+                    1
+                    for name in z.namelist()
+                    if name.startswith("samples/") and name.endswith(".json")
+                )
+            finally:
+                z.close()
+        else:
+            raise ValueError(f"Format must be 'json' or 'eval': {format}")
+
+
+def _runs_from_eval_file(
+    file: BinaryIO,
+) -> Tuple[dict[str, Any], Generator[AgentRun, None, None]]:
+    zip = ZipFile(file, mode="r")
+    header: dict[str, Any] = json.load(zip.open("header.json", "r"))
+    header_metadata = _run_metadata_from_header(header)
+
+    def _iter_runs() -> Generator[AgentRun, None, None]:
+        try:
+            for sample_file in zip.namelist():
+                if not (sample_file.startswith("samples/") and sample_file.endswith(".json")):
+                    continue
+                with zip.open(sample_file, "r") as f:
+                    data = json.load(f)
+                run: AgentRun = _read_sample_as_run(data, header_metadata)
+                yield run
+        finally:
+            zip.close()
+
+    return header_metadata, _iter_runs()
+
+
+def _runs_from_json_file(
+    file: BinaryIO,
+) -> Tuple[dict[str, Any], Generator[AgentRun, None, None]]:
+    data = json.load(file)
+    header_metadata = _run_metadata_from_header(data)
+
+    def _iter_runs() -> Generator[AgentRun, None, None]:
+        for sample in data["samples"]:
+            run: AgentRun = _read_sample_as_run(sample, header_metadata)
+            yield run
+
+    return header_metadata, _iter_runs()
+
+
+def runs_from_file(
+    file: BinaryIO, format: str = "json"
+) -> Tuple[dict[str, Any], Generator[AgentRun, None, None]]:
+    if format == "json":
+        result = _runs_from_json_file(file)
+    elif format == "eval":
+        result = _runs_from_eval_file(file)
+    else:
+        raise ValueError(f"Format must be 'json' or 'eval': {format}")
+    return result

docent/trace.py

@@ -12,6 +12,7 @@ from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar, Token
 from datetime import datetime, timezone
 from enum import Enum
+from importlib.metadata import Distribution, distributions
 from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Set, Union
 
 import requests
@@ -19,10 +20,6 @@ from opentelemetry import trace
 from opentelemetry.context import Context
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as GRPCExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as HTTPExporter
-from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor
-from opentelemetry.instrumentation.bedrock import BedrockInstrumentor
-from opentelemetry.instrumentation.langchain import LangchainInstrumentor
-from opentelemetry.instrumentation.openai import OpenAIInstrumentor
 from opentelemetry.instrumentation.threading import ThreadingInstrumentor
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
@@ -34,9 +31,8 @@ from opentelemetry.sdk.trace.export import (
 from opentelemetry.trace import Span
 
 # Configure logging
-logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-logger.disabled = True
+logger.setLevel(logging.ERROR)
 
 # Default configuration
 DEFAULT_ENDPOINT = "https://api.docent.transluce.org/rest/telemetry"
@@ -223,7 +219,7 @@ class DocentTracer:
                 exporters.append(exporter)
                 logger.info(f"Initialized exporter for endpoint: {endpoint}")
             else:
-                logger.warning(f"Failed to initialize exporter for endpoint: {endpoint}")
+                logger.critical(f"Failed to initialize exporter for endpoint: {endpoint}")
 
         return exporters
 
@@ -326,8 +322,6 @@ class DocentTracer:
                     logger.info(
                         f"Added {len(otlp_exporters)} OTLP exporters for {len(self.endpoints)} endpoints"
                     )
-                else:
-                    logger.warning("Failed to initialize OTLP exporter")
 
             if self.enable_console_export:
                 console_exporter: ConsoleSpanExporter = ConsoleSpanExporter()
@@ -355,32 +349,44 @@ class DocentTracer:
             # Instrument OpenAI with our isolated tracer provider
             if Instruments.OPENAI in enabled_instruments:
                 try:
-                    OpenAIInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                    logger.info("Instrumented OpenAI")
+                    if is_package_installed("openai"):
+                        from opentelemetry.instrumentation.openai import OpenAIInstrumentor
+
+                        OpenAIInstrumentor().instrument(tracer_provider=self._tracer_provider)
+                        logger.info("Instrumented OpenAI")
                 except Exception as e:
                     logger.warning(f"Failed to instrument OpenAI: {e}")
 
             # Instrument Anthropic with our isolated tracer provider
             if Instruments.ANTHROPIC in enabled_instruments:
                 try:
-                    AnthropicInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                    logger.info("Instrumented Anthropic")
+                    if is_package_installed("anthropic"):
+                        from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor
+
+                        AnthropicInstrumentor().instrument(tracer_provider=self._tracer_provider)
+                        logger.info("Instrumented Anthropic")
                 except Exception as e:
                     logger.warning(f"Failed to instrument Anthropic: {e}")
 
             # Instrument Bedrock with our isolated tracer provider
             if Instruments.BEDROCK in enabled_instruments:
                 try:
-                    BedrockInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                    logger.info("Instrumented Bedrock")
+                    if is_package_installed("boto3"):
+                        from opentelemetry.instrumentation.bedrock import BedrockInstrumentor
+
+                        BedrockInstrumentor().instrument(tracer_provider=self._tracer_provider)
+                        logger.info("Instrumented Bedrock")
                 except Exception as e:
                     logger.warning(f"Failed to instrument Bedrock: {e}")
 
             # Instrument LangChain with our isolated tracer provider
             if Instruments.LANGCHAIN in enabled_instruments:
                 try:
-                    LangchainInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                    logger.info("Instrumented LangChain")
+                    if is_package_installed("langchain") or is_package_installed("langgraph"):
+                        from opentelemetry.instrumentation.langchain import LangchainInstrumentor
+
+                        LangchainInstrumentor().instrument(tracer_provider=self._tracer_provider)
+                        logger.info("Instrumented LangChain")
                 except Exception as e:
                     logger.warning(f"Failed to instrument LangChain: {e}")
 
@@ -812,9 +818,19 @@ class DocentTracer:
             metadata: Optional metadata to send
         """
         collection_id = self.collection_id
+
+        # Get agent_run_id from current context
+        agent_run_id = self.get_current_agent_run_id()
+        if not agent_run_id:
+            logger.error(
+                f"Cannot send transcript group metadata for {transcript_group_id} - no agent_run_id in context"
+            )
+            return
+
         payload: Dict[str, Any] = {
             "collection_id": collection_id,
             "transcript_group_id": transcript_group_id,
+            "agent_run_id": agent_run_id,
             "timestamp": datetime.now(timezone.utc).isoformat(),
         }
 
@@ -1019,6 +1035,22 @@ def initialize_tracing(
     return _global_tracer
 
 
+def _get_package_name(dist: Distribution) -> str | None:
+    try:
+        return dist.name.lower()
+    except (KeyError, AttributeError):
+        return None
+
+
+installed_packages = {
+    name for dist in distributions() if (name := _get_package_name(dist)) is not None
+}
+
+
+def is_package_installed(package_name: str) -> bool:
+    return package_name.lower() in installed_packages
+
+
 def get_tracer() -> DocentTracer:
     """Get the global Docent tracer."""
     if _global_tracer is None:

{docent_python-0.1.5a0.dist-info → docent_python-0.1.6a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.5a0
+Version: 0.1.6a0
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues

{docent_python-0.1.5a0.dist-info → docent_python-0.1.6a0.dist-info}/RECORD

@@ -1,29 +1,29 @@
 docent/__init__.py,sha256=J2BbO6rzilfw9WXRUeolr439EGFezqbMU_kCpCCryRA,59
 docent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/trace.py,sha256=aPbV9fy_JpFq9y3yuqhlHKPm2QiNh9SLsKOL2azxghs,62611
+docent/trace.py,sha256=oGhNizXcOU-FZJZkmnd8WQjlCvlRKWI3PncMHUHKQ_4,63667
 docent/trace_temp.py,sha256=Z0lAPwVzXjFvxpiU-CuvfWIslq9Q4alNkZMoQ77Xudk,40711
 docent/_log_util/__init__.py,sha256=3HXXrxrSm8PxwG4llotrCnSnp7GuroK1FNHsdg6f7aE,73
 docent/_log_util/logger.py,sha256=kwM0yRW1IJd6-XTorjWn48B4l8qvD2ZM6VDjY5eskQI,4422
 docent/data_models/__init__.py,sha256=4JbTDVzRhS5VZgo8MALwd_YI17GaN7X9E3rOc4Xl7kw,327
 docent/data_models/_tiktoken_util.py,sha256=hC0EDDWItv5-0cONBnHWgZtQOflDU7ZNEhXPFo4DvPc,3057
-docent/data_models/agent_run.py,sha256=lw-odD2zzFi-RGvkAFjz9x8l6XWPrGT6uRGqTj9h8qU,9621
+docent/data_models/agent_run.py,sha256=bDRToWUlY52PugoHWU1D9hasr5t_fnTmRLpkzWP1s_k,9811
 docent/data_models/citation.py,sha256=WsVQZcBT2EJD24ysyeVOC5Xfo165RI7P5_cOnJBgHj0,10015
 docent/data_models/metadata.py,sha256=r0SYC4i2x096dXMLfw_rAMtcJQCsoV6EOMPZuEngbGA,9062
 docent/data_models/regex.py,sha256=0ciIerkrNwb91bY5mTcyO5nDWH67xx2tZYObV52fmBo,1684
 docent/data_models/shared_types.py,sha256=jjm-Dh5S6v7UKInW7SEqoziOsx6Z7Uu4e3VzgCbTWvc,225
-docent/data_models/transcript.py,sha256=NDcpvil4dJ8YhG_JJ0X-w0prkXhwhsdO-zoL-CZMipM,15446
+docent/data_models/transcript.py,sha256=0iF2ujcWhTss8WkkpNMeIKJyKOfMEsiMoAQMGwY4ing,15753
 docent/data_models/chat/__init__.py,sha256=O04XQ2NmO8GTWqkkB_Iydj8j_CucZuLhoyMVTxJN_cs,570
 docent/data_models/chat/content.py,sha256=Co-jO8frQa_DSP11wJuhPX0s-GpJk8yqtKqPeiAIZ_U,1672
 docent/data_models/chat/message.py,sha256=iAo38kbV6wYbFh8S23cxLy6HY4C_i3PzQ6RpSQG5dxM,3861
 docent/data_models/chat/tool.py,sha256=x7NKINswPe0Kqvcx4ubjHzB-n0-i4DbFodvaBb2vitk,3042
-docent/loaders/load_inspect.py,sha256=yK6LZgprT8kc0Jg4N_cnbhsGCq9lINmMcgALXA9AibY,2812
+docent/loaders/load_inspect.py,sha256=_cK2Qd6gyLQuJVzOlsvEZz7TrqzNmH6ZsLTkSCWAPqQ,6628
 docent/samples/__init__.py,sha256=roDFnU6515l9Q8v17Es_SpWyY9jbm5d6X9lV01V0MZo,143
 docent/samples/load.py,sha256=ZGE07r83GBNO4A0QBh5aQ18WAu3mTWA1vxUoHd90nrM,207
 docent/samples/log.eval,sha256=orrW__9WBfANq7NwKsPSq9oTsQRcG6KohG5tMr_X_XY,397708
 docent/samples/tb_airline.json,sha256=eR2jFFRtOw06xqbEglh6-dPewjifOk-cuxJq67Dtu5I,47028
 docent/sdk/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 docent/sdk/client.py,sha256=fLdniy8JzMLoZpaS9SP2pHban_ToavgtI8VeHZLMNZo,12773
-docent_python-0.1.5a0.dist-info/METADATA,sha256=XRzyvevWQ36Mk6UB6HIGIuD6vCQBX0cazhZylY25Q24,1037
-docent_python-0.1.5a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.5a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
-docent_python-0.1.5a0.dist-info/RECORD,,
+docent_python-0.1.6a0.dist-info/METADATA,sha256=ib_GqBFrOmPvacYb4uncrC5qLsoygIJ0wU852MOea_8,1037
+docent_python-0.1.6a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+docent_python-0.1.6a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
+docent_python-0.1.6a0.dist-info/RECORD,,