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

mirix/schemas/resource_memory.py

@@ -0,0 +1,134 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import Field, field_validator
+
+from mirix.client.constants import MAX_EMBEDDING_DIM
+from mirix.schemas.embedding_config import EmbeddingConfig
+from mirix.schemas.mirix_base import MirixBase
+from mirix.client.utils import get_utc_time
+
+
+class ResourceMemoryItemBase(MirixBase):
+    """
+    Base schema for resource memory items - storing docs, user files, references, etc.
+    """
+
+    __id_prefix__ = "res_item"
+    title: str = Field(..., description="Short name/title of the resource")
+    summary: str = Field(
+        ..., description="Short description or summary of the resource"
+    )
+    resource_type: str = Field(
+        ..., description="File type or format (e.g. 'doc', 'markdown', 'pdf_text')"
+    )
+    content: str = Field(
+        ..., description="Full or partial text content of the resource"
+    )
+
+
+class ResourceMemoryItem(ResourceMemoryItemBase):
+    """
+    Full schema for resource memory items with DB fields.
+    """
+
+    id: Optional[str] = Field(
+        None, description="Unique identifier for the resource memory item"
+    )
+    agent_id: Optional[str] = Field(
+        None, description="The id of the agent this resource memory item belongs to"
+    )
+    user_id: str = Field(
+        ..., description="The id of the user who generated the resource"
+    )
+    created_at: datetime = Field(
+        default_factory=get_utc_time, description="Creation timestamp"
+    )
+    updated_at: Optional[datetime] = Field(None, description="Last update timestamp")
+    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="The unique identifier of the organization"
+    )
+    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"
+            }
+        ]
+    )
+
+    @field_validator("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 ResourceMemoryItemUpdate(MirixBase):
+    """Schema for updating an existing resource memory item."""
+
+    id: str = Field(..., description="Unique ID for this resource memory entry")
+    agent_id: Optional[str] = Field(
+        None, description="The id of the agent this resource memory item belongs to"
+    )
+    title: Optional[str] = Field(None, description="Short name/title of the resource")
+    summary: Optional[str] = Field(
+        None, description="Short description or summary of the resource"
+    )
+    resource_type: Optional[str] = Field(
+        None, description="File type/format (e.g. 'doc', 'markdown')"
+    )
+    content: Optional[str] = Field(None, description="Full or partial text content")
+    organization_id: Optional[str] = Field(None, description="The organization ID")
+    updated_at: datetime = Field(
+        default_factory=get_utc_time, description="Update timestamp"
+    )
+    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"
+    )
+    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"
+    )
+
+class ResourceMemoryItemResponse(ResourceMemoryItem):
+    """Response schema for resource memory item with additional fields if needed."""
+
+    pass

mirix/schemas/sandbox_config.py

@@ -0,0 +1,132 @@
+import hashlib
+import json
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, Field, model_validator
+
+from mirix.schemas.agent import AgentState
+from mirix.schemas.mirix_base import MirixBase, OrmMetadataBase
+from mirix.settings import tool_settings
+from mirix.log import get_logger
+
+
+# Sandbox Config
+
+logger = get_logger(__name__)
+class SandboxType(str, Enum):
+    E2B = "e2b"
+    LOCAL = "local"
+
+
+class SandboxRunResult(BaseModel):
+    func_return: Optional[Any] = Field(None, description="The function return object")
+    agent_state: Optional[AgentState] = Field(None, description="The agent state")
+    stdout: Optional[List[str]] = Field(
+        None,
+        description="Captured stdout (e.g. prints, logs) from the function invocation",
+    )
+    stderr: Optional[List[str]] = Field(
+        None, description="Captured stderr from the function invocation"
+    )
+    status: Literal["success", "error"] = Field(
+        ..., description="The status of the tool execution and return object"
+    )
+    sandbox_config_fingerprint: str = Field(
+        None, description="The fingerprint of the config for the sandbox"
+    )
+
+
+class LocalSandboxConfig(BaseModel):
+    sandbox_dir: str = Field(..., description="Directory for the sandbox environment.")
+    use_venv: bool = Field(
+        False,
+        description="Whether or not to use the venv, or run directly in the same run loop.",
+    )
+    venv_name: str = Field(
+        "venv",
+        description="The name for the venv in the sandbox directory. We first search for an existing venv with this name, otherwise, we make it from the requirements.txt.",
+    )
+
+    @property
+    def type(self) -> "SandboxType":
+        return SandboxType.LOCAL
+
+
+class E2BSandboxConfig(BaseModel):
+    timeout: int = Field(5 * 60, description="Time limit for the sandbox (in seconds).")
+    template: Optional[str] = Field(
+        None, description="The E2B template id (docker image)."
+    )
+    pip_requirements: Optional[List[str]] = Field(
+        None, description="A list of pip packages to install on the E2B Sandbox"
+    )
+
+    @property
+    def type(self) -> "SandboxType":
+        return SandboxType.E2B
+
+    @model_validator(mode="before")
+    @classmethod
+    def set_default_template(cls, data: dict):
+        """
+        Assign a default template value if the template field is not provided.
+        """
+        if data.get("template") is None:
+            data["template"] = tool_settings.e2b_sandbox_template_id
+        return data
+
+
+class SandboxConfigBase(OrmMetadataBase):
+    __id_prefix__ = "sandbox"
+
+
+class SandboxConfig(SandboxConfigBase):
+    id: str = SandboxConfigBase.generate_id_field()
+    type: SandboxType = Field(None, description="The type of sandbox.")
+    organization_id: Optional[str] = Field(
+        None,
+        description="The unique identifier of the organization associated with the sandbox.",
+    )
+    config: Dict = Field(
+        default_factory=lambda: {}, description="The JSON sandbox settings data."
+    )
+
+    def get_e2b_config(self) -> E2BSandboxConfig:
+        return E2BSandboxConfig(**self.config)
+
+    def get_local_config(self) -> LocalSandboxConfig:
+        return LocalSandboxConfig(**self.config)
+
+    def fingerprint(self) -> str:
+        # Only take into account type, org_id, and the config items
+        # Canonicalize input data into JSON with sorted keys
+        hash_input = json.dumps(
+            {
+                "type": self.type.value,
+                "organization_id": self.organization_id,
+                "config": self.config,
+            },
+            sort_keys=True,  # Ensure stable ordering
+            separators=(",", ":"),  # Minimize serialization differences
+        )
+
+        # Compute SHA-256 hash
+        hash_digest = hashlib.sha256(hash_input.encode("utf-8")).digest()
+
+        # Convert the digest to an integer for compatibility with Python's hash requirements
+        return str(int.from_bytes(hash_digest, byteorder="big"))
+
+
+class SandboxConfigCreate(MirixBase):
+    config: Union[LocalSandboxConfig, E2BSandboxConfig] = Field(
+        ..., description="The configuration for the sandbox."
+    )
+
+
+class SandboxConfigUpdate(MirixBase):
+    """Pydantic model for updating SandboxConfig fields."""
+
+    config: Union[LocalSandboxConfig, E2BSandboxConfig] = Field(
+        None, description="The JSON configuration data for the sandbox."
+    )

