harborrag-engine 2.0.0a1__py3-none-any.whl

harborrag_engine/__init__.py

@@ -0,0 +1,3 @@
+"""HarborRAG business workflows for ingestion and retrieval."""
+
+__all__: list[str] = []

harborrag_engine/agent/__init__.py

@@ -0,0 +1,37 @@
+"""Provider-neutral agent orchestration."""
+
+from harborrag_core.ports.agent_runs import (
+    AgentCheckpoint,
+    AgentRunIdentity,
+    AgentRunRepository,
+    AgentRunStatus,
+    AgentStopReason,
+    AgentToolExecution,
+)
+
+from .events import AgentEvent, AgentEventSink
+from .guard import ExecutionGuard
+from .service import (
+    AgentRunOptions,
+    AgentRunResult,
+    AgentService,
+    AgentToolProvider,
+    AgentToolSpec,
+)
+
+__all__ = [
+    "AgentCheckpoint",
+    "AgentEvent",
+    "AgentEventSink",
+    "AgentRunIdentity",
+    "AgentRunOptions",
+    "AgentRunRepository",
+    "AgentRunResult",
+    "AgentRunStatus",
+    "AgentService",
+    "AgentStopReason",
+    "AgentToolExecution",
+    "AgentToolProvider",
+    "AgentToolSpec",
+    "ExecutionGuard",
+]

harborrag_engine/agent/events.py

