docent-python 0.1.59a0__tar.gz → 0.1.61a0__tar.gz

{docent_python-0.1.59a0 → docent_python-0.1.61a0}/PKG-INFO

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

{docent_python-0.1.59a0 → docent_python-0.1.61a0}/docent/data_models/reading.py

@@ -1,8 +1,8 @@
 from datetime import datetime
-from typing import Any, Literal, TypeAlias
+from typing import Annotated, Any, Literal, TypeAlias
 from uuid import uuid4
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 from docent._llm_util.providers.preference_types import ModelOption
 
@@ -41,6 +41,7 @@ prompt segments, model config, output schema, and user-supplied arguments.
   re-evaluation.
 """
 ReadingCacheMode = Literal["reading", "results", "none"]
+ReadingStatus = Literal["completed", "failed", "pending", "cached", "needs_approval", "unresolved"]
 
 
 class ContextFilterSection(BaseModel):
@@ -73,6 +74,7 @@ class ReadingPreset(BaseModel):
     collection_id: str
     name: str
     created_at: datetime | None = None
+    created_by: str | None = None
     updated_at: datetime | None = None
 
 
@@ -117,6 +119,7 @@ class Reading(BaseModel):
     user_metadata: dict[str, Any] | None = None
     source_reading_preset_id: str | None = None
     created_at: datetime | None = None
+    created_by: str | None = None
 
 
 class ReadingResult(BaseModel):
@@ -174,6 +177,7 @@ class ReadingStep(BaseModel):
     name: str | None = None
     reading_id: str | None = None
     dql_query: str | None = None
+    dql_step_alias: str | None = None
     prompt_template_segments: list[Any] | None = None
     context_config: dict[str, Any] | None = None
     model: ModelOption
@@ -188,7 +192,11 @@ class ReadingStep(BaseModel):
     def to_submission(self, *, dql_query: str | None = None) -> "ReadingStepSubmission":
         """Convert to a ReadingStepSubmission for resolve_reading_entry.
 
