Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-agents/src/lexigram/ai/agents/speculation/predictor.py: 29%

1"""KeywordToolCallPredictor — keyword-based tool call prediction.""" 

2 

3from __future__ import annotations 

4 

5import re 

6 

7from lexigram.contracts import ( 

8 ToolProtocol, 

9) 

10from lexigram.contracts.ai.llm import ChatMessage 

11from lexigram.logging import ( 

12 get_logger, 

13) 

14 

15logger = get_logger(__name__) 

16 

17 

18class KeywordToolCallPredictor: 

19 """Keyword-based tool call prediction heuristic. 

20 

21 Implements ToolCallPredictorProtocol. Scores each tool by keyword 

22 overlap between the query and the tool's name + description. 

23 

24 Heuristic algorithm: 

25 1. Tokenize query into lowercase words. 

26 2. For each tool, tokenize name + description into keywords. 

27 3. Score = |query_words ∩ tool_words| / |tool_words| 

28 4. Boost score by 1.5x if the tool was called in the last recency_window turns. 

29 5. Return tools sorted by score descending. 

30 """ 

31 

32 def __init__(self, recency_boost: float = 1.5, recency_window: int = 3) -> None: 

33 """Initialize the predictor. 

34 

35 Args: 

36 recency_boost: Multiplier applied to recently-used tools. 

37 recency_window: Number of recent turns to consider for recency boost. 

38 """ 

39 self._recency_boost = recency_boost 

40 self._recency_window = recency_window 

41 

42 def _tokenize(self, text: str) -> set[str]: 

43 """Split text on whitespace and punctuation, return lowercase tokens.""" 

44 return set(re.split(r"[\s\W]+", text.lower())) - {""} 

45 

46 def predict( 

47 self, 

48 query: str, 

49 available_tools: list[ToolProtocol], 

50 recent_history: list[ChatMessage] | None = None, 

51 ) -> list[ToolProtocol]: 

52 """Return tools ranked by likelihood of being called. 

53 

54 Args: 

55 query: The current user query. 

56 available_tools: All tools registered in the agent. 

57 recent_history: Recent conversation turns for context. 

58 

59 Returns: 

60 Tools sorted by predicted likelihood, most likely first. 

61 """ 

62 query_words = self._tokenize(query) 

63 

64 # Find recently-used tool names from history 

65 recently_used: set[str] = set() 

66 if recent_history: 

67 for msg in recent_history[-self._recency_window :]: 

68 content = getattr(msg, "content", "") or "" 

69 for tool in available_tools: 

70 tool_name = getattr(tool, "name", "") 

71 if tool_name and tool_name.lower() in content.lower(): 

72 recently_used.add(tool_name) 

73 

74 def score(tool: ToolProtocol) -> float: 

75 tool_name = getattr(tool, "name", "") or "" 

76 tool_desc = getattr(tool, "description", "") or "" 

77 tool_words = self._tokenize(f"{tool_name} {tool_desc}") 

78 if not tool_words: 

79 return 0.0 

80 overlap = len(query_words & tool_words) 

81 s = overlap / len(tool_words) 

82 if getattr(tool, "name", "") in recently_used: 

83 s *= self._recency_boost 

84 return s 

85 

86 return sorted(available_tools, key=score, reverse=True)