Skip to content

Orchestrator

The agent loop, tool dispatch, and provider discovery.

apab.agent.orchestrator

Agent orchestrator: LLM ↔ tool-calling loop.

AgentOrchestrator

Drives the agentic LLM ↔ tool loop.

Parameters:

Name Type Description Default
config ProjectConfig

Project configuration.

required
workspace Workspace | None

Workspace manager for run bundles.

None
provider LLMProvider | None

An already-instantiated LLM provider. If None one is created from config.llm.

None
Source code in src/apab/agent/orchestrator.py
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
class AgentOrchestrator:
    """Drives the agentic LLM ↔ tool loop.

    Parameters
    ----------
    config:
        Project configuration.
    workspace:
        Workspace manager for run bundles.
    provider:
        An already-instantiated LLM provider.  If ``None`` one is
        created from *config.llm*.
    """

    def __init__(
        self,
        config: ProjectConfig,
        workspace: Workspace | None = None,
        provider: LLMProvider | None = None,
    ) -> None:
        self.config = config
        self.workspace = workspace or Workspace(Path(config.project.workspace))
        self.provider = provider or get_provider(
            config.llm.provider,
            model=config.llm.model,
            base_url=config.llm.base_url,
        )
        mode = config.llm.redaction_mode
        self.dispatcher = ToolDispatcher(
            redaction_mode=mode.value if hasattr(mode, "value") else str(mode),
        )
        self._messages: list[dict[str, Any]] = []
        self._run_ctx: RunContext | None = None
        self._session_usage: dict[str, Any] = _empty_usage()
        self._trace_id: str | None = None

    # ── session lifecycle ─────────────────────────────────────────────

    def start_session(self, user_request: str) -> RunContext:
        """Begin a new agent session and return a :class:`RunContext`."""
        self.workspace.ensure_dirs()
        self._run_ctx = self.workspace.new_run()
        self._session_usage = _empty_usage()

        tool_schemas = self.dispatcher.get_tool_schemas()
        tool_names = [t["name"] for t in tool_schemas]
        system_prompt = build_system_prompt(
            self.config.model_dump(), tool_names=tool_names,
        )
        self._messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_request},
        ]

        logger.info(
            "Session started: run_id=%s, provider=%s",
            self._run_ctx.run_id,
            self.provider.name,
        )
        return self._run_ctx

    def step(self) -> dict[str, Any]:
        """Execute one LLM turn and return the raw response dict."""
        tools = self.dispatcher.get_tool_schemas()
        with span(
            "apab.llm.chat",
            **{
                "gen_ai.system": self.provider.name,
                "gen_ai.request.model": self.config.llm.model,
            },
        ) as s:
            try:
                response = self.provider.chat(
                    messages=self._messages,
                    tools=tools,
                )
            except Exception as exc:
                set_span_error(s, exc)
                raise

            usage = getattr(self.provider, "last_usage", None)
            if usage is not None:
                s.set_attribute("gen_ai.usage.input_tokens", usage.prompt_tokens)
                s.set_attribute("gen_ai.usage.output_tokens", usage.completion_tokens)
                s.set_attribute("apab.latency_s", usage.latency_s)
                s.set_attribute("apab.cost_estimate_usd", usage.cost_estimate_usd)
            s.set_attribute(
                "apab.tool_call_count", len(response.get("tool_calls") or []),
            )
            s.set_attribute(
                "apab.response.has_content", bool(response.get("content")),
            )

        # Apply redaction before logging
        self._log_egress(response)
        self._accumulate_usage()

        # Append assistant message
        self._messages.append({
            "role": "assistant",
            "content": response.get("content"),
            "tool_calls": response.get("tool_calls"),
        })

        return response

    def execute_tool_calls(self, response: dict[str, Any]) -> list[dict[str, Any]]:
        """Execute all tool calls from a response and append results to messages."""
        tool_calls = response.get("tool_calls") or []
        results = []

        capture_mode = self._capture_mode()

        for tc in tool_calls:
            tool_name = tc["name"]
            arguments = tc.get("arguments", {})

            logger.info("Calling tool: %s(%s)", tool_name, json.dumps(arguments, default=str))

            captured = capture_args(arguments, capture_mode)
            attrs = {
                "apab.tool.name": tool_name,
                "apab.tool.args_hash": captured["args_hash"],
            }
            if "args_json" in captured:
                attrs["apab.tool.args_json"] = captured["args_json"]
            if "arg_keys" in captured:
                attrs["apab.tool.arg_keys"] = captured["arg_keys"]

            with span(f"apab.tool.{tool_name}", **attrs) as s:
                result_str = self.dispatcher.dispatch(tool_name, arguments)
                s.set_attribute(
                    "apab.tool.status",
                    "error" if _is_error_result(result_str) else "ok",
                )
                summary = capture_text(result_str, capture_mode)
                if summary is not None:
                    s.set_attribute("apab.tool.result_summary", summary)

            result_msg = {
                "role": "tool",
                "content": result_str,
                "name": tool_name,
            }
            self._messages.append(result_msg)
            results.append({"tool": tool_name, "result": result_str})

        return results

    def run_to_completion(
        self,
        user_request: str,
        max_turns: int = 20,
        on_event: OnEvent | None = None,
    ) -> str:
        """Run the full agentic loop and return the final text response.

        Parameters
        ----------
        user_request:
            The user's natural-language request.
        max_turns:
            Maximum number of LLM turns before forcibly stopping.
        on_event:
            Optional callback invoked with ``(event_name, payload)`` as
            the loop progresses, for progress rendering. Exceptions
            raised by the callback are logged and ignored.
        """
        ctx = self.start_session(user_request)
        self._emit(on_event, "session_start", {"run_id": ctx.run_id})
        init_observability(self.config.observability, run_ctx=ctx)

        from apab.core.provenance import hash_config

        status = "error"
        try:
            with span(
                "apab.session",
                **{
                    "apab.run_id": ctx.run_id,
                    "gen_ai.system": self.provider.name,
                    "gen_ai.request.model": self.config.llm.model,
                    "apab.max_turns": max_turns,
                    "apab.config_hash": hash_config(self.config.model_dump()),
                },
            ) as session_span:
                ids = current_trace_ids()
                self._trace_id = ids[0] if ids else None

                try:
                    for turn in range(max_turns):
                        self._emit(on_event, "turn_start", {"turn": turn})

                        with span("apab.turn", **{"apab.turn.index": turn}):
                            response = self.step()
                            tool_calls = response.get("tool_calls")

                            if not tool_calls:
                                # No tool calls => final response
                                content = response.get("content") or ""
                                self._emit(
                                    on_event, "assistant_message",
                                    {"content": content},
                                )
                                status = "success"
                                return content

                            for tc in tool_calls:
                                self._emit(on_event, "tool_call", {
                                    "name": tc["name"],
                                    "arguments": tc.get("arguments", {}),
                                })

                            results = self.execute_tool_calls(response)

                            for r in results:
                                self._emit(on_event, "tool_result", r)

                    self._emit(on_event, "max_turns", {"max_turns": max_turns})
                    status = "max_turns"
                    return (
                        "Reached maximum number of turns. "
                        "The last response may be incomplete."
                    )
                except Exception as exc:
                    set_span_error(session_span, exc)
                    raise
                finally:
                    totals = self._session_usage
                    session_span.set_attribute("apab.status", status)
                    session_span.set_attribute(
                        "gen_ai.usage.input_tokens", totals["prompt_tokens"],
                    )
                    session_span.set_attribute(
                        "gen_ai.usage.output_tokens",
                        totals["completion_tokens"],
                    )
                    session_span.set_attribute(
                        "apab.cost_estimate_usd", totals["cost_estimate_usd"],
                    )
        finally:
            self._persist_audit_log()
            shutdown_observability()
            self._persist_manifest(status)

    # ── internals ─────────────────────────────────────────────────────

    @staticmethod
    def _emit(
        on_event: OnEvent | None,
        name: str,
        payload: dict[str, Any],
    ) -> None:
        """Invoke the progress callback, never letting it break the run."""
        if on_event is None:
            return
        try:
            on_event(name, payload)
        except Exception:
            logger.exception("on_event callback failed for %r", name)

    def _capture_mode(self) -> RedactionMode:
        """Redaction level for span attributes; inherits llm.redaction_mode."""
        mode = self.config.observability.capture_mode
        if mode is None:
            mode = self.config.llm.redaction_mode
        return RedactionMode(mode)

    def _accumulate_usage(self) -> None:
        """Add the provider's most recent call usage to session totals."""
        usage = getattr(self.provider, "last_usage", None)
        if usage is None:
            return
        totals = self._session_usage
        totals["prompt_tokens"] += usage.prompt_tokens
        totals["completion_tokens"] += usage.completion_tokens
        totals["cost_estimate_usd"] += usage.cost_estimate_usd
        totals["llm_calls"] += 1

    def _log_egress(self, response: dict[str, Any]) -> None:
        """Log outgoing LLM interactions with redaction applied."""
        mode = self.config.llm.redaction_mode

        if mode == RedactionMode.none:
            logger.debug("LLM response: %s", json.dumps(response, default=str)[:500])
        elif mode == RedactionMode.metadata_only:
            has_tools = bool(response.get("tool_calls"))
            logger.debug(
                "LLM response: has_content=%s, has_tool_calls=%s",
                bool(response.get("content")),
                has_tools,
            )
        elif mode == RedactionMode.strict:
            logger.debug("LLM response: [REDACTED]")

    def _persist_audit_log(self) -> None:
        """Save the audit log to the run directory if a session is active."""
        if self._run_ctx is not None and self.dispatcher.audit_log:
            audit_path = self._run_ctx.run_dir / "audit.json"
            try:
                self.dispatcher.save_audit_log(audit_path)
                logger.info("Audit log saved to %s", audit_path)
            except Exception:
                logger.exception("Failed to save audit log")

    def _persist_manifest(self, status: str) -> None:
        """Write a provenance manifest to the run directory."""
        if self._run_ctx is None:
            return
        try:
            from apab.core.provenance import build_manifest

            run_dir = self._run_ctx.run_dir
            artifacts = sorted(
                str(p.relative_to(run_dir))
                for p in self._run_ctx.artifacts_dir.rglob("*")
                if p.is_file()
            )
            manifest = build_manifest(
                self._run_ctx.run_id,
                config=self.config.model_dump(),
                artifacts=artifacts,
            )
            manifest["status"] = status
            manifest["usage"] = dict(self._session_usage)
            manifest["trace_id"] = self._trace_id or ""

            manifest_path = run_dir / "manifest.json"
            manifest_path.write_text(
                json.dumps(manifest, indent=2, default=str)
            )
            logger.info("Manifest saved to %s", manifest_path)
        except Exception:
            logger.exception("Failed to save manifest")

    @property
    def messages(self) -> list[dict[str, Any]]:
        """Return a copy of the conversation messages."""
        return list(self._messages)

    @property
    def run_context(self) -> RunContext | None:
        """Return the current run context, or ``None`` if not started."""
        return self._run_ctx

    @property
    def session_usage(self) -> dict[str, Any]:
        """Return a copy of the accumulated session usage totals."""
        return dict(self._session_usage)

