docent-python 0.1.0a1__py3-none-any.whl

docent/data_models/metadata.py

@@ -0,0 +1,219 @@
+import traceback
+from typing import Any, Optional
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    PrivateAttr,
+    SerializerFunctionWrapHandler,
+    model_serializer,
+    model_validator,
+)
+
+from docent._log_util import get_logger
+
+logger = get_logger(__name__)
+
+SINGLETONS = (int, float, str, bool)
+
+
+class BaseMetadata(BaseModel):
+    """Provides common functionality for accessing and validating metadata fields.
+    All metadata classes should inherit from this class.
+
+    Serialization Behavior:
+        - Field descriptions are highly recommended and stored in serialized versions of the object.
+        - When a subclass of BaseMetadata is uploaded to a server, all extra fields and their descriptions are retained.
+        - To recover the original structure with proper typing upon download, use:
+          `CustomMetadataClass.model_validate(obj.model_dump())`.
+
+    Attributes:
+        model_config: Pydantic configuration that allows extra fields.
+        allow_fields_without_descriptions: Boolean indicating whether to allow fields without descriptions.
+    """
+
+    model_config = ConfigDict(extra="allow")
+    allow_fields_without_descriptions: bool = False
+
+    # Private attribute to store field descriptions
+    _field_descriptions: dict[str, str | None] | None = PrivateAttr(default=None)
+    _internal_basemetadata_fields: set[str] = PrivateAttr(
+        default={
+            "allow_fields_without_descriptions",
+            "model_config",
+            "_field_descriptions",
+        }
+    )
+
+    @model_validator(mode="after")
+    def _validate_field_types_and_descriptions(self):
+        """Validates that all fields have descriptions and proper types.
+
+        Returns:
+            Self: The validated model instance.
+
+        Raises:
+            ValueError: If any field is missing a description or has an invalid type.
+        """
+        # Validate each field in the model
+        for field_name, field_info in self.__class__.model_fields.items():
+            if field_name in self._internal_basemetadata_fields:
+                continue
+
+            # Check that field has a description
+            if field_info.description is None:
+                if not self.allow_fields_without_descriptions:
+                    raise ValueError(
+                        f"Field `{field_name}` needs a description in the definition of `{self.__class__.__name__}`, like `{field_name}: T = Field(description=..., default=...)`. "
+                        "To allow un-described fields, set `allow_fields_without_descriptions = True` on the instance or in your metadata class definition."
+                    )
+
+        # Validate that the metadata is JSON serializable
+        try:
+            self.model_dump_json()
+        except Exception as e:
+            raise ValueError(
+                f"Metadata is not JSON serializable: {e}. Traceback: {traceback.format_exc()}"
+            )
+
+        return self
+
+    def model_post_init(self, __context: Any) -> None:
+        """Initializes field descriptions from extra data after model initialization.
+
+        Args:
+            __context: The context provided by Pydantic's post-initialization hook.
+        """
+        fd = self.model_extra.pop("_field_descriptions", None) if self.model_extra else None
+        if fd is not None:
+            self._field_descriptions = fd
+
+    @model_serializer(mode="wrap")
+    def _serialize_model(self, handler: SerializerFunctionWrapHandler):
+        # Call the default serializer
+        data = handler(self)
+
+        # Dump the field descriptions
+        if self._field_descriptions is None:
+            self._field_descriptions = self._compute_field_descriptions()
+        data["_field_descriptions"] = self._field_descriptions
+
+        return data
+
+    def model_dump(
+        self, *args: Any, strip_internal_fields: bool = False, **kwargs: Any
+    ) -> dict[str, Any]:
+        data = super().model_dump(*args, **kwargs)
+
+        # Remove internal fields if requested
+        if strip_internal_fields:
+            for field in self._internal_basemetadata_fields:
+                if field in data:
+                    data.pop(field)
+
+        return data
+
+    def get(self, key: str, default_value: Any = None) -> Any:
+        """Gets a value from the metadata by key.
+
+        Args:
+            key: The key to look up in the metadata.
+            default_value: Value to return if the key is not found. Defaults to None.
+
+        Returns:
+            Any: The value associated with the key, or the default value if not found.
+        """
+        # Check if the field exists in the model's fields
+        if key in self.__class__.model_fields or (
+            self.model_extra is not None and key in self.model_extra
+        ):
+            # Field exists, return its value (even if None)
+            return getattr(self, key)
+
+        logger.warning(f"Field '{key}' not found in {self.__class__.__name__}")
+        return default_value
+
+    def get_field_description(self, field_name: str) -> str | None:
+        """Gets the description of a field defined in the model schema.
+
+        Args:
+            field_name: The name of the field.
+
+        Returns:
+            str or None: The description string if the field is defined in the model schema
+                and has a description, otherwise None.
+        """
+        if self._field_descriptions is None:
+            self._field_descriptions = self._compute_field_descriptions()
+
+        if field_name in self._field_descriptions:
+            return self._field_descriptions[field_name]
+
+        logger.warning(
+            f"Field description for '{field_name}' not found in {self.__class__.__name__}"
+        )
+        return None
+
+    def get_all_field_descriptions(self) -> dict[str, str | None]:
+        """Gets descriptions for all fields defined in the model schema.
+
+        Returns:
+            dict: A dictionary mapping field names to their descriptions.
+                Only includes fields that have descriptions defined in the schema.
+        """
+        if self._field_descriptions is None:
+            self._field_descriptions = self._compute_field_descriptions()
+        return self._field_descriptions
+
+    def _compute_field_descriptions(self) -> dict[str, str | None]:
+        """Computes descriptions for all fields in the model.
+
+        Returns:
+            dict: A dictionary mapping field names to their descriptions.
+        """
+        field_descriptions: dict[str, Optional[str]] = {}
+        for field_name, field_info in self.__class__.model_fields.items():
+            if field_name not in self._internal_basemetadata_fields:
+                field_descriptions[field_name] = field_info.description
+        return field_descriptions
+
+
+class BaseAgentRunMetadata(BaseMetadata):
+    """Extends BaseMetadata with fields specific to agent evaluation runs.
+
+    Attributes:
+        scores: Dictionary of evaluation metrics.
+        default_score_key: The primary evaluation metric key.
+    """
+
+    scores: dict[str, int | float | bool | None] = Field(
+        description="A dict of score_key -> score_value. Use one key for each metric you're tracking."
+    )
+    default_score_key: str | None = Field(
+        description="The default score key for the transcript; one top-line metric"
+    )
+
+    def get_default_score(self) -> int | float | bool | None:
+        """Gets the default evaluation score.
+
+        Returns:
+            int, float, bool, or None: The value of the default score if a default score key is set,
+                otherwise None.
+        """
+        if self.default_score_key is None:
+            return None
+        return self.scores.get(self.default_score_key)
+
+
+class FrameDimension(BaseModel):
+    """A dimension for organizing agent runs."""
+
+    id: str
+    name: str
+    search_query: str | None = None
+    metadata_key: str | None = None
+    maintain_mece: bool | None = None
+    loading_clusters: bool = False
+    loading_bins: bool = False
+    binIds: list[dict[str, Any]] | None = None

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 []