@@ -0,0 +1,38 @@
+"""Transport-neutral progress events emitted by the agent loop.
+
+The agent loop knows nothing about SSE, WebSockets, or logging -- it emits
+``AgentEvent`` to an injected sink, and a transport layer decides how to
+render each one. Passing no sink (the default) costs nothing: ``run()``
+and ``resume()`` skip emission entirely when ``events`` is ``None``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class AgentEvent:
+    """One point-in-time occurrence during an agent run.
+
+    ``kind`` is a stable dotted name (``"run.started"``, ``"tool.completed"``,
+    ...) a transport can use as a discriminator; ``data`` carries only
+    JSON-serializable, already-public fields -- the same values that appear
+    in ``AgentRunResult``/``AgentToolExecution``.
+    """
+
+    kind: str
+    run_id: str
+    data: Mapping[str, object]
+
+
+AgentEventSink = Callable[[AgentEvent], Awaitable[None]]
+
+
+async def emit(sink: AgentEventSink | None, event: AgentEvent) -> None:
+    if sink is not None:
+        await sink(event)
+
+
+__all__ = ["AgentEvent", "AgentEventSink", "emit"]

harborrag_engine/agent/execution.py

@@ -0,0 +1,188 @@
+"""Model-turn and tool-call execution primitives used by the agent loop.
+
+Split out of ``loop.py`` so that module stays focused on step orchestration
+(guard checks, checkpointing, stop-reason bookkeeping) while this one owns
+the actual calls out to the chat model and the tool provider.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Sequence
+
+from harborrag_core.models.chat import (
+    HarborChatMessage,
+    HarborChatMetadata,
+    HarborChatRequest,
+    HarborChatResponse,
+    HarborChatTool,
+    HarborToolCall,
+)
+from harborrag_core.ports.agent_runs import AgentToolExecution
+from harborrag_engine.conversation import ConversationIdentity, ConversationMemory, ConversationTurn
+
+from .guard import ExecutionGuard, digest_arguments
+from .protocols import AgentChatModel, AgentToolProvider, AgentToolSpec
+from .schemas import AgentRunOptions
+from .tool_execution import VECTOR_SEARCH_TOOL, bounded_tool_result_content
+
+_BLOCKED_TOOL_NAMES = frozenset({"agent", "chat"})
+_GRAPH_TOOL_PREFIX = "graph_"
+
+
+class ChatAndToolExecutor:
+    """Bind one model turn or one batch of tool calls to their transports."""
+
+    def __init__(
+        self,
+        chat: AgentChatModel,
+        tools: AgentToolProvider,
+        *,
+        memory: ConversationMemory | None,
+    ) -> None:
+        self._chat = chat
+        self._tools = tools
+        self._memory = memory
+
+    def available_specs(
+        self,
+        tenant_id: str,
+        *,
+        graph_search: bool,
+    ) -> list[AgentToolSpec]:
+        return [
+            spec
+            for spec in self._tools.list_tools(tenant_id)
+            if spec.capability == "read"
+            and spec.name not in _BLOCKED_TOOL_NAMES
+            and (graph_search or not spec.name.startswith(_GRAPH_TOOL_PREFIX))
+        ]
+
+    async def complete(
+        self,
+        messages: Sequence[HarborChatMessage],
+        options: AgentRunOptions,
+        *,
+        tools: tuple[HarborChatTool, ...],
+        guard: ExecutionGuard | None = None,
+        completion_token_limit: int | None = None,
+    ) -> HarborChatResponse:
+        request = HarborChatRequest(
+            messages=tuple(messages),
+            tools=tools,
+            parallel_tool_calls=True if tools else None,
+            max_completion_tokens=completion_token_limit,
+            metadata=HarborChatMetadata(
+                tenant_id=options.tenant_id,
+                conversation_id=options.session_id,
+            ),
+            sensitive=True,
+        )
+        coro = self._chat.complete(request)
+        remaining = guard.remaining_seconds() if guard is not None else None
+        if remaining is None:
+            return await coro
+        return await asyncio.wait_for(coro, timeout=remaining)
+
+    async def execute_tool_calls(
+        self,
+        calls: Sequence[HarborToolCall],
+        *,
+        step: int,
+        options: AgentRunOptions,
+        allowed_names: set[str],
+        guard: ExecutionGuard,
+    ) -> tuple[tuple[HarborChatMessage, AgentToolExecution], ...]:
+        coro = asyncio.gather(
+            *(
+                self._execute(call, step=step, options=options, allowed_names=allowed_names)
+                for call in calls
+            )
+        )
+        remaining = guard.remaining_seconds()
+        if remaining is None:
+            return tuple(await coro)
+        return tuple(await asyncio.wait_for(coro, timeout=remaining))
+
+    async def remember(
+        self,
+        identity: ConversationIdentity | None,
+        current_user_message: HarborChatMessage | None,
+        response_text: str,
+    ) -> None:
+        if (
+            self._memory is not None
+            and identity is not None
+            and current_user_message is not None
+            and isinstance(current_user_message.content, str)
+        ):
+            await self._memory.append(
+                identity,
+                ConversationTurn(current_user_message.content, response_text),
+            )
+
+    async def _execute(
+        self,
+        call: HarborToolCall,
+        *,
+        step: int,
+        options: AgentRunOptions,
+        allowed_names: set[str],
+    ) -> tuple[HarborChatMessage, AgentToolExecution]:
+        name = call.function.name
+        arguments = call.function.parsed_arguments
+        digest = digest_arguments(
+            arguments if isinstance(arguments, dict) else {"__unparsed__": call.function.arguments}
+        )
+        result = await self._invoke(name, arguments, options=options, allowed_names=allowed_names)
+        content = bounded_tool_result_content(result)
+        return (
+            HarborChatMessage.tool(content, tool_call_id=call.id, name=name),
+            AgentToolExecution(
+                step=step,
+                call_id=call.id,
+                tool=name,
+                ok=result.get("ok") is True,
+                arguments_digest=digest,
+            ),
+        )
+
+    async def _invoke(
+        self,
+        name: str,
+        arguments: object,
+        *,
+        options: AgentRunOptions,
+        allowed_names: set[str],
+    ) -> dict[str, object]:
+        """Run one tool call, or explain why it can't run, without ever raising."""
+
+        if not isinstance(arguments, dict):
+            return {"ok": False, "error": "invalid tool arguments"}
+        if name not in allowed_names:
+            return {"ok": False, "error": "tool is not available to this agent"}
+        try:
+            return await self._tools.call_tool(
+                name,
+                self._scoped_arguments(name, arguments, options),
+                principal_id=options.principal_id,
+            )
+        except Exception:  # noqa: BLE001 - tool failures become model-visible data
+            return {"ok": False, "error": "tool call failed"}
+
+    @staticmethod
+    def _scoped_arguments(
+        name: str,
+        arguments: dict[str, object],
+        options: AgentRunOptions,
+    ) -> dict[str, object]:
+        """Bind a tool call to its caller's tenant, never trusting the model for it."""
+
+        scoped = dict(arguments)
+        scoped["tenant_id"] = options.tenant_id
+        if name == VECTOR_SEARCH_TOOL and not options.graph_search:
+            scoped["observe_graph"] = False
+        return scoped
+
+
+__all__ = ["ChatAndToolExecutor"]

harborrag_engine/agent/guard.py

