session_memory.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
from __future__ import annotations
 
import time
from typing import Any
 
from .registry import register
 
 
@register("write_session_memory", {
    "type": "function",
    "function": {
        "name": "write_session_memory",
        "description": "Persist a key/value pair to the session memory. Use to remember facts across turns (e.g. which features were interesting, user preferences, last findings).",
        "parameters": {
            "type": "object",
            "properties": {
                "key": {"type": "string"},
                "value": {},
            },
            "required": ["key", "value"],
        },
    },
})
def write_session_memory(args: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    key = args.get("key")
    value = args.get("value")
    if not key:
        return {"error": "key is required"}
 
    from aquin.engine.session_memory_store import patch_local_memory
 
    mem_key = ctx.get("session_id") or "local"
    patch_local_memory(mem_key, key, value)
 
    # Update local state so read_session_memory works within the same turn
    state: dict[str, Any] = ctx.get("state", {})
    mem: dict[str, Any] = state.get("memory") or {}
    mem[key] = value
    state["memory"] = mem
 
    return {"ok": True, "key": key, "written_at": int(time.time())}
 
 
@register("read_session_memory", {
    "type": "function",
    "function": {
        "name": "read_session_memory",
        "description": "Read a previously stored key from session memory.",
        "parameters": {
            "type": "object",
            "properties": {
                "key": {"type": "string"},
            },
            "required": ["key"],
        },
    },
})
def read_session_memory(args: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    key = args.get("key")
    if not key:
        return {"error": "key is required"}
 
    from aquin.engine.session_memory_store import load_local_memory
 
    mem_key = ctx.get("session_id") or "local"
    local_mem = dict(ctx.get("state", {}).get("memory") or {})
    mem = {**load_local_memory(mem_key), **local_mem}
 
    if key not in mem:
        return {"found": False, "key": key, "value": None}
 
    return {"found": True, "key": key, "value": mem[key]}