docent-python 0.1.4a0__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/sdk/client.py

@@ -197,75 +197,85 @@ class Docent:
         return response.json()
 
     def list_searches(self, collection_id: str) -> list[dict[str, Any]]:
-        """List all searches for a given collection.
+        """List all rubrics for a given collection.
 
         Args:
             collection_id: ID of the Collection.
 
         Returns:
-            list: List of dictionaries containing search query information.
+            list: List of dictionaries containing rubric information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{collection_id}/list_search_queries"
+        url = f"{self._server_url}/rubric/{collection_id}/rubrics"
         response = self._session.get(url)
         response.raise_for_status()
         return response.json()
 
-    def get_search_results(self, collection_id: str, search_query: str) -> list[dict[str, Any]]:
-        """Get search results for a given collection and search query.
-        Pass in either search_query or query_id.
+    def get_search_results(
+        self, collection_id: str, rubric_id: str, rubric_version: int
+    ) -> list[dict[str, Any]]:
+        """Get rubric results for a given collection, rubric and version.
 
         Args:
             collection_id: ID of the Collection.
-            search_query: The search query to get results for.
+            rubric_id: The ID of the rubric to get results for.
+            rubric_version: The version of the rubric to get results for.
 
         Returns:
-            list: List of dictionaries containing search result information.
+            list: List of dictionaries containing rubric result information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{collection_id}/get_search_results"
-        response = self._session.post(url, json={"search_query": search_query})
+        url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/results"
+        response = self._session.get(url, params={"rubric_version": rubric_version})
         response.raise_for_status()
         return response.json()
 
-    def list_search_clusters(self, collection_id: str, search_query: str) -> list[dict[str, Any]]:
-        """List all search clusters for a given collection.
-        Pass in either search_query or query_id.
+    def list_search_clusters(
+        self, collection_id: str, rubric_id: str, rubric_version: int | None = None
+    ) -> list[dict[str, Any]]:
+        """List all centroids for a given collection and rubric.
 
         Args:
             collection_id: ID of the Collection.
-            search_query: The search query to get clusters for.
+            rubric_id: The ID of the rubric to get centroids for.
+            rubric_version: Optional version of the rubric. If not provided, uses latest.
 
         Returns:
-            list: List of dictionaries containing search cluster information.
+            list: List of dictionaries containing centroid information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{collection_id}/list_search_clusters"
-        response = self._session.post(url, json={"search_query": search_query})
+        url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/centroids"
+        params: dict[str, int] = {}
+        if rubric_version is not None:
+            params["rubric_version"] = rubric_version
+        response = self._session.get(url, params=params)
         response.raise_for_status()
         return response.json()
 
-    def get_cluster_matches(self, collection_id: str, centroid: str) -> list[dict[str, Any]]:
-        """Get the matches for a given cluster.
+    def get_cluster_matches(
+        self, collection_id: str, rubric_id: str, rubric_version: int
+    ) -> list[dict[str, Any]]:
+        """Get centroid assignments for a given rubric.
 
         Args:
             collection_id: ID of the Collection.
-            cluster_id: The ID of the cluster to get matches for.
+            rubric_id: The ID of the rubric to get assignments for.
+            rubric_version: The version of the rubric to get assignments for.
 
         Returns:
-            list: List of dictionaries containing the search results that match the cluster.
+            list: List of dictionaries containing centroid assignment information.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/{collection_id}/get_cluster_matches"
-        response = self._session.post(url, json={"centroid": centroid})
+        url = f"{self._server_url}/rubric/{collection_id}/{rubric_id}/assignments"
+        response = self._session.get(url, params={"rubric_version": rubric_version})
         response.raise_for_status()
         return response.json()
 

docent/trace.py

@@ -11,17 +11,15 @@ from collections import defaultdict
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar, Token
 from datetime import datetime, timezone
-from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Union
+from enum import Enum
+from importlib.metadata import Distribution, distributions
+from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Set, Union
 
 import requests
 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