messages property

messages

Return a copy of the conversation messages.

run_context property

run_context

Return the current run context, or None if not started.

session_usage property

session_usage

Return a copy of the accumulated session usage totals.

start_session

start_session(user_request)

Begin a new agent session and return a :class:RunContext.

Source code in src/apab/agent/orchestrator.py
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def start_session(self, user_request: str) -> RunContext:
    """Begin a new agent session and return a :class:`RunContext`."""
    self.workspace.ensure_dirs()
    self._run_ctx = self.workspace.new_run()
    self._session_usage = _empty_usage()

    tool_schemas = self.dispatcher.get_tool_schemas()
    tool_names = [t["name"] for t in tool_schemas]
    system_prompt = build_system_prompt(
        self.config.model_dump(), tool_names=tool_names,
    )
    self._messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_request},
    ]

    logger.info(
        "Session started: run_id=%s, provider=%s",
        self._run_ctx.run_id,
        self.provider.name,
    )
    return self._run_ctx

step

step()

Execute one LLM turn and return the raw response dict.

Source code in src/apab/agent/orchestrator.py
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def step(self) -> dict[str, Any]:
    """Execute one LLM turn and return the raw response dict."""
    tools = self.dispatcher.get_tool_schemas()
    with span(
        "apab.llm.chat",
        **{
            "gen_ai.system": self.provider.name,
            "gen_ai.request.model": self.config.llm.model,
        },
    ) as s:
        try:
            response = self.provider.chat(
                messages=self._messages,
                tools=tools,
            )
        except Exception as exc:
            set_span_error(s, exc)
            raise

        usage = getattr(self.provider, "last_usage", None)
        if usage is not None:
            s.set_attribute("gen_ai.usage.input_tokens", usage.prompt_tokens)
            s.set_attribute("gen_ai.usage.output_tokens", usage.completion_tokens)
            s.set_attribute("apab.latency_s", usage.latency_s)
            s.set_attribute("apab.cost_estimate_usd", usage.cost_estimate_usd)
        s.set_attribute(
            "apab.tool_call_count", len(response.get("tool_calls") or []),
        )
        s.set_attribute(
            "apab.response.has_content", bool(response.get("content")),
        )

    # Apply redaction before logging
    self._log_egress(response)
    self._accumulate_usage()

    # Append assistant message
    self._messages.append({
        "role": "assistant",
        "content": response.get("content"),
        "tool_calls": response.get("tool_calls"),
    })

    return response

