376 lines
13 KiB
Python
376 lines
13 KiB
Python
"""
|
|
Pydantic models for ra_aid database entities.
|
|
|
|
This module defines Pydantic models that correspond to Peewee ORM models,
|
|
providing validation, serialization, and deserialization capabilities.
|
|
"""
|
|
|
|
import datetime
|
|
import json
|
|
from typing import Dict, List, Any, Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict, field_serializer, field_validator
|
|
|
|
|
|
class SessionModel(BaseModel):
|
|
"""
|
|
Pydantic model representing a Session.
|
|
|
|
This model corresponds to the Session Peewee ORM model and provides
|
|
validation and serialization capabilities. It handles the conversion
|
|
between JSON-encoded strings and Python dictionaries for the machine_info field.
|
|
|
|
Attributes:
|
|
id: Unique identifier for the session
|
|
created_at: When the session record was created
|
|
updated_at: When the session record was last updated
|
|
start_time: When the program session started
|
|
command_line: Command line arguments used to start the program
|
|
program_version: Version of the program
|
|
machine_info: Dictionary containing machine-specific metadata
|
|
"""
|
|
id: Optional[int] = None
|
|
created_at: datetime.datetime
|
|
updated_at: datetime.datetime
|
|
start_time: datetime.datetime
|
|
command_line: Optional[str] = None
|
|
program_version: Optional[str] = None
|
|
machine_info: Optional[Dict[str, Any]] = None
|
|
|
|
# Configure the model to work with ORM objects
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
@field_validator("machine_info", mode="before")
|
|
@classmethod
|
|
def parse_machine_info(cls, value: Any) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Parse the machine_info field from a JSON string to a dictionary.
|
|
|
|
Args:
|
|
value: The value to parse, can be a string, dict, or None
|
|
|
|
Returns:
|
|
Optional[Dict[str, Any]]: The parsed dictionary or None
|
|
|
|
Raises:
|
|
ValueError: If the JSON string is invalid
|
|
"""
|
|
if value is None:
|
|
return None
|
|
|
|
if isinstance(value, dict):
|
|
return value
|
|
|
|
if isinstance(value, str):
|
|
try:
|
|
return json.loads(value)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid JSON in machine_info: {e}")
|
|
|
|
raise ValueError(f"Unexpected type for machine_info: {type(value)}")
|
|
|
|
@field_serializer("machine_info")
|
|
def serialize_machine_info(self, machine_info: Optional[Dict[str, Any]]) -> Optional[str]:
|
|
"""
|
|
Serialize the machine_info dictionary to a JSON string for storage.
|
|
|
|
Args:
|
|
machine_info: Dictionary to serialize
|
|
|
|
Returns:
|
|
Optional[str]: JSON-encoded string or None
|
|
"""
|
|
if machine_info is None:
|
|
return None
|
|
|
|
return json.dumps(machine_info)
|
|
|
|
|
|
class HumanInputModel(BaseModel):
|
|
"""
|
|
Pydantic model representing a HumanInput.
|
|
|
|
This model corresponds to the HumanInput Peewee ORM model and provides
|
|
validation and serialization capabilities.
|
|
|
|
Attributes:
|
|
id: Unique identifier for the human input
|
|
created_at: When the record was created
|
|
updated_at: When the record was last updated
|
|
content: The text content of the input
|
|
source: The source of the input ('cli', 'chat', or 'hil')
|
|
session_id: Optional reference to the associated session
|
|
"""
|
|
id: Optional[int] = None
|
|
created_at: datetime.datetime
|
|
updated_at: datetime.datetime
|
|
content: str
|
|
source: str
|
|
session_id: Optional[int] = None
|
|
|
|
# Configure the model to work with ORM objects
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class KeyFactModel(BaseModel):
|
|
"""
|
|
Pydantic model representing a KeyFact.
|
|
|
|
This model corresponds to the KeyFact Peewee ORM model and provides
|
|
validation and serialization capabilities.
|
|
|
|
Attributes:
|
|
id: Unique identifier for the key fact
|
|
created_at: When the record was created
|
|
updated_at: When the record was last updated
|
|
content: The text content of the key fact
|
|
human_input_id: Optional reference to the associated human input
|
|
session_id: Optional reference to the associated session
|
|
"""
|
|
id: Optional[int] = None
|
|
created_at: datetime.datetime
|
|
updated_at: datetime.datetime
|
|
content: str
|
|
human_input_id: Optional[int] = None
|
|
session_id: Optional[int] = None
|
|
|
|
# Configure the model to work with ORM objects
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class KeySnippetModel(BaseModel):
|
|
"""
|
|
Pydantic model representing a KeySnippet.
|
|
|
|
This model corresponds to the KeySnippet Peewee ORM model and provides
|
|
validation and serialization capabilities.
|
|
|
|
Attributes:
|
|
id: Unique identifier for the key snippet
|
|
created_at: When the record was created
|
|
updated_at: When the record was last updated
|
|
filepath: Path to the source file
|
|
line_number: Line number where the snippet starts
|
|
snippet: The source code snippet text
|
|
description: Optional description of the significance
|
|
human_input_id: Optional reference to the associated human input
|
|
session_id: Optional reference to the associated session
|
|
"""
|
|
id: Optional[int] = None
|
|
created_at: datetime.datetime
|
|
updated_at: datetime.datetime
|
|
filepath: str
|
|
line_number: int
|
|
snippet: str
|
|
description: Optional[str] = None
|
|
human_input_id: Optional[int] = None
|
|
session_id: Optional[int] = None
|
|
|
|
# Configure the model to work with ORM objects
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class ResearchNoteModel(BaseModel):
|
|
"""
|
|
Pydantic model representing a ResearchNote.
|
|
|
|
This model corresponds to the ResearchNote Peewee ORM model and provides
|
|
validation and serialization capabilities.
|
|
|
|
Attributes:
|
|
id: Unique identifier for the research note
|
|
created_at: When the record was created
|
|
updated_at: When the record was last updated
|
|
content: The text content of the research note
|
|
human_input_id: Optional reference to the associated human input
|
|
session_id: Optional reference to the associated session
|
|
"""
|
|
id: Optional[int] = None
|
|
created_at: datetime.datetime
|
|
updated_at: datetime.datetime
|
|
content: str
|
|
human_input_id: Optional[int] = None
|
|
session_id: Optional[int] = None
|
|
|
|
# Configure the model to work with ORM objects
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class TrajectoryModel(BaseModel):
|
|
"""
|
|
Pydantic model representing a Trajectory.
|
|
|
|
This model corresponds to the Trajectory Peewee ORM model and provides
|
|
validation and serialization capabilities. It handles the conversion
|
|
between JSON-encoded strings and Python dictionaries for the tool_parameters,
|
|
tool_result, and step_data fields.
|
|
|
|
Attributes:
|
|
id: Unique identifier for the trajectory
|
|
created_at: When the record was created
|
|
updated_at: When the record was last updated
|
|
human_input_id: Optional reference to the associated human input
|
|
tool_name: Name of the tool that was executed
|
|
tool_parameters: Dictionary containing the parameters passed to the tool
|
|
tool_result: Dictionary containing the result returned by the tool
|
|
step_data: Dictionary containing UI rendering data
|
|
record_type: Type of trajectory record
|
|
cost: Optional cost of the tool execution
|
|
tokens: Optional token usage of the tool execution
|
|
is_error: Flag indicating if this record represents an error
|
|
error_message: The error message if is_error is True
|
|
error_type: The type/class of the error if is_error is True
|
|
error_details: Additional error details if is_error is True
|
|
session_id: Optional reference to the associated session
|
|
"""
|
|
id: Optional[int] = None
|
|
created_at: datetime.datetime
|
|
updated_at: datetime.datetime
|
|
human_input_id: Optional[int] = None
|
|
tool_name: Optional[str] = None
|
|
tool_parameters: Optional[Dict[str, Any]] = None
|
|
tool_result: Optional[Any] = None
|
|
step_data: Optional[Dict[str, Any]] = None
|
|
record_type: Optional[str] = None
|
|
cost: Optional[float] = None
|
|
tokens: Optional[int] = None
|
|
is_error: bool = False
|
|
error_message: Optional[str] = None
|
|
error_type: Optional[str] = None
|
|
error_details: Optional[str] = None
|
|
session_id: Optional[int] = None
|
|
|
|
# Configure the model to work with ORM objects
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
@field_validator("tool_parameters", mode="before")
|
|
@classmethod
|
|
def parse_tool_parameters(cls, value: Any) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Parse the tool_parameters field from a JSON string to a dictionary.
|
|
|
|
Args:
|
|
value: The value to parse, can be a string, dict, or None
|
|
|
|
Returns:
|
|
Optional[Dict[str, Any]]: The parsed dictionary or None
|
|
|
|
Raises:
|
|
ValueError: If the JSON string is invalid
|
|
"""
|
|
if value is None:
|
|
return None
|
|
|
|
if isinstance(value, dict):
|
|
return value
|
|
|
|
if isinstance(value, str):
|
|
try:
|
|
return json.loads(value)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid JSON in tool_parameters: {e}")
|
|
|
|
raise ValueError(f"Unexpected type for tool_parameters: {type(value)}")
|
|
|
|
@field_validator("tool_result", mode="before")
|
|
@classmethod
|
|
def parse_tool_result(cls, value: Any) -> Optional[Any]:
|
|
"""
|
|
Parse the tool_result field from a JSON string to a Python object.
|
|
|
|
Args:
|
|
value: The value to parse, can be a string, dict, list, or None
|
|
|
|
Returns:
|
|
Optional[Any]: The parsed object or None
|
|
|
|
Raises:
|
|
ValueError: If the JSON string is invalid
|
|
"""
|
|
if value is None:
|
|
return None
|
|
|
|
if not isinstance(value, str):
|
|
return value
|
|
|
|
try:
|
|
return json.loads(value)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid JSON in tool_result: {e}")
|
|
|
|
@field_validator("step_data", mode="before")
|
|
@classmethod
|
|
def parse_step_data(cls, value: Any) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Parse the step_data field from a JSON string to a dictionary.
|
|
|
|
Args:
|
|
value: The value to parse, can be a string, dict, or None
|
|
|
|
Returns:
|
|
Optional[Dict[str, Any]]: The parsed dictionary or None
|
|
|
|
Raises:
|
|
ValueError: If the JSON string is invalid
|
|
"""
|
|
if value is None:
|
|
return None
|
|
|
|
if isinstance(value, dict):
|
|
return value
|
|
|
|
if isinstance(value, str):
|
|
try:
|
|
return json.loads(value)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid JSON in step_data: {e}")
|
|
|
|
raise ValueError(f"Unexpected type for step_data: {type(value)}")
|
|
|
|
@field_serializer("tool_parameters")
|
|
def serialize_tool_parameters(self, tool_parameters: Optional[Dict[str, Any]]) -> Optional[str]:
|
|
"""
|
|
Serialize the tool_parameters dictionary to a JSON string for storage.
|
|
|
|
Args:
|
|
tool_parameters: Dictionary to serialize
|
|
|
|
Returns:
|
|
Optional[str]: JSON-encoded string or None
|
|
"""
|
|
if tool_parameters is None:
|
|
return None
|
|
|
|
return json.dumps(tool_parameters)
|
|
|
|
@field_serializer("tool_result")
|
|
def serialize_tool_result(self, tool_result: Optional[Any]) -> Optional[str]:
|
|
"""
|
|
Serialize the tool_result object to a JSON string for storage.
|
|
|
|
Args:
|
|
tool_result: Object to serialize
|
|
|
|
Returns:
|
|
Optional[str]: JSON-encoded string or None
|
|
"""
|
|
if tool_result is None:
|
|
return None
|
|
|
|
return json.dumps(tool_result)
|
|
|
|
@field_serializer("step_data")
|
|
def serialize_step_data(self, step_data: Optional[Dict[str, Any]]) -> Optional[str]:
|
|
"""
|
|
Serialize the step_data dictionary to a JSON string for storage.
|
|
|
|
Args:
|
|
step_data: Dictionary to serialize
|
|
|
|
Returns:
|
|
Optional[str]: JSON-encoded string or None
|
|
"""
|
|
if step_data is None:
|
|
return None
|
|
|
|
return json.dumps(step_data) |