search_personalization/agentic_memory/memory_client.py

"""
OpenSearch Agentic Memory client.
Wraps the /_plugins/_ml/memory_containers/ API (introduced in OpenSearch 3.3).
"""

from typing import Any, Optional
from opensearchpy import OpenSearch
from search_personalization.data_loader import get_opensearch_client
from search_personalization.agentic_memory.config import EMBEDDING_MODEL_ID, MEMORY_LLM_MODEL_ID, EMBEDDING_DIMENSION


def _client() -> OpenSearch:
    """Get the shared OpenSearch client."""
    return get_opensearch_client()


# --- Container Management ---


def create_container(
    name: str,
    description: str,
    strategies: list[dict],
) -> dict:
    """
    Create a memory container.

    Args:
        name: Human-readable container name.
        description: What this container stores.
        strategies: List of strategy dicts, e.g. [{"type": "USER_PREFERENCE", "namespace": ["user_id"]}]
    """
    client = _client()
    body = {
        "name": name,
        "description": description,
        "configuration": {
            "embedding_model_type": "TEXT_EMBEDDING",
            "embedding_model_id": EMBEDDING_MODEL_ID,
            "embedding_dimension": EMBEDDING_DIMENSION,
            "llm_id": MEMORY_LLM_MODEL_ID,
            "strategies": strategies,
        },
    }
    return client.transport.perform_request(
        "POST", "/_plugins/_ml/memory_containers/_create", body=body
    )


def list_containers() -> dict:
    """Search all memory containers."""
    client = _client()
    body = {"query": {"match_all": {}}}
    return client.transport.perform_request(
        "POST", "/_plugins/_ml/memory_containers/_search", body=body
    )


def get_container(container_id: str) -> dict:
    """Get a specific container by ID."""
    client = _client()
    return client.transport.perform_request(
        "GET", f"/_plugins/_ml/memory_containers/{container_id}"
    )


def delete_container(container_id: str) -> dict:
    """Delete a memory container."""
    client = _client()
    return client.transport.perform_request(
        "DELETE", f"/_plugins/_ml/memory_containers/{container_id}"
    )


# --- Memory Operations ---


def write_memory(
    container_id: str,
    namespace: dict,
    content: str,
    payload_type: str = "conversational",
    infer: bool = True,
) -> dict:
    """
    Write a memory to a container.

    Args:
        container_id: The memory container ID.
        namespace: Namespace dict, e.g. {"user_id": "user1", "session_id": "s1"}.
        content: The text content to store.
        payload_type: "conversational" or "data".
        infer: Whether to use LLM to extract knowledge.
    """
    client = _client()
    body: dict[str, Any] = {
        "messages": [
            {
                "role": "user",
                "content": [{"text": content, "type": "text"}],
            }
        ],
        "namespace": namespace,
        "payload_type": payload_type,
        "infer": infer,
    }
    return client.transport.perform_request(
        "POST", f"/_plugins/_ml/memory_containers/{container_id}/memories", body=body
    )


def read_memory(
    container_id: str,
    namespace: dict,
    memory_type: str = "long-term",
    size: int = 10,
) -> dict:
    """
    Search memories in a container by namespace.

    Args:
        container_id: The memory container ID.
        namespace: Namespace dict to filter by.
        memory_type: One of "sessions", "working", "long-term", "history".
        size: Max results.
    """
    client = _client()
    must_clauses = [{"term": {f"namespace.{k}": v}} for k, v in namespace.items()]
    body = {"query": {"bool": {"must": must_clauses}}, "size": size}
    return client.transport.perform_request(
        "GET",
        f"/_plugins/_ml/memory_containers/{container_id}/memories/{memory_type}/_search",
        body=body,
    )


def search_memory(
    container_id: str,
    namespace: dict,
    query: str,
    memory_type: str = "long-term",
    size: int = 5,
) -> dict:
    """
    Semantic search within a memory container namespace.

    Args:
        container_id: The memory container ID.
        namespace: Namespace dict to filter by.
        query: Natural language query for semantic search.
        memory_type: One of "sessions", "working", "long-term", "history".
        size: Max results.
    """
    client = _client()
    must_clauses = [{"term": {f"namespace.{k}": v}} for k, v in namespace.items()]
    body = {
        "query": {
            "bool": {
                "must": must_clauses + [{"match": {"content": query}}],
            }
        },
        "size": size,
    }
    return client.transport.perform_request(
        "GET",
        f"/_plugins/_ml/memory_containers/{container_id}/memories/{memory_type}/_search",
        body=body,
    )