docent-python 0.1.21a0__py3-none-any.whl → 0.1.23a0__py3-none-any.whl

docent/sdk/client.py

@@ -8,7 +8,8 @@ from tqdm import tqdm
 
 from docent._log_util.logger import get_logger
 from docent.data_models.agent_run import AgentRun
-from docent.data_models.judge import JudgeRunLabel
+from docent.data_models.judge import Label
+from docent.judges.util.meta_schema import validate_judge_result_schema
 from docent.loaders import load_inspect
 
 logger = get_logger(__name__)
@@ -50,9 +51,15 @@ class Docent:
         self._login(api_key)
 
     def _handle_response_errors(self, response: requests.Response):
-        """Handle API response and raise informative errors.
-        TODO: make this more informative."""
-        response.raise_for_status()
+        """Handle API response and raise informative errors."""
+        if response.status_code >= 400:
+            try:
+                error_data = response.json()
+                detail = error_data.get("detail", response.text)
+            except Exception:
+                detail = response.text
+
+            raise requests.HTTPError(f"HTTP {response.status_code}: {detail}", response=response)
 
     def _login(self, api_key: str):
         """Login with email/password to establish session."""
@@ -253,30 +260,59 @@ class Docent:
         clustering_state = self.get_clustering_state(collection_id, rubric_id)
         return clustering_state.get("assignments", {})
 
+    def create_label_set(
+        self,
+        collection_id: str,
+        name: str,
+        label_schema: dict[str, Any],
+        description: str | None = None,
+    ) -> str:
+        """Create a new label set with a JSON schema.
+
+        Args:
+            collection_id: ID of the collection.
+            name: Name of the label set.
+            label_schema: JSON schema for validating labels in this set.
+            description: Optional description of the label set.
+
+        Returns:
+            str: The ID of the created label set.
+
+        Raises:
+            ValueError: If the response is missing the label_set_id.
+            jsonschema.ValidationError: If the label schema is invalid.
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        validate_judge_result_schema(label_schema)
+
+        url = f"{self._server_url}/label/{collection_id}/label_set"
+        payload = {
+            "name": name,
+            "label_schema": label_schema,
+            "description": description,
+        }
+        response = self._session.post(url, json=payload)
+        self._handle_response_errors(response)
+        return response.json()["label_set_id"]
+
     def add_label(
         self,
         collection_id: str,
-        rubric_id: str,
-        label: JudgeRunLabel,
-    ) -> dict[str, Any]:
-        """Attach a manual label to an agent run for a rubric.
+        label: Label,
+    ) -> dict[str, str]:
+        """Create a label in a label set.
 
         Args:
-            collection_id: ID of the Collection that owns the rubric.
-            rubric_id: ID of the rubric the label applies to.
-            label: A `JudgeRunLabel` that must comply with the rubric's output schema.
+            collection_id: ID of the Collection.
+            label: A `Label` object that must comply with the label set's schema.
 
         Returns:
-            dict: API response containing a status message.
+            dict: API response containing the label_id.
 
         Raises:
-            ValueError: If the label does not target the rubric specified in the path.
             requests.exceptions.HTTPError: If the API request fails or validation errors occur.
         """
-        if label.rubric_id != rubric_id:
-            raise ValueError("Label rubric_id must match the rubric_id argument")
-
-        url = f"{self._server_url}/rubric/{collection_id}/rubric/{rubric_id}/label"
+        url = f"{self._server_url}/label/{collection_id}/label"
         payload = {"label": label.model_dump(mode="json")}
         response = self._session.post(url, json=payload)
         self._handle_response_errors(response)
@@ -285,55 +321,50 @@ class Docent:
     def add_labels(
         self,
         collection_id: str,
-        rubric_id: str,
-        labels: list[JudgeRunLabel],
+        labels: list[Label],
     ) -> dict[str, Any]:
-        """Attach multiple manual labels to a rubric.
+        """Create multiple labels.
 
         Args:
-            collection_id: ID of the Collection that owns the rubric.
-            rubric_id: ID of the rubric the labels apply to.
-            labels: List of `JudgeRunLabel` objects.
+            collection_id: ID of the Collection.
+            labels: List of `Label` objects.
 
         Returns:
-            dict: API response containing status information.
+            dict: API response containing label_ids list and optional errors list.
 
         Raises:
             ValueError: If no labels are provided.
