docent-python 0.1.41a0__py3-none-any.whl

docent/data_models/chat/message.py

@@ -0,0 +1,191 @@
+from logging import getLogger
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Discriminator
+
+from docent.data_models.chat.content import Content
+from docent.data_models.chat.tool import ToolCall
+from docent.data_models.citation import InlineCitation
+
+logger = getLogger(__name__)
+
+
+class BaseChatMessage(BaseModel):
+    """Base class for all chat message types.
+
+    Attributes:
+        id: Optional unique identifier for the message.
+        content: The message content, either as a string or list of Content objects.
+        role: The role of the message sender (system, user, assistant, tool).
+        metadata: Additional structured metadata about the message.
+    """
+
+    id: str | None = None
+    content: str | list[Content]
+    role: Literal["system", "user", "assistant", "tool"]
+    metadata: dict[str, Any] | None = None
+
+    @property
+    def text(self) -> str:
+        """Get the text content of the message.
+
+        Returns:
+            str: The text content of the message. If content is a list,
+                 concatenates all text content elements with newlines.
+        """
+        if isinstance(self.content, str):
+            return self.content
+        else:
+            all_text = [content.text for content in self.content if content.type == "text"]
+            return "\n".join(all_text)
+
+
+class SystemMessage(BaseChatMessage):
+    """System message in a chat conversation.
+
+    Attributes:
+        role: Always set to "system".
+    """
+
+    role: Literal["system"] = "system"  # type: ignore
+
+
+class UserMessage(BaseChatMessage):
+    """User message in a chat conversation.
+
+    Attributes:
+        role: Always set to "user".
+        tool_call_id: Optional list of tool call IDs this message is responding to.
+    """
+
+    role: Literal["user"] = "user"  # type: ignore
+    tool_call_id: list[str] | None = None
+
+
+class AssistantMessage(BaseChatMessage):
+    """Assistant message in a chat conversation.
+
+    Attributes:
+        role: Always set to "assistant".
+        model: Optional identifier for the model that generated this message.
+        tool_calls: Optional list of tool calls made by the assistant.
+    """
+
+    role: Literal["assistant"] = "assistant"  # type: ignore
+    model: str | None = None
+    tool_calls: list[ToolCall] | None = None
+
+
+class DocentAssistantMessage(AssistantMessage):
+    """Assistant message in a chat session with additional chat-specific metadata.
+
+    This extends AssistantMessage with fields that are only relevant in Docent chat contexts
+
+    Attributes:
+        citations: Optional list of citations referenced in the message content.
+        suggested_messages: Optional list of suggested followup messages.
+    """
+
+    citations: list[InlineCitation] | None = None
+    suggested_messages: list[str] | None = None
+
+
+class ToolMessage(BaseChatMessage):
+    """Tool message in a chat conversation.
+
+    Attributes:
+        role: Always set to "tool".
+        tool_call_id: Optional ID of the tool call this message is responding to.
+        function: Optional name of the function that was called.
+        error: Optional error information if the tool call failed.
+    """
+
+    role: Literal["tool"] = "tool"  # type: ignore
+
+    tool_call_id: str | None = None
+    function: str | None = None
+    error: dict[str, Any] | None = None
+
+
+ChatMessage = Annotated[
+    SystemMessage | UserMessage | AssistantMessage | ToolMessage,
+    Discriminator("role"),
+]
+"""Type alias for any chat message type, discriminated by the role field.
+
+This is the base message union used in Transcript and AgentRun contexts.
+For chat sessions, use ChatSessionMessage instead.
+"""
+
+DocentChatMessage = Annotated[
+    SystemMessage | UserMessage | DocentAssistantMessage | ToolMessage,
+    Discriminator("role"),
+]
+"""Type alias for chat session messages with chat-specific assistant metadata."""
+
+
+def parse_chat_message(message_data: dict[str, Any] | ChatMessage) -> ChatMessage:
+    """Parse a message dictionary or object into the appropriate ChatMessage subclass.
+
+    This parses base messages without chat-specific fields. For chat sessions,
+    use parse_chat_session_message instead.
+
+    Args:
+        message_data: A dictionary or ChatMessage object representing a chat message.
+
+    Returns:
+        ChatMessage: An instance of a ChatMessage subclass based on the role.
+
+    Raises:
+        ValueError: If the message role is unknown.
+    """
+    if isinstance(message_data, (SystemMessage, UserMessage, AssistantMessage, ToolMessage)):
+        return message_data
+
+    role = message_data.get("role")
+    if role == "system":
+        return SystemMessage.model_validate(message_data)
+    elif role == "user":
+        return UserMessage.model_validate(message_data)
+    elif role == "assistant":
+        return AssistantMessage.model_validate(message_data)
+    elif role == "tool":
+        return ToolMessage.model_validate(message_data)
+    else:
+        raise ValueError(f"Unknown message role: {role}")
+
+
+def parse_docent_chat_message(
+    message_data: dict[str, Any] | DocentChatMessage,
+) -> DocentChatMessage:
+    """Parse a message dictionary or object into the appropriate ChatSessionMessage subclass.
+
+    This handles chat session messages which may include ChatAssistantMessage with
+    citations and suggested_messages fields.
+
+    Args:
+        message_data: A dictionary or ChatSessionMessage object representing a chat session message.
+
+    Returns:
+        ChatSessionMessage: An instance of a ChatSessionMessage subclass based on the role.
+
+    Raises:
+        ValueError: If the message role is unknown.
+    """
+    if isinstance(
+        message_data,
+        (SystemMessage, UserMessage, DocentAssistantMessage, AssistantMessage, ToolMessage),
+    ):
+        return message_data
+
+    role = message_data.get("role")
+    if role == "system":
+        return SystemMessage.model_validate(message_data)
+    elif role == "user":
+        return UserMessage.model_validate(message_data)
+    elif role == "assistant":
+        return DocentAssistantMessage.model_validate(message_data)
+    elif role == "tool":
+        return ToolMessage.model_validate(message_data)
+    else:
+        raise ValueError(f"Unknown message role: {role}")

