MemoryOS 2.0.3__py3-none-any.whl

memos/types/openai_chat_completion_types/chat_completion_content_part_param.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import Literal, TypeAlias
+
+from typing_extensions import Required, TypedDict
+
+from .chat_completion_content_part_image_param import ChatCompletionContentPartImageParam
+from .chat_completion_content_part_input_audio_param import ChatCompletionContentPartInputAudioParam
+from .chat_completion_content_part_text_param import ChatCompletionContentPartTextParam
+
+
+__all__ = ["ChatCompletionContentPartParam", "File", "FileFile"]
+
+
+class FileFile(TypedDict, total=False):
+    file_data: str
+    """
+    The base64 encoded file data, used when passing the file to the model as a
+    string.
+    or a url.
+    or just string which is the content of the file.
+    """
+
+    file_id: str
+    """The ID of an uploaded file to use as input."""
+
+    filename: str
+    """The name of the file, used when passing the file to the model as a string."""
+
+
+class File(TypedDict, total=False):
+    file: Required[FileFile]
+
+    type: Required[Literal["file"]]
+    """The type of the content part. Always `file`."""
+
+
+ChatCompletionContentPartParam: TypeAlias = (
+    ChatCompletionContentPartTextParam
+    | ChatCompletionContentPartImageParam
+    | ChatCompletionContentPartInputAudioParam
+    | File
+)

memos/types/openai_chat_completion_types/chat_completion_content_part_refusal_param.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from typing_extensions import Required, TypedDict
+
+
+__all__ = ["ChatCompletionContentPartRefusalParam"]
+
+
+class ChatCompletionContentPartRefusalParam(TypedDict, total=False):
+    refusal: Required[str]
+    """The refusal message generated by the model."""
+
+    type: Required[Literal["refusal"]]
+    """The type of the content part."""

memos/types/openai_chat_completion_types/chat_completion_content_part_text_param.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from typing_extensions import Required, TypedDict
+
+
+__all__ = ["ChatCompletionContentPartTextParam"]
+
+
+class ChatCompletionContentPartTextParam(TypedDict, total=False):
+    text: Required[str]
+    """The text content."""
+
+    type: Required[Literal["text"]]
+    """The type of the content part."""

memos/types/openai_chat_completion_types/chat_completion_message_custom_tool_call_param.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from typing_extensions import Required, TypedDict
+
+
+__all__ = ["ChatCompletionMessageCustomToolCallParam", "Custom"]
+
+
+class Custom(TypedDict, total=False):
+    input: Required[str]
+    """The input for the custom tool call generated by the model."""
+
+    name: Required[str]
+    """The name of the custom tool to call."""
+
+
+class ChatCompletionMessageCustomToolCallParam(TypedDict, total=False):
+    id: Required[str]
+    """The ID of the tool call."""
+
+    custom: Required[Custom]
+    """The custom tool that the model called."""
+
+    type: Required[Literal["custom"]]
+    """The type of the tool. Always `custom`."""

memos/types/openai_chat_completion_types/chat_completion_message_function_tool_call_param.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from typing_extensions import Required, TypedDict
+
+
+__all__ = ["ChatCompletionMessageFunctionToolCallParam", "Function"]
+
+
+class Function(TypedDict, total=False):
+    arguments: Required[str]
+    """
+    The arguments to call the function with, as generated by the model in JSON
+    format. Note that the model does not always generate valid JSON, and may
+    hallucinate parameters not defined by your function schema. Validate the
+    arguments in your code before calling your function.
+    """
+
+    name: Required[str]
+    """The name of the function to call."""
+
+
+class ChatCompletionMessageFunctionToolCallParam(TypedDict, total=False):
+    id: Required[str]
+    """The ID of the tool call."""
+
+    function: Required[Function]
+    """The function that the model called."""
+
+    type: Required[Literal["function"]]
+    """The type of the tool. Currently, only `function` is supported."""

memos/types/openai_chat_completion_types/chat_completion_message_param.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import TypeAlias
+
+from .chat_completion_assistant_message_param import ChatCompletionAssistantMessageParam
+from .chat_completion_system_message_param import ChatCompletionSystemMessageParam
+from .chat_completion_tool_message_param import ChatCompletionToolMessageParam
+from .chat_completion_user_message_param import ChatCompletionUserMessageParam
+
+
+__all__ = ["ChatCompletionMessageParam"]
+
+ChatCompletionMessageParam: TypeAlias = (
+    ChatCompletionSystemMessageParam
+    | ChatCompletionUserMessageParam
+    | ChatCompletionAssistantMessageParam
+    | ChatCompletionToolMessageParam
+)