-            ValueError: If any label targets a different rubric.
             requests.exceptions.HTTPError: If the API request fails.
         """
         if not labels:
             raise ValueError("labels must contain at least one entry")
 
-        rubric_ids = {label.rubric_id for label in labels}
-        if rubric_ids != {rubric_id}:
-            raise ValueError(
-                "All labels must specify the same rubric_id that is provided to add_labels"
-            )
-
-        payload = {"labels": [l.model_dump(mode="json") for l in labels]}
-
-        url = f"{self._server_url}/rubric/{collection_id}/rubric/{rubric_id}/labels"
+        url = f"{self._server_url}/label/{collection_id}/labels"
+        payload = {"labels": [label.model_dump(mode="json") for label in labels]}
         response = self._session.post(url, json=payload)
         self._handle_response_errors(response)
         return response.json()
 
-    def get_labels(self, collection_id: str, rubric_id: str) -> list[dict[str, Any]]:
-        """Retrieve all manual labels for a rubric.
+    def get_labels(
+        self, collection_id: str, label_set_id: str, filter_valid_labels: bool = False
+    ) -> list[dict[str, Any]]:
+        """Retrieve all labels in a label set.
 
         Args:
-            collection_id: ID of the Collection that owns the rubric.
-            rubric_id: ID of the rubric to fetch labels for.
+            collection_id: ID of the Collection.
+            label_set_id: ID of the label set to fetch labels for.
+            filter_valid_labels: If True, only return labels that match the label set schema
+                INCLUDING requirements. Default is False (returns all labels).
 
         Returns:
-            list: List of label dictionaries. Each includes agent_run_id and label content.
+            list: List of label dictionaries.
 
         Raises:
             requests.exceptions.HTTPError: If the API request fails.
         """
-        url = f"{self._server_url}/rubric/{collection_id}/rubric/{rubric_id}/labels"
-        response = self._session.get(url)
+        url = f"{self._server_url}/label/{collection_id}/label_set/{label_set_id}/labels"
+        params = {"filter_valid_labels": filter_valid_labels}
+        response = self._session.get(url, params=params)
         self._handle_response_errors(response)
         return response.json()
 

docent/trace.py

@@ -43,6 +43,7 @@ class Instruments(Enum):
     ANTHROPIC = "anthropic"
     BEDROCK = "bedrock"
     LANGCHAIN = "langchain"
+    GOOGLE_GENERATIVEAI = "google_generativeai"
 
 
 class DocentTracer:
@@ -392,6 +393,23 @@ class DocentTracer:
                 except Exception as e:
                     logger.warning(f"Failed to instrument LangChain: {e}")
 
+            # Instrument Google Generative AI with our isolated tracer provider
+            if Instruments.GOOGLE_GENERATIVEAI in enabled_instruments:
+                try:
+                    if is_package_installed("google-generativeai") or is_package_installed(
+                        "google-genai"
+                    ):
+                        from opentelemetry.instrumentation.google_generativeai import (
+                            GoogleGenerativeAiInstrumentor,
+                        )
+
+                        GoogleGenerativeAiInstrumentor().instrument(
+                            tracer_provider=self._tracer_provider
+                        )
+                        logger.info("Instrumented Google Generative AI")
+                except Exception as e:
+                    logger.warning(f"Failed to instrument Google Generative AI: {e}")
+
             # Register cleanup handlers
             self._register_cleanup()
 