@@ -0,0 +1,91 @@
+"""Runaway-loop protection for the bounded agent loop.
+
+``AgentRunOptions.max_steps`` alone does not stop a pathological run: a model
+can spend its entire step budget calling the same tool with the same
+arguments, or one step can hang indefinitely on a slow tool. ``ExecutionGuard``
+adds the two checks a step-count bound cannot express: a wall-clock deadline
+and repeated-identical-tool-call detection.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import time
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from json import dumps
+
+from harborrag_core.ports.agent_runs import AgentToolExecution
+
+
+def digest_arguments(arguments: dict[str, object]) -> str:
+    """Stable hash of tool-call arguments, used as a repeat-detection key.
+
+    Only the digest is ever persisted or compared -- never the raw
+    arguments -- so this is safe to store in a public execution trace.
+    """
+
+    canonical = dumps(arguments, sort_keys=True, separators=(",", ":"), default=str)
+    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
+
+
+@dataclass(slots=True)
+class ExecutionGuard:
+    """Bound wall-clock time and repeated tool calls within one run attempt.
+
+    A fresh guard is armed for each call to ``run()`` or ``resume()`` --
+    ``timeout_seconds`` is a budget for *this* execution attempt, not the
+    run's total lifetime, since a resumed run may follow an arbitrarily long
+    gap after a crash.
+    """
+
+    timeout_seconds: float | None = None
+    max_repeated_tool_calls: int = 2
+    _deadline: float | None = field(default=None, init=False, repr=False)
+    _call_counts: dict[tuple[str, str], int] = field(
+        default_factory=dict,
+        init=False,
+        repr=False,
+    )
+
+    def __post_init__(self) -> None:
+        if self.timeout_seconds is not None and self.timeout_seconds <= 0:
+            raise ValueError("agent timeout_seconds must be positive")
+        if self.max_repeated_tool_calls < 1:
+            raise ValueError("agent max_repeated_tool_calls must be at least 1")
+
+    def start(self) -> None:
+        """Arm the deadline. Call exactly once at the start of an execution attempt."""
+
+        self._deadline = (
+            None if self.timeout_seconds is None else time.monotonic() + self.timeout_seconds
+        )
+
+    def remaining_seconds(self) -> float | None:
+        """Seconds left before the deadline, or ``None`` when unbounded."""
+
+        if self._deadline is None:
+            return None
+        return max(0.0, self._deadline - time.monotonic())
+
+    def timed_out(self) -> bool:
+        remaining = self.remaining_seconds()
+        return remaining is not None and remaining <= 0
+
+    def observe_tool_call(self, tool: str, arguments_digest: str) -> bool:
+        """Record one tool call; return ``True`` once its repeat limit is exceeded."""
+
+        key = (tool, arguments_digest)
+        count = self._call_counts.get(key, 0) + 1
+        self._call_counts[key] = count
+        return count > self.max_repeated_tool_calls
+
+    def replay(self, executions: Iterable[AgentToolExecution]) -> None:
+        """Rebuild repeated-call state from a checkpoint's executions on resume."""
+
+        for execution in executions:
+            key = (execution.tool, execution.arguments_digest)
+            self._call_counts[key] = self._call_counts.get(key, 0) + 1
+
+
+__all__ = ["ExecutionGuard", "digest_arguments"]

harborrag_engine/agent/helpers.py

@@ -0,0 +1,68 @@
+"""Small pure helpers shared across the agent loop."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from harborrag_core.models.chat import HarborChatMessage, HarborChatUsage, MessageRole
+from harborrag_engine.conversation import ConversationTurn
+
+from .schemas import AgentRunOptions
+
+
+def validate_options(options: AgentRunOptions) -> None:
+    for field_name, value in (
+        ("tenant_id", options.tenant_id),
+        ("principal_id", options.principal_id),
+        ("session_id", options.session_id),
+    ):
+        if not value.strip() or len(value) > 128:
+            raise ValueError(f"agent {field_name} must contain 1 to 128 non-whitespace characters")
+    if not 1 <= options.max_steps <= 8:
+        raise ValueError("agent max_steps must be between 1 and 8")
+    if options.timeout_seconds is not None and options.timeout_seconds <= 0:
+        raise ValueError("agent timeout_seconds must be positive")
+    if options.max_repeated_tool_calls < 1:
+        raise ValueError("agent max_repeated_tool_calls must be at least 1")
+    if options.synthesis_timeout_seconds is not None and options.synthesis_timeout_seconds <= 0:
+        raise ValueError("agent synthesis_timeout_seconds must be positive")
+    if options.max_total_tokens is not None and options.max_total_tokens < 1:
+        raise ValueError("agent max_total_tokens must be at least 1")
+
+
+def add_usage(left: HarborChatUsage, right: HarborChatUsage) -> HarborChatUsage:
+    def optional_sum(name: str) -> int | None:
+        first = getattr(left, name)
+        second = getattr(right, name)
+        return None if first is None and second is None else (first or 0) + (second or 0)
+
+    return HarborChatUsage(
+        prompt_tokens=left.prompt_tokens + right.prompt_tokens,
+        completion_tokens=left.completion_tokens + right.completion_tokens,
+        total_tokens=left.total_tokens + right.total_tokens,
+        cache_read_input_tokens=optional_sum("cache_read_input_tokens"),
+        cache_creation_input_tokens=optional_sum("cache_creation_input_tokens"),
+        reasoning_tokens=optional_sum("reasoning_tokens"),
+    )
+
+
+def turn_messages(turns: Sequence[ConversationTurn]) -> list[HarborChatMessage]:
+    messages: list[HarborChatMessage] = []
+    for turn in turns:
+        messages.extend(
+            (
+                HarborChatMessage.user(turn.user_content),
+                HarborChatMessage.assistant(turn.assistant_content),
+            )
+        )
+    return messages
+
+
+def last_user_message(messages: Sequence[HarborChatMessage]) -> HarborChatMessage | None:
+    return next(
+        (message for message in reversed(messages) if message.role is MessageRole.USER),
+        None,
+    )
+
+
+__all__ = ["add_usage", "last_user_message", "turn_messages", "validate_options"]