memos/types/openai_chat_completion_types/chat_completion_message_tool_call_union_param.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from typing import TypeAlias
+
+from .chat_completion_message_custom_tool_call_param import ChatCompletionMessageCustomToolCallParam
+from .chat_completion_message_function_tool_call_param import (
+    ChatCompletionMessageFunctionToolCallParam,
+)
+
+
+__all__ = ["ChatCompletionMessageToolCallUnionParam"]
+
+ChatCompletionMessageToolCallUnionParam: TypeAlias = (
+    ChatCompletionMessageFunctionToolCallParam | ChatCompletionMessageCustomToolCallParam
+)

memos/types/openai_chat_completion_types/chat_completion_system_message_param.py

@@ -0,0 +1,36 @@
+# ruff: noqa: TC001, TC003
+
+from __future__ import annotations
+
+from typing import Literal
+
+from typing_extensions import Required, TypedDict
+
+from .chat_completion_content_part_text_param import ChatCompletionContentPartTextParam
+
+
+__all__ = ["ChatCompletionSystemMessageParam"]
+
+
+class ChatCompletionSystemMessageParam(TypedDict, total=False):
+    content: Required[
+        str | list[ChatCompletionContentPartTextParam] | ChatCompletionContentPartTextParam
+    ]
+    """The contents of the system message."""
+
+    role: Required[Literal["system"]]
+    """The role of the messages author, in this case `system`."""
+
+    name: str
+    """An optional name for the participant.
+
+    Provides the model information to differentiate between participants of the same
+    role.
+    """
+
+    chat_time: str | None
+    """Optional timestamp for the message, format is not
+    restricted, it can be any vague or precise time string."""
+
+    message_id: str | None
+    """Optional unique identifier for the message"""

memos/types/openai_chat_completion_types/chat_completion_tool_message_param.py

@@ -0,0 +1,30 @@
+# ruff: noqa: TC001, TC003
+
+from __future__ import annotations
+
+from typing import Literal
+
+from typing_extensions import Required, TypedDict
+
+from .chat_completion_content_part_param import ChatCompletionContentPartParam
+
+
+__all__ = ["ChatCompletionToolMessageParam"]
+
+
+class ChatCompletionToolMessageParam(TypedDict, total=False):
+    content: Required[str | list[ChatCompletionContentPartParam] | ChatCompletionContentPartParam]
+    """The contents of the tool message."""
+
+    role: Required[Literal["tool"]]
+    """The role of the messages author, in this case `tool`."""
+
+    tool_call_id: Required[str]
+    """Tool call that this message is responding to."""
+
+    chat_time: str | None
+    """Optional timestamp for the message, format is not
+    restricted, it can be any vague or precise time string."""
+
+    message_id: str | None
+    """Optional unique identifier for the message"""

memos/types/openai_chat_completion_types/chat_completion_user_message_param.py

@@ -0,0 +1,34 @@
+# ruff: noqa: TC001, TC003
+
+from __future__ import annotations
+
+from typing import Literal
+
+from typing_extensions import Required, TypedDict
+
+from .chat_completion_content_part_param import ChatCompletionContentPartParam
+
+
+__all__ = ["ChatCompletionUserMessageParam"]
+
+
+class ChatCompletionUserMessageParam(TypedDict, total=False):
+    content: Required[str | list[ChatCompletionContentPartParam] | ChatCompletionContentPartParam]
+    """The contents of the user message."""
+
+    role: Required[Literal["user"]]
+    """The role of the messages author, in this case `user`."""
+
+    name: str
+    """An optional name for the participant.
+
+    Provides the model information to differentiate between participants of the same
+    role.
+    """
+
+    chat_time: str | None
+    """Optional timestamp for the message, format is not
+    restricted, it can be any vague or precise time string."""
+
+    message_id: str | None
+    """Optional unique identifier for the message"""

memos/utils.py

