docent-python 0.1.58a0__tar.gz → 0.1.60a0__tar.gz

{docent_python-0.1.58a0 → docent_python-0.1.60a0}/PKG-INFO

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

{docent_python-0.1.58a0 → docent_python-0.1.60a0}/docent/_llm_util/model_registry.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass
 from functools import lru_cache
 from typing import Optional
@@ -9,6 +10,8 @@ from docent._log_util import get_logger
 
 logger = get_logger(__name__)
 
+_CLAUDE_VERSION_PATTERN = re.compile(r"(claude-(?:haiku|sonnet|opus)-4)[.-](\d+)\b")
+
 
 """
 Values are USD per million tokens
@@ -34,6 +37,14 @@ _REGISTRY: list[tuple[str, ModelInfo]] = [
         "gpt-5-chat-latest",
         ModelInfo(rate={"input": 1.25, "output": 10.0}, context_window=128_000),
     ),
+    (
+        "gpt-5.4-mini",
+        ModelInfo(rate={"input": 0.75, "output": 4.50}, context_window=400_000),
+    ),
+    (
+        "gpt-5.4",
+        ModelInfo(rate={"input": 2.50, "output": 15.0}, context_window=1_050_000),
+    ),
     (
         "gpt-5-nano",
         ModelInfo(rate={"input": 0.05, "output": 0.40}, context_window=400_000),
@@ -62,18 +73,9 @@ _REGISTRY: list[tuple[str, ModelInfo]] = [
         "claude-sonnet-4-5",
         ModelInfo(rate={"input": 3.0, "output": 15.0}, context_window=200_000),
     ),
-    (
-        "claude-sonnet-4-6",
-        ModelInfo(rate={"input": 3.0, "output": 15.0}, context_window=1_000_000),
-    ),
-    (
-        "claude-opus-4-6",
-        ModelInfo(rate={"input": 5.0, "output": 25.0}, context_window=200_000),
-    ),
-    (
-        "claude-haiku-4-5",
-        ModelInfo(rate={"input": 1.0, "output": 5.0}, context_window=200_000),
-    ),
+    ("claude-sonnet-4-6", ModelInfo(rate={"input": 3.0, "output": 15.0}, context_window=1_000_000)),
+    ("claude-opus-4-6", ModelInfo(rate={"input": 5.0, "output": 25.0}, context_window=1_000_000)),
+    ("claude-haiku-4-5", ModelInfo(rate={"input": 1.0, "output": 5.0}, context_window=200_000)),
     (
         "claude-opus-4-5-20251101",
         ModelInfo(rate={"input": 5.0, "output": 25.0}, context_window=200_000),
@@ -108,6 +110,13 @@ _REGISTRY: list[tuple[str, ModelInfo]] = [
             context_window=1_048_576,
         ),
     ),
+    (
+        "gemini-3.1-pro-preview",
+        ModelInfo(
+            rate={"input": 2.00, "output": 12.00},
+            context_window=1_048_576,
+        ),
+    ),
     (
         "gemini-3-flash-preview",
         ModelInfo(
@@ -146,12 +155,27 @@ _REGISTRY: list[tuple[str, ModelInfo]] = [
 ]
 
 
+def normalize_model_name(model_name: str) -> str:
+    """Normalize provider-specific naming differences before registry lookup."""
+
+    return _CLAUDE_VERSION_PATTERN.sub(r"\1-\2", model_name)
+
+
+def model_names_match(expected: str, actual: str) -> bool:
+    """Match a configured model name against a provider-reported model string."""
+
+    normalized_expected = normalize_model_name(expected)
+    normalized_actual = normalize_model_name(actual)
+    return normalized_expected in normalized_actual
+
+
 @lru_cache(maxsize=None)
 def get_model_info(model_name: str) -> Optional[ModelInfo]:
+    normalized_model_name = normalize_model_name(model_name)
     for registry_model_name, info in sorted(
         _REGISTRY, key=lambda entry: len(entry[0]), reverse=True
     ):
-        if registry_model_name in model_name:
+        if registry_model_name in normalized_model_name:
             return info
     return None
 

docent_python-0.1.60a0/docent/_llm_util/providers/preference_types.py

@@ -0,0 +1,268 @@
+"""Provides preferences of which LLM models to use for different Docent functions."""
+
+from functools import cached_property
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+from docent._llm_util.model_registry import get_context_window
+from docent._log_util import get_logger
+
+logger = get_logger(__name__)
+
+
+class ModelOption(BaseModel):
+    """Configuration for a specific model from a provider. Not to be confused with ModelInfo.
+
+    Attributes:
+        provider: The name of the LLM provider (e.g., "openai", "anthropic").
+        model_name: The specific model to use from the provider.
+        reasoning_effort: Optional indication of computational effort to use.
+    """
+
+    model_config = ConfigDict(extra="ignore")
+
+    provider: str
+    model_name: str
+    reasoning_effort: Literal["minimal", "low", "medium", "high"] | None = None
+
+
+class ModelOptionWithContext(BaseModel):
+    """Enhanced model option that includes context window information for frontend use.
+    Not to be confused with ModelInfo or ModelOption.
+
+    Attributes:
+        provider: The name of the LLM provider (e.g., "openai", "anthropic").
+        model_name: The specific model to use from the provider.
+        reasoning_effort: Optional indication of computational effort to use.
+        context_window: The context window size in tokens.
+        uses_byok: Whether this model would use the user's own API key.
+    """
+
+    provider: str
+    model_name: str
+    reasoning_effort: Literal["minimal", "low", "medium", "high"] | None = None
+    context_window: int
+    uses_byok: bool
+
+    @classmethod
+    def from_model_option(
+        cls, model_option: ModelOption, uses_byok: bool = False
+    ) -> "ModelOptionWithContext":
+        """Create a ModelOptionWithContext from a ModelOption.
+
+        Args:
+            model_option: The base model option
+            uses_byok: Whether this model requires bring-your-own-key
+
+        Returns:
+            ModelOptionWithContext with context window looked up from global mapping
+        """
+        context_window = get_context_window(model_option.model_name)
+
+        return cls(
+            provider=model_option.provider,
+            model_name=model_option.model_name,
+            reasoning_effort=model_option.reasoning_effort,
+            context_window=context_window,
+            uses_byok=uses_byok,
+        )
+
+
+def merge_models_with_byok(
+    defaults: list[ModelOption],
+    byok: list[ModelOption],
+    api_keys: dict[str, str] | None,
+) -> list[ModelOptionWithContext]:
+    user_keys = api_keys or {}
+
+    merged: list[ModelOption] = list(defaults)
+    if user_keys:
+        merged.extend([m for m in byok if m.provider in user_keys])
+
+    return [ModelOptionWithContext.from_model_option(m, m.provider in user_keys) for m in merged]
+
+
+class PublicProviderPreferences(BaseModel):
+    @cached_property
+    def default_judge_models(self) -> list[ModelOption]:
+        """Judge models that any user can access without providing their own API key"""
+
+        return [
+            ModelOption(provider="openai", model_name="gpt-5.4-mini", reasoning_effort="low"),
+            ModelOption(provider="openai", model_name="gpt-5.4-mini", reasoning_effort="medium"),
+            ModelOption(provider="openai", model_name="gpt-5.4-mini", reasoning_effort="high"),
+            ModelOption(provider="openai", model_name="gpt-5.4", reasoning_effort="low"),
+            ModelOption(provider="openai", model_name="gpt-5.4", reasoning_effort="medium"),
+            ModelOption(provider="openai", model_name="gpt-5.4", reasoning_effort="high"),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-haiku-4-5",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-haiku-4-5",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-haiku-4-5",
+                reasoning_effort="high",
+            ),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-opus-4-6",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-opus-4-6",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-opus-4-6",
+                reasoning_effort="high",
+            ),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-sonnet-4-6",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-sonnet-4-6",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-sonnet-4-6",
+                reasoning_effort="high",
+            ),
+            ModelOption(
+                provider="google",
+                model_name="gemini-3-flash-preview",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="google",
+                model_name="gemini-3-flash-preview",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="google",
+                model_name="gemini-3-flash-preview",
+                reasoning_effort="high",
+            ),
+            ModelOption(
+                provider="google",
+                model_name="gemini-3.1-pro-preview",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="google",
+                model_name="gemini-3.1-pro-preview",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="google",
+                model_name="gemini-3.1-pro-preview",
+                reasoning_effort="high",
+            ),
+            # Open Router equivalents
+            ModelOption(
+                provider="openrouter", model_name="openai/gpt-5.4-mini", reasoning_effort="low"
+            ),
+            ModelOption(
+                provider="openrouter", model_name="openai/gpt-5.4-mini", reasoning_effort="medium"
+            ),
+            ModelOption(
+                provider="openrouter", model_name="openai/gpt-5.4-mini", reasoning_effort="high"
+            ),
+            ModelOption(provider="openrouter", model_name="openai/gpt-5.4", reasoning_effort="low"),
+            ModelOption(
+                provider="openrouter", model_name="openai/gpt-5.4", reasoning_effort="medium"
+            ),
+            ModelOption(
+                provider="openrouter", model_name="openai/gpt-5.4", reasoning_effort="high"
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-haiku-4.5",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-haiku-4.5",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-haiku-4.5",
+                reasoning_effort="high",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-opus-4.6",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-opus-4.6",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-opus-4.6",
+                reasoning_effort="high",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-sonnet-4.6",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-sonnet-4.6",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="anthropic/claude-sonnet-4.6",
+                reasoning_effort="high",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="google/gemini-3-flash-preview",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="google/gemini-3-flash-preview",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="google/gemini-3-flash-preview",
+                reasoning_effort="high",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="google/gemini-3.1-pro-preview",
+                reasoning_effort="low",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="google/gemini-3.1-pro-preview",
+                reasoning_effort="medium",
+            ),
+            ModelOption(
+                provider="openrouter",
+                model_name="google/gemini-3.1-pro-preview",
+                reasoning_effort="high",
+            ),
+        ]
+
+
+PUBLIC_PROVIDER_PREFERENCES = PublicProviderPreferences()

{docent_python-0.1.58a0 → docent_python-0.1.60a0}/docent/data_models/reading.py

@@ -1,5 +1,5 @@
 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
@@ -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):
@@ -226,6 +229,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
 
 
@@ -306,23 +310,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 +456,8 @@ __all__ = [
     "ReadingCacheMode",
     "ReadingParamPlaceholder",
     "ReadingParamType",
+    "ReadingResultPreview",
+    "ReadingStatus",
     "ReadingStep",
     "ReadingStepSubmission",
     "ReadingTemplateSegment",
@@ -346,4 +469,17 @@ __all__ = [
     "ReadingResult",
     "StepGroupSubmission",
     "PresetReadingStepSubmission",
+    "PlanStreamEvent",
+    "PlanStreamStepStatus",
+    "PlanSnapshotEvent",
+    "PlanStepStartedEvent",
+    "PlanStepCompletedEvent",
+    "PlanStepError",
+    "PlanStepFailedEvent",
+    "PlanStepsUpdatedEvent",
+    "PlanJobStartedEvent",
+    "PlanJobCompletedEvent",
+    "PlanJobFailedEvent",
+    "PlanJobCancelledEvent",
+    "PlanSupersededEvent",
 ]

{docent_python-0.1.58a0 → docent_python-0.1.60a0}/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