execute_tool_calls

execute_tool_calls(response)

Execute all tool calls from a response and append results to messages.

Source code in src/apab/agent/orchestrator.py
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def execute_tool_calls(self, response: dict[str, Any]) -> list[dict[str, Any]]:
    """Execute all tool calls from a response and append results to messages."""
    tool_calls = response.get("tool_calls") or []
    results = []

    capture_mode = self._capture_mode()

    for tc in tool_calls:
        tool_name = tc["name"]
        arguments = tc.get("arguments", {})

        logger.info("Calling tool: %s(%s)", tool_name, json.dumps(arguments, default=str))

        captured = capture_args(arguments, capture_mode)
        attrs = {
            "apab.tool.name": tool_name,
            "apab.tool.args_hash": captured["args_hash"],
        }
        if "args_json" in captured:
            attrs["apab.tool.args_json"] = captured["args_json"]
        if "arg_keys" in captured:
            attrs["apab.tool.arg_keys"] = captured["arg_keys"]

        with span(f"apab.tool.{tool_name}", **attrs) as s:
            result_str = self.dispatcher.dispatch(tool_name, arguments)
            s.set_attribute(
                "apab.tool.status",
                "error" if _is_error_result(result_str) else "ok",
            )
            summary = capture_text(result_str, capture_mode)
            if summary is not None:
                s.set_attribute("apab.tool.result_summary", summary)

        result_msg = {
            "role": "tool",
            "content": result_str,
            "name": tool_name,
        }
        self._messages.append(result_msg)
        results.append({"tool": tool_name, "result": result_str})

    return results