docent/data_models/chat/tool.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+@dataclass
+class ToolCall:
+    """Tool call information.
+
+    Attributes:
+        id: Unique identifier for tool call.
+        type: Type of tool call. Can only be "function" or None.
+        function: Function called.
+        arguments: Arguments to function.
+        parse_error: Error which occurred parsing tool call.
+        view: Custom view of tool call input.
+    """
+
+    id: str
+    function: str
+    arguments: dict[str, Any]
+    type: Literal["function"] | None = None
+    parse_error: str | None = None
+    view: ToolCallContent | None = None
+
+
+class ToolCallContent(BaseModel):
+    """Content to include in tool call view.
+
+    Attributes:
+        title: Optional (plain text) title for tool call content.
+        format: Format (text or markdown).
+        content: Text or markdown content.
+    """
+
+    title: str | None = None
+    format: Literal["text", "markdown"]
+    content: str
+
+
+class ToolParam(BaseModel):
+    """A parameter for a tool function.
+
+    Args:
+        name: The name of the parameter.
+        description: A description of what the parameter does.
+        input_schema: JSON Schema describing the parameter's type and validation rules.
+    """
+
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+
+class ToolParams(BaseModel):
+    """Description of tool parameters object in JSON Schema format.
+
+    Args:
+        type: The type of the parameters object, always 'object'.
+        properties: Dictionary mapping parameter names to their ToolParam definitions.
+        required: List of required parameter names.
+        additionalProperties: Whether additional properties are allowed beyond those
+            specified. Always False.
+    """
+
+    type: Literal["object"] = "object"
+    properties: dict[str, ToolParam] = Field(default_factory=dict)
+    required: list[str] = Field(default_factory=list)
+    additionalProperties: bool = False
+
+
+class ToolInfo(BaseModel):
+    """Specification of a tool (JSON Schema compatible).
+
+    If you are implementing a ModelAPI, most LLM libraries can
+    be passed this object (dumped to a dict) directly as a function
+    specification. For example, in the OpenAI provider:
+
+    ```python
+    ChatCompletionToolParam(
+        type="function",
+        function=tool.model_dump(exclude_none=True),
+    )
+    ```
+
+    In some cases the field names don't match up exactly. In that case
+    call `model_dump()` on the `parameters` field. For example, in the
+    Anthropic provider:
+
+    ```python
+    ToolParam(
+        name=tool.name,
+        description=tool.description,
+        input_schema=tool.parameters.model_dump(exclude_none=True),
+    )
+    ```
+
+    Attributes:
+        name: Name of tool.
+        description: Short description of tool.
+        parameters: JSON Schema of tool parameters object.
+    """
+
+    name: str
+    description: str
+    parameters: ToolParams = Field(default_factory=ToolParams)

