giers10/Heimgeist
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Implement core agent tooling infrastructure (chat, knowledge search) and update tool registration

Browse Source
This commit is contained in:
Victor Giers
2026-06-15 15:03:22 +02:00
parent 0c04d5bf7d
commit 97366bb8e0
4 changed files with 378 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
backend/agent/ollama_runtime.py Normal file
View File

@@ -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,
}

13
backend/agent/tools/__init__.py Normal file
View File

@@ -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)

190
backend/agent/tools/chat.py Normal file
View File

@@ -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,
))

74
backend/agent/tools/knowledge.py Normal file
View File

@@ -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,
))