@@ -0,0 +1,123 @@
+import functools
+import time
+import traceback
+
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def timed_with_status(
+    func=None,
+    *,
+    log_prefix="",
+    log_args=None,
+    log_extra_args=None,
+    fallback=None,
+):
+    """
+    Parameters:
+    - log: enable timing logs (default True)
+    - log_prefix: prefix; falls back to function name
+    - log_args: names to include in logs (str or list/tuple of str), values are taken from kwargs by name.
+    - log_extra_args:
+        - can be a dict: fixed contextual fields that are always attached to logs;
+        - or a callable: like `fn(*args, **kwargs) -> dict`, used to dynamically generate contextual fields at runtime.
+    """
+
+    if isinstance(log_args, str):
+        effective_log_args = [log_args]
+    else:
+        effective_log_args = list(log_args) if log_args else []
+
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            exc_type = None
+            exc_message = None
+            result = None
+            success_flag = False
+
+            try:
+                result = fn(*args, **kwargs)
+                success_flag = True
+                return result
+            except Exception as e:
+                exc_type = type(e)
+                stack_info = "".join(traceback.format_stack()[:-1])
+                exc_message = f"{stack_info}{traceback.format_exc()}"
+                success_flag = False
+
+                if fallback is not None and callable(fallback):
+                    result = fallback(e, *args, **kwargs)
+                    return result
+            finally:
+                elapsed_ms = (time.perf_counter() - start) * 1000.0
+
+                ctx_parts = []
+                # 1) Collect parameters from kwargs by name
+                for key in effective_log_args:
+                    val = kwargs.get(key)
+                    ctx_parts.append(f"{key}={val}")
+
+                # 2) Support log_extra_args as dict or callable, so we can dynamically
+                #    extract values from self or other runtime context
+                extra_items = {}
+                try:
+                    if callable(log_extra_args):
+                        extra_items = log_extra_args(*args, **kwargs) or {}
+                    elif isinstance(log_extra_args, dict):
+                        extra_items = log_extra_args
+                except Exception as e:
+                    logger.warning(f"[TIMER_WITH_STATUS] log_extra_args callback error: {e!r}")
+
+                if extra_items:
+                    ctx_parts.extend(f"{key}={val}" for key, val in extra_items.items())
+
+                ctx_str = f" [{', '.join(ctx_parts)}]" if ctx_parts else ""
+
+                status = "SUCCESS" if success_flag else "FAILED"
+                status_info = f", status: {status}"
+                if not success_flag and exc_type is not None:
+                    status_info += (
+                        f", error_type: {exc_type.__name__}, error_message: {exc_message}"
+                    )
+
+                msg = (
+                    f"[TIMER_WITH_STATUS] {log_prefix or fn.__name__} "
+                    f"took {elapsed_ms:.0f} ms{status_info}, args: {ctx_str}"
+                )
+
+                logger.info(msg)
+
+        return wrapper
+
+    if func is None:
+        return decorator
+    return decorator(func)
+
+
+def timed(func=None, *, log=True, log_prefix=""):
+    def decorator(fn):
+        def wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            result = fn(*args, **kwargs)
+            elapsed_ms = (time.perf_counter() - start) * 1000.0
+
+            if log is not True:
+                return result
+
+            # 100ms threshold
+            if elapsed_ms >= 100.0:
+                logger.info(f"[TIMER] {log_prefix or fn.__name__} took {elapsed_ms:.0f} ms")
+
+            return result
+
+        return wrapper
+
+    # Handle both @timed and @timed(log=True) cases
+    if func is None:
+        return decorator
+    return decorator(func)

memos/vec_dbs/base.py

