agent_search_gateway.llm.stages
Resolved prompt-level LLM stages.
1"""Resolved prompt-level LLM stages.""" 2 3import asyncio 4import logging 5import time 6from collections.abc import Awaitable, Callable, Mapping 7from typing import TypeVar 8 9from ..errors import ErrorCode, ExecutionFailure 10from ..models import LLMInvocation, StageDecision 11from ..observability import elapsed_ms, log_event 12from ..providers.contracts import LLMClient 13from .prompts import ( 14 content_clean_messages, 15 focus_summary_messages, 16 judge_messages, 17 llm_paper_search_messages, 18 llm_search_messages, 19 safety_messages, 20) 21 22T = TypeVar("T") 23 24 25def cheap_check(candidate: str) -> bool: 26 return bool(candidate.strip()) 27 28 29class LLMStages: 30 def __init__( 31 self, 32 clients: Mapping[str, LLMClient], 33 *, 34 judge: LLMInvocation, 35 safety: LLMInvocation, 36 content_clean: LLMInvocation, 37 focus_summary: LLMInvocation, 38 logger: logging.Logger | None = None, 39 monotonic: Callable[[], float] = time.monotonic, 40 ) -> None: 41 self._clients = dict(clients) 42 self._judge = judge 43 self._safety = safety 44 self._content_clean = content_clean 45 self._focus_summary = focus_summary 46 self._logger = logger or logging.getLogger(__name__) 47 self._monotonic = monotonic 48 49 @property 50 def judge_provider(self) -> str: 51 return self._judge.provider 52 53 async def judge(self, candidate: str) -> StageDecision: 54 return await self._run_decision_stage( 55 self._judge, 56 "judge", 57 input_chars=len(candidate), 58 operation=lambda: self._client(self._judge.provider).complete_json( 59 self._judge, 60 judge_messages(candidate), 61 ), 62 ) 63 64 async def safety(self, content: str) -> StageDecision: 65 return await self._run_decision_stage( 66 self._safety, 67 "safety", 68 input_chars=len(content), 69 operation=lambda: self._client(self._safety.provider).complete_json( 70 self._safety, 71 safety_messages(content), 72 ), 73 ) 74 75 async def content_clean(self, raw_content: str) -> str: 76 return await self._run_text_stage( 77 self._content_clean, 78 "content_clean", 79 input_chars=len(raw_content), 80 operation=lambda: self._client(self._content_clean.provider).complete_text( 81 self._content_clean, 82 content_clean_messages(raw_content), 83 ), 84 ) 85 86 async def focus_summary(self, content: str, focus: str) -> str: 87 normalized_focus = focus.strip() 88 if not normalized_focus: 89 raise ExecutionFailure(ErrorCode.LLM_STAGE_FAILED, "focus must be non-empty") 90 return await self._run_text_stage( 91 self._focus_summary, 92 "focus_summary", 93 input_chars=len(content), 94 focus_chars=len(normalized_focus), 95 operation=lambda: self._client(self._focus_summary.provider).complete_text( 96 self._focus_summary, 97 focus_summary_messages(content, normalized_focus), 98 ), 99 ) 100 101 async def llm_search_markdown(self, invocation: LLMInvocation, prompt: str) -> str: 102 return await self._run_text_stage( 103 invocation, 104 "llm_search", 105 input_chars=len(prompt), 106 operation=lambda: self._client(invocation.provider).complete_text( 107 invocation, 108 llm_search_messages(prompt), 109 ), 110 ) 111 112 async def llm_paper_search_markdown( 113 self, 114 invocation: LLMInvocation, 115 prompt: str, 116 ) -> str: 117 return await self._run_text_stage( 118 invocation, 119 "llm_paper_search", 120 input_chars=len(prompt), 121 operation=lambda: self._client(invocation.provider).complete_text( 122 invocation, 123 llm_paper_search_messages(prompt), 124 ), 125 ) 126 127 async def _run_decision_stage( 128 self, 129 invocation: LLMInvocation, 130 stage: str, 131 *, 132 input_chars: int, 133 operation: Callable[[], Awaitable[Mapping[str, object]]], 134 ) -> StageDecision: 135 started = self._stage_started(invocation, stage, input_chars=input_chars) 136 try: 137 decision = self._parse_decision(await operation()) 138 except asyncio.CancelledError: 139 self._stage_cancelled(invocation, stage, started) 140 raise 141 except Exception as exc: 142 self._stage_failed(invocation, stage, started, exc) 143 raise 144 log_event( 145 self._logger, 146 logging.DEBUG, 147 "llm_stage_completed", 148 provider=invocation.provider, 149 stage=stage, 150 model=invocation.model, 151 ok=decision.ok, 152 reason_present=bool(decision.reason), 153 elapsed_ms=elapsed_ms(self._monotonic, started), 154 ) 155 return decision 156 157 async def _run_text_stage( 158 self, 159 invocation: LLMInvocation, 160 stage: str, 161 *, 162 input_chars: int, 163 operation: Callable[[], Awaitable[str]], 164 focus_chars: int = 0, 165 ) -> str: 166 started = self._stage_started( 167 invocation, 168 stage, 169 input_chars=input_chars, 170 focus_chars=focus_chars, 171 ) 172 try: 173 text = self._require_non_empty(await operation(), stage.replace("_", "-")) 174 except asyncio.CancelledError: 175 self._stage_cancelled(invocation, stage, started) 176 raise 177 except Exception as exc: 178 self._stage_failed(invocation, stage, started, exc) 179 raise 180 log_event( 181 self._logger, 182 logging.DEBUG, 183 "llm_stage_completed", 184 provider=invocation.provider, 185 stage=stage, 186 model=invocation.model, 187 output_chars=len(text), 188 elapsed_ms=elapsed_ms(self._monotonic, started), 189 ) 190 return text 191 192 def _stage_started( 193 self, 194 invocation: LLMInvocation, 195 stage: str, 196 *, 197 input_chars: int, 198 focus_chars: int = 0, 199 ) -> float: 200 started = self._monotonic() 201 log_event( 202 self._logger, 203 logging.DEBUG, 204 "llm_stage_started", 205 provider=invocation.provider, 206 stage=stage, 207 model=invocation.model, 208 input_chars=input_chars, 209 focus_present=focus_chars > 0, 210 focus_chars=focus_chars, 211 ) 212 return started 213 214 def _stage_failed( 215 self, 216 invocation: LLMInvocation, 217 stage: str, 218 started: float, 219 exc: Exception, 220 ) -> None: 221 log_event( 222 self._logger, 223 logging.DEBUG, 224 "llm_stage_failed", 225 provider=invocation.provider, 226 stage=stage, 227 model=invocation.model, 228 error_type=type(exc).__name__, 229 elapsed_ms=elapsed_ms(self._monotonic, started), 230 ) 231 232 def _stage_cancelled(self, invocation: LLMInvocation, stage: str, started: float) -> None: 233 log_event( 234 self._logger, 235 logging.DEBUG, 236 "llm_stage_cancelled", 237 provider=invocation.provider, 238 stage=stage, 239 model=invocation.model, 240 elapsed_ms=elapsed_ms(self._monotonic, started), 241 ) 242 243 def _client(self, provider: str) -> LLMClient: 244 client = self._clients.get(provider) 245 if client is None: 246 raise ExecutionFailure( 247 ErrorCode.LLM_STAGE_FAILED, 248 f"LLM provider is not initialized: {provider}", 249 ) 250 return client 251 252 @staticmethod 253 def _require_non_empty(text: str, stage: str) -> str: 254 normalized = text.strip() 255 if not normalized: 256 raise ExecutionFailure( 257 ErrorCode.LLM_STAGE_FAILED, 258 f"{stage} returned empty text", 259 ) 260 return normalized 261 262 @staticmethod 263 def _parse_decision(payload: Mapping[str, object]) -> StageDecision: 264 ok = payload.get("ok") 265 if not isinstance(ok, bool): 266 raise ExecutionFailure( 267 ErrorCode.LLM_STAGE_FAILED, 268 "LLM decision response requires boolean ok", 269 ) 270 reason = payload.get("reason", "") 271 if not isinstance(reason, str): 272 raise ExecutionFailure( 273 ErrorCode.LLM_STAGE_FAILED, 274 "LLM decision reason must be a string", 275 ) 276 return StageDecision(ok=ok, reason=reason.strip())
def
cheap_check(candidate: str) -> bool:
class
LLMStages:
30class LLMStages: 31 def __init__( 32 self, 33 clients: Mapping[str, LLMClient], 34 *, 35 judge: LLMInvocation, 36 safety: LLMInvocation, 37 content_clean: LLMInvocation, 38 focus_summary: LLMInvocation, 39 logger: logging.Logger | None = None, 40 monotonic: Callable[[], float] = time.monotonic, 41 ) -> None: 42 self._clients = dict(clients) 43 self._judge = judge 44 self._safety = safety 45 self._content_clean = content_clean 46 self._focus_summary = focus_summary 47 self._logger = logger or logging.getLogger(__name__) 48 self._monotonic = monotonic 49 50 @property 51 def judge_provider(self) -> str: 52 return self._judge.provider 53 54 async def judge(self, candidate: str) -> StageDecision: 55 return await self._run_decision_stage( 56 self._judge, 57 "judge", 58 input_chars=len(candidate), 59 operation=lambda: self._client(self._judge.provider).complete_json( 60 self._judge, 61 judge_messages(candidate), 62 ), 63 ) 64 65 async def safety(self, content: str) -> StageDecision: 66 return await self._run_decision_stage( 67 self._safety, 68 "safety", 69 input_chars=len(content), 70 operation=lambda: self._client(self._safety.provider).complete_json( 71 self._safety, 72 safety_messages(content), 73 ), 74 ) 75 76 async def content_clean(self, raw_content: str) -> str: 77 return await self._run_text_stage( 78 self._content_clean, 79 "content_clean", 80 input_chars=len(raw_content), 81 operation=lambda: self._client(self._content_clean.provider).complete_text( 82 self._content_clean, 83 content_clean_messages(raw_content), 84 ), 85 ) 86 87 async def focus_summary(self, content: str, focus: str) -> str: 88 normalized_focus = focus.strip() 89 if not normalized_focus: 90 raise ExecutionFailure(ErrorCode.LLM_STAGE_FAILED, "focus must be non-empty") 91 return await self._run_text_stage( 92 self._focus_summary, 93 "focus_summary", 94 input_chars=len(content), 95 focus_chars=len(normalized_focus), 96 operation=lambda: self._client(self._focus_summary.provider).complete_text( 97 self._focus_summary, 98 focus_summary_messages(content, normalized_focus), 99 ), 100 ) 101 102 async def llm_search_markdown(self, invocation: LLMInvocation, prompt: str) -> str: 103 return await self._run_text_stage( 104 invocation, 105 "llm_search", 106 input_chars=len(prompt), 107 operation=lambda: self._client(invocation.provider).complete_text( 108 invocation, 109 llm_search_messages(prompt), 110 ), 111 ) 112 113 async def llm_paper_search_markdown( 114 self, 115 invocation: LLMInvocation, 116 prompt: str, 117 ) -> str: 118 return await self._run_text_stage( 119 invocation, 120 "llm_paper_search", 121 input_chars=len(prompt), 122 operation=lambda: self._client(invocation.provider).complete_text( 123 invocation, 124 llm_paper_search_messages(prompt), 125 ), 126 ) 127 128 async def _run_decision_stage( 129 self, 130 invocation: LLMInvocation, 131 stage: str, 132 *, 133 input_chars: int, 134 operation: Callable[[], Awaitable[Mapping[str, object]]], 135 ) -> StageDecision: 136 started = self._stage_started(invocation, stage, input_chars=input_chars) 137 try: 138 decision = self._parse_decision(await operation()) 139 except asyncio.CancelledError: 140 self._stage_cancelled(invocation, stage, started) 141 raise 142 except Exception as exc: 143 self._stage_failed(invocation, stage, started, exc) 144 raise 145 log_event( 146 self._logger, 147 logging.DEBUG, 148 "llm_stage_completed", 149 provider=invocation.provider, 150 stage=stage, 151 model=invocation.model, 152 ok=decision.ok, 153 reason_present=bool(decision.reason), 154 elapsed_ms=elapsed_ms(self._monotonic, started), 155 ) 156 return decision 157 158 async def _run_text_stage( 159 self, 160 invocation: LLMInvocation, 161 stage: str, 162 *, 163 input_chars: int, 164 operation: Callable[[], Awaitable[str]], 165 focus_chars: int = 0, 166 ) -> str: 167 started = self._stage_started( 168 invocation, 169 stage, 170 input_chars=input_chars, 171 focus_chars=focus_chars, 172 ) 173 try: 174 text = self._require_non_empty(await operation(), stage.replace("_", "-")) 175 except asyncio.CancelledError: 176 self._stage_cancelled(invocation, stage, started) 177 raise 178 except Exception as exc: 179 self._stage_failed(invocation, stage, started, exc) 180 raise 181 log_event( 182 self._logger, 183 logging.DEBUG, 184 "llm_stage_completed", 185 provider=invocation.provider, 186 stage=stage, 187 model=invocation.model, 188 output_chars=len(text), 189 elapsed_ms=elapsed_ms(self._monotonic, started), 190 ) 191 return text 192 193 def _stage_started( 194 self, 195 invocation: LLMInvocation, 196 stage: str, 197 *, 198 input_chars: int, 199 focus_chars: int = 0, 200 ) -> float: 201 started = self._monotonic() 202 log_event( 203 self._logger, 204 logging.DEBUG, 205 "llm_stage_started", 206 provider=invocation.provider, 207 stage=stage, 208 model=invocation.model, 209 input_chars=input_chars, 210 focus_present=focus_chars > 0, 211 focus_chars=focus_chars, 212 ) 213 return started 214 215 def _stage_failed( 216 self, 217 invocation: LLMInvocation, 218 stage: str, 219 started: float, 220 exc: Exception, 221 ) -> None: 222 log_event( 223 self._logger, 224 logging.DEBUG, 225 "llm_stage_failed", 226 provider=invocation.provider, 227 stage=stage, 228 model=invocation.model, 229 error_type=type(exc).__name__, 230 elapsed_ms=elapsed_ms(self._monotonic, started), 231 ) 232 233 def _stage_cancelled(self, invocation: LLMInvocation, stage: str, started: float) -> None: 234 log_event( 235 self._logger, 236 logging.DEBUG, 237 "llm_stage_cancelled", 238 provider=invocation.provider, 239 stage=stage, 240 model=invocation.model, 241 elapsed_ms=elapsed_ms(self._monotonic, started), 242 ) 243 244 def _client(self, provider: str) -> LLMClient: 245 client = self._clients.get(provider) 246 if client is None: 247 raise ExecutionFailure( 248 ErrorCode.LLM_STAGE_FAILED, 249 f"LLM provider is not initialized: {provider}", 250 ) 251 return client 252 253 @staticmethod 254 def _require_non_empty(text: str, stage: str) -> str: 255 normalized = text.strip() 256 if not normalized: 257 raise ExecutionFailure( 258 ErrorCode.LLM_STAGE_FAILED, 259 f"{stage} returned empty text", 260 ) 261 return normalized 262 263 @staticmethod 264 def _parse_decision(payload: Mapping[str, object]) -> StageDecision: 265 ok = payload.get("ok") 266 if not isinstance(ok, bool): 267 raise ExecutionFailure( 268 ErrorCode.LLM_STAGE_FAILED, 269 "LLM decision response requires boolean ok", 270 ) 271 reason = payload.get("reason", "") 272 if not isinstance(reason, str): 273 raise ExecutionFailure( 274 ErrorCode.LLM_STAGE_FAILED, 275 "LLM decision reason must be a string", 276 ) 277 return StageDecision(ok=ok, reason=reason.strip())
LLMStages( clients: Mapping[str, agent_search_gateway.providers.contracts.LLMClient], *, judge: agent_search_gateway.models.LLMInvocation, safety: agent_search_gateway.models.LLMInvocation, content_clean: agent_search_gateway.models.LLMInvocation, focus_summary: agent_search_gateway.models.LLMInvocation, logger: logging.Logger | None = None, monotonic: Callable[[], float] = <built-in function monotonic>)
31 def __init__( 32 self, 33 clients: Mapping[str, LLMClient], 34 *, 35 judge: LLMInvocation, 36 safety: LLMInvocation, 37 content_clean: LLMInvocation, 38 focus_summary: LLMInvocation, 39 logger: logging.Logger | None = None, 40 monotonic: Callable[[], float] = time.monotonic, 41 ) -> None: 42 self._clients = dict(clients) 43 self._judge = judge 44 self._safety = safety 45 self._content_clean = content_clean 46 self._focus_summary = focus_summary 47 self._logger = logger or logging.getLogger(__name__) 48 self._monotonic = monotonic
async def
content_clean(self, raw_content: str) -> str:
76 async def content_clean(self, raw_content: str) -> str: 77 return await self._run_text_stage( 78 self._content_clean, 79 "content_clean", 80 input_chars=len(raw_content), 81 operation=lambda: self._client(self._content_clean.provider).complete_text( 82 self._content_clean, 83 content_clean_messages(raw_content), 84 ), 85 )
async def
focus_summary(self, content: str, focus: str) -> str:
87 async def focus_summary(self, content: str, focus: str) -> str: 88 normalized_focus = focus.strip() 89 if not normalized_focus: 90 raise ExecutionFailure(ErrorCode.LLM_STAGE_FAILED, "focus must be non-empty") 91 return await self._run_text_stage( 92 self._focus_summary, 93 "focus_summary", 94 input_chars=len(content), 95 focus_chars=len(normalized_focus), 96 operation=lambda: self._client(self._focus_summary.provider).complete_text( 97 self._focus_summary, 98 focus_summary_messages(content, normalized_focus), 99 ), 100 )
async def
llm_search_markdown( self, invocation: agent_search_gateway.models.LLMInvocation, prompt: str) -> str:
102 async def llm_search_markdown(self, invocation: LLMInvocation, prompt: str) -> str: 103 return await self._run_text_stage( 104 invocation, 105 "llm_search", 106 input_chars=len(prompt), 107 operation=lambda: self._client(invocation.provider).complete_text( 108 invocation, 109 llm_search_messages(prompt), 110 ), 111 )
async def
llm_paper_search_markdown( self, invocation: agent_search_gateway.models.LLMInvocation, prompt: str) -> str:
113 async def llm_paper_search_markdown( 114 self, 115 invocation: LLMInvocation, 116 prompt: str, 117 ) -> str: 118 return await self._run_text_stage( 119 invocation, 120 "llm_paper_search", 121 input_chars=len(prompt), 122 operation=lambda: self._client(invocation.provider).complete_text( 123 invocation, 124 llm_paper_search_messages(prompt), 125 ), 126 )