run_to_completion

run_to_completion(user_request, max_turns=20, on_event=None)

Run the full agentic loop and return the final text response.

Parameters:

Name Type Description Default
user_request str

The user's natural-language request.

required
max_turns int

Maximum number of LLM turns before forcibly stopping.

20
on_event OnEvent | None

Optional callback invoked with (event_name, payload) as the loop progresses, for progress rendering. Exceptions raised by the callback are logged and ignored.

None
Source code in src/apab/agent/orchestrator.py
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
def run_to_completion(
    self,
    user_request: str,
    max_turns: int = 20,
    on_event: OnEvent | None = None,
) -> str:
    """Run the full agentic loop and return the final text response.

    Parameters
    ----------
    user_request:
        The user's natural-language request.
    max_turns:
        Maximum number of LLM turns before forcibly stopping.
    on_event:
        Optional callback invoked with ``(event_name, payload)`` as
        the loop progresses, for progress rendering. Exceptions
        raised by the callback are logged and ignored.
    """
    ctx = self.start_session(user_request)
    self._emit(on_event, "session_start", {"run_id": ctx.run_id})
    init_observability(self.config.observability, run_ctx=ctx)

    from apab.core.provenance import hash_config

    status = "error"
    try:
        with span(
            "apab.session",
            **{
                "apab.run_id": ctx.run_id,
                "gen_ai.system": self.provider.name,
                "gen_ai.request.model": self.config.llm.model,
                "apab.max_turns": max_turns,
                "apab.config_hash": hash_config(self.config.model_dump()),
            },
        ) as session_span:
            ids = current_trace_ids()
            self._trace_id = ids[0] if ids else None

            try:
                for turn in range(max_turns):
                    self._emit(on_event, "turn_start", {"turn": turn})

                    with span("apab.turn", **{"apab.turn.index": turn}):
                        response = self.step()
                        tool_calls = response.get("tool_calls")

                        if not tool_calls:
                            # No tool calls => final response
                            content = response.get("content") or ""
                            self._emit(
                                on_event, "assistant_message",
                                {"content": content},
                            )
                            status = "success"
                            return content

                        for tc in tool_calls:
                            self._emit(on_event, "tool_call", {
                                "name": tc["name"],
                                "arguments": tc.get("arguments", {}),
                            })

                        results = self.execute_tool_calls(response)

                        for r in results:
                            self._emit(on_event, "tool_result", r)

                self._emit(on_event, "max_turns", {"max_turns": max_turns})
                status = "max_turns"
                return (
                    "Reached maximum number of turns. "
                    "The last response may be incomplete."
                )
            except Exception as exc:
                set_span_error(session_span, exc)
                raise
            finally:
                totals = self._session_usage
                session_span.set_attribute("apab.status", status)
                session_span.set_attribute(
                    "gen_ai.usage.input_tokens", totals["prompt_tokens"],
                )
                session_span.set_attribute(
                    "gen_ai.usage.output_tokens",
                    totals["completion_tokens"],
                )
                session_span.set_attribute(
                    "apab.cost_estimate_usd", totals["cost_estimate_usd"],
                )
    finally:
        self._persist_audit_log()
        shutdown_observability()
        self._persist_manifest(status)

