jl-ecms-client 0.2.8__py3-none-any.whl

mirix/schemas/block.py

@@ -0,0 +1,188 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
+
+from mirix.client.constants import CORE_MEMORY_BLOCK_CHAR_LIMIT
+from mirix.schemas.mirix_base import MirixBase
+
+# block of the LLM context
+
+
+class BaseBlock(MirixBase, validate_assignment=True):
+    """Base block of the LLM context"""
+
+    __id_prefix__ = "block"
+
+    # data value
+    value: Optional[str] = Field(None, description="Value of the block.")
+    limit: int = Field(
+        CORE_MEMORY_BLOCK_CHAR_LIMIT, description="Character limit of the block."
+    )
+
+    # context window label
+    label: Optional[str] = Field(
+        None,
+        description="Label of the block (e.g. 'human', 'persona') in the context window.",
+    )
+
+    class Config:
+        extra = "ignore"  # Ignores extra fields
+
+    @model_validator(mode="after")
+    def verify_char_limit(self) -> Self:
+        if self.value and len(self.value) > self.limit:
+            error_msg = f"Edit failed: Exceeds {self.limit} character limit (requested {len(self.value)}) - {str(self)}."
+            raise ValueError(error_msg)
+
+        return self
+
+    def __setattr__(self, name, value):
+        """Run validation if self.value is updated"""
+        super().__setattr__(name, value)
+        if name == "value":
+            # run validation
+            self.__class__.model_validate(self.model_dump(exclude_unset=True))
+
+
+class Block(BaseBlock):
+    """
+    A Block represents a reserved section of the LLM's context window which is editable. `Block` objects contained in the `Memory` object, which is able to edit the Block values.
+
+    Parameters:
+        label (str): The label of the block (e.g. 'human', 'persona'). This defines a category for the block.
+        value (str): The value of the block. This is the string that is represented in the context window.
+        limit (int): The character limit of the block.
+        user_id (str): The unique identifier of the user associated with the block.
+        agent_id (str): The unique identifier of the agent associated with the block.
+    """
+
+    id: str = BaseBlock.generate_id_field()
+
+    # associated user/agent/organization
+    user_id: Optional[str] = Field(
+        None, description="The unique identifier of the user associated with the block."
+    )
+    agent_id: Optional[str] = Field(
+        None, description="The unique identifier of the agent associated with the block."
+    )
+    organization_id: Optional[str] = Field(
+        None,
+        description="The unique identifier of the organization associated with the block.",
+    )
+
+    # default orm fields
+    created_by_id: Optional[str] = Field(
+        None, description="The id of the user that made this Block."
+    )
+    last_updated_by_id: Optional[str] = Field(
+        None, description="The id of the user that last updated this Block."
+    )
+
+
+class Human(Block):
+    """Human block of the LLM context"""
+
+    label: str = "human"
+
+
+class Persona(Block):
+    """Persona block of the LLM context"""
+
+    label: str = "persona"
+
+
+# class CreateBlock(BaseBlock):
+#    """Create a block"""
+#
+#    label: str = Field(..., description="Label of the block.")
+
+
+class BlockLabelUpdate(BaseModel):
+    """Update the label of a block"""
+
+    current_label: str = Field(..., description="Current label of the block.")
+    new_label: str = Field(..., description="New label of the block.")
+
+
+# class CreatePersona(CreateBlock):
+#    """Create a persona block"""
+#
+#    label: str = "persona"
+#
+#
+# class CreateHuman(CreateBlock):
+#    """Create a human block"""
+#
+#    label: str = "human"
+
+
+class BlockUpdate(BaseBlock):
+    """Update a block"""
+
+    limit: Optional[int] = Field(
+        CORE_MEMORY_BLOCK_CHAR_LIMIT, description="Character limit of the block."
+    )
+    value: Optional[str] = Field(None, description="Value of the block.")
+
+    class Config:
+        extra = "ignore"  # Ignores extra fields
+
+
+class BlockLimitUpdate(BaseModel):
+    """Update the limit of a block"""
+
+    label: str = Field(..., description="Label of the block.")
+    limit: int = Field(..., description="New limit of the block.")
+
+
+# class UpdatePersona(BlockUpdate):
+#    """Update a persona block"""
+#
+#    label: str = "persona"
+#
+#
+# class UpdateHuman(BlockUpdate):
+#    """Update a human block"""
+#
+#    label: str = "human"
+
+
+class CreateBlock(BaseBlock):
+    """Create a block"""
+
+    label: str = Field(..., description="Label of the block.")
+    limit: int = Field(
+        CORE_MEMORY_BLOCK_CHAR_LIMIT, description="Character limit of the block."
+    )
+    value: str = Field(..., description="Value of the block.")
+
+
+class CreateHuman(CreateBlock):
+    """Create a human block"""
+
+    label: str = "human"
+
+
+class CreatePersona(CreateBlock):
+    """Create a persona block"""
+
+    label: str = "persona"
+
+
+class CreateBlockTemplate(CreateBlock):
+    """Create a block template"""
+
+    pass
+
+
+class CreateHumanBlockTemplate(CreateHuman):
+    """Create a human block template"""
+
+    label: str = "human"
+
+
+class CreatePersonaBlockTemplate(CreatePersona):
+    """Create a persona block template"""
+
+    label: str = "persona"