docent/data_models/citation.py

@@ -0,0 +1,187 @@
+from typing import Annotated, Literal, Union
+
+from pydantic import BaseModel, Discriminator
+
+
+class CitationTargetTextRange(BaseModel):
+    start_pattern: str | None = None
+    end_pattern: str | None = None
+
+
+class ResolvedCitationItem(BaseModel):
+    pass
+
+
+class CitationTarget(BaseModel):
+    item: "ResolvedCitationItemUnion"
+    text_range: CitationTargetTextRange | None = None
+
+
+class ParsedCitation(BaseModel):
+    start_idx: int
+    end_idx: int
+    item_alias: str
+    text_range: CitationTargetTextRange | None = None
+
+
+class InlineCitation(BaseModel):
+    start_idx: int
+    end_idx: int
+    target: CitationTarget
+
+
+class AgentRunMetadataItem(ResolvedCitationItem):
+    item_type: Literal["agent_run_metadata"] = "agent_run_metadata"
+    agent_run_id: str
+    collection_id: str
+    metadata_key: str
+
+
+class TranscriptMetadataItem(ResolvedCitationItem):
+    item_type: Literal["transcript_metadata"] = "transcript_metadata"
+    agent_run_id: str
+    collection_id: str
+    transcript_id: str
+    metadata_key: str
+
+
+class TranscriptBlockMetadataItem(ResolvedCitationItem):
+    item_type: Literal["block_metadata"] = "block_metadata"
+    agent_run_id: str
+    collection_id: str
+    transcript_id: str
+    block_idx: int
+    metadata_key: str
+
+
+class TranscriptBlockContentItem(ResolvedCitationItem):
+    item_type: Literal["block_content"] = "block_content"
+    agent_run_id: str
+    collection_id: str
+    transcript_id: str
+    block_idx: int
+
+
+ResolvedCitationItemUnion = Annotated[
+    Union[
+        AgentRunMetadataItem,
+        TranscriptMetadataItem,
+        TranscriptBlockMetadataItem,
+        TranscriptBlockContentItem,
+    ],
+    Discriminator("item_type"),
+]
+
+RANGE_BEGIN = "<RANGE>"
+RANGE_END = "</RANGE>"
+
+
+def scan_brackets(text: str) -> list[tuple[int, int, str]]:
+    """Scan text for bracketed segments, respecting RANGE markers and nested brackets.
+
+    Returns a list of (start_index, end_index_exclusive, inner_content).
+    """
+    matches: list[tuple[int, int, str]] = []
+    i = 0
+    while i < len(text):
+        if text[i] == "[":
+            start = i
+            bracket_count = 1
+            j = i + 1
+            in_range = False
+
+            while j < len(text) and bracket_count > 0:
+                if text[j : j + len(RANGE_BEGIN)] == RANGE_BEGIN:
+                    in_range = True
+                elif text[j : j + len(RANGE_END)] == RANGE_END:
+                    in_range = False
+                elif text[j] == "[" and not in_range:
+                    bracket_count += 1
+                elif text[j] == "]" and not in_range:
+                    bracket_count -= 1
+                j += 1
+
+            if bracket_count == 0:
+                end_exclusive = j
+                bracket_content = text[start + 1 : end_exclusive - 1]
+                matches.append((start, end_exclusive, bracket_content))
+                i = j
+            else:
+                i += 1
+        else:
+            i += 1
+    return matches
+
+
+def _extract_range_pattern(range_part: str) -> CitationTargetTextRange | None:
+    if RANGE_BEGIN in range_part and RANGE_END in range_part:
+        range_begin_idx = range_part.find(RANGE_BEGIN)
+        range_end_idx = range_part.find(RANGE_END)
+        if range_begin_idx != -1 and range_end_idx != -1:
+            range_content = range_part[range_begin_idx + len(RANGE_BEGIN) : range_end_idx]
+            start_pattern = range_content if range_content else None
+            return CitationTargetTextRange(start_pattern=start_pattern)
+
+    return None
+
+
+def parse_single_citation(part: str) -> tuple[str, CitationTargetTextRange | None] | None:
+    """
+    Parse a single citation token inside a bracket and return its components.
+
+    Returns ParsedCitation or None if invalid.
+    For metadata citations, transcript_idx may be None (for agent run metadata).
+    Supports optional text range for all valid citation kinds.
+    """
+    token = part.strip()
+    if not token:
+        return None
+
+    # Extract optional range part
+    item_alias = token
+    text_range: CitationTargetTextRange | None = None
+    if ":" in token:
+        left, right = token.split(":", 1)
+        item_alias = left.strip()
+        text_range = _extract_range_pattern(right)
+
+    return item_alias, text_range
+
+
+def parse_citations(text: str) -> tuple[str, list[ParsedCitation]]:
+    """
+    Parse citations from text in the format described by TEXT_RANGE_CITE_INSTRUCTION.
+
+    Supported formats:
+    - Single block: [T<key>B<idx>]
+    - Text range with start pattern: [T<key>B<idx>:<RANGE>start_pattern</RANGE>]
+    - Agent run metadata: [M.key]
+    - Transcript metadata: [T<key>M.key]
+    - Message metadata: [T<key>B<idx>M.key]
+    - Message metadata with text range: [T<key>B<idx>M.key:<RANGE>start_pattern</RANGE>]
+
+    Args:
+        text: The text to parse citations from
+
+    Returns:
+        A tuple of (cleaned_text, citations) where cleaned_text has brackets and range markers removed
+        and citations have start_idx and end_idx representing character positions
+        in the cleaned text
+    """
+    citations: list[ParsedCitation] = []
+
+    bracket_matches = scan_brackets(text)
+
+    for start, end, bracket_content in bracket_matches:
+        # Parse a single citation token inside the bracket
+        parsed = parse_single_citation(bracket_content)
+        if not parsed:
+            continue
+        label, text_range = parsed
+
+        citations.append(
+            ParsedCitation(start_idx=start, end_idx=end, item_alias=label, text_range=text_range)
+        )
+
+    # We're not cleaning the text right now but may do that later
+    return text, citations