mirix/schemas/semantic_memory.py

@@ -0,0 +1,162 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import Field, field_validator
+
+from mirix.client.constants import MAX_EMBEDDING_DIM
+from mirix.schemas.embedding_config import EmbeddingConfig
+from mirix.schemas.mirix_base import MirixBase
+from mirix.client.utils import get_utc_time
+
+
+class SemanticMemoryItemBase(MirixBase):
+    """
+    Base schema for storing semantic memory items (e.g., general knowledge, concepts, facts).
+    """
+
+    __id_prefix__ = "sem_item"
+    name: str = Field(
+        ..., description="The name or main concept/object for the knowledge entry"
+    )
+    summary: str = Field(
+        ..., description="A concise explanation or summary of the concept"
+    )
+    details: str = Field(
+        ..., description="Detailed explanation or additional context for the concept"
+    )
+    source: str = Field(
+        ...,
+        description="Reference or origin of this information (e.g., book, article, movie)",
+    )
+
+
+class SemanticMemoryItem(SemanticMemoryItemBase):
+    """
+    Full semantic memory item schema, including database-related fields.
+    """
+
+    id: Optional[str] = Field(
+        None, description="Unique identifier for the semantic memory item"
+    )
+    agent_id: Optional[str] = Field(
+        None, description="The id of the agent this semantic memory item belongs to"
+    )
+    user_id: str = Field(
+        ..., description="The id of the user who generated the semantic memory"
+    )
+    created_at: datetime = Field(
+        default_factory=get_utc_time, description="Creation timestamp"
+    )
+    updated_at: Optional[datetime] = Field(None, description="Last update timestamp")
+    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="The unique identifier of the organization"
+    )
+    details_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the details"
+    )
+    name_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the name"
+    )
+    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", "name_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 SemanticMemoryItemUpdate(MirixBase):
+    """
+    Schema for updating an existing semantic memory item.
+    """
+
+    id: str = Field(..., description="Unique ID for this semantic memory entry")
+    agent_id: Optional[str] = Field(
+        None, description="The id of the agent this semantic memory item belongs to"
+    )
+    name: Optional[str] = Field(
+        None, description="The name or main concept for the knowledge entry"
+    )
+    summary: Optional[str] = Field(
+        None, description="A concise explanation or summary of the concept"
+    )
+    details: Optional[str] = Field(
+        None, description="Detailed explanation or additional context for the concept"
+    )
+    source: Optional[str] = Field(
+        None,
+        description="Reference or origin of this information (e.g., book, article, movie)",
+    )
+    actor: Optional[str] = Field(
+        None,
+        description="The actor who generated the semantic memory (user or assistant)",
+    )
+    organization_id: Optional[str] = Field(None, description="The organization ID")
+    updated_at: datetime = Field(
+        default_factory=get_utc_time, description="Update timestamp"
+    )
+    last_modify: Optional[Dict[str, Any]] = Field(
+        None,
+        description="Last modification info including timestamp and operation type",
+    )
+    details_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the details"
+    )
+    name_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the name"
+    )
+    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"
+    )
+    filter_tags: Optional[Dict[str, Any]] = Field(
+        None, description="Custom filter tags for filtering and categorization"
+    )
+
+
+class SemanticMemoryItemResponse(SemanticMemoryItem):
+    """
+    Response schema for semantic memory item.
+    """
+
+    pass

