115 lines
4.5 KiB
Python
115 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import multiprocessing
|
|
import time
|
|
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 {},
|
|
}
|
|
|
|
|
|
def _knowledge_search_worker(connection, arguments: Dict[str, Any]) -> None:
|
|
try:
|
|
request = LibraryContextRequest(
|
|
prompt=arguments["prompt"],
|
|
top_k=arguments.get("top_k", 5),
|
|
embed_model=arguments.get("embedding_model"),
|
|
)
|
|
connection.send(("ok", library_context(arguments["library_slug"], request)))
|
|
except BaseException as exc:
|
|
detail = getattr(exc, "detail", None) or str(exc)
|
|
connection.send(("error", f"{type(exc).__name__}: {detail}"))
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
async def _isolated_library_context(arguments: Dict[str, Any], context: ToolExecutionContext) -> Dict[str, Any]:
|
|
process_context = multiprocessing.get_context("spawn")
|
|
parent_connection, child_connection = process_context.Pipe(duplex=False)
|
|
process = process_context.Process(target=_knowledge_search_worker, args=(child_connection, arguments), daemon=True)
|
|
process.start()
|
|
child_connection.close()
|
|
deadline = time.monotonic() + 170
|
|
try:
|
|
while True:
|
|
if context.cancellation_event.is_set():
|
|
raise asyncio.CancelledError()
|
|
if parent_connection.poll():
|
|
status, value = parent_connection.recv()
|
|
if status == "ok":
|
|
return value
|
|
raise RuntimeError(value)
|
|
if not process.is_alive():
|
|
raise RuntimeError(f"Local retrieval process exited unexpectedly (exit code {process.exitcode}).")
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError("Local retrieval exceeded its process timeout.")
|
|
await asyncio.sleep(0.05)
|
|
finally:
|
|
parent_connection.close()
|
|
if process.is_alive():
|
|
process.terminate()
|
|
process.join(timeout=2)
|
|
|
|
|
|
async def knowledge_search_handler(arguments: Dict[str, Any], context: ToolExecutionContext) -> Dict[str, Any]:
|
|
payload = await _isolated_library_context(arguments, context)
|
|
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 Heimgeist's indexed local Knowledge store.",
|
|
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,
|
|
))
|