docent/data_models/formatted_objects.py

@@ -0,0 +1,84 @@
+from uuid import uuid4
+
+from pydantic import Field, model_validator
+
+from docent.data_models.agent_run import AgentRun
+from docent.data_models.transcript import Transcript
+
+
+class FormattedTranscript(Transcript):
+    """A Transcript that preserves original message indices during edits.
+
+    This class extends Transcript to support customization while maintaining accurate
+    citations. Each message retains its original index from the source transcript,
+    even if messages are added, removed, or reordered.
+
+    Use this class when you need to customize which parts of a transcript are visible
+    to an LLM while ensuring citations remain valid.
+    """
+
+    id_to_original_index: dict[str, int]
+
+    @classmethod
+    def from_transcript(cls, transcript: Transcript) -> "FormattedTranscript":
+        """Create a FormattedTranscript from a regular Transcript."""
+        # Ensure all messages have IDs and build id_to_original_index
+        id_to_original_index: dict[str, int] = {}
+        for idx, msg in enumerate(transcript.messages):
+            if msg.id is None:
+                msg.id = str(uuid4())
+            id_to_original_index[msg.id] = idx
+
+        return cls(
+            id=transcript.id,
+            name=transcript.name,
+            description=transcript.description,
+            transcript_group_id=transcript.transcript_group_id,
+            created_at=transcript.created_at,
+            messages=transcript.messages,
+            metadata=transcript.metadata,
+            id_to_original_index=id_to_original_index,
+        )
+
+    @model_validator(mode="after")
+    def _validate_id_to_original_index(self) -> "FormattedTranscript":
+        """Ensure id_to_original_index covers all messages."""
+        for msg in self.messages:
+            if msg.id not in self.id_to_original_index:
+                raise ValueError(
+                    f"Message {msg.id} missing from id_to_original_index. "
+                    "Use FormattedTranscript.from_transcript() to create a new instance."
+                )
+        return self
+
+    def _enumerate_messages(self):
+        """Yield (original index, message) for each message."""
+        for message in self.messages:
+            assert message.id is not None
+            original_idx = self.id_to_original_index[message.id]
+            yield (original_idx, message)
+
+
+class FormattedAgentRun(AgentRun):
+    """An AgentRun that allows customization while tracking original identifiers.
+
+    This class extends AgentRun to support modifications to what an LLM sees
+    while maintaining accurate citations back to the original agent run.
+
+    Use this class when you need to customize which parts of an agent run are visible
+    to an LLM (e.g., hiding metadata, truncating long outputs).
+    """
+
+    transcripts: list[FormattedTranscript] = Field(default_factory=list)  # type: ignore[assignment]
+
+    @classmethod
+    def from_agent_run(cls, agent_run: AgentRun) -> "FormattedAgentRun":
+        """Create a FormattedAgentRun from a regular AgentRun."""
+        return cls(
+            id=agent_run.id,
+            name=agent_run.name,
+            description=agent_run.description,
+            transcripts=[FormattedTranscript.from_transcript(t) for t in agent_run.transcripts],
+            transcript_groups=agent_run.transcript_groups,
+            metadata=agent_run.metadata,
+        )