mirix/schemas/cloud_file_mapping.py

@@ -0,0 +1,29 @@
+from datetime import datetime
+
+from pydantic import Field
+
+from mirix.client.utils import get_utc_time
+from mirix.schemas.mirix_base import MirixBase
+
+
+class CloudFileMappingBase(MirixBase):
+    __id_prefix__ = "cloud_map"
+
+
+class CloudFileMapping(CloudFileMappingBase):
+    """
+    Schema for the mapping between cloud file and the local file
+    """
+
+    id: str = CloudFileMappingBase.generate_id_field()
+    cloud_file_id: str = Field(..., description="The ID of the cloud file")
+    local_file_id: str = Field(..., description="The ID of the local file")
+    status: str = Field(
+        ..., description="whether it has been processed into the memory system."
+    )
+    # TODO: change timestamp from str to datetime to accommodate timezone changes.
+    timestamp: str = Field(..., description="timestamp of the screenshot")
+    created_at: datetime = Field(
+        default_factory=get_utc_time,
+        description="Timestamp when this memory record was created",
+    )

mirix/schemas/embedding_config.py

@@ -0,0 +1,114 @@
+from typing import Literal, Optional
+
+from pydantic import BaseModel, Field
+
+from mirix.log import get_logger
+
+logger = get_logger(__name__)
+
+class EmbeddingConfig(BaseModel):
+    """
+
+    Embedding model configuration. This object specifies all the information necessary to access an embedding model to usage with Mirix, except for secret keys.
+
+    Attributes:
+        embedding_endpoint_type (str): The endpoint type for the model.
+        embedding_endpoint (str): The endpoint for the model.
+        embedding_model (str): The model for the embedding.
+        embedding_dim (int): The dimension of the embedding.
+        embedding_chunk_size (int): The chunk size of the embedding.
+        azure_endpoint (:obj:`str`, optional): The Azure endpoint for the model (Azure only).
+        azure_version (str): The Azure version for the model (Azure only).
+        azure_deployment (str): The Azure deployment for the model (Azure only).
+
+    """
+
+    embedding_endpoint_type: Literal[
+        "openai",
+        "anthropic",
+        "bedrock",
+        "cohere",
+        "google_ai",
+        "azure",
+        "groq",
+        "ollama",
+        "webui",
+        "webui-legacy",
+        "lmstudio",
+        "lmstudio-legacy",
+        "llamacpp",
+        "koboldcpp",
+        "vllm",
+        "hugging-face",
+        "mistral",
+        "together",  # completions endpoint
+    ] = Field(..., description="The endpoint type for the model.")
+    embedding_endpoint: Optional[str] = Field(
+        None, description="The endpoint for the model (`None` if local)."
+    )
+    embedding_model: str = Field(..., description="The model for the embedding.")
+    embedding_dim: int = Field(..., description="The dimension of the embedding.")
+    embedding_chunk_size: Optional[int] = Field(
+        300, description="The chunk size of the embedding."
+    )
+    handle: Optional[str] = Field(
+        None,
+        description="The handle for this config, in the format provider/model-name.",
+    )
+
+    # azure only
+    azure_endpoint: Optional[str] = Field(
+        None, description="The Azure endpoint for the model."
+    )
+    azure_version: Optional[str] = Field(
+        None, description="The Azure version for the model."
+    )
+    azure_deployment: Optional[str] = Field(
+        None, description="The Azure deployment for the model."
+    )
+
+    @classmethod
+    def default_config(
+        cls, model_name: Optional[str] = None, provider: Optional[str] = None
+    ):
+        if model_name == "text-embedding-3-small" or (
+            not model_name and provider == "openai"
+        ):
+            return cls(
+                embedding_model="text-embedding-3-small",
+                embedding_endpoint_type="openai",
+                embedding_endpoint="https://api.openai.com/v1",
+                embedding_dim=1536,
+                embedding_chunk_size=8191,
+            )
+        elif model_name == "text-embedding-004" or (
+            not model_name and provider == "google_ai"
+        ):
+            return cls(
+                embedding_model="text-embedding-004",
+                embedding_endpoint_type="google_ai",
+                embedding_endpoint="https://generativelanguage.googleapis.com",
+                embedding_dim=768,
+                embedding_chunk_size=2048,
+            )
+        elif model_name == "mirix":
+            return cls(
+                embedding_endpoint="https://embeddings.memgpt.ai",
+                embedding_model="BAAI/bge-large-en-v1.5",
+                embedding_dim=1024,
+                embedding_chunk_size=300,
+                embedding_endpoint_type="hugging-face",
+            )
+        else:
+            raise ValueError(f"Model {model_name} not supported.")
+
+    def pretty_print(self) -> str:
+        return (
+            f"{self.embedding_model}"
+            + (
+                f" [type={self.embedding_endpoint_type}]"
+                if self.embedding_endpoint_type
+                else ""
+            )
+            + (f" [ip={self.embedding_endpoint}]" if self.embedding_endpoint else "")
+        )

