Add tools for knowledge persistence (saving messages/websites) and visual attachment analysis.

2026-06-15 15:04:56 +02:00
parent b2f1aeba42
commit 48cf63d535
2 changed files with 135 additions and 0 deletions
--- a/backend/agent/tools/saving.py
+++ b/backend/agent/tools/saving.py
@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from ...local_rag import CreateWebsiteRequest, create_library_website
+from ..registry import NativeToolProvider, ToolDefinition, ToolExecutionContext
+
+
+async def save_message_handler(arguments: Dict[str, Any], context: ToolExecutionContext) -> Dict[str, Any]:
+    from ... import main as main_module
+
+    db = context.db_factory()
+    try:
+        result = await main_module._save_message_to_knowledge_impl(
+            arguments["library_slug"],
+            arguments["message_id"],
+            db,
+            title=arguments.get("title"),
+            content=arguments.get("edited_content"),
+            saved_by="agent" if context.selection_mode == "automatic" else "user",
+        )
+    finally:
+        db.close()
+    return {"status": "saved", "confirmed": True, "result": result}
+
+
+async def save_website_handler(arguments: Dict[str, Any], context: ToolExecutionContext) -> Dict[str, Any]:
+    result = await create_library_website(
+        arguments["library_slug"],
+        CreateWebsiteRequest(url=arguments["url"], title=arguments.get("title")),
+    )
+    return {"status": "saved", "confirmed": True, "result": result}
+
+
+SAVE_OUTPUT = {
+    "type": "object",
+    "properties": {
+        "status": {"type": "string"},
+        "confirmed": {"type": "boolean"},
+        "result": {"type": "object"},
+    },
+    "required": ["status", "confirmed", "result"],
+    "additionalProperties": False,
+}
+
+
+def register_saving_tools(registry: NativeToolProvider) -> None:
+    registry.register(ToolDefinition(
+        name="heimgeist.save_message_to_knowledge",
+        description="Save a stable chat message snapshot to a Heimgeist knowledge library.",
+        input_schema={
+            "type": "object",
+            "properties": {
+                "message_id": {"type": "string", "minLength": 1},
+                "library_slug": {"type": "string", "minLength": 1},
+                "title": {"type": ["string", "null"]},
+                "edited_content": {"type": ["string", "null"]},
+            },
+            "required": ["message_id", "library_slug"],
+            "additionalProperties": False,
+        },
+        output_schema=SAVE_OUTPUT,
+        handler=save_message_handler,
+        risk_level="write",
+        requires_confirmation=True,
+        timeout_seconds=180,
+    ))
+    registry.register(ToolDefinition(
+        name="heimgeist.save_website_to_knowledge",
+        description="Fetch and save a durable website snapshot to a Heimgeist knowledge library.",
+        input_schema={
+            "type": "object",
+            "properties": {
+                "url": {"type": "string", "minLength": 1},
+                "library_slug": {"type": "string", "minLength": 1},
+                "title": {"type": ["string", "null"]},
+            },
+            "required": ["url", "library_slug"],
+            "additionalProperties": False,
+        },
+        output_schema=SAVE_OUTPUT,
+        handler=save_website_handler,
+        risk_level="write",
+        requires_confirmation=True,
+        timeout_seconds=180,
+    ))
--- a/backend/agent/tools/vision.py
+++ b/backend/agent/tools/vision.py
@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from ..registry import NativeToolProvider, ToolDefinition, ToolExecutionContext
+
+
+async def vision_analyze_handler(arguments: Dict[str, Any], _context: ToolExecutionContext) -> Dict[str, Any]:
+    from ... import main as main_module
+
+    prepared, _images, context_block = await main_module._prepare_chat_message_attachments(
+        arguments.get("attachments") or [],
+        request_model_supports_vision=False,
+        vision_model=arguments.get("vision_model"),
+        transcription_model=arguments.get("transcription_model"),
+        persist_file_text=False,
+    )
+    return {
+        "context_block": context_block,
+        "attachments": [main_module._attachment_history_payload(item) for item in prepared],
+        "sources": [],
+    }
+
+
+def register_vision_tools(registry: NativeToolProvider) -> None:
+    registry.register(ToolDefinition(
+        name="heimgeist.vision_analyze",
+        description="Analyze chat attachments through Heimgeist's existing vision and extraction path.",
+        input_schema={
+            "type": "object",
+            "properties": {
+                "attachments": {"type": "array"},
+                "vision_model": {"type": "string", "minLength": 1},
+                "transcription_model": {"type": ["string", "null"]},
+            },
+            "required": ["attachments", "vision_model"],
+            "additionalProperties": False,
+        },
+        output_schema={
+            "type": "object",
+            "properties": {"context_block": {"type": "string"}, "attachments": {"type": "array"}, "sources": {"type": "array"}},
+            "required": ["context_block", "attachments", "sources"],
+            "additionalProperties": False,
+        },
+        handler=vision_analyze_handler,
+        timeout_seconds=600,
+        result_size_limit=140_000,
+        llm_call=True,
+    ))