« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""SpeculativeToolPreFetcher — parallel tool pre-fetching during LLM decision-making."""
3from __future__ import annotations
5import asyncio
7from lexigram.ai.agents.speculation.predictor import KeywordToolCallPredictor
8from lexigram.ai.agents.speculation.protocols import ToolCallPredictorProtocol
9from lexigram.contracts import (
10 AgentError,
11 LLMClientProtocol,
12 ToolProtocol,
13 ToolRegistryProtocol,
14)
15from lexigram.contracts.ai.llm import ChatMessage, Completion
16from lexigram.logging import (
17 get_logger,
18)
19from lexigram.result import Err, Ok, Result
21logger = get_logger(__name__)
24class SpeculativeToolPreFetcher:
25 """Parallel tool pre-fetching during LLM decision-making.
27 While the LLM is deciding which tool to call, speculatively executes
28 the top-N most likely tools in parallel. If the LLM picks a pre-fetched
29 tool, the result is returned instantly. Unused speculative tasks are
30 cancelled.
32 Opt-in strategy — not enabled by default in the agent loop.
33 """
35 def __init__(
36 self,
37 tool_registry: ToolRegistryProtocol,
38 max_speculative: int = 3,
39 predictor: ToolCallPredictorProtocol | None = None,
40 ) -> None:
41 """Initialize the prefetcher.
43 Args:
44 tool_registry: Registry of available tools.
45 max_speculative: Maximum number of tools to pre-fetch in parallel.
46 predictor: Tool call predictor. Defaults to KeywordToolCallPredictor.
47 """
48 self._registry = tool_registry
49 self._max_speculative = max_speculative
50 self._predictor = predictor or KeywordToolCallPredictor()
51 self._background_tasks: set[asyncio.Task] = set()
53 async def execute_with_speculation(
54 self,
55 query: str,
56 tools: list[ToolProtocol],
57 llm_client: LLMClientProtocol,
58 messages: list[ChatMessage],
59 ) -> Result[Completion, AgentError]:
60 """Execute LLM call with parallel speculative tool pre-fetching.
62 1. Predict likely tool calls from query + history.
63 2. Fire LLM call AND top-N tool calls in parallel.
64 3. If LLM picks a pre-fetched tool, return pre-fetched result.
65 4. If LLM picks a non-predicted tool, execute normally.
66 5. Cancel all unused speculative tasks.
67 6. Store task references per RUF006.
69 Args:
70 query: Current user query for tool prediction.
71 tools: All available tools.
72 llm_client: LLM client to use for the main call.
73 messages: Conversation messages to send to LLM.
75 Returns:
76 Result containing the LLM Completion or an AgentError.
77 """
78 predicted = self._predictor.predict(query, tools)[: self._max_speculative]
80 # Speculatively execute top-N tools in parallel while LLM decides
81 speculative_tasks: dict[str, asyncio.Task] = {}
82 for tool in predicted:
83 tool_name = getattr(tool, "name", "")
84 if not tool_name:
85 continue
86 task = asyncio.create_task(tool.execute({})) # type: ignore[call-arg]
87 self._background_tasks.add(task)
88 task.add_done_callback(self._background_tasks.discard)
89 speculative_tasks[tool_name] = task
91 # Run LLM call
92 llm_task = asyncio.create_task(llm_client.complete(messages))
93 self._background_tasks.add(llm_task)
94 llm_task.add_done_callback(self._background_tasks.discard)
96 llm_result = await llm_task
98 # Cancel all unused speculative tasks
99 for task in speculative_tasks.values():
100 if not task.done():
101 task.cancel()
103 if llm_result.is_ok():
104 return Ok(llm_result.unwrap()) # type: ignore[arg-type]
105 err = llm_result.unwrap_err()
106 return Err(AgentError(f"LLM call failed: {err}"))