docent/data_models/shared_types.py

@@ -0,0 +1,10 @@
+from typing import List, TypedDict
+
+from docent.data_models.citation import Citation
+
+
+class EvidenceWithCitation(TypedDict):
+    """A piece of evidence with its citations."""
+
+    evidence: str
+    citations: List[Citation]

docent/data_models/transcript.py

@@ -0,0 +1,347 @@
+import sys
+from typing import Any
+from uuid import uuid4
+
+import yaml
+from pydantic import BaseModel, Field, PrivateAttr, field_serializer, field_validator
+
+from docent.data_models._tiktoken_util import (
+    get_token_count,
+    group_messages_into_ranges,
+    truncate_to_token_limit,
+)
+from docent.data_models.chat import AssistantMessage, ChatMessage, ContentReasoning
+from docent.data_models.metadata import BaseMetadata
+
+# Template for formatting individual transcript blocks
+TRANSCRIPT_BLOCK_TEMPLATE = """
+<{index_label} | role: {role}>
+{content}
+</{index_label}>
+""".strip()
+
+# Instructions for citing single transcript blocks
+SINGLE_RUN_CITE_INSTRUCTION = "Each transcript and each block has a unique index. Cite the relevant indices in brackets when relevant, like [T<idx>B<idx>]. Use multiple tags to cite multiple blocks, like [T<idx1>B<idx1>][T<idx2>B<idx2>]. Use an inner dash to cite a range of blocks, like [T<idx1>B<idx1>-T<idx2>B<idx2>]. Remember to cite specific blocks and NOT action units."
+
+# Instructions for citing multiple transcript blocks
+MULTI_RUN_CITE_INSTRUCTION = "Each run, each transcript, and each block has a unique index. Cite the relevant indices in brackets when relevant, like [R<idx>T<idx>B<idx>]. Use multiple tags to cite multiple blocks, like [R<idx1>T<idx1>B<idx1>][R<idx2>T<idx2>B<idx2>]. Use an inner dash to cite a range of blocks, like [R<idx1>T<idx1>B<idx1>-R<idx2>T<idx2>B<idx2>]. Remember to cite specific blocks and NOT action units."
+
+
+def format_chat_message(
+    message: ChatMessage,
+    block_idx: int,
+    transcript_idx: int = 0,
+    agent_run_idx: int | None = None,
+) -> str:
+    if agent_run_idx is not None:
+        index_label = f"R{agent_run_idx}T{transcript_idx}B{block_idx}"
+    else:
+        index_label = f"T{transcript_idx}B{block_idx}"
+
+    cur_content = ""
+
+    # Add reasoning at beginning if applicable
+    if isinstance(message, AssistantMessage) and message.content:
+        for content in message.content:
+            if isinstance(content, ContentReasoning):
+                cur_content = f"<reasoning>\n{content.reasoning}\n</reasoning>\n"
+
+    # Main content text
+    cur_content += message.text
+
+    # Update content in case there's a view
+    if isinstance(message, AssistantMessage) and message.tool_calls:
+        for tool_call in message.tool_calls:
+            if tool_call.view:
+                cur_content += f"\n<tool call>\n{tool_call.view.content}\n</tool call>"
+            else:
+                args = ", ".join([f"{k}={v}" for k, v in tool_call.arguments.items()])
+                cur_content += f"\n<tool call>\n{tool_call.function}({args})\n</tool call>"
+
+    return TRANSCRIPT_BLOCK_TEMPLATE.format(
+        index_label=index_label, role=message.role, content=cur_content
+    )
+
+
+class Transcript(BaseModel):
+    """Represents a transcript of messages in a conversation with an AI agent.
+
+    A transcript contains a sequence of messages exchanged between different roles
+    (system, user, assistant, tool) and provides methods to organize these messages
+    into logical units of action.
+
+    Attributes:
+        id: Unique identifier for the transcript, auto-generated by default.
+        name: Optional human-readable name for the transcript.
+        description: Optional description of the transcript.
+        messages: List of chat messages in the transcript.
+        metadata: Additional structured metadata about the transcript.
+    """
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str | None = None
+    description: str | None = None
+
+    messages: list[ChatMessage]
+    metadata: BaseMetadata = Field(default_factory=BaseMetadata)
+
+    _units_of_action: list[list[int]] | None = PrivateAttr(default=None)
+
+    @field_serializer("metadata")
+    def serialize_metadata(self, metadata: BaseMetadata, _info: Any) -> dict[str, Any]:
+        """
+        Custom serializer for the metadata field so the internal fields are explicitly preserved.
+        """
+        return metadata.model_dump(strip_internal_fields=False)
+
+    @field_validator("metadata", mode="before")
+    @classmethod
+    def _validate_metadata_type(cls, v: Any) -> Any:
+        if v is not None and not isinstance(v, BaseMetadata):
+            raise ValueError(
+                f"metadata must be an instance of BaseMetadata, got {type(v).__name__}"
+            )
+        return v
+
+    @property
+    def units_of_action(self) -> list[list[int]]:
+        """Get the units of action in the transcript.
+
+        A unit of action represents a logical group of messages, such as a system message
+        on its own or a user message followed by assistant responses and tool outputs.
+
+        Returns:
+            list[list[int]]: List of units of action, where each unit is a list of message indices.
+        """
+        if self._units_of_action is None:
+            self._units_of_action = self._compute_units_of_action()
+        return self._units_of_action
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self._units_of_action = self._compute_units_of_action()
+
+    def _compute_units_of_action(self) -> list[list[int]]:
+        """Compute the units of action in the transcript.
+
+        A unit of action is defined as:
+        - A system prompt by itself
+        - A group consisting of a user message, assistant response, and any associated tool outputs
+
+        Returns:
+            list[list[int]]: A list of units of action, where each unit is a list of message indices.
+        """
+        if not self.messages:
+            return []
+
+        units: list[list[int]] = []
+        current_unit: list[int] = []
+
+        def _start_new_unit():
+            nonlocal current_unit
+            if current_unit:
+                units.append(current_unit.copy())
+            current_unit = []
+
+        for i, message in enumerate(self.messages):
+            role = message.role
+            prev_message = self.messages[i - 1] if i > 0 else None
+
+            # System messages are their own unit
+            if role == "system":
+                assert not current_unit, "System message should be the first message"
+                units.append([i])
+
+            # User message always starts a new unit UNLESS the previous message was a user message
+            elif role == "user":
+                if current_unit and prev_message and prev_message.role != "user":
+                    _start_new_unit()
+                current_unit.append(i)
+
+            # Start a new unit if the previous message was not a user or assistant message
+            elif role == "assistant":
+                if (
+                    current_unit
+                    and prev_message
+                    and prev_message.role != "user"
+                    and prev_message.role != "assistant"
+                ):
+                    _start_new_unit()
+                current_unit.append(i)
+
+            # Tool messages are part of the current unit
+            elif role == "tool":
+                current_unit.append(i)
+
+            else:
+                raise ValueError(f"Unknown message role: {role}")
+
+        # Add the last unit if it exists
+        _start_new_unit()
+
+        return units
+
+    def get_first_block_in_action_unit(self, action_unit_idx: int) -> int | None:
+        """Get the index of the first message in a given action unit.
+
+        Args:
+            action_unit_idx: The index of the action unit.
+
+        Returns:
+            int | None: The index of the first message in the action unit,
+                        or None if the action unit doesn't exist.
+
+        Raises:
+            IndexError: If the action unit index is out of range.
+        """
+        if not self._units_of_action:
+            self._units_of_action = self._compute_units_of_action()
+
+        if 0 <= action_unit_idx < len(self._units_of_action):
+            unit = self._units_of_action[action_unit_idx]
+            return unit[0] if unit else None
+        return None
+
+    def get_action_unit_for_block(self, block_idx: int) -> int | None:
+        """Find the action unit that contains the specified message block.
+
+        Args:
+            block_idx: The index of the message block to find.
+
+        Returns:
+            int | None: The index of the action unit containing the block,
+                        or None if no action unit contains the block.
+        """
+        if not self._units_of_action:
+            self._units_of_action = self._compute_units_of_action()
+
+        for unit_idx, unit in enumerate(self._units_of_action):
+            if block_idx in unit:
+                return unit_idx
+        return None
+
+    def set_messages(self, messages: list[ChatMessage]):
+        """Set the messages in the transcript and recompute units of action.
+
+        Args:
+            messages: The new list of chat messages to set.
+        """
+        self.messages = messages
+        self._units_of_action = self._compute_units_of_action()
+
+    def to_str(
+        self,
+        transcript_idx: int = 0,
+        agent_run_idx: int | None = None,
+        highlight_action_unit: int | None = None,
+    ) -> str:
+        return self.to_str_with_token_limit(
+            token_limit=sys.maxsize,
+            agent_run_idx=agent_run_idx,
+            transcript_idx=transcript_idx,
+            highlight_action_unit=highlight_action_unit,
+        )[0]
+
+    def to_str_with_token_limit(
+        self,
+        token_limit: int,
+        transcript_idx: int = 0,
+        agent_run_idx: int | None = None,
+        highlight_action_unit: int | None = None,
+    ) -> list[str]:
+        """Represents the transcript as a list of strings, each of which is at most token_limit tokens
+        under the GPT-4 tokenization scheme.
+
+        We'll try to split up long transcripts along message boundaries and include metadata.
+        For very long messages, we'll have to truncate them and remove metadata.
+
+        Returns:
+            list[str]: A list of strings, each of which is at most token_limit tokens
+            under the GPT-4 tokenization scheme.
+        """
+        if highlight_action_unit is not None and not (
+            0 <= highlight_action_unit < len(self._units_of_action or [])
+        ):
+            raise ValueError(f"Invalid action unit index: {highlight_action_unit}")
+
+        # Format blocks by units of action
+        au_blocks: list[str] = []
+        for unit_idx, unit in enumerate(self._units_of_action or []):
+            unit_blocks: list[str] = []
+            for msg_idx in unit:
+                unit_blocks.append(
+                    format_chat_message(
+                        self.messages[msg_idx],
+                        msg_idx,
+                        transcript_idx,
+                        agent_run_idx,
+                    )
+                )
+
+            unit_content = "\n".join(unit_blocks)
+
+            # Add highlighting if requested
+            if highlight_action_unit and unit_idx == highlight_action_unit:
+                blocks_str_template = "<HIGHLIGHTED>\n{}\n</HIGHLIGHTED>"
+            else:
+                blocks_str_template = "{}"
+            au_blocks.append(
+                blocks_str_template.format(
+                    f"<action unit {unit_idx}>\n{unit_content}\n</action unit {unit_idx}>"
+                )
+            )
+        blocks_str = "\n".join(au_blocks)
+
+        # Gather metadata
+        metadata_obj = self.metadata.model_dump(strip_internal_fields=True)
+        # Add the field descriptions if they exist
+        metadata_obj = {
+            (f"{k} ({d})" if (d := self.metadata.get_field_description(k)) is not None else k): v
+            for k, v in metadata_obj.items()
+        }
+
+        yaml_width = float("inf")
+        block_str = f"<blocks>\n{blocks_str}\n</blocks>\n"
+        metadata_str = f"<metadata>\n{yaml.dump(metadata_obj, width=yaml_width)}\n</metadata>"
+
+        if token_limit == sys.maxsize:
+            return [f"{block_str}" f"{metadata_str}"]
+
+        metadata_token_count = get_token_count(metadata_str)
+        block_token_count = get_token_count(block_str)
+
+        if metadata_token_count + block_token_count <= token_limit:
+            return [f"{block_str}" f"{metadata_str}"]
+        else:
+            results: list[str] = []
+            block_token_counts = [get_token_count(block) for block in au_blocks]
+            ranges = group_messages_into_ranges(
+                block_token_counts, metadata_token_count, token_limit
+            )
+            for msg_range in ranges:
+                if msg_range.include_metadata:
+                    cur_au_blocks = "\n".join(au_blocks[msg_range.start : msg_range.end])
+                    results.append(f"<blocks>\n{cur_au_blocks}\n</blocks>\n" f"{metadata_str}")
+                else:
+                    assert (
+                        msg_range.end == msg_range.start + 1
+                    ), "Ranges without metadata should be a single message"
+                    result = str(au_blocks[msg_range.start])
+                    if msg_range.num_tokens > token_limit - 10:
+                        result = truncate_to_token_limit(result, token_limit - 10)
+                    results.append(f"<blocks>\n{result}\n</blocks>\n")
+
+            return results
+
+
+class TranscriptWithoutMetadataValidator(Transcript):
+    """
+    A version of Transcript that doesn't have the model_validator on metadata.
+    Needed for sending/receiving transcripts via JSON, since they incorrectly trip the existing model_validator.
+    """
+
+    @field_validator("metadata", mode="before")
+    @classmethod
+    def _validate_metadata_type(cls, v: Any) -> Any:
+        # Bypass the model_validator
+        return v