@@ -33,15 +31,23 @@ 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"
 DEFAULT_COLLECTION_NAME = "default-collection-name"
 
 
+class Instruments(Enum):
+    """Enumeration of available instrument types."""
+
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    BEDROCK = "bedrock"
+    LANGCHAIN = "langchain"
+
+
 def _is_notebook() -> bool:
     """Check if we're running in a Jupyter notebook."""
     try:
@@ -64,6 +70,8 @@ class DocentTracer:
         enable_console_export: bool = False,
         enable_otlp_export: bool = True,
         disable_batch: bool = False,
+        instruments: Optional[Set[Instruments]] = None,
+        block_instruments: Optional[Set[Instruments]] = None,
     ):
         """
         Initialize Docent tracing manager.
@@ -78,6 +86,8 @@ class DocentTracer:
             enable_console_export: Whether to export to console
             enable_otlp_export: Whether to export to OTLP endpoint
             disable_batch: Whether to disable batch processing (use SimpleSpanProcessor)
+            instruments: Set of instruments to enable (None = all instruments)
+            block_instruments: Set of instruments to explicitly disable
         """
         self.collection_name: str = collection_name
         self.collection_id: str = collection_id if collection_id else str(uuid.uuid4())
@@ -105,6 +115,9 @@ class DocentTracer:
         self.enable_console_export = enable_console_export
         self.enable_otlp_export = enable_otlp_export
         self.disable_batch = disable_batch
+        self.disabled_instruments: Set[Instruments] = {Instruments.LANGCHAIN}
+        self.instruments = instruments or (set(Instruments) - self.disabled_instruments)
+        self.block_instruments = block_instruments or set()
 
         # Use separate tracer provider to avoid interfering with existing OTEL setup
         self._tracer_provider: Optional[TracerProvider] = None
@@ -206,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
 
@@ -309,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()
@@ -333,33 +344,51 @@ class DocentTracer:
             except Exception as e:
                 logger.warning(f"Failed to instrument threading: {e}")
 
+            enabled_instruments = self.instruments - self.block_instruments
+
             # Instrument OpenAI with our isolated tracer provider
-            try:
-                OpenAIInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                logger.info("Instrumented OpenAI")
-            except Exception as e:
-                logger.warning(f"Failed to instrument OpenAI: {e}")
+            if Instruments.OPENAI in enabled_instruments:
+                try:
+                    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
-            try:
-                AnthropicInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                logger.info("Instrumented Anthropic")
-            except Exception as e:
-                logger.warning(f"Failed to instrument Anthropic: {e}")
+            if Instruments.ANTHROPIC in enabled_instruments:
+                try:
+                    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
-            try:
-                BedrockInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                logger.info("Instrumented Bedrock")
-            except Exception as e:
-                logger.warning(f"Failed to instrument Bedrock: {e}")
+            if Instruments.BEDROCK in enabled_instruments:
+                try:
+                    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
-            try:
-                LangchainInstrumentor().instrument(tracer_provider=self._tracer_provider)
-                logger.info("Instrumented LangChain")
-            except Exception as e:
-                logger.warning(f"Failed to instrument LangChain: {e}")
+            if Instruments.LANGCHAIN in enabled_instruments:
+                try:
+                    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}")
 
             # Register cleanup handlers
             self._register_cleanup()
@@ -789,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(),
         }
 
@@ -942,6 +981,8 @@ def initialize_tracing(
     enable_console_export: bool = False,
     enable_otlp_export: bool = True,
     disable_batch: bool = False,
+    instruments: Optional[Set[Instruments]] = None,
+    block_instruments: Optional[Set[Instruments]] = None,
 ) -> DocentTracer:
     """
     Initialize the global Docent tracer.
@@ -958,6 +999,8 @@ def initialize_tracing(
         enable_console_export: Whether to export spans to console
         enable_otlp_export: Whether to export spans to OTLP endpoint
         disable_batch: Whether to disable batch processing (use SimpleSpanProcessor)
+        instruments: Set of instruments to enable (None = all instruments).
+        block_instruments: Set of instruments to explicitly disable.
 
     Returns:
         The initialized Docent tracer
@@ -966,6 +1009,7 @@ def initialize_tracing(
         # Basic setup
         initialize_tracing("my-collection")
     """