-        Optionally overrides dql_query (e.g. after alias substitution).
+        Optionally overrides dql_query (e.g. after alias substitution). The
+        stored step's own dql_query may be None when dql_step_alias is set;
+        callers are expected to pass the resolved DQL explicitly in that case.
+        When a concrete DQL is supplied, clear dql_step_alias so the
+        submission continues to satisfy the "exactly one DQL source" contract.
         """
         return ReadingStepSubmission(
             alias=self.alias,
@@ -200,6 +208,7 @@ class ReadingStep(BaseModel):
             prompt_template_segments=self.prompt_template_segments,
             context_config=self.context_config,
             dql_query=dql_query if dql_query is not None else self.dql_query,
+            dql_step_alias=None if dql_query is not None else self.dql_step_alias,
             source_reading_preset_id=self.source_reading_preset_id,
             cache_mode=self.cache_mode,
         )
@@ -226,6 +235,7 @@ class ReadingPlan(BaseModel):
     name: str | None = None
     steps: list[PlanStep] = Field(default_factory=list)  # type: ignore[reportUnknownVariableType]
     created_at: datetime | None = None
+    created_by: str | None = None
     updated_at: datetime | None = None
 
 
@@ -268,10 +278,29 @@ class ReadingStepSubmission(BaseModel):
     prompt_template_segments: list[Any] | None = None
     context_config: dict[str, ParameterContextConfig] | None = None
     dql_query: str | None = None
+    # References a DqlOnlyStep in the same plan whose rows feed this reading.
+    # Mutually exclusive with dql_query for template entries.
+    dql_step_alias: str | None = None
 
     # Scripted reading fields (mutually exclusive with template fields)
     requests: list[ScriptedRequest] | None = None
 
+    @model_validator(mode="after")
+    def _validate_dql_source(self) -> "ReadingStepSubmission":
+        if self.requests is not None:
+            if self.dql_query is not None or self.dql_step_alias is not None:
+                raise ValueError(
+                    "Scripted reading submissions must not set dql_query or dql_step_alias"
+                )
+            return self
+        if self.dql_query is not None and self.dql_step_alias is not None:
+            raise ValueError("ReadingStepSubmission: set exactly one of dql_query / dql_step_alias")
+        if self.dql_query is None and self.dql_step_alias is None:
+            raise ValueError(
+                "ReadingStepSubmission: template entries must set one of dql_query / dql_step_alias"
+            )
+        return self
+
 
 class PresetReadingStepSubmission(BaseModel):
     entry_type: Literal["preset_reading"] = "preset_reading"
@@ -280,8 +309,21 @@ class PresetReadingStepSubmission(BaseModel):
     source_reading_preset_id: str
     user_metadata: dict[str, Any] | None = None
     dql_query: str | None = None
+    dql_step_alias: str | None = None
     cache_mode: ReadingCacheMode = "reading"
 
+    @model_validator(mode="after")
+    def _validate_dql_source(self) -> "PresetReadingStepSubmission":
+        if self.dql_query is not None and self.dql_step_alias is not None:
+            raise ValueError(
+                "PresetReadingStepSubmission: set exactly one of dql_query / dql_step_alias"
+            )
+        if self.dql_query is None and self.dql_step_alias is None:
+            raise ValueError(
+                "PresetReadingStepSubmission: must set one of dql_query / dql_step_alias"
+            )
+        return self
+
 
 class DqlOnlyStepSubmission(BaseModel):
     entry_type: Literal["dql_only"] = "dql_only"
@@ -306,23 +348,140 @@ class PlanSubmissionRequest(BaseModel):
     entries: list[PlanStepSubmission]
 
 
+class DqlPreview(BaseModel):
+    columns: list[str]
+    rows: list[list[Any]]
+    truncated: bool
+    row_count: int
+
+
+class ReadingResultPreview(BaseModel):
+    id: str
+    output: dict[str, Any] | None = None
+    error: dict[str, Any] | None = None
+
+
 class PlanStepSubmissionStatus(BaseModel):
     alias: str
-    status: Literal["cached", "needs_approval", "unresolved"]
+    entry_type: str
+    status: ReadingStatus
     reading_id: str | None = None
+    result_count: int | None = None
+    dql_preview: DqlPreview | None = None
+    result_preview: list[ReadingResultPreview] | None = None
 
 
 class PlanSubmissionResponse(BaseModel):
     plan_id: str
+    plan_name: str | None = None
+    previous_latest_plan_id: str | None = None
+    has_active_listeners: bool = False
     entry_statuses: list[PlanStepSubmissionStatus]
 
 
+# ── Plan SSE stream events (server → SDK) ────────────────────────────
+
+
+class PlanStreamStepStatus(BaseModel):
+    """Minimal step shape carried inside a snapshot event."""
+
+    alias: str
+    reading_id: str | None = None
+    derived_status: str
+
+
+class PlanSnapshotEvent(BaseModel):
+    type: Literal["snapshot"] = "snapshot"
+    steps: list[PlanStreamStepStatus]
+
+
+class PlanStepError(BaseModel):
+    message: str
+
+
+class PlanStepStartedEvent(BaseModel):
+    type: Literal["step_started"] = "step_started"
+    plan_id: str
+    step_alias: str
+    job_id: str
+    reading_id: str
+
+
+class PlanStepCompletedEvent(BaseModel):
+    type: Literal["step_completed"] = "step_completed"
+    plan_id: str
+    step_alias: str
+    job_id: str
+    reading_id: str
+    result_count: int | None = None
+
+
+class PlanStepFailedEvent(BaseModel):
+    type: Literal["step_failed"] = "step_failed"
+    plan_id: str
+    step_alias: str
+    job_id: str
+    error: PlanStepError | None = None
+
+
+class PlanStepsUpdatedEvent(BaseModel):
+    type: Literal["steps_updated"] = "steps_updated"
+    plan_id: str
+
+
+class PlanJobStartedEvent(BaseModel):
+    type: Literal["job_started"] = "job_started"
+    plan_id: str
+    job_id: str
+
+
+class PlanJobCompletedEvent(BaseModel):
+    type: Literal["job_completed"] = "job_completed"
+    plan_id: str
+    job_id: str
+
+
+class PlanJobFailedEvent(BaseModel):
+    type: Literal["job_failed"] = "job_failed"
+    plan_id: str
+    job_id: str
+    error: PlanStepError | None = None
+
+
+class PlanSupersededEvent(BaseModel):
+    type: Literal["plan_superseded"] = "plan_superseded"
+    plan_id: str
+    superseded_by_plan_id: str
+    name: str | None = None
+
+
+class PlanJobCancelledEvent(BaseModel):
+    type: Literal["job_cancelled"] = "job_cancelled"
+    plan_id: str
+
+
+PlanStreamEvent: TypeAlias = Annotated[
+    PlanSnapshotEvent
+    | PlanStepStartedEvent
+    | PlanStepCompletedEvent
+    | PlanStepFailedEvent
+    | PlanStepsUpdatedEvent
+    | PlanJobStartedEvent
+    | PlanJobCompletedEvent
+    | PlanJobFailedEvent
+    | PlanJobCancelledEvent
+    | PlanSupersededEvent,
+    Field(discriminator="type"),
+]
+
+
 __all__ = [
     "AnnotatableReadingParamType",
     "BeginGroupStep",
     "ContextFilterSection",
     "DqlOnlyStep",
     "DqlOnlyStepSubmission",
+    "DqlPreview",
     "EndGroupStep",
     "EndStepGroupSubmission",
     "ScriptedRequest",
@@ -335,6 +494,8 @@ __all__ = [
     "ReadingCacheMode",
     "ReadingParamPlaceholder",
     "ReadingParamType",
+    "ReadingResultPreview",
+    "ReadingStatus",
     "ReadingStep",
     "ReadingStepSubmission",
     "ReadingTemplateSegment",
@@ -346,4 +507,17 @@ __all__ = [
     "ReadingResult",
     "StepGroupSubmission",
     "PresetReadingStepSubmission",
+    "PlanStreamEvent",
+    "PlanStreamStepStatus",
+    "PlanSnapshotEvent",
+    "PlanStepStartedEvent",
+    "PlanStepCompletedEvent",
+    "PlanStepError",
+    "PlanStepFailedEvent",
+    "PlanStepsUpdatedEvent",
+    "PlanJobStartedEvent",
+    "PlanJobCompletedEvent",
+    "PlanJobFailedEvent",
+    "PlanJobCancelledEvent",
+    "PlanSupersededEvent",
 ]

{docent_python-0.1.59a0 → docent_python-0.1.61a0}/docent/mcp/server.py

@@ -33,10 +33,13 @@ def get_metadata_fields(collection_id: str) -> str:
     """
     client = get_client()
     try:
