Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-prompt/src/lexigram/ai/prompt/assembly/assembler.py: 30%

1"""Cache-aware prompt assembler implementing PromptAssemblerProtocol.""" 

2 

3from __future__ import annotations 

4 

5from lexigram.ai.prompt.assembly.cache_strategies import ProviderCacheStrategyRegistry 

6from lexigram.contracts.ai import ToolDefinition 

7from lexigram.contracts.ai.llm import ( 

8 ChatMessage, 

9 PromptAssemblerProtocol, 

10 TokenCounterProtocol, 

11) 

12from lexigram.logging import ( 

13 get_logger, 

14) 

15 

16logger = get_logger(__name__) 

17 

18 

19class CacheAwarePromptAssembler: 

20 """Prompt assembler that enforces static-before-dynamic ordering and 

21 applies provider-specific cache annotations. 

22 

23 Implements PromptAssemblerProtocol and enforces the 7-layer ordering: 

24 1. System instructions (STATIC) 

25 2. Tool/function definitions (STATIC) 

26 3. Reference documents (SEMI-STATIC) 

27 4. Few-shot examples (STATIC) 

28 ── CACHE BOUNDARY ── 

29 5. Chat history (DYNAMIC) 

30 6. Current user query (DYNAMIC) 

31 7. Dynamic metadata (DYNAMIC) 

32 

33 Args: 

34 strategy_registry: Provider cache strategy registry. 

35 token_counter: Optional token counter for cache-size validation. 

36 """ 

37 

38 def __init__( 

39 self, 

40 strategy_registry: ProviderCacheStrategyRegistry, 

41 token_counter: TokenCounterProtocol | None = None, 

42 ) -> None: 

43 """Initialize with strategy registry and optional token counter.""" 

44 self._registry = strategy_registry 

45 self._token_counter = token_counter 

46 

47 def set_token_counter(self, counter: TokenCounterProtocol) -> None: 

48 """Inject token counter after construction (called from DI boot phase). 

49 

50 Args: 

51 counter: Token counter to use for cache size validation. 

52 """ 

53 self._token_counter = counter 

54 

55 def assemble( 

56 self, 

57 system: str, 

58 tools: list[ToolDefinition] | None, 

59 reference_docs: list[str] | None, 

60 few_shot: list[ChatMessage] | None, 

61 history: list[ChatMessage], 

62 query: str, 

63 provider: str, 

64 dynamic_metadata: str | None = None, 

65 ) -> list[ChatMessage]: 

66 """Assemble messages in canonical static-before-dynamic order. 

67 

68 This method matches the PromptAssemblerProtocol.assemble() signature 

69 exactly. It enforces the 7-layer static-to-dynamic ordering: 

70 

71 Layers 1-4 (STATIC) come before cache boundary: 

72 - Layer 1: System instructions 

73 - Layer 2: Tool definitions 

74 - Layer 3: Reference documents 

75 - Layer 4: Few-shot examples 

76 

77 Layers 5-7 (DYNAMIC) come after cache boundary: 

78 - Layer 5: Chat history 

79 - Layer 6: Current user query 

80 - Layer 7: Dynamic metadata 

81 

82 Args: 

83 system: System instructions (Layer 1, always static). 

84 tools: Tool definitions as list of ToolDefinition objects. 

85 reference_docs: Reference document strings (Layer 3). 

86 few_shot: Few-shot example messages (Layer 4, static). 

87 history: Chat history messages (Layer 5, dynamic). 

88 query: Current user query text (Layer 6, dynamic). 

89 provider: Provider key for cache strategy selection. 

90 dynamic_metadata: Dynamic metadata string appended after query. 

91 

92 Returns: 

93 Assembled and cache-annotated list of ChatMessage objects. 

94 """ 

95 messages: list[ChatMessage] = [] 

96 

97 # Layer 1: System instructions (STATIC) 

98 if system: 

99 messages.append(ChatMessage(role="system", content=system)) 

100 

101 # Layer 2: Tool definitions (STATIC) — append as system context 

102 if tools: 

103 tool_lines: list[str] = [] 

104 for tool in tools: 

105 if isinstance(tool, str): 

106 tool_lines.append(tool) 

107 elif hasattr(tool, "name") and hasattr(tool, "description"): 

108 tool_lines.append(f"Tool: {tool.name}\n{tool.description}") 

109 else: 

110 tool_lines.append(str(tool)) 

111 if tool_lines: 

112 tool_text = "\n".join(tool_lines) 

113 messages.append( 

114 ChatMessage(role="system", content=f"Available tools:\n{tool_text}") 

115 ) 

116 

117 # Layer 3: Reference documents (SEMI-STATIC) 

118 if reference_docs: 

119 doc_text = "\n\n".join(reference_docs) 

120 messages.append( 

121 ChatMessage(role="system", content=f"Reference documents:\n{doc_text}") 

122 ) 

123 

124 # Layer 4: Few-shot examples (STATIC) 

125 if few_shot: 

126 messages.extend(few_shot) 

127 

128 # Track static boundary (layers 1-4) 

129 static_count = len(messages) 

130 

131 # Layer 5: Chat history (DYNAMIC) 

132 if history: 

133 messages.extend(history) 

134 

135 # Layer 6: Current user query (DYNAMIC) 

136 messages.append(ChatMessage(role="user", content=query)) 

137 

138 # Layer 7: Dynamic metadata (DYNAMIC) — appended last 

139 if dynamic_metadata: 

140 messages.append( 

141 ChatMessage(role="system", content=f"Metadata: {dynamic_metadata}") 

142 ) 

143 

144 # Apply provider-specific cache annotations to static layers 

145 strategy = self._registry.for_provider(provider) 

146 messages = strategy.annotate(messages, static_count, self._token_counter) 

147 

148 logger.debug( 

149 "prompt_assembled", 

150 provider=provider, 

151 static_layers=static_count, 

152 total_messages=len(messages), 

153 ) 

154 return messages 

155 

156 

157# Runtime structural protocol verification for PromptAssemblerProtocol 

158_: PromptAssemblerProtocol = CacheAwarePromptAssembler.__new__( 

159 CacheAwarePromptAssembler 

160)