docent/data_models/judge.py

@@ -0,0 +1,17 @@
+"""Judge-related data models shared across Docent components."""
+
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+
+class Label(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+
+    label_set_id: str
+    label_value: dict[str, Any]
+    agent_run_id: str
+
+
+__all__ = ["Label"]

docent/data_models/metadata_util.py

@@ -0,0 +1,16 @@
+import json
+from typing import Any
+
+from pydantic_core import to_jsonable_python
+
+
+def dump_metadata(metadata: dict[str, Any]) -> str | None:
+    """
+    Dump metadata to a JSON string.
+    We used to use YAML to save tokens, but JSON makes it easier to find cited ranges on the frontend because the frontend uses JSON.
+    """
+    if not metadata:
+        return None
+    metadata_obj = to_jsonable_python(metadata)
+    text = json.dumps(metadata_obj, indent=2)
+    return text.strip()

docent/data_models/regex.py

@@ -0,0 +1,56 @@
+import re
+
+from pydantic import BaseModel
+
+from docent._log_util import get_logger
+
+logger = get_logger(__name__)
+
+
+class RegexSnippet(BaseModel):
+    snippet: str
+    match_start: int
+    match_end: int
+
+
+def get_regex_snippets(text: str, pattern: str, window_size: int = 50) -> list[RegexSnippet]:
+    """Extracts snippets from text that match a regex pattern, with surrounding context.
+
+    Args:
+        text: The text to search in.
+        pattern: The regex pattern to match.
+        window_size: The number of characters to include before and after the match.
+
+    Returns:
+        A list of RegexSnippet objects containing the snippets and match positions.
+    """
+    # Find all matches
+    try:
+        matches = list(re.compile(pattern, re.IGNORECASE | re.DOTALL).finditer(text))
+        if not matches:
+            logger.warning(f"No regex matches found for {pattern}: this shouldn't happen!")
+
+        if not matches:
+            return []
+
+        snippets: list[RegexSnippet] = []
+        for match in matches:
+            start, end = match.span()
+
+            # Calculate window around the match
+            snippet_start = max(0, start - window_size)
+            snippet_end = min(len(text), end + window_size)
+
+            # Create the snippet with the match indices adjusted for the window
+            snippets.append(
+                RegexSnippet(
+                    snippet=text[snippet_start:snippet_end],
+                    match_start=start - snippet_start,
+                    match_end=end - snippet_start,
+                )
+            )
+
+        return snippets
+    except re.error as e:
+        logger.error(f"Got regex error: {e}")
+        return []