{docent_python-0.1.21a0.dist-info → docent_python-0.1.23a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docent-python
-Version: 0.1.21a0
+Version: 0.1.23a0
 Summary: Docent SDK
 Project-URL: Homepage, https://github.com/TransluceAI/docent
 Project-URL: Issues, https://github.com/TransluceAI/docent/issues
@@ -20,6 +20,7 @@ Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.34.1
 Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.34.1
 Requires-Dist: opentelemetry-instrumentation-anthropic>=0.40.14
 Requires-Dist: opentelemetry-instrumentation-bedrock>=0.40.14
+Requires-Dist: opentelemetry-instrumentation-google-generativeai>=0.40.14
 Requires-Dist: opentelemetry-instrumentation-langchain>=0.40.14
 Requires-Dist: opentelemetry-instrumentation-openai>=0.40.14
 Requires-Dist: opentelemetry-instrumentation-threading>=0.55b1

{docent_python-0.1.21a0.dist-info → docent_python-0.1.23a0.dist-info}/RECORD

@@ -1,31 +1,29 @@
 docent/__init__.py,sha256=fuhETwJPcesiB76Zxa64HBJxeaaTyRalIH-fs77TWsU,112
 docent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/trace.py,sha256=_JvDmtWVFARPYvXsNx8-RKRdev4mMxNK6iq9AARzoJE,66362
-docent/trace_2.py,sha256=-OxzXF2kOFkhto1UGXHWVM797EN_BT_uwDSbzgMme8o,67145
+docent/trace.py,sha256=u_1M_B1ncCR5a0Yy5ppQY_0k0AFYe4UhSIGeTuZCH0w,67271
 docent/trace_temp.py,sha256=Z0lAPwVzXjFvxpiU-CuvfWIslq9Q4alNkZMoQ77Xudk,40711
 docent/_llm_util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docent/_llm_util/llm_cache.py,sha256=p8pJ-B1vyJQlNn48ce1Pqv2gLocBVD6fZLPsd2VC5LA,6386
+docent/_llm_util/llm_cache.py,sha256=nGrvfFikFbEnfmzZRvWvZ60gfVSTvW1iC8-ciCXwbAk,6430
+docent/_llm_util/llm_svc.py,sha256=PQ-96UDJrnPa9csTKL_JDO8jzOrLzysVBqUHywuij0w,18046
 docent/_llm_util/model_registry.py,sha256=8Y4VwrA2f2EX78cG1VBIBHVvT_p4qqBTdu9a9zJpfTo,3382
-docent/_llm_util/prod_llms.py,sha256=HuGOg5Bhnpk_TijC3mOH8CTRIBy2C8w0_SebiEouNoE,16859
 docent/_llm_util/data_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 docent/_llm_util/data_models/exceptions.py,sha256=IW4BVMVp8r5TufNXyrhy3acgwJiQQQPQjB9VA4RVXw8,1489
-docent/_llm_util/data_models/llm_output.py,sha256=fuYPJ-SwxZjB4XGATA6XpLyc42Ix-kXHgBqFr_jPhK8,10123
-docent/_llm_util/data_models/simple_svc.py,sha256=0twuXP6aEU-jYY0obDSEgjT2lDSJCuZG_NgiqEzZIPM,2881
+docent/_llm_util/data_models/llm_output.py,sha256=ZAIIcgfxMZtTft8bXTPAhUcXEO48GLG3epkul_4gQNQ,10239
 docent/_llm_util/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 docent/_llm_util/providers/anthropic.py,sha256=-1oPd5FB4aFwKSmNvXzG8PVewjhgsogLRX1SCpnCxoA,18720
 docent/_llm_util/providers/common.py,sha256=dgcTuU4XkCKoAaM48UW8zMgRYUzj7TDBhvWqtnxBO7g,1166
 docent/_llm_util/providers/google.py,sha256=2D9mDgenZW0pt0_V7koX-aoZzpl8jo8xE5EWOLK7I0k,20314
 docent/_llm_util/providers/openai.py,sha256=4niQV9CNaJ-iiEwYG0BSFxCwcsCAWZz0JuUs4wBKu9M,25904
 docent/_llm_util/providers/openrouter.py,sha256=sT2onpeQ1gAwJLjkQbzD2RodJWTm013Q-siTXezca10,11958
-docent/_llm_util/providers/preference_types.py,sha256=z-TOxj_es1_cs5DzknZaganGyjMkmh5NgtiDYKMRI1I,3751
+docent/_llm_util/providers/preference_types.py,sha256=K1IH15YyU8E8F7Jk7LgLKiCJeDNk9uASTdnZGk5hSDc,3773
 docent/_llm_util/providers/provider_registry.py,sha256=EPYGQlegYPtg4ogEusCftm_5PZP-_XVKH1qg3xjPFTU,6337
 docent/_log_util/__init__.py,sha256=3HXXrxrSm8PxwG4llotrCnSnp7GuroK1FNHsdg6f7aE,73
 docent/_log_util/logger.py,sha256=kwM0yRW1IJd6-XTorjWn48B4l8qvD2ZM6VDjY5eskQI,4422
-docent/data_models/__init__.py,sha256=bE_Wy4Ql-9-0ZPcolMCPHhYvaE_Ug6h-jV7wOJ_DAi0,399
+docent/data_models/__init__.py,sha256=vEcFppE6wtKFp37KF_hUv00Ncn6fK_qUbVGZE5ltz-o,383
 docent/data_models/_tiktoken_util.py,sha256=hC0EDDWItv5-0cONBnHWgZtQOflDU7ZNEhXPFo4DvPc,3057
 docent/data_models/agent_run.py,sha256=7_37I9aS9rhDTkAvMPwoJGssQldvvKte8qVb93EnAiY,19329
 docent/data_models/citation.py,sha256=2_M1-_olVOJtjCGGFx1GIwGYWl0ILHxRsW8-EFDS9j0,7844
-docent/data_models/judge.py,sha256=zPbTqztn-yWu6tgD3R5JTyGnNiDhY6cWQ-gz3e_eM5k,340
+docent/data_models/judge.py,sha256=BOKAfZmNoLPclJNz_b7NvH8G8FzfR7kc6OpIv91GMDQ,336
 docent/data_models/metadata_util.py,sha256=E-EClAP5vVm9xbfTlPSz0tUyCalOfN9Jujd6JGoRnBg,487
 docent/data_models/regex.py,sha256=0ciIerkrNwb91bY5mTcyO5nDWH67xx2tZYObV52fmBo,1684
 docent/data_models/remove_invalid_citation_ranges.py,sha256=3RSMsOzFO2cSjkxI549TAo12qdvD-AGHd05Jxu0amvs,6282
@@ -36,14 +34,17 @@ docent/data_models/chat/__init__.py,sha256=ws77P3raDiOv6XesAMycUwu-uT75D5f9aNgjF
 docent/data_models/chat/content.py,sha256=Co-jO8frQa_DSP11wJuhPX0s-GpJk8yqtKqPeiAIZ_U,1672
 docent/data_models/chat/message.py,sha256=_72xeTdgv8ogQd4WLl1P3yXfIDkIEQrHlWgdvObeQxY,4291
 docent/data_models/chat/tool.py,sha256=MMglNHzkwHqUoK0xDWqs2FtelPsgHqwVpGpI1F8KZyw,3049
-docent/judges/__init__.py,sha256=Sob1uxJRgmr2S2sz4J6skHP8iqcVoiUq7Jlh8S5Sj9Y,462
-docent/judges/impl.py,sha256=qiItNKWPvB0KlB5b0rQoIfT-7m1xzyI028WtgvgvRhU,8864
-docent/judges/types.py,sha256=NlLv42iLDORbPAHppCz-YWZ6ksR4QYDWAweGw75izJ0,8439
+docent/judges/__init__.py,sha256=aTsQ2mIQnZt8HEMau02KrEA4m5w-lGC3U9Dirkj3to4,500
+docent/judges/analysis.py,sha256=bn7XIT7mj77LjFHMh1PqjALknq3nN-fRXqgg8cfJF8o,2486
+docent/judges/impl.py,sha256=JOq2tEBTqNbWIG2gRuI8OmEW2dHdx7nfnJnHeGwdyOk,24035
+docent/judges/runner.py,sha256=ANUVrrfgT61_zTV9pErLXoerMiD6x_RIJQGpwxWIIMg,1928
+docent/judges/stats.py,sha256=zejJle583xHG2G3gcYHiWcHoIOkeKwpSkl8lfeKQhFs,7805
+docent/judges/types.py,sha256=goNaKs3PF5wMHWLnFerYCEjUjPR0IVI9cVrxCK2TfjI,11539
 docent/judges/util/forgiving_json.py,sha256=zSh0LF3UVHdSjuMNvEiqUmSxpxPaqK1rSLiI6KCNihg,3549
-docent/judges/util/meta_schema.json,sha256=g3MUa_6e38I3GqZryy8b1w_Y9Krx2xSiWIuaG8Zpszc,2055
+docent/judges/util/meta_schema.json,sha256=7VHCGQUM0PbMIiwWDar15Sqaodi2y2Ty7yIW3PDL268,2105
 docent/judges/util/meta_schema.py,sha256=6IrIRHERJ6tkRcUtUShJ84I68yUJgkwfFeBjgt42qEA,930
-docent/judges/util/parse_output.py,sha256=qvqt7TEnrAqvzYHqip48boMQSUcoGa-1PA1gIGn-w4s,3381
-docent/judges/util/voting.py,sha256=cAty9b4w7M1OWeW-j8t6vxpZn7VXyE3aBL9Ex2ERKcU,3071
+docent/judges/util/parse_output.py,sha256=XMdU-hd3iozlGa72bpTrOADRpRze2sg68RSkQAw-yD4,2975
+docent/judges/util/voting.py,sha256=IRYXXYLsdc8MsgdzBBNVI1nnsx4kxnnWLPeQyLrDhwc,5152
 docent/loaders/load_inspect.py,sha256=VLrtpvcVZ44n2DIPMwUivXqbvOWjaooGw6moY8UQ0VE,6789
 docent/samples/__init__.py,sha256=roDFnU6515l9Q8v17Es_SpWyY9jbm5d6X9lV01V0MZo,143
 docent/samples/load.py,sha256=ZGE07r83GBNO4A0QBh5aQ18WAu3mTWA1vxUoHd90nrM,207
@@ -51,8 +52,8 @@ docent/samples/log.eval,sha256=orrW__9WBfANq7NwKsPSq9oTsQRcG6KohG5tMr_X_XY,39770
 docent/samples/tb_airline.json,sha256=eR2jFFRtOw06xqbEglh6-dPewjifOk-cuxJq67Dtu5I,47028
 docent/sdk/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 docent/sdk/agent_run_writer.py,sha256=0AWdxejoqZyuj9JSA39WlEwGcMSYTWNqnzIuluySY-M,11043
-docent/sdk/client.py,sha256=K1NVkj_CFj0q-2mSFvWfh8NTqXqosED--dv5aLD7yOE,18239
-docent_python-0.1.21a0.dist-info/METADATA,sha256=H0iEQ39cv90MW0lRZ94XhER6C_znvNU3DBPc6M72i9g,1277
-docent_python-0.1.21a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-docent_python-0.1.21a0.dist-info/licenses/LICENSE.md,sha256=QIMv2UiT6MppRasso4ymaA0w7ltkqmlL0HCt8CLD7Rc,580
-docent_python-0.1.21a0.dist-info/RECORD,,
+docent/sdk/client.py,sha256=ZRJDz1RKcV7_gA0zHbr3bIdvUqFrSU99Zes__2kVnKo,19184
+docent_python-0.1.23a0.dist-info/METADATA,sha256=sxm-OrFzMKg2U2G_EmrHlGeceqhpY7cGzKa0H989FOo,1351
+docent_python-0.1.23a0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+docent_python-0.1.23a0.dist-info/licenses/LICENSE.md,sha256=QIMv2UiT6MppRasso4ymaA0w7ltkqmlL0HCt8CLD7Rc,580
+docent_python-0.1.23a0.dist-info/RECORD,,

docent/_llm_util/data_models/simple_svc.py

@@ -1,79 +0,0 @@
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from typing import Literal
-
-from docent._llm_util.data_models.llm_output import (
-    AsyncLLMOutputStreamingCallback,
-    LLMOutput,
-)
-from docent._llm_util.prod_llms import MessagesInput, get_llm_completions_async
-from docent._llm_util.providers.preference_types import ModelOption
-from docent.data_models.chat import ToolInfo
-
-__all__ = ["BaseLLMService"]
-
-
-class BaseLLMService(ABC):
-    """Common interface for LLM services."""
-
-    @abstractmethod
-    async def get_completions(
-        self,
-        *,
-        inputs: list[MessagesInput],
-        model_options: list[ModelOption],
-        tools: list[ToolInfo] | None = None,
-        tool_choice: Literal["auto", "required"] | None = None,
-        max_new_tokens: int = 1024,
-        temperature: float = 1.0,
-        logprobs: bool = False,
-        top_logprobs: int | None = None,
-        max_concurrency: int = 100,
-        timeout: float = 120.0,
-        streaming_callback: AsyncLLMOutputStreamingCallback | None = None,
-        validation_callback: AsyncLLMOutputStreamingCallback | None = None,
-        completion_callback: AsyncLLMOutputStreamingCallback | None = None,
-        use_cache: bool = False,
-    ) -> list[LLMOutput]:
-        """Request completions from a configured LLM provider."""
-
-
-class SimpleLLMService(BaseLLMService):
-    """Lightweight LLM service that simply forwards completion requests.
-    Does not support cost tracking, usage limits, global scheduling or rate limiting."""
-
-    async def get_completions(
-        self,
-        *,
-        inputs: list[MessagesInput],
-        model_options: list[ModelOption],
-        tools: list[ToolInfo] | None = None,
-        tool_choice: Literal["auto", "required"] | None = None,
-        max_new_tokens: int = 1024,
-        temperature: float = 1.0,
-        logprobs: bool = False,
-        top_logprobs: int | None = None,
-        max_concurrency: int = 100,
-        timeout: float = 120.0,
-        streaming_callback: AsyncLLMOutputStreamingCallback | None = None,
-        validation_callback: AsyncLLMOutputStreamingCallback | None = None,
-        completion_callback: AsyncLLMOutputStreamingCallback | None = None,
-        use_cache: bool = False,
-    ) -> list[LLMOutput]:
-        return await get_llm_completions_async(
-            inputs=inputs,
-            model_options=model_options,
-            tools=tools,
-            tool_choice=tool_choice,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            logprobs=logprobs,
-            top_logprobs=top_logprobs,
-            max_concurrency=max_concurrency,
-            timeout=timeout,
-            streaming_callback=streaming_callback,
-            validation_callback=validation_callback,
-            completion_callback=completion_callback,
-            use_cache=use_cache,
-        )