+
     global _global_tracer
 
     # Check for API key in environment variable if not provided as parameter
@@ -983,12 +1027,30 @@ def initialize_tracing(
             enable_console_export=enable_console_export,
             enable_otlp_export=enable_otlp_export,
             disable_batch=disable_batch,
+            instruments=instruments,
+            block_instruments=block_instruments,
         )
         _global_tracer.initialize()
 
     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.4a0.dist-info → docent_python-0.1.6a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.4a0
+Version: 0.1.6a0
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues
@@ -22,4 +22,3 @@ Requires-Dist: pydantic>=2.11.7
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: tiktoken>=0.7.0
 Requires-Dist: tqdm>=4.67.1
-Requires-Dist: traceloop-sdk>=0.44.1

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

@@ -1,30 +1,29 @@
 docent/__init__.py,sha256=J2BbO6rzilfw9WXRUeolr439EGFezqbMU_kCpCCryRA,59
 docent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/trace.py,sha256=T6F9jbj0JcnGBxlcR65iDPF0lmeFR6laTlhDBcw0Mh0,61167
-docent/trace_alt.py,sha256=k0EZ_jyvVDqQ3HExDPZzdZYj6tKFL_cgbNRnOq4oGyg,17747
+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=uyhTisb9bHk7Hd2G4UKLdfvuiAmYOOqJiwEPbYWN9IE,12371
-docent_python-0.1.4a0.dist-info/METADATA,sha256=AUptJVGzZtABJ8V1Hpzhi2EOaMdBJB_nSaldpN6J8Bg,1074
-docent_python-0.1.4a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.4a0.dist-info/licenses/LICENSE.md,sha256=vOHzq3K4Ndu0UV9hPrtXvlD7pHOjyDQmGjHuLSIkRQY,1087
-docent_python-0.1.4a0.dist-info/RECORD,,
+docent/sdk/client.py,sha256=fLdniy8JzMLoZpaS9SP2pHban_ToavgtI8VeHZLMNZo,12773
+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,,

docent/trace_alt.py