mirix/schemas/source.py

@@ -0,0 +1,96 @@
+from datetime import datetime
+from typing import Optional
+
+from pydantic import Field
+
+from mirix.schemas.embedding_config import EmbeddingConfig
+from mirix.schemas.mirix_base import MirixBase
+
+
+class BaseSource(MirixBase):
+    """
+    Shared attributes accourss all source schemas.
+    """
+
+    __id_prefix__ = "source"
+
+
+class Source(BaseSource):
+    """
+    Representation of a source, which is a collection of files and passages.
+
+    Parameters:
+        id (str): The ID of the source
+        name (str): The name of the source.
+        embedding_config (EmbeddingConfig): The embedding configuration used by the source.
+        user_id (str): The ID of the user that created the source.
+        metadata_ (dict): Metadata associated with the source.
+        description (str): The description of the source.
+    """
+
+    id: str = BaseSource.generate_id_field()
+    name: str = Field(..., description="The name of the source.")
+    description: Optional[str] = Field(
+        None, description="The description of the source."
+    )
+    embedding_config: EmbeddingConfig = Field(
+        ..., description="The embedding configuration used by the source."
+    )
+    organization_id: Optional[str] = Field(
+        None, description="The ID of the organization that created the source."
+    )
+    metadata_: Optional[dict] = Field(
+        None, description="Metadata associated with the source."
+    )
+
+    # metadata fields
+    created_by_id: Optional[str] = Field(
+        None, description="The id of the user that made this Tool."
+    )
+    last_updated_by_id: Optional[str] = Field(
+        None, description="The id of the user that made this Tool."
+    )
+    created_at: Optional[datetime] = Field(
+        None, description="The timestamp when the source was created."
+    )
+    updated_at: Optional[datetime] = Field(
+        None, description="The timestamp when the source was last updated."
+    )
+
+
+class SourceCreate(BaseSource):
+    """
+    Schema for creating a new Source.
+    """
+
+    # required
+    name: str = Field(..., description="The name of the source.")
+    # TODO: @matt, make this required after shub makes the FE changes
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None, description="The embedding configuration used by the source."
+    )
+
+    # optional
+    description: Optional[str] = Field(
+        None, description="The description of the source."
+    )
+    metadata_: Optional[dict] = Field(
+        None, description="Metadata associated with the source."
+    )
+
+
+class SourceUpdate(BaseSource):
+    """
+    Schema for updating an existing Source.
+    """
+
+    name: Optional[str] = Field(None, description="The name of the source.")
+    description: Optional[str] = Field(
+        None, description="The description of the source."
+    )
+    metadata_: Optional[dict] = Field(
+        None, description="Metadata associated with the source."
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None, description="The embedding configuration used by the source."
+    )

mirix/schemas/step.py

@@ -0,0 +1,53 @@
+from typing import Dict, List, Optional
+
+from pydantic import Field
+
+from mirix.schemas.message import Message
+from mirix.schemas.mirix_base import MirixBase
+
+
+class StepBase(MirixBase):
+    __id_prefix__ = "step"
+
+
+class Step(StepBase):
+    id: str = Field(..., description="The id of the step. Assigned by the database.")
+    origin: Optional[str] = Field(
+        None, description="The surface that this agent step was initiated from."
+    )
+    organization_id: Optional[str] = Field(
+        None,
+        description="The unique identifier of the organization associated with the step.",
+    )
+    provider_name: Optional[str] = Field(
+        None, description="The name of the provider used for this step."
+    )
+    model: Optional[str] = Field(
+        None, description="The name of the model used for this step."
+    )
+    context_window_limit: Optional[int] = Field(
+        None, description="The context window limit configured for this step."
+    )
+    completion_tokens: Optional[int] = Field(
+        None,
+        description="The number of tokens generated by the agent during this step.",
+    )
+    prompt_tokens: Optional[int] = Field(
+        None, description="The number of tokens in the prompt during this step."
+    )
+    total_tokens: Optional[int] = Field(
+        None,
+        description="The total number of tokens processed by the agent during this step.",
+    )
+    completion_tokens_details: Optional[Dict] = Field(
+        None, description="Metadata for the agent."
+    )
+
+    tags: List[str] = Field([], description="Metadata tags.")
+    tid: Optional[str] = Field(
+        None,
+        description="The unique identifier of the transaction that processed this step.",
+    )
+    messages: List[Message] = Field(
+        [], description="The messages generated during this step."
+    )