Текстовый файл из учебного комплекта, не свидетельство установленного runtime. Вернуться к учебнику · Исходный файл
1 from __future__ import annotations 2 3 import json 4 import logging 5 import time 6 from dataclasses import dataclass 7 from typing import Any 8 from uuid import UUID, uuid4 9 10 from nats import errors as nats_errors 11 from nats.aio.client import Client as NATS 12 13 from src_shared.contracts import ( 14 ErrorInfo, 15 LlmRequest, 16 LlmResponse, 17 ToolCallRequest, 18 ToolCallResponse, 19 WorkflowRunRequest, 20 now_ts_ms, 21 ) 22 23 logger = logging.getLogger("src_langgraph.runtime_io") 24 25 26 def _child_trace(parent: WorkflowRunRequest, *, request_id: UUID | None = None) -> dict[str, Any]: 27 """Создает trace-блок для дочернего вызова LLM/tool на базе входного запроса.""" 28 # Preserve trace/correlation across nested LLM/tool calls while generating per-call request_id. 29 return { 30 "trace_id": parent.trace_id, 31 "correlation_id": parent.correlation_id, 32 "request_id": request_id or uuid4(), 33 "session_id": parent.session_id, 34 "ts_ms": now_ts_ms(), 35 } 36 37 38 @dataclass(slots=True) 39 class RuntimeIO: 40 """Транспортный слой для request/reply вызовов LLM и tools по NATS.""" 41 42 nc: NATS 43 llm_subject_prefix: str 44 tools_subject_prefix: str 45 user_id: str 46 request_timeout_s: float = 120.0 47 48 async def call_llm( 49 self, 50 *, 51 parent: WorkflowRunRequest, 52 mode: str, 53 input_data: dict[str, Any], 54 constraints: dict[str, Any] | None = None, 55 ) -> LlmResponse: 56 """Отправляет запрос в LLM-сервис и валидирует структуру ответа.""" 57 req = LlmRequest( 58 **_child_trace(parent), 59 mode=mode, # type: ignore[arg-type] 60 input=input_data, 61 constraints=constraints or {}, 62 ) 63 subject = f"{self.llm_subject_prefix}{self.user_id}" 64 started_at = time.monotonic() 65 logger.info( 66 "call_llm start subject=%s mode=%s request_id=%s session_id=%s", 67 subject, 68 mode, 69 str(req.request_id), 70 str(req.session_id) if req.session_id else "<none>", 71 ) 72 try: 73 raw = await self._request_json(subject, req.model_dump_json().encode("utf-8")) 74 resp = LlmResponse.model_validate(raw) 75 duration_ms = int((time.monotonic() - started_at) * 1000) 76 logger.info( 77 "call_llm done subject=%s mode=%s request_id=%s ok=%s duration_ms=%s", 78 subject, 79 mode, 80 str(req.request_id), 81 resp.ok, 82 duration_ms, 83 ) 84 return resp 85 except nats_errors.TimeoutError: 86 duration_ms = int((time.monotonic() - started_at) * 1000) 87 logger.warning( 88 "call_llm timeout subject=%s mode=%s request_id=%s timeout_s=%s duration_ms=%s", 89 subject, 90 mode, 91 str(req.request_id), 92 self.request_timeout_s, 93 duration_ms, 94 ) 95 return LlmResponse( 96 trace_id=req.trace_id, 97 correlation_id=req.correlation_id, 98 request_id=req.request_id, 99 session_id=req.session_id, 100 ts_ms=now_ts_ms(), 101 ok=False, 102 data=None, 103 error=ErrorInfo( 104 code="nats_timeout", 105 message=f"Timeout waiting for LLM response from '{subject}'", 106 ), 107 ) 108 except Exception as exc: 109 duration_ms = int((time.monotonic() - started_at) * 1000) 110 logger.exception( 111 "call_llm transport error subject=%s mode=%s request_id=%s duration_ms=%s", 112 subject, 113 mode, 114 str(req.request_id), 115 duration_ms, 116 ) 117 return LlmResponse( 118 trace_id=req.trace_id, 119 correlation_id=req.correlation_id, 120 request_id=req.request_id, 121 session_id=req.session_id, 122 ts_ms=now_ts_ms(), 123 ok=False, 124 data=None, 125 error=ErrorInfo( 126 code="llm_transport_error", 127 message=f"LLM request failed for '{subject}': {exc}", 128 ), 129 ) 130 131 async def call_tool( 132 self, 133 *, 134 parent: WorkflowRunRequest, 135 tool_name: str, 136 args: dict[str, Any], 137 ) -> ToolCallResponse: 138 """Вызывает tool-сервис по имени инструмента и валидирует ответ.""" 139 req = ToolCallRequest( 140 **_child_trace(parent), 141 tool_name=tool_name, 142 args=args, 143 ) 144 subject = f"{self.tools_subject_prefix}{tool_name}" 145 started_at = time.monotonic() 146 logger.info( 147 "call_tool start subject=%s tool=%s request_id=%s session_id=%s", 148 subject, 149 tool_name, 150 str(req.request_id), 151 str(req.session_id) if req.session_id else "<none>", 152 ) 153 try: 154 raw = await self._request_json(subject, req.model_dump_json().encode("utf-8")) 155 resp = ToolCallResponse.model_validate(raw) 156 duration_ms = int((time.monotonic() - started_at) * 1000) 157 logger.info( 158 "call_tool done subject=%s tool=%s request_id=%s ok=%s duration_ms=%s", 159 subject, 160 tool_name, 161 str(req.request_id), 162 resp.ok, 163 duration_ms, 164 ) 165 return resp 166 except nats_errors.TimeoutError: 167 duration_ms = int((time.monotonic() - started_at) * 1000) 168 logger.warning( 169 "call_tool timeout subject=%s tool=%s request_id=%s timeout_s=%s duration_ms=%s", 170 subject, 171 tool_name, 172 str(req.request_id), 173 self.request_timeout_s, 174 duration_ms, 175 ) 176 return ToolCallResponse( 177 trace_id=req.trace_id, 178 correlation_id=req.correlation_id, 179 request_id=req.request_id, 180 session_id=req.session_id, 181 ts_ms=now_ts_ms(), 182 ok=False, 183 artifact_key=None, 184 data=None, 185 error=ErrorInfo( 186 code="nats_timeout", 187 message=f"Timeout waiting for tool response from '{subject}'", 188 ), 189 ) 190 except Exception as exc: 191 duration_ms = int((time.monotonic() - started_at) * 1000) 192 logger.exception( 193 "call_tool transport error subject=%s tool=%s request_id=%s duration_ms=%s", 194 subject, 195 tool_name, 196 str(req.request_id), 197 duration_ms, 198 ) 199 return ToolCallResponse( 200 trace_id=req.trace_id, 201 correlation_id=req.correlation_id, 202 request_id=req.request_id, 203 session_id=req.session_id, 204 ts_ms=now_ts_ms(), 205 ok=False, 206 artifact_key=None, 207 data=None, 208 error=ErrorInfo( 209 code="tool_transport_error", 210 message=f"Tool request failed for '{subject}': {exc}", 211 ), 212 ) 213 214 async def _request_json(self, subject: str, payload: bytes) -> dict[str, Any]: 215 """Базовый request/reply helper: отправляет payload и требует JSON object в ответе.""" 216 msg = await self.nc.request(subject, payload, timeout=self.request_timeout_s) 217 data = json.loads(msg.data.decode("utf-8")) 218 if not isinstance(data, dict): 219 raise RuntimeError(f"Invalid JSON response from {subject}") 220 return data