Текстовый файл из учебного комплекта, не свидетельство установленного runtime. Вернуться к учебнику · Исходный файл
1 from __future__ import annotations 2 3 import asyncio 4 import json 5 import logging 6 import time 7 from dataclasses import dataclass 8 from typing import Any 9 from uuid import UUID, uuid4 10 11 from nats.aio.client import Client as NATS 12 from nats.errors import NoRespondersError 13 from nats.errors import TimeoutError as NatsTimeoutError 14 15 from src_agent.repositories.runtime_state import ( 16 RuntimeStateConflict, 17 load_runtime_state, 18 save_runtime_state, 19 ) 20 from src_agent.utils.db import ( 21 Database, 22 chat_turn_exists, 23 create_session, 24 persist_chat_turn, 25 ) 26 from src_shared.contracts import ( 27 AgentInboundRequest, 28 ErrorInfo, 29 WorkflowRunRequest, 30 WorkflowRunResponse, 31 WorkflowRuntimeState, 32 now_ts_ms, 33 ) 34 35 logger = logging.getLogger("src_agent.runner") 36 37 38 @dataclass(frozen=True) 39 class RunnerResult: 40 session_id: UUID 41 workflow: WorkflowRunResponse 42 43 44 @dataclass 45 class AgentRunner: 46 nc: NATS 47 db: Database 48 workflow_subject: str 49 user_id: str 50 workflow_timeout_s: int 51 max_runtime_conflict_retries: int = 3 52 workflow_no_responders_retries: int = 4 53 workflow_no_responders_retry_delay_s: float = 1.0 54 55 async def run(self, req: AgentInboundRequest) -> RunnerResult: 56 # Ensure the `sessions` row exists even when the session_id is provided externally. 57 session_id = req.session_id or uuid4() 58 session_id = UUID( 59 await create_session(self.db, user_id=self.user_id, session_id=str(session_id)) 60 ) 61 62 runtime = await load_runtime_state(self.db, session_id=session_id) 63 workflow = await self._run_with_optimistic_lock(req=req, session_id=session_id, runtime=runtime) 64 return RunnerResult(session_id=session_id, workflow=workflow) 65 66 @staticmethod 67 def _failed_workflow_response( 68 *, 69 trace_id, 70 correlation_id, 71 request_id, 72 session_id, 73 ts_ms: int, 74 code: str, 75 message: str, 76 next_runtime: WorkflowRuntimeState, 77 details: dict[str, Any] | None = None, 78 ) -> WorkflowRunResponse: 79 return WorkflowRunResponse( 80 trace_id=trace_id, 81 correlation_id=correlation_id, 82 request_id=request_id, 83 session_id=session_id, 84 ts_ms=ts_ms, 85 status="FAILED", 86 result="", 87 errors=[ErrorInfo(code=code, message=message, details=details)], 88 next_runtime=next_runtime, 89 ) 90 91 async def _run_with_optimistic_lock( 92 self, 93 *, 94 req: AgentInboundRequest, 95 session_id: UUID, 96 runtime: WorkflowRuntimeState, 97 ) -> WorkflowRunResponse: 98 correlation_id = req.correlation_id or req.trace_id 99 edit_turn_id = self._edit_turn_id(req.edit) 100 if edit_turn_id: 101 exists = await chat_turn_exists( 102 self.db, 103 session_id=str(session_id), 104 turn_id=edit_turn_id, 105 ) 106 if not exists: 107 return self._failed_workflow_response( 108 trace_id=req.trace_id, 109 correlation_id=correlation_id, 110 request_id=req.request_id, 111 session_id=session_id, 112 ts_ms=req.ts_ms, 113 code="edit_turn_not_found", 114 message="Edited turn was not found in session history.", 115 next_runtime=runtime, 116 ) 117 # Edits rebuild the branch from the edited user turn, so pending workflow lock must be reset. 118 runtime = runtime.model_copy(update={"active_workflow_id": None, "pending": None}) 119 user_turn_id = edit_turn_id or (req.turn_id.strip() if isinstance(req.turn_id, str) and req.turn_id.strip() else str(uuid4())) 120 121 last_error: Exception | None = None 122 for attempt in range(self.max_runtime_conflict_retries + 1): 123 inner_request_id = req.request_id if attempt == 0 else uuid4() 124 inner_ts_ms = req.ts_ms if attempt == 0 else now_ts_ms() 125 126 workflow_req = WorkflowRunRequest( 127 trace_id=req.trace_id, 128 correlation_id=correlation_id, 129 request_id=inner_request_id, 130 session_id=session_id, 131 ts_ms=inner_ts_ms, 132 text=req.text, 133 turn_id=user_turn_id, 134 edit=req.edit, 135 runtime=runtime, 136 ) 137 138 workflow_resp = await self._call_workflow(workflow_req) 139 140 try: 141 assistant_turn_id = self._assistant_turn_id(workflow_resp) or str(uuid4()) 142 assistant_text = self._assistant_text(workflow_resp) 143 async with self.db.transaction() as connection: 144 await save_runtime_state( 145 self.db, 146 session_id=session_id, 147 expected_version=runtime.version, 148 next_state=workflow_resp.next_runtime, 149 connection=connection, 150 ) 151 if assistant_text: 152 await persist_chat_turn( 153 self.db, 154 session_id=str(session_id), 155 user_turn_id=user_turn_id, 156 user_text=req.text, 157 assistant_turn_id=assistant_turn_id, 158 assistant_text=assistant_text, 159 edit_from_turn_id=edit_turn_id, 160 user_meta={"request_id": str(req.request_id)}, 161 assistant_meta={ 162 "request_id": str(workflow_resp.request_id), 163 "status": workflow_resp.status, 164 "user_turn_id": user_turn_id, 165 }, 166 connection=connection, 167 ) 168 return self._inject_turn_ids(workflow_resp, user_turn_id=user_turn_id, assistant_turn_id=assistant_turn_id) 169 except ValueError as exc: 170 if str(exc) == "edit_turn_not_found": 171 return self._failed_workflow_response( 172 trace_id=req.trace_id, 173 correlation_id=correlation_id, 174 request_id=req.request_id, 175 session_id=session_id, 176 ts_ms=req.ts_ms, 177 code="edit_turn_not_found", 178 message="Edited turn was not found in session history.", 179 next_runtime=runtime, 180 ) 181 raise 182 except RuntimeStateConflict as exc: 183 last_error = exc 184 if attempt >= self.max_runtime_conflict_retries: 185 break 186 runtime = await load_runtime_state(self.db, session_id=session_id) 187 188 logger.warning("runtime_state update failed after retries. session_id=%s", str(session_id)) 189 return self._failed_workflow_response( 190 trace_id=req.trace_id, 191 correlation_id=correlation_id, 192 request_id=req.request_id, 193 session_id=session_id, 194 ts_ms=req.ts_ms, 195 code="runtime_state_conflict", 196 message="Failed to update runtime state due to concurrent updates.", 197 details={"error": str(last_error)} if last_error else None, 198 next_runtime=runtime, 199 ) 200 201 async def _call_workflow(self, req: WorkflowRunRequest) -> WorkflowRunResponse: 202 payload = req.model_dump_json().encode("utf-8") 203 total_attempts = self.workflow_no_responders_retries + 1 204 last_exc: Exception | None = None 205 for attempt in range(total_attempts): 206 try: 207 started_at = time.monotonic() 208 logger.info( 209 "workflow request start subject=%s attempt=%s/%s request_id=%s session_id=%s timeout_s=%s", 210 self.workflow_subject, 211 attempt + 1, 212 total_attempts, 213 str(req.request_id), 214 str(req.session_id) if req.session_id else "<none>", 215 self.workflow_timeout_s, 216 ) 217 msg = await self.nc.request(self.workflow_subject, payload, timeout=self.workflow_timeout_s) 218 duration_ms = int((time.monotonic() - started_at) * 1000) 219 raw = json.loads(msg.data.decode("utf-8")) 220 resp = WorkflowRunResponse.model_validate(raw) 221 logger.info( 222 "workflow request done subject=%s request_id=%s status=%s duration_ms=%s", 223 self.workflow_subject, 224 str(req.request_id), 225 resp.status, 226 duration_ms, 227 ) 228 return resp 229 except NatsTimeoutError: 230 logger.error( 231 "Workflow request timeout subject=%s timeout=%ss request_id=%s session_id=%s", 232 self.workflow_subject, 233 self.workflow_timeout_s, 234 str(req.request_id), 235 str(req.session_id) if req.session_id else "<none>", 236 ) 237 return self._failed_workflow_response( 238 trace_id=req.trace_id, 239 correlation_id=req.correlation_id, 240 request_id=req.request_id, 241 session_id=req.session_id, 242 ts_ms=now_ts_ms(), 243 code="workflow_timeout", 244 message="Сервис сценариев не успел ответить вовремя. Попробуйте повторить запрос.", 245 details={ 246 "subject": self.workflow_subject, 247 "timeout_s": self.workflow_timeout_s, 248 }, 249 next_runtime=req.runtime, 250 ) 251 except NoRespondersError as exc: 252 last_exc = exc 253 if attempt >= total_attempts - 1: 254 break 255 logger.warning( 256 "No responders for subject=%s request_id=%s (attempt=%s/%s), retrying in %.1fs", 257 self.workflow_subject, 258 str(req.request_id), 259 attempt + 1, 260 total_attempts, 261 self.workflow_no_responders_retry_delay_s, 262 ) 263 await asyncio.sleep(self.workflow_no_responders_retry_delay_s) 264 265 waited_s = self.workflow_no_responders_retries * self.workflow_no_responders_retry_delay_s 266 return self._failed_workflow_response( 267 trace_id=req.trace_id, 268 correlation_id=req.correlation_id, 269 request_id=req.request_id, 270 session_id=req.session_id, 271 ts_ms=now_ts_ms(), 272 code="workflow_unavailable", 273 message=( 274 "Сервис сценариев временно недоступен " 275 f"(нет responder для NATS subject `{self.workflow_subject}` ~{waited_s:.0f}с). " 276 "Повторите запрос через несколько секунд." 277 ), 278 details={ 279 "subject": self.workflow_subject, 280 "error": str(last_exc) if last_exc else None, 281 }, 282 next_runtime=req.runtime, 283 ) 284 285 @staticmethod 286 def _edit_turn_id(edit: Any) -> str | None: 287 if not isinstance(edit, dict): 288 return None 289 value = edit.get("turn_id") 290 if not isinstance(value, str): 291 return None 292 value = value.strip() 293 return value or None 294 295 @staticmethod 296 def _assistant_text(resp: WorkflowRunResponse) -> str: 297 if isinstance(resp.result, str) and resp.result.strip(): 298 return resp.result.strip() 299 if resp.errors and isinstance(resp.errors[0].message, str): 300 return resp.errors[0].message.strip() 301 return "" 302 303 @staticmethod 304 def _assistant_turn_id(resp: WorkflowRunResponse) -> str | None: 305 if not isinstance(resp.client_handler, dict): 306 return None 307 payload = resp.client_handler.get("payload") 308 if not isinstance(payload, dict): 309 return None 310 value = payload.get("assistant_turn_id") 311 if not isinstance(value, str): 312 return None 313 value = value.strip() 314 return value or None 315 316 @staticmethod 317 def _inject_turn_ids(resp: WorkflowRunResponse, *, user_turn_id: str, assistant_turn_id: str) -> WorkflowRunResponse: 318 handler = dict(resp.client_handler or {}) 319 payload = handler.get("payload") if isinstance(handler.get("payload"), dict) else {} 320 payload = dict(payload or {}) 321 payload.setdefault("user_turn_id", user_turn_id) 322 payload.setdefault("assistant_turn_id", assistant_turn_id) 323 if handler: 324 handler["payload"] = payload 325 return resp.model_copy(update={"client_handler": handler})