276 lines
11 KiB
Python
276 lines
11 KiB
Python
from typing import Literal, List, Optional, Union, Any, Dict, Set
|
|
from pydantic import BaseModel, Field, model_validator, field_validator, RootModel
|
|
from datetime import datetime
|
|
import hashlib
|
|
|
|
|
|
def _content_item_type(value: Any) -> Optional[str]:
|
|
if isinstance(value, dict):
|
|
item_type = value.get("type")
|
|
return item_type if isinstance(item_type, str) else None
|
|
|
|
item_type = getattr(value, "type", None)
|
|
return item_type if isinstance(item_type, str) else None
|
|
|
|
|
|
def _content_item_text(value: Any) -> Optional[str]:
|
|
if isinstance(value, dict):
|
|
text = value.get("text")
|
|
return text if isinstance(text, str) else None
|
|
|
|
text = getattr(value, "text", None)
|
|
return text if isinstance(text, str) else None
|
|
|
|
|
|
def _content_item_image_payload(value: Any) -> Any:
|
|
if isinstance(value, dict):
|
|
return value.get("image_url")
|
|
return getattr(value, "image_url", None)
|
|
|
|
|
|
def _image_payload_string_field(payload: Any, field_name: str) -> Optional[str]:
|
|
if isinstance(payload, dict):
|
|
value = payload.get(field_name)
|
|
return value if isinstance(value, str) else None
|
|
|
|
value = getattr(payload, field_name, None)
|
|
return value if isinstance(value, str) else None
|
|
|
|
|
|
def _normalize_multimodal_content_item(value: Any) -> Dict[str, Any]:
|
|
if hasattr(value, "model_dump"):
|
|
value = value.model_dump()
|
|
|
|
item_type = _content_item_type(value)
|
|
if item_type == "text":
|
|
text = _content_item_text(value)
|
|
if text is None:
|
|
raise ValueError("Text content item is missing a valid 'text' field")
|
|
return {"type": "text", "text": text}
|
|
|
|
if item_type == "image_url":
|
|
image_payload = _content_item_image_payload(value)
|
|
url = _image_payload_string_field(image_payload, "url")
|
|
if url is None:
|
|
raise ValueError(
|
|
"Image content item is missing a valid 'image_url.url' field"
|
|
)
|
|
|
|
normalized_payload: Dict[str, Any] = {"url": url}
|
|
|
|
detail = _image_payload_string_field(image_payload, "detail")
|
|
if detail in {"auto", "low", "high"}:
|
|
normalized_payload["detail"] = detail
|
|
|
|
local_path = _image_payload_string_field(image_payload, "local_path")
|
|
if local_path:
|
|
normalized_payload["local_path"] = local_path
|
|
|
|
return {"type": "image_url", "image_url": normalized_payload}
|
|
|
|
raise ValueError(f"Unsupported multimodal content item: {type(value).__name__}")
|
|
|
|
|
|
def _normalize_message_content(value: Any) -> Any:
|
|
if value is None or isinstance(value, str):
|
|
return value
|
|
|
|
if isinstance(value, list):
|
|
return [_normalize_multimodal_content_item(item) for item in value]
|
|
|
|
return value
|
|
|
|
|
|
class TextContent(BaseModel):
|
|
type: Literal["text"] = Field(..., description="Content block type: plain text")
|
|
text: str = Field(..., description="Text content of the message")
|
|
|
|
|
|
class ImageURL(BaseModel):
|
|
url: str = Field(..., description="Public access URL for the image")
|
|
detail: Optional[Literal["auto", "low", "high"]] = Field(
|
|
None, description="Optional image detail level"
|
|
)
|
|
local_path: Optional[str] = Field(None, description="Local path of the image in the workspace")
|
|
|
|
|
|
class ImageContent(BaseModel):
|
|
type: Literal["image_url"] = Field(..., description="Content block type: image URL")
|
|
image_url: ImageURL = Field(..., description="Detailed information for the image content")
|
|
|
|
|
|
MessageContent = Union[str, None, List[Union[TextContent, ImageContent]]]
|
|
|
|
|
|
class FunctionCall(BaseModel):
|
|
name: str = Field(..., description="Name of the function to call")
|
|
arguments: str = Field(..., description="JSON-formatted string of arguments to pass")
|
|
|
|
|
|
class ToolCall(BaseModel):
|
|
id: str = Field(..., description="Unique ID of this tool call")
|
|
type: Literal["function"] = Field(..., description="Type of called tool (function)")
|
|
function: FunctionCall = Field(..., description="Function call specification")
|
|
|
|
|
|
class Message(BaseModel):
|
|
role: Literal["system", "user", "assistant", "tool"] = Field(
|
|
..., description="Role of the message sender"
|
|
)
|
|
content: MessageContent = Field(
|
|
...,
|
|
description=(
|
|
"Message content. It can be a string, null (when calling tools), or a list of structured multimodal blocks."
|
|
),
|
|
)
|
|
name: Optional[str] = Field(
|
|
default=None,
|
|
description="Optional sender name, required when the role is 'user' or 'tool'",
|
|
max_length=64,
|
|
pattern=r"^[a-zA-Z0-9_]*$",
|
|
)
|
|
tool_calls: Optional[List[ToolCall]] = Field(
|
|
default=None, description="List of tool calls the assistant wants to invoke"
|
|
)
|
|
tool_call_id: Optional[str] = Field(
|
|
default=None, description="Tool call ID that this tool message responds to"
|
|
)
|
|
timestamp: Optional[str] = Field(default=None, description="Message timestamp (ISO format)")
|
|
|
|
@field_validator("content", mode="before")
|
|
@classmethod
|
|
def normalize_content(cls, value: Any):
|
|
return _normalize_message_content(value)
|
|
|
|
@model_validator(mode="after")
|
|
def validate_tool_message_consistency(cls, values):
|
|
role = values.role
|
|
content = values.content
|
|
tool_calls = values.tool_calls
|
|
tool_call_id = values.tool_call_id
|
|
|
|
if role == "assistant" and tool_calls and content is not None:
|
|
raise ValueError(
|
|
"When role is 'assistant' and tool_calls exist, content must be None."
|
|
)
|
|
if role == "tool" and not tool_call_id:
|
|
raise ValueError("When role is 'tool', tool_call_id must be provided.")
|
|
return values
|
|
|
|
|
|
class ChatMessages(RootModel[List[Message]]):
|
|
"""Chat message list ordered chronologically."""
|
|
|
|
|
|
class SketchPadItem(BaseModel):
|
|
"""Data structure for a SketchPad storage item."""
|
|
|
|
value: Any = Field(..., description="Stored value")
|
|
timestamp: datetime = Field(default_factory=datetime.now, description="Creation time")
|
|
summary: Optional[str] = Field(default=None, description="Content summary")
|
|
expires_at: Optional[datetime] = Field(default=None, description="Expiration time")
|
|
access_count: int = Field(default=0, description="Access count")
|
|
last_accessed: Optional[datetime] = Field(default=None, description="Last access time")
|
|
tags: Set[str] = Field(default_factory=set, description="Tag set")
|
|
content_type: str = Field(default="text", description="Content type")
|
|
content_hash: Optional[str] = Field(default=None, description="Content hash value")
|
|
|
|
@field_validator("last_accessed", mode="before")
|
|
@classmethod
|
|
def set_last_accessed(cls, v):
|
|
"""If last_accessed is None, set it to the current time."""
|
|
if v is None:
|
|
return datetime.now()
|
|
return v
|
|
|
|
@field_validator("content_hash", mode="before")
|
|
@classmethod
|
|
def set_content_hash(cls, v, info):
|
|
"""If content_hash is None, compute the hash value."""
|
|
if v is None:
|
|
value = info.data.get("value")
|
|
if value is not None:
|
|
content_str = str(value)
|
|
return hashlib.md5(content_str.encode()).hexdigest()[:8]
|
|
return v
|
|
|
|
def is_expired(self) -> bool:
|
|
"""Check whether the item has expired."""
|
|
return self.expires_at is not None and datetime.now() > self.expires_at
|
|
|
|
def update_access(self):
|
|
"""Update access information for LRU caching."""
|
|
self.access_count += 1
|
|
self.last_accessed = datetime.now()
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to a dictionary for serialization."""
|
|
return {
|
|
"value": self.value,
|
|
"timestamp": self.timestamp.isoformat(),
|
|
"summary": self.summary,
|
|
"expires_at": self.expires_at.isoformat() if self.expires_at else None,
|
|
"access_count": self.access_count,
|
|
"last_accessed": (
|
|
self.last_accessed.isoformat() if self.last_accessed else None
|
|
),
|
|
"tags": list(self.tags),
|
|
"content_type": self.content_type,
|
|
"content_hash": self.content_hash,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> "SketchPadItem":
|
|
"""Create an instance from a dictionary."""
|
|
# Process time fields.
|
|
if isinstance(data.get("timestamp"), str):
|
|
data["timestamp"] = datetime.fromisoformat(data["timestamp"])
|
|
if data.get("expires_at") and isinstance(data["expires_at"], str):
|
|
data["expires_at"] = datetime.fromisoformat(data["expires_at"])
|
|
if data.get("last_accessed") and isinstance(data["last_accessed"], str):
|
|
data["last_accessed"] = datetime.fromisoformat(data["last_accessed"])
|
|
|
|
# Process tags.
|
|
if data.get("tags") and isinstance(data["tags"], list):
|
|
data["tags"] = set(data["tags"])
|
|
|
|
return cls(**data)
|
|
|
|
|
|
class SketchPadStatistics(BaseModel):
|
|
"""SketchPad statistics."""
|
|
|
|
total_items: int = Field(..., description="Total number of items")
|
|
max_items: int = Field(..., description="Maximum number of items")
|
|
items_with_summary: int = Field(..., description="Number of items with summaries")
|
|
total_accesses: int = Field(..., description="Total number of accesses")
|
|
popular_tags: Dict[str, int] = Field(..., description="Popular tag statistics")
|
|
content_types: Dict[str, int] = Field(..., description="Content type statistics")
|
|
avg_access_per_item: float = Field(..., description="Average accesses per item")
|
|
memory_usage_percent: float = Field(..., description="Memory usage percentage")
|
|
|
|
|
|
class SketchPadSearchResult(BaseModel):
|
|
"""SketchPad search result."""
|
|
|
|
key: str = Field(..., description="Item key")
|
|
value: Any = Field(..., description="Item value")
|
|
summary: Optional[str] = Field(default=None, description="Item summary")
|
|
timestamp: str = Field(..., description="Creation time (ISO format)")
|
|
tags: List[str] = Field(default_factory=list, description="Tag list")
|
|
content_type: str = Field(..., description="Content type")
|
|
access_count: int = Field(..., description="Access count")
|
|
|
|
|
|
class SketchPadListItem(BaseModel):
|
|
"""SketchPad list item."""
|
|
|
|
key: str = Field(..., description="Item key")
|
|
summary: Optional[str] = Field(default=None, description="Item summary")
|
|
timestamp: str = Field(..., description="Creation time (ISO format)")
|
|
tags: List[str] = Field(default_factory=list, description="Tag list")
|
|
content_type: str = Field(..., description="Content type")
|
|
access_count: int = Field(..., description="Access count")
|
|
content_hash: Optional[str] = Field(default=None, description="Content hash value")
|
|
value: Optional[Any] = Field(default=None, description="Item value, only present when content is included")
|