Add tools for knowledge persistence (saving messages/websites) and visual attachment analysis.
This commit is contained in:
86
backend/agent/tools/saving.py
Normal file
86
backend/agent/tools/saving.py
Normal file
@@ -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,
|
||||||
|
))
|
||||||
49
backend/agent/tools/vision.py
Normal file
49
backend/agent/tools/vision.py
Normal file
@@ -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,
|
||||||
|
))
|
||||||
Reference in New Issue
Block a user