apab.agent.tool_dispatch

Tool dispatch: extract MCP tool schemas and execute tool calls.

ToolDispatcher

Bridges the LLM's tool calls to the MCP server's tool implementations.

Source code in src/apab/agent/tool_dispatch.py
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class ToolDispatcher:
    """Bridges the LLM's tool calls to the MCP server's tool implementations."""

    def __init__(self, redaction_mode: str = "none") -> None:
        self._audit_log: list[dict[str, Any]] = []
        self._redaction_mode = redaction_mode

    def get_tool_schemas(self) -> list[dict[str, Any]]:
        """Extract tool schemas from the MCP server in LLM-friendly format."""
        from apab.mcp.server import get_mcp

        server = get_mcp()
        tools = server._tool_manager._tools

        schemas = []
        for name, tool in tools.items():
            schema = {
                "name": name,
                "description": tool.description or "",
                "inputSchema": tool.parameters if hasattr(tool, "parameters") else {},
            }
            schemas.append(schema)
        return schemas

    def dispatch(self, tool_name: str, arguments: dict[str, Any]) -> str:
        """Execute a tool call and return the JSON result string.

        Uses asyncio to call the async tool function from a sync context.
        """
        from apab.mcp.server import get_mcp

        server = get_mcp()
        tools = server._tool_manager._tools

        if tool_name not in tools:
            error = {"error": f"Unknown tool: {tool_name}"}
            self._log_call(tool_name, arguments, error)
            return json.dumps(error)

        tool = tools[tool_name]
        fn = tool.fn

        try:
            # Run the async tool function
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = None

            if loop and loop.is_running():
                # We're inside an async context — use the shared executor
                result = _executor.submit(
                    asyncio.run, fn(**arguments)
                ).result()
            else:
                result = asyncio.run(fn(**arguments))

            self._log_call(tool_name, arguments, result)
            return json.dumps(result, default=str)

        except Exception as e:
            error = {"error": str(e), "tool": tool_name}
            self._log_call(tool_name, arguments, error)
            logger.exception("Tool dispatch failed: %s", tool_name)
            return json.dumps(error)

    def _log_call(
        self,
        tool_name: str,
        arguments: dict[str, Any],
        result: Any,
    ) -> None:
        """Append an entry to the audit log, respecting redaction mode."""
        entry: dict[str, Any] = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "tool": tool_name,
        }

        # Correlate with the active trace when observability is on.
        from apab.observability import current_trace_ids

        ids = current_trace_ids()
        if ids:
            entry["trace_id"], entry["span_id"] = ids

        if self._redaction_mode == "strict":
            entry["arguments"] = "[REDACTED]"
            entry["result_summary"] = "[REDACTED]"
        elif self._redaction_mode == "metadata_only":
            entry["arguments"] = list(arguments.keys())
            entry["result_summary"] = _summarise(result)
        else:
            entry["arguments"] = arguments
            entry["result_summary"] = _summarise(result)

        self._audit_log.append(entry)

    def save_audit_log(self, path: str | Any) -> None:
        """Persist the audit log to a JSON file."""
        from pathlib import Path

        out = Path(path)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps(self._audit_log, indent=2, default=str))

    @property
    def audit_log(self) -> list[dict[str, Any]]:
        """Return a copy of the audit log."""
        return list(self._audit_log)