mirix/schemas/enums.py

@@ -0,0 +1,69 @@
+from enum import Enum
+
+
+class ProviderType(str, Enum):
+    anthropic = "anthropic"
+
+
+class MessageRole(str, Enum):
+    assistant = "assistant"
+    user = "user"
+    tool = "tool"
+    function = "function"
+    system = "system"
+
+
+class OptionState(str, Enum):
+    """Useful for kwargs that are bool + default option"""
+
+    YES = "yes"
+    NO = "no"
+    DEFAULT = "default"
+
+
+class JobStatus(str, Enum):
+    """
+    Status of the job.
+    """
+
+    not_started = "not_started"
+    created = "created"
+    running = "running"
+    completed = "completed"
+    failed = "failed"
+    pending = "pending"
+    cancelled = "cancelled"
+    expired = "expired"
+
+
+class AgentStepStatus(str, Enum):
+    """
+    Status of the job.
+    """
+
+    paused = "paused"
+    resumed = "resumed"
+    completed = "completed"
+
+
+class MessageStreamStatus(str, Enum):
+    done = "[DONE]"
+
+    def model_dump_json(self):
+        return "[DONE]"
+
+
+class ToolRuleType(str, Enum):
+    """
+    Type of tool rule.
+    """
+
+    # note: some of these should be renamed when we do the data migration
+
+    run_first = "run_first"
+    exit_loop = "exit_loop"  # reasoning loop should exit
+    continue_loop = "continue_loop"
+    conditional = "conditional"
+    constrain_child_tools = "constrain_child_tools"
+    max_count_per_step = "max_count_per_step"
+    parent_last_tool = "parent_last_tool"

