retrieve_memory_tool.py

"""RetrieveMemoryTool — mid-iteration Viking memory retrieval for GroundingAgent.

Registered as an internal tool alongside ``retrieve_skill`` so the LLM
can pull Viking memories (tool knowledge, patterns, past cases,
anti-patterns) during execution when the initial enrichment is
insufficient.

The pre-execution enrichment in ``tool_layer.execute()`` injects L0
abstracts into iter-1 and strips them at iter-2. This tool gives the
agent a way to re-query dynamically if it discovers mid-execution that
it needs additional cross-session context — e.g. when a tool fails
unexpectedly and the agent wants to check for known issues.

Unlike the pre-execution path, this tool accepts an optional
``category`` argument so the agent can target the most relevant
memory bucket (``tools`` for fail reasons, ``antipatterns`` to check
if an approach has been tried before, etc.).
"""
from __future__ import annotations

from typing import TYPE_CHECKING, List, Optional

from openspace.grounding.core.tool.local_tool import LocalTool
from openspace.grounding.core.types import BackendType
from openspace.utils.logging import Logger

if TYPE_CHECKING:
    from .client import OpenVikingClient

logger = Logger.get_logger(__name__)


_VALID_CATEGORIES = {
    "tools", "patterns", "skills", "cases", "preferences", "antipatterns",
}


class RetrieveMemoryTool(LocalTool):
    """Internal tool: mid-iteration cross-session memory retrieval.

    Thin wrapper around :class:`OpenVikingClient` — formats the result
    as a plain-text block suitable for direct injection into the LLM
    context. Fails gracefully (returns an explanatory string, never
    raises) so the iteration loop is never broken by Viking issues.
    """

    _name = "retrieve_memory"
    _description = (
        "Query the cross-session memory store (OpenViking) for relevant "
        "knowledge when the current approach isn't working. Returns "
        "abstracts from prior tasks: tool knowledge, patterns, past "
        "cases, known failure modes (anti-patterns), and user preferences. "
        "Use this when encountering an unexpected error, before trying a "
        "risky approach, or when you need user preference context mid-task. "
        "Arguments: query (required), category (optional: tools, patterns, "
        "skills, cases, preferences, antipatterns — omit to search all)."
    )
    backend_type = BackendType.SYSTEM

    def __init__(
        self,
        viking_client: "OpenVikingClient",
        score_threshold: float = 0.0,
    ):
        super().__init__()
        self._client = viking_client
        self._score_threshold = score_threshold

    async def _arun(
        self,
        query: str,
        category: Optional[str] = None,
        limit: int = 5,
    ) -> str:
        if not query or not isinstance(query, str):
            return "retrieve_memory error: query is required."

        limit = max(1, min(int(limit or 5), 10))

        try:
            if not await self._client.is_available():
                return "retrieve_memory: OpenViking server is not available."
        except Exception as exc:
            return f"retrieve_memory: health check failed: {exc}"

        # Normalize category → target_uri
        norm: Optional[str] = None
        target_uri = ""
        if category:
            raw = category.strip().lower().replace("-", "_")
            if raw in _VALID_CATEGORIES:
                norm = raw
                if raw == "preferences":
                    target_uri = self._client.user_memory_uri("preferences")
                else:
                    target_uri = self._client.agent_memory_uri(raw)
            else:
                return (
                    f"retrieve_memory error: invalid category '{category}'. "
                    f"Valid: {sorted(_VALID_CATEGORIES)}."
                )

        try:
            results = await self._client.find_memories(
                query,
                target_uri=target_uri,
                limit=limit,
                score_threshold=self._score_threshold,
            )
        except Exception as exc:
            return f"retrieve_memory error: {exc}"

        if not results:
            scope = norm or "all categories"
            return f"retrieve_memory: no matches found for '{query[:80]}' in {scope}."

        lines: List[str] = [
            f"# Retrieved from OpenViking — query: {query[:80]!r}",
            f"# Category: {norm or 'all'}, {len(results)} result(s)",
            "",
        ]
        for i, m in enumerate(results, 1):
            abstract = (
                m.get("abstract") or m.get("summary")
                or m.get("content") or m.get("uri", "")
            )
            score = m.get("score", 0.0)
            uri = m.get("uri", "")
            lines.append(f"{i}. [score={score:.2f}] {abstract}")
            if uri:
                lines.append(f"   uri: {uri}")

        if norm == "antipatterns":
            lines.insert(
                3,
                "## WARNING: these are FAILED approaches. Do NOT repeat them "
                "unless you have evidence they will work now.\n"
            )

        return "\n".join(lines)