audit_log property

audit_log

Return a copy of the audit log.

get_tool_schemas

get_tool_schemas()

Extract tool schemas from the MCP server in LLM-friendly format.

Source code in src/apab/agent/tool_dispatch.py
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def get_tool_schemas(self) -> list[dict[str, Any]]:
    """Extract tool schemas from the MCP server in LLM-friendly format."""
    from apab.mcp.server import get_mcp

    server = get_mcp()
    tools = server._tool_manager._tools

    schemas = []
    for name, tool in tools.items():
        schema = {
            "name": name,
            "description": tool.description or "",
            "inputSchema": tool.parameters if hasattr(tool, "parameters") else {},
        }
        schemas.append(schema)
    return schemas

dispatch

dispatch(tool_name, arguments)

Execute a tool call and return the JSON result string.

Uses asyncio to call the async tool function from a sync context.

Source code in src/apab/agent/tool_dispatch.py
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def dispatch(self, tool_name: str, arguments: dict[str, Any]) -> str:
    """Execute a tool call and return the JSON result string.

    Uses asyncio to call the async tool function from a sync context.
    """
    from apab.mcp.server import get_mcp

    server = get_mcp()
    tools = server._tool_manager._tools

    if tool_name not in tools:
        error = {"error": f"Unknown tool: {tool_name}"}
        self._log_call(tool_name, arguments, error)
        return json.dumps(error)

    tool = tools[tool_name]
    fn = tool.fn

    try:
        # Run the async tool function
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if loop and loop.is_running():
            # We're inside an async context — use the shared executor
            result = _executor.submit(
                asyncio.run, fn(**arguments)
            ).result()
        else:
            result = asyncio.run(fn(**arguments))

        self._log_call(tool_name, arguments, result)
        return json.dumps(result, default=str)

    except Exception as e:
        error = {"error": str(e), "tool": tool_name}
        self._log_call(tool_name, arguments, error)
        logger.exception("Tool dispatch failed: %s", tool_name)
        return json.dumps(error)

save_audit_log

save_audit_log(path)

Persist the audit log to a JSON file.

Source code in src/apab/agent/tool_dispatch.py
115
116
117
118
119
120
121
def save_audit_log(self, path: str | Any) -> None:
    """Persist the audit log to a JSON file."""
    from pathlib import Path

    out = Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(self._audit_log, indent=2, default=str))

apab.agent.provider_registry

LLM provider protocol and registry.

ProviderUsage dataclass

Token and cost tracking for a single LLM call.

Providers expose the most recent call's usage via a last_usage property. Local providers (e.g. Ollama) report a zero cost estimate.

Source code in src/apab/providers/usage.py
 8
 9
10
11
12
13
14
15
16
17
18
19
@dataclass
class ProviderUsage:
    """Token and cost tracking for a single LLM call.

    Providers expose the most recent call's usage via a ``last_usage``
    property. Local providers (e.g. Ollama) report a zero cost estimate.
    """

    prompt_tokens: int = 0
    completion_tokens: int = 0
    latency_s: float = 0.0
    cost_estimate_usd: float = 0.0

LLMProvider

Bases: Protocol

Protocol for LLM provider backends.

Providers may additionally expose an optional last_usage property returning an :class:apab.providers.usage.ProviderUsage (or None) for the most recent chat() call. It is not part of the protocol so that third-party providers registered via entry points keep working; consumers must read it with getattr(provider, "last_usage", None).