@@ -0,0 +1,117 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from memos.configs.vec_db import BaseVecDBConfig
+from memos.vec_dbs.item import VecDBItem
+
+
+class BaseVecDB(ABC):
+    """Base class for all vector databases."""
+
+    @abstractmethod
+    def __init__(self, config: BaseVecDBConfig):
+        """Initialize the vector database with the given configuration."""
+
+    # Collection management methods
+
+    @abstractmethod
+    def create_collection(self) -> None:
+        """Create a new collection/index with specified parameters."""
+
+    @abstractmethod
+    def list_collections(self) -> list[str]:
+        """List all collections/indexes."""
+
+    @abstractmethod
+    def delete_collection(self, name: str) -> None:
+        """Delete a collection/index."""
+
+    @abstractmethod
+    def collection_exists(self, name: str) -> bool:
+        """Check if a collection/index exists."""
+
+    # Vector management methods
+
+    @abstractmethod
+    def search(
+        self,
+        query_vector: list[float],
+        top_k: int,
+        filter: dict[str, Any] | None = None,
+    ) -> list[VecDBItem]:
+        """
+        Search for similar items in the vector database.
+
+        Args:
+            query_vector: Single vector to search
+            top_k: Number of results to return
+            filter: payload filters (may not be supported by all implementations)
+
+        Returns:
+            List of search results with distance scores and payloads.
+        """
+
+    @abstractmethod
+    def get_by_id(self, id: str) -> VecDBItem | None:
+        """Get an item from the vector database."""
+
+    @abstractmethod
+    def get_by_ids(self, ids: list[str]) -> list[VecDBItem]:
+        """Get multiple items by their IDs."""
+
+    @abstractmethod
+    def get_by_filter(self, filter: dict[str, Any]) -> list[VecDBItem]:
+        """
+        Retrieve all items that match the given filter criteria.
+
+        Args:
+            filter: Payload filters to match against stored items
+
+        Returns:
+            List of items including vectors and payloads that match the filter
+        """
+
+    @abstractmethod
+    def get_all(self) -> list[VecDBItem]:
+        """Retrieve all items in the vector database."""
+
+    @abstractmethod
+    def count(self, filter: dict[str, Any] | None = None) -> int:
+        """Count items in the database, optionally with filter."""
+
+    @abstractmethod
+    def add(self, data: list[VecDBItem | dict[str, Any]]) -> None:
+        """
+        Add data to the vector database.
+
+        Args:
+            data: List of VecDBItem objects or dictionaries containing:
+                - 'id': unique identifier
+                - 'vector': embedding vector
+                - 'payload': additional fields for filtering/retrieval
+        """
+
+    @abstractmethod
+    def update(self, id: str, data: VecDBItem | dict[str, Any]) -> None:
+        """Update an item in the vector database."""
+
+    @abstractmethod
+    def upsert(self, data: list[VecDBItem | dict[str, Any]]) -> None:
+        """
+        Add or update data in the vector database.
+
+        If an item with the same ID exists, it will be updated.
+        Otherwise, it will be added as a new item.
+        """
+
+    @abstractmethod
+    def delete(self, ids: list[str]) -> None:
+        """Delete items from the vector database."""
+
+    @abstractmethod
+    def ensure_payload_indexes(self, fields: list[str]) -> None:
+        """
+        Create payload indexes for specified fields in the collection.
+        Args:
+            fields (list[str]): List of field names to index (as keyword).
+        """

memos/vec_dbs/factory.py

@@ -0,0 +1,23 @@
+from typing import Any, ClassVar
+
+from memos.configs.vec_db import VectorDBConfigFactory
+from memos.vec_dbs.base import BaseVecDB
+from memos.vec_dbs.milvus import MilvusVecDB
+from memos.vec_dbs.qdrant import QdrantVecDB
+
+
+class VecDBFactory(BaseVecDB):
+    """Factory class for creating Vector Database instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "qdrant": QdrantVecDB,
+        "milvus": MilvusVecDB,
+    }
+
+    @classmethod
+    def from_config(cls, config_factory: VectorDBConfigFactory) -> BaseVecDB:
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        vec_db_class = cls.backend_to_class[backend]
+        return vec_db_class(config_factory.config)

memos/vec_dbs/item.py

@@ -0,0 +1,50 @@
+"""Defines vector database item types."""
+
+import uuid
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+class VecDBItem(BaseModel):
+    """Represents a single item in the vector database.
+
+    This serves as a standardized format for vector database items across different
+    vector database implementations (Qdrant, FAISS, Weaviate, etc.).
+    """
+
+    id: str = Field(default=str(uuid.uuid4()), description="Unique identifier for the item")
+    vector: list[float] | None = Field(default=None, description="Embedding vector")
+    payload: dict[str, Any] | None = Field(
+        default=None, description="Additional payload for filtering/retrieval"
+    )
+    score: float | None = Field(
+        default=None, description="Similarity score (used in search results)"
+    )
+
+    model_config = ConfigDict(extra="forbid")
+
+    @field_validator("id")
+    @classmethod
+    def validate_id(cls, v):
+        """Validate that ID is a valid UUID."""
+        if not isinstance(v, str) or not uuid.UUID(v, version=4):
+            raise ValueError("ID must be a valid UUID string")
+        return v
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "VecDBItem":
+        """Create VecDBItem from dictionary."""
+        return cls(**data)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format."""
+        return self.model_dump(exclude_none=True)
+
+
+class MilvusVecDBItem(VecDBItem):
+    """Represents a single item in the Milvus vector database."""
+
+    memory: str | None = Field(default=None, description="Memory string")
+    original_text: str | None = Field(default=None, description="Original text content")