@@ -1,513 +0,0 @@
-import asyncio
-import atexit
-import functools
-import io
-import logging
-import os
-import uuid
-from contextlib import asynccontextmanager, contextmanager, redirect_stdout
-from contextvars import ContextVar, Token
-from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, Set
-
-import requests
-from opentelemetry.context import Context
-from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
-from opentelemetry.trace import Span
-from traceloop.sdk import Traceloop
-
-# Configure logging
-logger = logging.getLogger(__name__)
-logger.disabled = True
-
-DEFAULT_ENDPOINT = "https://api.docent.transluce.org/rest/telemetry"
-
-# Context variables for tracking current agent run and collection
-_current_agent_run_id: ContextVar[Optional[str]] = ContextVar("current_agent_run_id", default=None)
-_current_collection_id: ContextVar[Optional[str]] = ContextVar(
-    "current_collection_id", default=None
-)
-
-# Global configuration
-_tracing_initialized = False
-_collection_name: Optional[str] = None
-_collection_id: Optional[str] = None
-_default_agent_run_id: Optional[str] = None
-_endpoint: Optional[str] = None
-_api_key: Optional[str] = None
-_enable_console_export = False
-_disable_batch = False
-_instruments: Optional[Set[Any]] = None
-_block_instruments: Optional[Set[Any]] = None
-
-
-class DocentSpanProcessor(SpanProcessor):
-    """Custom span processor to add Docent metadata to spans.
-
-    This processor integrates cleanly with Traceloop's existing span processing
-    and adds Docent-specific attributes to all spans.
-    """
-
-    def __init__(self, collection_id: str, enable_console_export: bool = False):
-        self.collection_id = collection_id
-        self.enable_console_export = enable_console_export
-
-    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
-        """Add Docent metadata when a span starts."""
-        # Always add collection_id
-        span.set_attribute("collection_id", self.collection_id)
-
-        # Add agent_run_id if available
-        agent_run_id = _get_current_agent_run_id()
-        if agent_run_id:
-            span.set_attribute("agent_run_id", agent_run_id)
-        else:
-            span.set_attribute("agent_run_id", _get_default_agent_run_id())
-            span.set_attribute("agent_run_id_default", True)
-
-        # Add service name for better integration with existing OTEL setups
-        span.set_attribute("service.name", _collection_name or "docent-trace")
-
-        if self.enable_console_export:
-            logging.debug(
-                f"Span started - collection_id: {self.collection_id}, agent_run_id: {agent_run_id}"
-            )
-
-    def on_end(self, span: ReadableSpan) -> None:
-        pass
-
-    def shutdown(self) -> None:
-        """Called when the processor is shut down."""
-
-    def force_flush(self, timeout_millis: float = 30000) -> bool:
-        """Force flush any pending spans."""
-        return True
-
-
-def initialize_tracing(
-    collection_name: str,
-    collection_id: Optional[str] = None,
-    endpoint: Optional[str] = None,
-    api_key: Optional[str] = None,
-    enable_console_export: bool = False,
-    disable_batch: bool = False,
-    instruments: Optional[Set[Any]] = None,
-    block_instruments: Optional[Set[Any]] = None,
-) -> None:
-    """Initialize Docent tracing with the specified configuration.
-
-    This function provides a comprehensive initialization that integrates cleanly
-    with existing OpenTelemetry setups and provides extensive configuration options.
-
-    Args:
-        collection_name: Name for your application/collection
-        collection_id: Optional collection ID (auto-generated if not provided)
-        endpoint: Optional OTLP endpoint URL (defaults to Docent's hosted service)
-        api_key: Optional API key (uses DOCENT_API_KEY environment variable if not provided)
-        enable_console_export: Whether to also export traces to console for debugging
-        disable_batch: Whether to disable batch processing (use SimpleSpanProcessor)
-        instruments: Set of instruments to enable (None = all instruments)
-        block_instruments: Set of instruments to explicitly disable
-    """
-    global _tracing_initialized, _collection_name, _collection_id, _default_agent_run_id, _endpoint, _api_key
-    global _enable_console_export, _disable_batch, _instruments, _block_instruments
-
-    if _tracing_initialized:
-        logging.warning("Docent tracing already initialized")
-        return
-
-    _collection_name = collection_name
-    _collection_id = collection_id or _generate_id()
-    _default_agent_run_id = _get_default_agent_run_id()  # Generate default ID if not set
-    _endpoint = endpoint or DEFAULT_ENDPOINT
-    _api_key = api_key or os.getenv("DOCENT_API_KEY")
-    _enable_console_export = enable_console_export
-    _disable_batch = disable_batch
-    _instruments = instruments
-    _block_instruments = block_instruments
-
-    _set_current_collection_id(_collection_id)
-
-    if not _api_key:
-        raise ValueError(
-            "API key is required. Set DOCENT_API_KEY environment variable or pass api_key parameter."
-        )
-
-    # Initialize Traceloop with comprehensive configuration
-
-    # Get Traceloop's default span processor
-    from traceloop.sdk.tracing.tracing import get_default_span_processor
-
-    # Create our custom context span processor (only adds metadata, doesn't export)
-    docent_processor = DocentSpanProcessor(_collection_id, enable_console_export)
-
-    # Get Traceloop's default span processor for export
-    export_processor = get_default_span_processor(
-        disable_batch=_disable_batch,
-        api_endpoint=_endpoint,
-        headers={"Authorization": f"Bearer {_api_key}"},
-    )
-
-    # Combine both processors
-    processors = [docent_processor, export_processor]
-
-    os.environ["TRACELOOP_METRICS_ENABLED"] = "false"
-    os.environ["TRACELOOP_TRACE_ENABLED"] = "true"
-
-    # Temporarily redirect stdout to suppress print statements
-    with redirect_stdout(io.StringIO()):
-        Traceloop.init(  # type: ignore
-            app_name=collection_name,
-            api_endpoint=_endpoint,
-            api_key=_api_key,
-            telemetry_enabled=False,  # don't send analytics to traceloop's backend
-            disable_batch=_disable_batch,
-            instruments=_instruments,
-            block_instruments=_block_instruments,
-            processor=processors,  # Add both our context processor and export processor
-        )
-
-    _tracing_initialized = True
-    logging.info(
-        f"Docent tracing initialized for collection: {collection_name} with collection_id: {_collection_id}"
-    )
-
-    # Register cleanup handlers
-    atexit.register(_cleanup_tracing)
-
-
-def _cleanup_tracing() -> None:
-    """Clean up tracing resources on shutdown."""
-    global _tracing_initialized
-    if _tracing_initialized:
-        try:
-            # Notify API that the trace is over
-            _notify_trace_done()
-
-            logging.info("Docent tracing cleanup completed")
-        except Exception as e:
-            logging.warning(f"Error during tracing cleanup: {e}")
-        finally:
-            _tracing_initialized = False
-
-
-def _ensure_tracing_initialized():
-    """Ensure tracing has been initialized before use."""
-    if not _tracing_initialized:
-        raise RuntimeError("Docent tracing not initialized. Call initialize_tracing() first.")
-
-
-def _generate_id() -> str:
-    """Generate a unique ID for agent runs or collections."""
-    return str(uuid.uuid4())
-
-
-def _get_current_agent_run_id() -> Optional[str]:
-    """Get the current agent run ID from context."""
-    return _current_agent_run_id.get()
-
-
-def _get_current_collection_id() -> Optional[str]:
-    """Get the current collection ID from context."""
-    return _current_collection_id.get()
-
-
-def _get_default_agent_run_id() -> str:
-    """Get the default agent run ID, generating it if not set."""
-    global _default_agent_run_id
-    if _default_agent_run_id is None:
-        _default_agent_run_id = _generate_id()
-    return _default_agent_run_id
-
-
-def _set_current_agent_run_id(agent_run_id: Optional[str]) -> Token[Optional[str]]:
-    """Set the current agent run ID in context."""
-    return _current_agent_run_id.set(agent_run_id)
-
-
-def _set_current_collection_id(collection_id: Optional[str]) -> Token[Optional[str]]:
-    """Set the current collection ID in context."""
-    return _current_collection_id.set(collection_id)
-
-
-def _send_to_api(endpoint: str, data: Dict[str, Any]) -> None:
-    """Send data to the Docent API endpoint.
-
-    Args:
-        endpoint: The API endpoint URL
-        data: The data to send
-    """
-    try:
-        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {_api_key}"}
-
-        response = requests.post(endpoint, json=data, headers=headers, timeout=10)
-        response.raise_for_status()
-
-        logging.debug(f"Successfully sent data to {endpoint}")
-    except requests.exceptions.RequestException as e:
-        logging.error(f"Failed to send data to {endpoint}: {e}")
-    except Exception as e:
-        logging.error(f"Unexpected error sending data to {endpoint}: {e}")
-
-
-def _notify_trace_done() -> None:
-    """Notify the Docent API that the trace is done."""
-    collection_id = _get_current_collection_id()
-    if collection_id and _endpoint:
-        data = {"collection_id": collection_id, "status": "completed"}
-        _send_to_api(f"{_endpoint}/v1/trace-done", data)
-
-
-def agent_run_score(name: str, score: float, attributes: Optional[Dict[str, Any]] = None) -> None:
-    """
-    Record a score event on the current span.
-    Automatically works in both sync and async contexts.
-
-    Args:
-        name: Name of the score metric
-        score: Numeric score value
-        attributes: Optional additional attributes for the score event
-    """
-    _ensure_tracing_initialized()
-
-    agent_run_id = _get_current_agent_run_id()
-    if not agent_run_id:
-        logging.warning("No active agent run context. Score will not be sent.")
-        return
-
-    collection_id = _get_current_collection_id() or _collection_id
-    if not collection_id:
-        logging.warning("No collection ID available. Score will not be sent.")
-        return
-
-    # Send score directly to API
-    score_data = {
-        "collection_id": collection_id,
-        "agent_run_id": agent_run_id,
-        "score_name": name,
-        "score_value": score,
-    }
-
-    # Add additional attributes if provided
-    if attributes:
-        score_data.update(attributes)
-
-    _send_to_api(f"{_endpoint}/v1/scores", score_data)
-
-
-def agent_run_metadata(metadata: Dict[str, Any]) -> None:
-    """Attach metadata to the current agent run.
-
-    Args:
-        metadata: Dictionary of metadata to attach
-    """
-    _ensure_tracing_initialized()
-
-    agent_run_id = _get_current_agent_run_id()
-    if not agent_run_id:
-        logging.warning("No active agent run context. Metadata will not be sent.")
-        return
-
-    collection_id = _get_current_collection_id() or _collection_id
-    if not collection_id:
-        logging.warning("No collection ID available. Metadata will not be sent.")
-        return
-
-    # Send metadata directly to API
-    metadata_data = {
-        "collection_id": collection_id,
-        "agent_run_id": agent_run_id,
-        "metadata": metadata,
-    }
-
-    _send_to_api(f"{_endpoint}/v1/metadata", metadata_data)
-
-
-@contextmanager
-def _agent_run_context_sync(
-    agent_run_id: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-) -> Iterator[tuple[str, Optional[str]]]:
-    """Synchronous context manager for creating and managing agent runs."""
-    _ensure_tracing_initialized()
-
-    # Generate IDs if not provided
-    current_agent_run_id = agent_run_id or _generate_id()
-
-    # Set up context
-    agent_run_token = _set_current_agent_run_id(current_agent_run_id)
-
-    try:
-        # Send metadata to API if provided
-        if metadata:
-            agent_run_metadata(metadata)
-
-        # Yield the agent run ID and None for transcript_id (handled by backend)
-        # Traceloop will automatically create spans for any instrumented operations
-        # and our DocentSpanProcessor will add the appropriate metadata
-        yield (current_agent_run_id, None)
-    finally:
-        # Restore context
-        _current_agent_run_id.reset(agent_run_token)
-
-
-@asynccontextmanager
-async def _agent_run_context_async(
-    agent_run_id: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-) -> AsyncIterator[tuple[str, Optional[str]]]:
-    """Asynchronous context manager for creating and managing agent runs."""
-    _ensure_tracing_initialized()
-
-    # Generate IDs if not provided
-    current_agent_run_id = agent_run_id or _generate_id()
-
-    # Set up context
-    agent_run_token = _set_current_agent_run_id(current_agent_run_id)
-
-    try:
-        # Send metadata to API if provided
-        if metadata:
-            agent_run_metadata(metadata)
-
-        # Yield the agent run ID and None for transcript_id (handled by backend)
-        # Traceloop will automatically create spans for any instrumented operations
-        # and our DocentSpanProcessor will add the appropriate metadata
-        yield (current_agent_run_id, None)
-    finally:
-        # Restore context
-        _current_agent_run_id.reset(agent_run_token)
-
-
-def agent_run_context(
-    agent_run_id: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-):
-    """Context manager for creating and managing agent runs.
-
-    This context manager can be used in both synchronous and asynchronous contexts.
-    In async contexts, use it with `async with agent_run_context()`.
-    In sync contexts, use it with `with agent_run_context()`.
-
-    Args:
-        agent_run_id: Optional agent run ID (auto-generated if not provided)
-        metadata: Optional metadata to attach to the agent run
-
-    Returns:
-        A context manager that yields a tuple of (agent_run_id, transcript_id)
-        where transcript_id is None for now as it's handled by backend
-    """
-    # Check if we're in an async context by looking at the current frame
-    import inspect
-
-    frame = inspect.currentframe()
-    try:
-        # Look for async context indicators in the call stack
-        while frame:
-            if frame.f_code.co_flags & 0x80:  # CO_COROUTINE flag
-                return _agent_run_context_async(agent_run_id=agent_run_id, metadata=metadata)
-            frame = frame.f_back
-    finally:
-        # Clean up the frame reference
-        del frame
-
-    # Default to sync context manager
-    return _agent_run_context_sync(agent_run_id=agent_run_id, metadata=metadata)
-
-
-def agent_run(
-    func: Optional[Callable[..., Any]] = None,
-    *,
-    agent_run_id: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-) -> Callable[..., Any]:
-    """Decorator for creating agent runs around functions.
-
-    Args:
-        func: Function to decorate
-        agent_run_id: Optional agent run ID (auto-generated if not provided)
-        metadata: Optional metadata to attach to the agent run
-
-    Returns:
-        Decorated function
-    """
-
-    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-        @functools.wraps(func)
-        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
-            with _agent_run_context_sync(agent_run_id=agent_run_id, metadata=metadata) as (
-                run_id,
-                _,
-            ):
-                result = func(*args, **kwargs)
-                # Store agent run ID as an attribute for access
-                setattr(sync_wrapper, "docent", type("DocentInfo", (), {"agent_run_id": run_id})())  # type: ignore
-                return result
-
-        @functools.wraps(func)
-        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
-            async with _agent_run_context_async(agent_run_id=agent_run_id, metadata=metadata) as (
-                run_id,
-                _,
-            ):
-                result = await func(*args, **kwargs)
-                # Store agent run ID as an attribute for access
-                setattr(async_wrapper, "docent", type("DocentInfo", (), {"agent_run_id": run_id})())  # type: ignore
-                return result
-
-        # Return appropriate wrapper based on function type
-        if asyncio.iscoroutinefunction(func):
-            return async_wrapper
-        else:
-            return sync_wrapper
-
-    # Handle both @agent_run and @agent_run(agent_run_id=..., metadata=...)
-    if func is None:
-        return decorator
-    else:
-        return decorator(func)
-
-
-# Additional utility functions for better integration
-
-
-def get_current_agent_run_id() -> Optional[str]:
-    """Get the current agent run ID from context.
-
-    Returns:
-        The current agent run ID if available, None otherwise
-    """
-    return _get_current_agent_run_id()
-
-
-def get_current_collection_id() -> Optional[str]:
-    """Get the current collection ID from context.
-
-    Returns:
-        The current collection ID if available, None otherwise
-    """
-    return _get_current_collection_id()
-
-
-def is_tracing_initialized() -> bool:
-    """Check if tracing has been initialized.
-
-    Returns:
-        True if tracing is initialized, False otherwise
-    """
-    return _tracing_initialized
-
-
-def flush_spans() -> None:
-    """Force flush any pending spans to the backend.
-
-    This is useful for ensuring all spans are sent before shutdown
-    or for debugging purposes.
-    """
-    if _tracing_initialized:
-        try:
-            traceloop_instance = Traceloop.get()
-            if hasattr(traceloop_instance, "flush"):
-                traceloop_instance.flush()  # type: ignore
-            logging.debug("Spans flushed successfully")
-        except Exception as e:
-            logging.warning(f"Error flushing spans: {e}")