Implement chat memory tool and integrate it into agent capabilities

This commit is contained in:
2026-06-16 19:52:19 +02:00
parent be7e1c7c49
commit 05e8b0296f
3 changed files with 99 additions and 1 deletions

View File

@@ -76,7 +76,7 @@ def _workflow_result(item: Dict[str, Any], *, confidence: float, reason: str, in
def _capability_enabled(capability: str, *, library_slug: Optional[str], web_search_enabled: bool, has_attachments: bool) -> bool:
if capability in {"chat"}:
if capability in {"chat", "chat_memory"}:
return True
if capability == "web":
return True

View File

@@ -1,5 +1,6 @@
from .chat import register_chat_tools
from .knowledge import register_knowledge_tools
from .memory import register_memory_tools
from .saving import register_saving_tools
from .vision import register_vision_tools
from .web import register_web_tools
@@ -8,6 +9,7 @@ from .web import register_web_tools
def register_native_tools(registry) -> None:
register_chat_tools(registry)
register_knowledge_tools(registry)
register_memory_tools(registry)
register_web_tools(registry)
register_vision_tools(registry)
register_saving_tools(registry)

View File

@@ -0,0 +1,96 @@
from __future__ import annotations
from typing import Any, Dict
from ...chat_memory import build_chat_memory_context, ensure_chat_memory_schema
from ..registry import NativeToolProvider, ToolDefinition, ToolExecutionContext
def _int_argument(arguments: Dict[str, Any], key: str, default: int, minimum: int, maximum: int) -> int:
try:
value = int(arguments.get(key) or default)
except Exception:
value = default
return max(minimum, min(value, maximum))
def _memory_hit(item: Dict[str, Any]) -> Dict[str, Any]:
return {
"session_id": item.get("session_id"),
"session_name": item.get("session_name"),
"user_message_id": item.get("user_message_id"),
"assistant_message_id": item.get("assistant_message_id"),
"created_at": item.get("created_at"),
"score": round(float(item.get("_score") or item.get("score") or 0), 4),
"user_excerpt": item.get("user_content"),
"assistant_excerpt": item.get("assistant_content"),
}
async def chat_memory_search_handler(arguments: Dict[str, Any], context: ToolExecutionContext) -> Dict[str, Any]:
top_k = _int_argument(arguments, "top_k", 6, 1, 20)
context_character_budget = _int_argument(arguments, "context_character_budget", 10_000, 500, 60_000)
mode = str(arguments.get("mode") or "search").strip().lower()
if mode not in {"search", "recent"}:
mode = "search"
exclude_current_session = bool(arguments.get("exclude_current_session", True))
exclude_session_id = context.session_id if exclude_current_session else None
db = context.db_factory()
try:
ensure_chat_memory_schema(db.get_bind())
payload = build_chat_memory_context(
db,
str(arguments.get("prompt") or ""),
exclude_session_id=exclude_session_id,
top_k=top_k,
context_character_budget=context_character_budget,
mode=mode,
)
finally:
db.close()
sources = payload.get("sources") or []
return {
"hits": [_memory_hit(item) for item in (payload.get("hits") or [])],
"context_block": str(payload.get("context_block") or ""),
"sources": sources,
"scores": [source.get("score") for source in sources if isinstance(source, dict)],
}
def register_memory_tools(registry: NativeToolProvider) -> None:
registry.register(ToolDefinition(
name="heimgeist.chat_memory_search",
description=(
"Retrieve relevant completed turns from past Heimgeist chats. Use only when a workflow or agent "
"decides previous chats are needed, such as questions about earlier conversations, prior decisions, "
"or remembered discussion context; it does not run automatically."
),
input_schema={
"type": "object",
"properties": {
"prompt": {"type": "string", "minLength": 1},
"top_k": {"type": "integer", "minimum": 1, "maximum": 20},
"context_character_budget": {"type": "integer", "minimum": 500, "maximum": 60000},
"mode": {"type": "string", "enum": ["search", "recent"]},
"exclude_current_session": {"type": "boolean"},
},
"required": ["prompt"],
"additionalProperties": False,
},
output_schema={
"type": "object",
"properties": {
"hits": {"type": "array"},
"context_block": {"type": "string"},
"sources": {"type": "array"},
"scores": {"type": "array"},
},
"required": ["hits", "context_block", "sources", "scores"],
"additionalProperties": False,
},
handler=chat_memory_search_handler,
timeout_seconds=20,
result_size_limit=100_000,
))