docent-python 0.1.0a1__py3-none-any.whl

docent/data_models/chat/message.py

@@ -0,0 +1,125 @@
+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
+
+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).
+    """
+
+    id: str | None = None
+    content: str | list[Content]
+    role: Literal["system", "user", "assistant", "tool"]
+
+    @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 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."""
+
+
+def parse_chat_message(message_data: dict[str, Any] | ChatMessage) -> ChatMessage:
+    """Parse a message dictionary or object into the appropriate ChatMessage subclass.
+
+    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}")

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
+    type: Literal["function"] | None
+    function: str
+    arguments: dict[str, Any]
+    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,223 @@
+import re
+from typing import TypedDict
+
+
+class Citation(TypedDict):
+    start_idx: int
+    end_idx: int
+    agent_run_idx: int | None
+    transcript_idx: int | None
+    block_idx: int
+    action_unit_idx: int | None
+
+
+def parse_citations_single_run(text: str) -> list[Citation]:
+    """
+    Parse citations from text in the format described by SINGLE_BLOCK_CITE_INSTRUCTION.
+
+    Supported formats:
+    - Single block: [T<key>B<idx>]
+    - Multiple blocks: [T<key1>B<idx1>, T<key2>B<idx2>, ...]
+    - Dash-separated blocks: [T<key1>B<idx1>-T<key2>B<idx2>]
+
+    Args:
+        text: The text to parse citations from
+
+    Returns:
+        A list of Citation objects with start_idx and end_idx representing
+        the character positions in the text (excluding brackets)
+    """
+    citations: list[Citation] = []
+
+    # Find all bracketed content first
+    bracket_pattern = r"\[(.*?)\]"
+    bracket_matches = re.finditer(bracket_pattern, text)
+
+    for bracket_match in bracket_matches:
+        bracket_content = bracket_match.group(1)
+        # Starting position of the bracket content (excluding '[')
+        content_start_pos = bracket_match.start() + 1
+
+        # Split by commas if present
+        parts = [part.strip() for part in bracket_content.split(",")]
+
+        for part in parts:
+            # Check if this part contains a dash (range citation)
+            if "-" in part:
+                # Split by dash and process each sub-part
+                dash_parts = [dash_part.strip() for dash_part in part.split("-")]
+                for dash_part in dash_parts:
+                    # Check for single block citation: T<key>B<idx>
+                    single_match = re.match(r"T(\d+)B(\d+)", dash_part)
+                    if single_match:
+                        transcript_idx = int(single_match.group(1))
+                        block_idx = int(single_match.group(2))
+
+                        # Find position within the original text
+                        citation_text = f"T{transcript_idx}B{block_idx}"
+                        part_pos_in_content = bracket_content.find(dash_part)
+                        ref_pos = content_start_pos + part_pos_in_content
+                        ref_end = ref_pos + len(citation_text)
+
+                        # Check if this citation overlaps with any existing citation
+                        if not any(
+                            citation["start_idx"] <= ref_pos < citation["end_idx"]
+                            or citation["start_idx"] < ref_end <= citation["end_idx"]
+                            for citation in citations
+                        ):
+                            citations.append(
+                                Citation(
+                                    start_idx=ref_pos,
+                                    end_idx=ref_end,
+                                    agent_run_idx=None,
+                                    transcript_idx=transcript_idx,
+                                    block_idx=block_idx,
+                                    action_unit_idx=None,
+                                )
+                            )
+            else:
+                # Check for single block citation: T<key>B<idx>
+                single_match = re.match(r"T(\d+)B(\d+)", part)
+                if single_match:
+                    transcript_idx = int(single_match.group(1))
+                    block_idx = int(single_match.group(2))
+
+                    # Find position within the original text
+                    citation_text = f"T{transcript_idx}B{block_idx}"
+                    part_pos_in_content = bracket_content.find(part)
+                    ref_pos = content_start_pos + part_pos_in_content
+                    ref_end = ref_pos + len(citation_text)
+
+                    # Check if this citation overlaps with any existing citation
+                    if not any(
+                        citation["start_idx"] <= ref_pos < citation["end_idx"]
+                        or citation["start_idx"] < ref_end <= citation["end_idx"]
+                        for citation in citations
+                    ):
+                        citations.append(
+                            Citation(
+                                start_idx=ref_pos,
+                                end_idx=ref_end,
+                                agent_run_idx=None,
+                                transcript_idx=transcript_idx,
+                                block_idx=block_idx,
+                                action_unit_idx=None,
+                            )
+                        )
+
+    return citations
+
+
+def parse_citations_multi_run(text: str) -> list[Citation]:
+    """
+    Parse citations from text in the format described by MULTI_BLOCK_CITE_INSTRUCTION.
+
+    Supported formats:
+    - Single block in transcript: [R<idx>T<key>B<idx>] or ([R<idx>T<key>B<idx>])
+    - Multiple blocks: [R<idx1>T<key1>B<idx1>][R<idx2>T<key2>B<idx2>]
+    - Comma-separated blocks: [R<idx1>T<key1>B<idx1>, R<idx2>T<key2>B<idx2>, ...]
+    - Dash-separated blocks: [R<idx1>T<key1>B<idx1>-R<idx2>T<key2>B<idx2>]
+
+    Args:
+        text: The text to parse citations from
+
+    Returns:
+        A list of Citation objects with start_idx and end_idx representing
+        the character positions in the text (excluding brackets)
+    """
+    citations: list[Citation] = []
+
+    # Find all content within brackets - this handles nested brackets too
+    bracket_pattern = r"\[([^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*)\]"
+    # Also handle optional parentheses around the brackets
+    paren_bracket_pattern = r"\(\[([^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*)\]\)"
+
+    # Single citation pattern
+    single_pattern = r"R(\d+)T(\d+)B(\d+)"
+
+    # Find all bracket matches
+    for pattern in [bracket_pattern, paren_bracket_pattern]:
+        matches = re.finditer(pattern, text)
+        for match in matches:
+            # Get the content inside brackets
+            if pattern == bracket_pattern:
+                content = match.group(1)
+                start_pos = match.start() + 1  # +1 to skip the opening bracket
+            else:
+                content = match.group(1)
+                start_pos = match.start() + 2  # +2 to skip the opening parenthesis and bracket
+
+            # Split by comma if present
+            items = [item.strip() for item in content.split(",")]
+
+            for item in items:
+                # Check if this item contains a dash (range citation)
+                if "-" in item:
+                    # Split by dash and process each sub-item
+                    dash_items = [dash_item.strip() for dash_item in item.split("-")]
+                    for dash_item in dash_items:
+                        # Check for single citation
+                        single_match = re.match(single_pattern, dash_item)
+                        if single_match:
+                            agent_run_idx = int(single_match.group(1))
+                            transcript_idx = int(single_match.group(2))
+                            block_idx = int(single_match.group(3))
+
+                            # Calculate position in the original text
+                            citation_text = f"R{agent_run_idx}T{transcript_idx}B{block_idx}"
+                            citation_start = text.find(citation_text, start_pos)
+                            citation_end = citation_start + len(citation_text)
+
+                            # Move start_pos for the next item if there are more items
+                            start_pos = citation_end
+
+                            # Avoid duplicate citations
+                            if not any(
+                                citation["start_idx"] == citation_start
+                                and citation["end_idx"] == citation_end
+                                for citation in citations
+                            ):
+                                citations.append(
+                                    Citation(
+                                        start_idx=citation_start,
+                                        end_idx=citation_end,
+                                        agent_run_idx=agent_run_idx,
+                                        transcript_idx=transcript_idx,
+                                        block_idx=block_idx,
+                                        action_unit_idx=None,
+                                    )
+                                )
+                else:
+                    # Check for single citation
+                    single_match = re.match(single_pattern, item)
+                    if single_match:
+                        agent_run_idx = int(single_match.group(1))
+                        transcript_idx = int(single_match.group(2))
+                        block_idx = int(single_match.group(3))
+
+                        # Calculate position in the original text
+                        citation_text = f"R{agent_run_idx}T{transcript_idx}B{block_idx}"
+                        citation_start = text.find(citation_text, start_pos)
+                        citation_end = citation_start + len(citation_text)
+
+                        # Move start_pos for the next item if there are more items
+                        start_pos = citation_end
+
+                        # Avoid duplicate citations
+                        if not any(
+                            citation["start_idx"] == citation_start
+                            and citation["end_idx"] == citation_end
+                            for citation in citations
+                        ):
+                            citations.append(
+                                Citation(
+                                    start_idx=citation_start,
+                                    end_idx=citation_end,
+                                    agent_run_idx=agent_run_idx,
+                                    transcript_idx=transcript_idx,
+                                    block_idx=block_idx,
+                                    action_unit_idx=None,
+                                )
+                            )
+
+    return citations

docent/data_models/filters.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Type, Union
+from uuid import uuid4
+
+from pydantic import BaseModel, Discriminator, Field, field_validator
+from sqlalchemy import ColumnElement, and_, or_
+
+from docent._log_util import get_logger
+
+if TYPE_CHECKING:
+    from docent_core._db_service.schemas.tables import SQLAAgentRun
+
+logger = get_logger(__name__)
+
+
+class BaseFrameFilter(BaseModel):
+    """Base class for all frame filters."""
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str | None = None
+    supports_sql: bool = True  # All filters must support SQL
+
+    def to_sqla_where_clause(self, table: Type["SQLAAgentRun"]) -> ColumnElement[bool] | None:
+        """Convert this filter to a SQLAlchemy WHERE clause.
+
+        All filters must implement this method to support SQL execution.
+        """
+        raise NotImplementedError(
+            f"Filter {self.__class__.__name__} must implement to_sqla_where_clause"
+        )
+
+
+class PrimitiveFilter(BaseFrameFilter):
+    """Filter that applies a primitive operation to a metadata field."""
+
+    type: Literal["primitive"] = "primitive"
+    key_path: list[str]
+    value: Any
+    op: Literal[">", ">=", "<", "<=", "==", "!=", "~*", "!~*"]
+
+    def to_sqla_where_clause(self, table: Type["SQLAAgentRun"]) -> ColumnElement[bool] | None:
+        """Convert this filter to a SQLAlchemy WHERE clause."""
+
+        mode = self.key_path[0]
+
+        # Extract value from JSONB using the table parameter
+        if mode == "text":
+            sqla_value = table.text_for_search  # type: ignore
+        elif mode == "metadata":
+            sqla_value = table.metadata_json  # type: ignore
+        else:
+            raise ValueError(f"Unsupported mode: {mode}")
+
+        for key in self.key_path[1:]:
+            sqla_value = sqla_value[key]
+
+        # Cast the extracted value to the correct type
+        # This is only necessary for metadata which is JSONB
+        if mode == "metadata":
+            if isinstance(self.value, str):
+                sqla_value = sqla_value.as_string()
+            elif isinstance(self.value, bool):
+                sqla_value = sqla_value.as_boolean()
+            elif isinstance(self.value, float) or isinstance(self.value, int):  # type: ignore warning about unnecessary comparison
+                # if self.value is an int, we may still need to do sql comparisons with floats
+                sqla_value = sqla_value.as_float()
+            else:
+                raise ValueError(f"Unsupported value type: {type(self.value)}")
+
+        # Handle different operations using SQLAlchemy expressions
+        if self.op == "==":
+            return sqla_value == self.value
+        elif self.op == "!=":
+            return sqla_value != self.value
+        elif self.op == ">":
+            return sqla_value > self.value
+        elif self.op == ">=":
+            return sqla_value >= self.value
+        elif self.op == "<":
+            return sqla_value < self.value
+        elif self.op == "<=":
+            return sqla_value <= self.value
+        elif self.op == "~*":
+            return sqla_value.op("~*")(self.value)
+        else:
+            raise ValueError(f"Unsupported operation: {self.op}")
+
+
+class ComplexFilter(BaseFrameFilter):
+    """Filter that combines multiple filters with AND/OR/NOT logic."""
+
+    type: Literal["complex"] = "complex"
+    filters: list[FrameFilter]
+    op: Literal["and", "or"] = "and"
+
+    @field_validator("filters")
+    @classmethod
+    def validate_filters(cls, v: list[FrameFilter]) -> list[FrameFilter]:
+        if not v:
+            raise ValueError("ComplexFilter must have at least one filter")
+        return v
+
+    def to_sqla_where_clause(self, table: Type["SQLAAgentRun"]) -> ColumnElement[bool] | None:
+        """Convert this filter to a SQLAlchemy WHERE clause."""
+
+        # Get WHERE clauses for all sub-filters
+        where_clauses: list[ColumnElement[bool]] = []
+        for filter_obj in self.filters:
+            where_clause = filter_obj.to_sqla_where_clause(table)
+            if where_clause is not None:
+                where_clauses.append(where_clause)
+
+        if not where_clauses:
+            return None
+
+        # Apply the operation
+        if self.op == "and":
+            result = and_(*where_clauses)
+        elif self.op == "or":
+            result = or_(*where_clauses)
+        else:
+            raise ValueError(f"Unsupported operation: {self.op}")
+
+        return result
+
+
+class AgentRunIdFilter(BaseFrameFilter):
+    """Filter that matches specific agent run IDs."""
+
+    type: Literal["agent_run_id"] = "agent_run_id"
+    agent_run_ids: list[str]
+
+    @field_validator("agent_run_ids")
+    @classmethod
+    def validate_agent_run_ids(cls, v: list[str]) -> list[str]:
+        if not v:
+            raise ValueError("AgentRunIdFilter must have at least one agent run ID")
+        return v
+
+    def to_sqla_where_clause(self, table: Type["SQLAAgentRun"]) -> ColumnElement[bool] | None:
+        """Convert to SQLAlchemy WHERE clause for agent run ID filtering."""
+        return table.id.in_(self.agent_run_ids)
+
+
+class SearchResultPredicateFilter(BaseFrameFilter):
+    """Filter that applies a predicate to search results."""
+
+    type: Literal["search_result_predicate"] = "search_result_predicate"
+    predicate: str
+    search_query: str
+
+    def to_sqla_where_clause(self, table: Type["SQLAAgentRun"]) -> ColumnElement[bool] | None:
+        """Convert to SQLAlchemy WHERE clause for search result filtering."""
+        # This filter requires joining with search results table
+        # For now, we'll return None to indicate it needs special handling
+        # In practice, this would join with the search_results table
+        return None
+
+
+class SearchResultExistsFilter(BaseFrameFilter):
+    """Filter that checks if search results exist."""
+
+    type: Literal["search_result_exists"] = "search_result_exists"
+    search_query: str
+
+    def to_sqla_where_clause(self, table: Type["SQLAAgentRun"]) -> ColumnElement[bool] | None:
+        """Convert to SQLAlchemy WHERE clause for search result existence filtering."""
+        # This filter requires joining with search results table
+        # For now, we'll return None to indicate it needs special handling
+        # In practice, this would join with the search_results table
+        return None
+
+
+FrameFilter = Annotated[
+    Union[
+        PrimitiveFilter,
+        ComplexFilter,
+        AgentRunIdFilter,
+        SearchResultPredicateFilter,
+        SearchResultExistsFilter,
+    ],
+    Discriminator("type"),
+]
+
+
+def parse_filter_dict(filter_dict: dict[str, Any]) -> FrameFilter:
+    """Parse a filter dictionary into a FrameFilter object."""
+    filter_type = filter_dict.get("type")
+
+    if filter_type == "primitive":
+        return PrimitiveFilter(**filter_dict)
+    elif filter_type == "complex":
+        # Recursively parse nested filters
+        nested_filters = [parse_filter_dict(f) for f in filter_dict.get("filters", [])]
+        filter_dict["filters"] = nested_filters
+        return ComplexFilter(**filter_dict)
+    elif filter_type == "agent_run_id":
+        return AgentRunIdFilter(**filter_dict)
+    elif filter_type == "search_result_predicate":
+        return SearchResultPredicateFilter(**filter_dict)
+    elif filter_type == "search_result_exists":
+        return SearchResultExistsFilter(**filter_dict)
+    else:
+        raise ValueError(f"Unknown filter type: {filter_type}")