123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
from app.tools.base import BaseTool
|
|
from app.core.contracts import ToolResult, UserTask
|
|
from app.tools.sandbox import ToolSandbox
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MemoryInsertTool(BaseTool):
|
|
_name = "memory_insert"
|
|
_description = "Store information in memory"
|
|
|
|
def __init__(self, sandbox: ToolSandbox, memory_interface=None) -> None:
|
|
super().__init__()
|
|
self._sandbox = sandbox
|
|
self._memory = memory_interface
|
|
|
|
def execute(self, task: UserTask, args: dict[str, Any]) -> ToolResult:
|
|
text = args.get("text", "")
|
|
kind = args.get("kind", "fact")
|
|
source = args.get("source", "user")
|
|
weight = args.get("weight", 0.5)
|
|
|
|
if not text:
|
|
return ToolResult(tool="memory_insert", ok=False, output="", error="text is required")
|
|
if not self._memory:
|
|
return ToolResult(tool="memory_insert", ok=False, output="", error="Memory not available")
|
|
|
|
try:
|
|
entry = self._memory.insert(
|
|
text=text,
|
|
kind=kind,
|
|
source=source,
|
|
task_id=task.task_id,
|
|
session_id=task.session_id,
|
|
weight=weight,
|
|
)
|
|
return ToolResult(
|
|
tool="memory_insert",
|
|
ok=True,
|
|
output=f"Stored: {entry.id}",
|
|
metadata={"entry_id": entry.id},
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Memory insert failed: {e}")
|
|
return ToolResult(tool="memory_insert", ok=False, output="", error=str(e))
|
|
|
|
|
|
class MemorySearchTool(BaseTool):
|
|
_name = "memory_search"
|
|
_description = "Search memory for information"
|
|
|
|
def __init__(self, sandbox: ToolSandbox, memory_interface=None) -> None:
|
|
super().__init__()
|
|
self._sandbox = sandbox
|
|
self._memory = memory_interface
|
|
|
|
def execute(self, task: UserTask, args: dict[str, Any]) -> ToolResult:
|
|
query = args.get("query", "")
|
|
top_k = args.get("top_k", 5)
|
|
|
|
if not query:
|
|
return ToolResult(tool="memory_search", ok=False, output="", error="query is required")
|
|
if not self._memory:
|
|
return ToolResult(tool="memory_search", ok=False, output="", error="Memory not available")
|
|
|
|
try:
|
|
results = self._memory.search(query, top_k=top_k)
|
|
if not results:
|
|
return ToolResult(tool="memory_search", ok=True, output="No results found", metadata={"count": 0})
|
|
|
|
output_lines = []
|
|
for entry, score in results:
|
|
output_lines.append(f"[{score:.2f}] {entry.text[:100]}")
|
|
|
|
return ToolResult(
|
|
tool="memory_search",
|
|
ok=True,
|
|
output="\n".join(output_lines),
|
|
metadata={"count": len(results)},
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Memory search failed: {e}")
|
|
return ToolResult(tool="memory_search", ok=False, output="", error=str(e))
|
|
|
|
|
|
class MemoryListTool(BaseTool):
|
|
_name = "memory_list"
|
|
_description = "List recent memories"
|
|
|
|
def __init__(self, sandbox: ToolSandbox, memory_interface=None) -> None:
|
|
super().__init__()
|
|
self._sandbox = sandbox
|
|
self._memory = memory_interface
|
|
|
|
def execute(self, task: UserTask, args: dict[str, Any]) -> ToolResult:
|
|
limit = args.get("limit", 10)
|
|
|
|
if not self._memory:
|
|
return ToolResult(tool="memory_list", ok=False, output="", error="Memory not available")
|
|
|
|
try:
|
|
entries = self._memory.get_recent(limit=limit)
|
|
if not entries:
|
|
return ToolResult(tool="memory_list", ok=True, output="No memories", metadata={"count": 0})
|
|
|
|
output_lines = []
|
|
for entry in entries:
|
|
output_lines.append(f"{entry.kind}: {entry.text[:80]}")
|
|
|
|
return ToolResult(
|
|
tool="memory_list",
|
|
ok=True,
|
|
output="\n".join(output_lines),
|
|
metadata={"count": len(entries)},
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Memory list failed: {e}")
|
|
return ToolResult(tool="memory_list", ok=False, output="", error=str(e)) |