mirix/schemas/environment_variables.py

@@ -0,0 +1,82 @@
+from typing import Optional
+
+from pydantic import Field
+
+from mirix.schemas.mirix_base import MirixBase, OrmMetadataBase
+
+
+# Base Environment Variable
+class EnvironmentVariableBase(OrmMetadataBase):
+    id: str = Field(
+        ..., description="The unique identifier for the environment variable."
+    )
+    key: str = Field(..., description="The name of the environment variable.")
+    value: str = Field(..., description="The value of the environment variable.")
+    description: Optional[str] = Field(
+        None, description="An optional description of the environment variable."
+    )
+    organization_id: Optional[str] = Field(
+        None,
+        description="The ID of the organization this environment variable belongs to.",
+    )
+
+
+class EnvironmentVariableCreateBase(MirixBase):
+    key: str = Field(..., description="The name of the environment variable.")
+    value: str = Field(..., description="The value of the environment variable.")
+    description: Optional[str] = Field(
+        None, description="An optional description of the environment variable."
+    )
+
+
+class EnvironmentVariableUpdateBase(MirixBase):
+    key: Optional[str] = Field(
+        None, description="The name of the environment variable."
+    )
+    value: Optional[str] = Field(
+        None, description="The value of the environment variable."
+    )
+    description: Optional[str] = Field(
+        None, description="An optional description of the environment variable."
+    )
+
+
+# Environment Variable
+class SandboxEnvironmentVariableBase(EnvironmentVariableBase):
+    __id_prefix__ = "sandbox-env"
+    sandbox_config_id: str = Field(
+        ...,
+        description="The ID of the sandbox config this environment variable belongs to.",
+    )
+
+
+class SandboxEnvironmentVariable(SandboxEnvironmentVariableBase):
+    id: str = SandboxEnvironmentVariableBase.generate_id_field()
+
+
+class SandboxEnvironmentVariableCreate(EnvironmentVariableCreateBase):
+    pass
+
+
+class SandboxEnvironmentVariableUpdate(EnvironmentVariableUpdateBase):
+    pass
+
+
+# Agent-Specific Environment Variable
+class AgentEnvironmentVariableBase(EnvironmentVariableBase):
+    __id_prefix__ = "agent-env"
+    agent_id: str = Field(
+        ..., description="The ID of the agent this environment variable belongs to."
+    )
+
+
+class AgentEnvironmentVariable(AgentEnvironmentVariableBase):
+    id: str = AgentEnvironmentVariableBase.generate_id_field()
+
+
+class AgentEnvironmentVariableCreate(EnvironmentVariableCreateBase):
+    pass
+
+
+class AgentEnvironmentVariableUpdate(EnvironmentVariableUpdateBase):
+    pass

mirix/schemas/episodic_memory.py