Source code in src/apab/agent/provider_registry.py
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
@runtime_checkable
class LLMProvider(Protocol):
    """Protocol for LLM provider backends.

    Providers may additionally expose an optional ``last_usage`` property
    returning an :class:`apab.providers.usage.ProviderUsage` (or ``None``)
    for the most recent ``chat()`` call. It is not part of the protocol so
    that third-party providers registered via entry points keep working;
    consumers must read it with ``getattr(provider, "last_usage", None)``.
    """

    @property
    def name(self) -> str:
        """Short identifier for the provider (e.g. ``"ollama"``)."""
        ...

    def supports_tool_calling(self) -> bool:
        """Whether this provider supports tool-calling."""
        ...

    def supports_streaming(self) -> bool:
        """Whether this provider supports streaming responses."""
        ...

    def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Send a chat request and return a normalised response.

        The returned dict has at minimum:
        - ``"role"`` — always ``"assistant"``
        - ``"content"`` — text content (may be ``None`` if tool calls)
        - ``"tool_calls"`` — list of ``{"name": str, "arguments": dict}``
          dicts, or ``None`` if no tool calls.
        """
        ...

name property

name

Short identifier for the provider (e.g. "ollama").

supports_tool_calling

supports_tool_calling()

Whether this provider supports tool-calling.

Source code in src/apab/agent/provider_registry.py
39
40
41
def supports_tool_calling(self) -> bool:
    """Whether this provider supports tool-calling."""
    ...

supports_streaming

supports_streaming()

Whether this provider supports streaming responses.

Source code in src/apab/agent/provider_registry.py
43
44
45
def supports_streaming(self) -> bool:
    """Whether this provider supports streaming responses."""
    ...

chat

chat(messages, tools=None, **kwargs)

Send a chat request and return a normalised response.

The returned dict has at minimum: - "role" — always "assistant" - "content" — text content (may be None if tool calls) - "tool_calls" — list of {"name": str, "arguments": dict} dicts, or None if no tool calls.

Source code in src/apab/agent/provider_registry.py
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def chat(
    self,
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]] | None = None,
    **kwargs: Any,
) -> dict[str, Any]:
    """Send a chat request and return a normalised response.

    The returned dict has at minimum:
    - ``"role"`` — always ``"assistant"``
    - ``"content"`` — text content (may be ``None`` if tool calls)
    - ``"tool_calls"`` — list of ``{"name": str, "arguments": dict}``
      dicts, or ``None`` if no tool calls.
    """
    ...

discover_providers

discover_providers()

Discover LLM provider classes registered via entry points.

Source code in src/apab/agent/provider_registry.py
64
65
66
67
68
69
70
71
72
73
74
75
76
def discover_providers() -> dict[str, type]:
    """Discover LLM provider classes registered via entry points."""
    providers: dict[str, type] = {}
    # entry_points(group=...) is available on all supported versions (3.10+).
    eps = importlib.metadata.entry_points(group="apab.llm_providers")

    for ep in eps:
        try:
            cls = ep.load()
            providers[ep.name] = cls
        except Exception:
            pass
    return providers

get_provider

get_provider(provider_name, **kwargs)

Instantiate and return an LLM provider by name.

Source code in src/apab/agent/provider_registry.py
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def get_provider(provider_name: str, **kwargs: Any) -> LLMProvider:
    """Instantiate and return an LLM provider by name."""
    discovered = discover_providers()
    if provider_name in discovered:
        return discovered[provider_name](**kwargs)  # type: ignore[no-any-return]

    module_path = _BUILTINS.get(provider_name)
    if module_path is None:
        raise ValueError(
            f"Unknown LLM provider '{provider_name}'. "
            f"Available: {sorted(set(discovered) | set(_BUILTINS))}"
        )

    mod = importlib.import_module(module_path)
    for attr_name in dir(mod):
        if attr_name.startswith("_"):
            continue
        attr = getattr(mod, attr_name)
        if isinstance(attr, type) and attr is not LLMProvider:
            # Check if it looks like a provider (has name property and chat method)
            if hasattr(attr, "chat") and hasattr(attr, "name"):
                return attr(**kwargs)  # type: ignore[no-any-return]

    raise ValueError(f"No LLMProvider found in module '{module_path}'")

validate_provider

validate_provider(provider_name)

Check if a provider name is available without instantiating it.

Returns True if the provider can be found; logs a warning and returns False otherwise.

Source code in src/apab/agent/provider_registry.py
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def validate_provider(provider_name: str) -> bool:
    """Check if a provider name is available without instantiating it.

    Returns ``True`` if the provider can be found; logs a warning and
    returns ``False`` otherwise.
    """
    discovered = discover_providers()
    if provider_name in discovered:
        return True
    if provider_name in _BUILTINS:
        return True
    logger.warning(
        "LLM provider %r not found. Available: %s",
        provider_name,
        sorted(set(discovered) | set(_BUILTINS)),
    )
    return False