-        fields = client.get_metadata_fields(
+        response = client.get_metadata_fields(
             collection_id, include_sample_values=True, sample_limit=10
         )
 
+        fields = response.get("fields", [])
+        total_runs = response.get("total_runs")
+
         if not fields:
             return f"No metadata fields found for collection {collection_id}"
 
@@ -71,7 +74,10 @@ def get_metadata_fields(collection_id: str) -> str:
             lines.append(line)
 
         field_list = "\n".join(lines)
-        return f"Metadata fields for collection {collection_id}:\n{field_list}"
+        tool_output = f"Metadata fields for collection {collection_id}:\n{field_list}"
+        if total_runs is not None:
+            tool_output += f"\n\nTotal runs: {total_runs}"
+        return tool_output
     except Exception as e:
         error_msg = str(e)
         if "404" in error_msg:
@@ -321,7 +327,11 @@ def get_reading_plan_results(
             for step in steps:
                 if step.get("type") == "reading" and step.get("reading_id"):
                     try:
-                        results = client.get_reading_results(collection_id, step["reading_id"])
+                        results = client.get_reading_results(
+                            collection_id,
+                            step["reading_id"],
+                            include_output=False,
+                        )
                         result_counts[step["alias"]] = len(results)
                     except Exception:
                         pass

docent_python-0.1.61a0/docent/sdk/_agent_runs.py

@@ -0,0 +1,217 @@
+"""Agent run fetch, metadata, transcript groups, and chat sessions."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from docent.data_models.agent_run import AgentRun
+from docent.sdk._base import DocentBase
+
+
+class DocentAgentRunsMixin(DocentBase):
+    """Agent run and transcript-group operations."""
+
+    def get_agent_run(self, collection_id: str, agent_run_id: str) -> AgentRun | None:
+        """Get a specific agent run by its ID.
+
+        Args:
+            collection_id: ID of the Collection.
+            agent_run_id: The ID of the agent run to retrieve.
+
+        Returns:
+            dict: Dictionary containing the agent run information.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._api_url}/{collection_id}/agent_run"
+        response = self._session.get(url, params={"agent_run_id": agent_run_id})
+        self._handle_response_errors(response)
+        if response.json() is None:
+            return None
+        else:
+            # We do this to avoid metadata validation failing
+            # TODO(mengk): kinda hacky
+            return AgentRun.model_validate(response.json())
+
+    def update_agent_run_metadata(
+        self,
+        collection_id: str,
+        agent_run_id: str,
+        metadata: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Merge metadata into an agent run's existing metadata.
+
+        Uses a deep merge: nested dictionaries are merged recursively so
+        existing keys are preserved, while non-dict values are overwritten.
+        Keys not present in ``metadata`` are left unchanged.
+
+        Requires WRITE permission on the collection.
+
+        Args:
+            collection_id: ID of the Collection containing the agent run.
+            agent_run_id: ID of the agent run to update.
+            metadata: Dictionary of metadata fields to merge.
+
+        Returns:
+            The full merged metadata dictionary after the update.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails (e.g., 404 if the agent run is not found).
+        """
+        url = f"{self._server_url}/{collection_id}/agent_run/{agent_run_id}/metadata"
+        response = self._session.put(url, json={"metadata": metadata})
+        self._handle_response_errors(response)
+        data: dict[str, Any] = response.json()
+        return data
+
+    def delete_agent_run_metadata_keys(
+        self,
+        collection_id: str,
+        agent_run_id: str,
+        keys: list[str],
+    ) -> tuple[dict[str, Any], list[str]]:
+        """Remove keys from an agent run's metadata.
+
+        Supports dot-delimited paths for nested deletion. For example,
+        ``"config.model"`` removes the ``model`` key inside ``config``
+        without affecting other keys in that dict.
+
+        Requires WRITE permission on the collection.
+
+        Args:
+            collection_id: ID of the Collection containing the agent run.
+            agent_run_id: ID of the agent run to modify.
+            keys: Metadata keys to remove. Use dot-delimited paths for nested
+                keys (e.g. ``["top_level_key", "nested.child_key"]``).
+
+        Returns:
+            A tuple of (metadata after deletion, list of keys that were not found).
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails (e.g., 404 if the agent run is not found).
+        """
+        url = f"{self._server_url}/{collection_id}/agent_run/{agent_run_id}/metadata/delete"
+        response = self._session.post(url, json={"keys": keys})
+        self._handle_response_errors(response)
+        data: dict[str, Any] = response.json()
+        metadata: dict[str, Any] = data["metadata"]
+        not_found: list[str] = data["not_found"]
+        return metadata, not_found
+
+    def get_agent_run_metadata(
+        self,
+        collection_id: str,
+        agent_run_id: str,
+    ) -> dict[str, Any]:
+        """Get an agent run's metadata.
+
+        Args:
+            collection_id: ID of the Collection containing the agent run.
+            agent_run_id: ID of the agent run.
+
+        Returns:
+            The agent run's metadata dict.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._server_url}/{collection_id}/agent_run/{agent_run_id}/metadata"
+        response = self._session.get(url)
+        self._handle_response_errors(response)
+        return response.json()
+
+    # ──────────────────────────────────────────
+    # Transcript group metadata
+    # ──────────────────────────────────────────
+
+    def get_transcript_group_metadata(
+        self,
+        collection_id: str,
+        transcript_group_id: str,
+    ) -> dict[str, Any]:
+        """Get a transcript group's metadata.
+
+        Args:
+            collection_id: ID of the Collection containing the transcript group.
+            transcript_group_id: ID of the transcript group.
+
+        Returns:
+            The transcript group's metadata dict.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._server_url}/{collection_id}/transcript_group/{transcript_group_id}/metadata"
+        response = self._session.get(url)
+        self._handle_response_errors(response)
+        return response.json()
+
+    def update_transcript_group_metadata(
+        self,
+        collection_id: str,
+        transcript_group_id: str,
+        metadata: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Deep-merge metadata into a transcript group's existing metadata.
+
+        Args:
+            collection_id: ID of the Collection containing the transcript group.
+            transcript_group_id: ID of the transcript group.
+            metadata: Metadata dict to merge into the existing metadata.
+
+        Returns:
+            The full merged metadata dict.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._server_url}/{collection_id}/transcript_group/{transcript_group_id}/metadata"
+        response = self._session.put(url, json={"metadata": metadata})
+        self._handle_response_errors(response)
+        return response.json()
+
+    def delete_transcript_group_metadata_keys(
+        self,
+        collection_id: str,
+        transcript_group_id: str,
+        keys: list[str],
+    ) -> tuple[dict[str, Any], list[str]]:
+        """Remove keys from a transcript group's metadata.
+
+        Supports dot-delimited paths for nested deletion.
+
+        Args:
+            collection_id: ID of the Collection containing the transcript group.
+            transcript_group_id: ID of the transcript group.
+            keys: Keys to remove. Use dot notation for nested keys.
+
+        Returns:
+            Tuple of (metadata after deletion, list of keys that were not found).
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._server_url}/{collection_id}/transcript_group/{transcript_group_id}/metadata/delete"
+        response = self._session.post(url, json={"keys": keys})
+        self._handle_response_errors(response)
+        data = response.json()
+        return data["metadata"], data["not_found"]
+
+    def get_chat_sessions(self, collection_id: str, agent_run_id: str) -> list[dict[str, Any]]:
+        """Get all chat sessions for an agent run, excluding judge result sessions.
+
+        Args:
+            collection_id: ID of the Collection.
+            agent_run_id: The ID of the agent run to retrieve chat sessions for.
+
+        Returns:
+            list: List of chat session dictionaries.
+
+        Raises:
+            requests.exceptions.HTTPError: If the API request fails.
+        """
+        url = f"{self._api_url}/chat/{collection_id}/{agent_run_id}/sessions"
+        response = self._session.get(url)
+        self._handle_response_errors(response)
+        return response.json()