@@ -0,0 +1,170 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import Field, field_validator
+
+from mirix.client.utils import get_utc_time
+from mirix.client.constants import MAX_EMBEDDING_DIM
+from mirix.schemas.embedding_config import EmbeddingConfig
+from mirix.schemas.mirix_base import MirixBase
+
+
+class EpisodicEventBase(MirixBase):
+    """
+    Base schema for episodic memory events containing common fields.
+    """
+
+    __id_prefix__ = "ep_mem"
+    event_type: str = Field(
+        ...,
+        description="Type/category of the episodic event (e.g., user_message, inference, system_notification)",
+    )
+    summary: str = Field(..., description="Short textual summary of the event")
+    details: str = Field(..., description="Detailed description or text for the event")
+    actor: str = Field(
+        ..., description="The actor who generated the event (user or assistant)"
+    )
+
+
+class EpisodicEventForLLM(EpisodicEventBase):
+    """
+    Schema for creating a new episodic memory record.
+    """
+
+    # TODO: make `occurred_at` optional
+    occurred_at: str = Field(
+        ...,
+        description="When the event happened (it should be mentioned in the user's response and it should be in the format of 'YYYY-MM-DD HH:MM:SS')",
+    )
+
+
+class EpisodicEvent(EpisodicEventBase):
+    """
+    Representation of a single episodic memory event in the system.
+
+    Additional Parameters:
+        id (str): Unique identifier for this memory item
+        occurred_at (datetime): When the event occurred or was recorded
+        created_at (datetime): When the memory record was created in the system
+        updated_at (Optional[datetime]): Last update timestamp
+    """
+
+    id: Optional[str] = Field(
+        None, description="Unique identifier for the episodic event"
+    )
+
+    agent_id: Optional[str] = Field(
+        None, description="The id of the agent this episodic event belongs to"
+    )
+
+    user_id: str = Field(
+        ..., description="The id of the user who generated the episodic event"
+    )
+
+    occurred_at: datetime = Field(
+        default_factory=get_utc_time,
+        description="When the event actually happened (recorded or user-labeled).",
+    )
+    created_at: datetime = Field(
+        default_factory=get_utc_time,
+        description="Timestamp when this memory record was created",
+    )
+    updated_at: Optional[datetime] = Field(
+        None, description="When this memory record was last updated"
+    )
+    last_modify: Dict[str, Any] = Field(
+        default_factory=lambda: {
+            "timestamp": get_utc_time().isoformat(),
+            "operation": "created",
+        },
+        description="Last modification info including timestamp and operation type",
+    )
+    organization_id: str = Field(
+        ..., description="Unique identifier of the organization"
+    )
+    details_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the event"
+    )
+    summary_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the summary"
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None, description="The embedding configuration used by the event"
+    )
+    
+    # NEW: Filter tags for flexible filtering and categorization
+    filter_tags: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="Custom filter tags for filtering and categorization",
+        examples=[
+            {
+                "project_id": "proj-abc",
+                "session_id": "sess-xyz",
+                "tags": ["important", "work"],
+                "priority": "high"
+            }
+        ]
+    )
+
+    # need to validate both details_embedding and summary_embedding to ensure they are the same size
+    @field_validator("details_embedding", "summary_embedding")
+    @classmethod
+    def pad_embeddings(cls, embedding: List[float]) -> List[float]:
+        """Pad embeddings to `MAX_EMBEDDING_SIZE`. This is necessary to ensure all stored embeddings are the same size."""
+        import numpy as np
+
+        if embedding and len(embedding) != MAX_EMBEDDING_DIM:
+            np_embedding = np.array(embedding)
+            padded_embedding = np.pad(
+                np_embedding,
+                (0, MAX_EMBEDDING_DIM - np_embedding.shape[0]),
+                mode="constant",
+            )
+            return padded_embedding.tolist()
+        return embedding
+
+
+class EpisodicEventUpdate(MirixBase):
+    """
+    Schema for updating an existing episodic memory record.
+
+    All fields (except id) are optional so that only provided fields are updated.
+    """
+
+    id: str = Field(..., description="Unique ID for this episodic memory record")
+    agent_id: Optional[str] = Field(
+        None, description="The id of the agent this episodic event belongs to"
+    )
+    event_type: Optional[str] = Field(None, description="Type/category of the event")
+    summary: Optional[str] = Field(
+        None, description="Short textual summary of the event"
+    )
+    details: Optional[str] = Field(
+        None, description="Detailed text describing the event"
+    )
+    organization_id: Optional[str] = Field(
+        None, description="Unique identifier of the organization"
+    )
+    occurred_at: Optional[datetime] = Field(
+        None, description="If the event's time is updated"
+    )
+    updated_at: datetime = Field(
+        default_factory=get_utc_time,
+        description="Timestamp when this memory record was last updated",
+    )
+    last_modify: Optional[Dict[str, Any]] = Field(
+        None,
+        description="Last modification info including timestamp and operation type",
+    )
+    summary_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the summary"
+    )
+    details_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the event"
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None, description="The embedding configuration used by the event"
+    )
+    filter_tags: Optional[Dict[str, Any]] = Field(
+        None, description="Custom filter tags for filtering and categorization"
+    )