From 97366bb8e036eb0fbbdc75c18acc97ac31b9de98 Mon Sep 17 00:00:00 2001 From: Victor Giers Date: Mon, 15 Jun 2026 15:03:22 +0200 Subject: [PATCH] Implement core agent tooling infrastructure (chat, knowledge search) and update tool registration --- backend/agent/ollama_runtime.py | 101 ++++++++++++++++ backend/agent/tools/__init__.py | 13 +++ backend/agent/tools/chat.py | 190 +++++++++++++++++++++++++++++++ backend/agent/tools/knowledge.py | 74 ++++++++++++ 4 files changed, 378 insertions(+) create mode 100644 backend/agent/ollama_runtime.py create mode 100644 backend/agent/tools/__init__.py create mode 100644 backend/agent/tools/chat.py create mode 100644 backend/agent/tools/knowledge.py diff --git a/backend/agent/ollama_runtime.py b/backend/agent/ollama_runtime.py new file mode 100644 index 0000000..9847a5b --- /dev/null +++ b/backend/agent/ollama_runtime.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import json +from typing import Any, Dict, List + +from ..ollama_client import chat_typed +from .registry import NativeToolProvider, ToolExecutionContext + + +def _usage_dict(usage: Any) -> Dict[str, int]: + return { + "prompt_eval_count": usage.prompt_eval_count, + "eval_count": usage.eval_count, + "total_duration": usage.total_duration, + "load_duration": usage.load_duration, + "prompt_eval_duration": usage.prompt_eval_duration, + "eval_duration": usage.eval_duration, + } + + +async def run_agent_loop(arguments: Dict[str, Any], context: ToolExecutionContext) -> Dict[str, Any]: + registry = context.registry + if registry is None: + raise RuntimeError("Agent tool registry is unavailable.") + + allowed_names = [] + for name in arguments.get("allowed_tool_names") or []: + if name == "heimgeist.agent": + continue + registry.get_tool(name) + allowed_names.append(name) + tools = [registry.get_tool(name).ollama_manifest() for name in allowed_names] + messages: List[Dict[str, Any]] = [dict(item) for item in arguments.get("messages") or []] + maximum_rounds = max(1, min(int(arguments.get("maximum_rounds") or 4), 8)) + maximum_output_chars = max(1000, min(int(arguments.get("maximum_output_chars") or 60_000), 120_000)) + tool_summary: List[Dict[str, Any]] = [] + usage_totals: Dict[str, int] = {} + + for round_index in range(maximum_rounds): + if context.cancellation_event.is_set(): + raise RuntimeError("Agent execution was cancelled.") + if context.consume_llm_call: + context.consume_llm_call() + result = await chat_typed( + arguments["model"], + messages, + options=arguments.get("generation_options") or {}, + tools=tools or None, + think=bool(arguments.get("reasoning")), + cancellation_event=context.cancellation_event, + ) + for key, value in _usage_dict(result.usage).items(): + usage_totals[key] = usage_totals.get(key, 0) + int(value or 0) + if result.thinking and context.show_thinking: + await context.emit("model_thinking", {"text": result.thinking, "round": round_index + 1}) + + if not result.tool_calls: + content = result.content[:maximum_output_chars] + return { + "content": content, + "tool_summary": tool_summary, + "usage": usage_totals, + "rounds": round_index + 1, + } + + assistant_message: Dict[str, Any] = { + "role": "assistant", + "content": result.content or "", + "tool_calls": [call.raw for call in result.tool_calls], + } + messages.append(assistant_message) + for call in result.tool_calls: + if call.name not in allowed_names: + tool_result = {"error": f"Tool is not allowed: {call.name}"} + else: + definition = registry.get_tool(call.name) + approved = True + if definition.requires_confirmation and context.request_confirmation: + approved = await context.request_confirmation(definition, call.arguments) + if approved: + await context.emit("tool_started", {"tool": call.name, "arguments": call.arguments, "agent_round": round_index + 1}) + try: + tool_result = await registry.call_tool(call.name, call.arguments, context) + except Exception as exc: + tool_result = {"error": f"{type(exc).__name__}: {exc}"} + await context.emit("tool_result", {"tool": call.name, "result": tool_result, "agent_round": round_index + 1}) + else: + tool_result = {"status": "rejected", "confirmed": False} + tool_summary.append({"tool": call.name, "arguments": call.arguments, "result": tool_result}) + messages.append({ + "role": "tool", + "tool_name": call.name, + "content": json.dumps(tool_result, ensure_ascii=False, default=str), + }) + + return { + "content": "I could not complete the task within the configured agent round limit.", + "tool_summary": tool_summary, + "usage": usage_totals, + "rounds": maximum_rounds, + } diff --git a/backend/agent/tools/__init__.py b/backend/agent/tools/__init__.py new file mode 100644 index 0000000..185cfd1 --- /dev/null +++ b/backend/agent/tools/__init__.py @@ -0,0 +1,13 @@ +from .chat import register_chat_tools +from .knowledge import register_knowledge_tools +from .saving import register_saving_tools +from .vision import register_vision_tools +from .web import register_web_tools + + +def register_native_tools(registry) -> None: + register_chat_tools(registry) + register_knowledge_tools(registry) + register_web_tools(registry) + register_vision_tools(registry) + register_saving_tools(registry) diff --git a/backend/agent/tools/chat.py b/backend/agent/tools/chat.py new file mode 100644 index 0000000..f5839c0 --- /dev/null +++ b/backend/agent/tools/chat.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +from typing import Any, Dict, List + +from ... import models +from ...ollama_client import chat_stream_typed +from ..ollama_runtime import run_agent_loop +from ..registry import NativeToolProvider, ToolDefinition, ToolExecutionContext + + +MESSAGE_SCHEMA = { + "type": "object", + "properties": { + "role": {"type": "string"}, + "content": {"type": "string"}, + "images": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["role", "content"], + "additionalProperties": True, +} + + +def _usage_dict(usage: Any) -> Dict[str, int]: + return { + "prompt_eval_count": usage.prompt_eval_count, + "eval_count": usage.eval_count, + "total_duration": usage.total_duration, + "load_duration": usage.load_duration, + "prompt_eval_duration": usage.prompt_eval_duration, + "eval_duration": usage.eval_duration, + } + + +def _context_text(blocks: List[Any]) -> str: + parts: List[str] = [] + for block in blocks or []: + if isinstance(block, str): + text = block + elif isinstance(block, dict): + text = str(block.get("context_block") or block.get("text") or block.get("content") or "") + else: + text = str(block or "") + if text.strip(): + parts.append(text.strip()) + return "\n\n".join(parts) + + +def _collect_sources(blocks: List[Any]) -> List[Any]: + sources: List[Any] = [] + seen = set() + for block in blocks or []: + if not isinstance(block, dict): + continue + for source in block.get("sources") or []: + key = str(source.get("url") if isinstance(source, dict) else source) + if key and key not in seen: + seen.add(key) + sources.append(source) + return sources + + +async def chat_handler(arguments: Dict[str, Any], context: ToolExecutionContext) -> Dict[str, Any]: + messages = [dict(item) for item in arguments.get("messages") or []] + if context.session_id: + db = context.db_factory() + try: + session = db.query(models.ChatSession).filter(models.ChatSession.session_id == context.session_id).first() + rows = [] if not session else ( + db.query(models.ChatMessage) + .filter(models.ChatMessage.session_pk == session.id) + .order_by(models.ChatMessage.created_at.asc(), models.ChatMessage.id.asc()) + .all() + )[-20:] + finally: + db.close() + if rows: + from ... import main as main_module + supports_vision = await main_module._model_supports_vision(arguments["model"]) + override_index = len(rows) - 1 if rows[-1].role == "user" else None + prompt = rows[-1].content if override_index is not None else "" + context_text = _context_text(arguments.get("context_blocks") or []) + override_content = f"{prompt}\n\n{context_text}".strip() if context_text else None + messages = await main_module._build_ollama_messages_from_rows( + rows, + request_model_supports_vision=supports_vision, + vision_model=arguments.get("vision_model"), + transcription_model=arguments.get("transcription_model"), + override_user_row_index=override_index, + override_user_content=override_content, + ) + system_prompt = str(arguments.get("system_prompt") or "").strip() + if system_prompt: + messages.insert(0, {"role": "system", "content": system_prompt}) + + full_content = "" + full_thinking = "" + usage: Dict[str, int] = {} + async for chunk in chat_stream_typed( + arguments["model"], + messages, + options=arguments.get("generation_options") or {}, + think=bool(arguments.get("reasoning")), + cancellation_event=context.cancellation_event, + ): + if chunk.thinking: + full_thinking += chunk.thinking + if context.show_thinking: + await context.emit("model_thinking", {"text": chunk.thinking}) + if chunk.content: + full_content += chunk.content + await context.emit("model_token", {"text": chunk.content}) + if chunk.done: + usage = _usage_dict(chunk.usage) + return { + "content": full_content, + "sources": _collect_sources(arguments.get("context_blocks") or []), + "usage": usage, + "thinking_available": bool(full_thinking), + } + + +def register_chat_tools(registry: NativeToolProvider) -> None: + registry.register(ToolDefinition( + name="heimgeist.chat", + description="Generate an Ollama chat response from conversation history and supplied context blocks.", + input_schema={ + "type": "object", + "properties": { + "model": {"type": "string", "minLength": 1}, + "messages": {"type": "array", "items": MESSAGE_SCHEMA}, + "system_prompt": {"type": "string"}, + "context_blocks": {"type": "array"}, + "attachments": {"type": "array"}, + "vision_model": {"type": ["string", "null"]}, + "transcription_model": {"type": ["string", "null"]}, + "generation_options": {"type": "object"}, + "stream": {"type": "boolean"}, + "reasoning": {"type": "boolean"}, + }, + "required": ["model", "messages"], + "additionalProperties": False, + }, + output_schema={ + "type": "object", + "properties": { + "content": {"type": "string"}, + "sources": {"type": "array"}, + "usage": {"type": "object"}, + "thinking_available": {"type": "boolean"}, + }, + "required": ["content", "sources", "usage", "thinking_available"], + "additionalProperties": False, + }, + handler=chat_handler, + timeout_seconds=600, + result_size_limit=180_000, + llm_call=True, + )) + registry.register(ToolDefinition( + name="heimgeist.agent", + description="Run a bounded Ollama tool-calling loop over an approved native-tool allowlist.", + input_schema={ + "type": "object", + "properties": { + "model": {"type": "string", "minLength": 1}, + "messages": {"type": "array", "items": MESSAGE_SCHEMA}, + "allowed_tool_names": {"type": "array", "items": {"type": "string"}, "maxItems": 12}, + "maximum_rounds": {"type": "integer", "minimum": 1, "maximum": 8}, + "maximum_output_chars": {"type": "integer", "minimum": 1000, "maximum": 120000}, + "reasoning": {"type": "boolean"}, + "generation_options": {"type": "object"}, + }, + "required": ["model", "messages", "allowed_tool_names", "maximum_rounds"], + "additionalProperties": False, + }, + output_schema={ + "type": "object", + "properties": { + "content": {"type": "string"}, + "tool_summary": {"type": "array"}, + "usage": {"type": "object"}, + "rounds": {"type": "integer"}, + }, + "required": ["content", "tool_summary", "usage", "rounds"], + "additionalProperties": False, + }, + handler=run_agent_loop, + timeout_seconds=600, + result_size_limit=180_000, + )) diff --git a/backend/agent/tools/knowledge.py b/backend/agent/tools/knowledge.py new file mode 100644 index 0000000..6a61fb7 --- /dev/null +++ b/backend/agent/tools/knowledge.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import asyncio +from typing import Any, Dict + +from ...local_rag import LibraryContextRequest, library_context +from ..registry import NativeToolProvider, ToolDefinition, ToolExecutionContext + + +def _knowledge_source(source: Dict[str, Any], library_slug: str) -> Dict[str, Any]: + return { + "type": "knowledge", + "library_slug": library_slug, + "doc_id": source.get("doc_id"), + "title": source.get("title"), + "url": source.get("url"), + "record_type": source.get("record_type"), + "mime": source.get("mime"), + "lang": source.get("lang"), + "snippet": source.get("snippet"), + "scores": source.get("scores") or {}, + } + + +async def knowledge_search_handler(arguments: Dict[str, Any], _context: ToolExecutionContext) -> Dict[str, Any]: + request = LibraryContextRequest( + prompt=arguments["prompt"], + top_k=arguments.get("top_k", 5), + embed_model=arguments.get("embedding_model"), + ) + payload = await asyncio.to_thread(library_context, arguments["library_slug"], request) + raw_hits = (payload.get("result") or {}).get("sources") or [] + budget = int(arguments.get("context_character_budget") or 12_000) + context_block = str(payload.get("context_block") or "")[:budget] + sources = [_knowledge_source(item, arguments["library_slug"]) for item in raw_hits] + return { + "hits": sources, + "context_block": context_block, + "sources": sources, + "scores": [item.get("scores") or {} for item in sources], + } + + +def register_knowledge_tools(registry: NativeToolProvider) -> None: + registry.register(ToolDefinition( + name="heimgeist.knowledge_search", + description="Retrieve structured evidence from an indexed Heimgeist knowledge library.", + input_schema={ + "type": "object", + "properties": { + "prompt": {"type": "string", "minLength": 1}, + "library_slug": {"type": "string", "minLength": 1}, + "top_k": {"type": "integer", "minimum": 1, "maximum": 20}, + "context_character_budget": {"type": "integer", "minimum": 500, "maximum": 60000}, + "embedding_model": {"type": ["string", "null"]}, + }, + "required": ["prompt", "library_slug"], + "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=knowledge_search_handler, + timeout_seconds=180, + result_size_limit=100_000, + ))