"""OpenAI-compatible LLM client for Modal vLLM with local tool execution.""" import json import logging import re import time from typing import cast from openai import OpenAI from openai.types.chat import ChatCompletionMessageParam from vital_types.llm import ChatMessage, LlmResponse, ToolCallRequest from llm.chat_helpers import ( extract_food_description, infer_meal_type_from_text, turn_had_successful_tool, user_requests_food_log, recent_user_text, ) from llm.config import LlmConfig, get_llm_config from llm.system_prompt import build_system_prompt, trim_prompt_to_budget from llm.tool_runner import execute_tool, tool_result_to_json from llm.tools import TOOL_SCHEMAS logger = logging.getLogger(__name__) _JSON_BLOCK_PATTERN = re.compile(r"\{[\s\S]*\}") def _json_schema_response_format(schema: dict[str, object]) -> dict[str, object]: """Build a schema-constrained response_format for vLLM structured outputs.""" return { "type": "json_schema", "json_schema": {"name": "structured_output", "schema": schema}, } _PROMPT_TOOL_PATTERN = re.compile( r'\{\s*"tool"\s*:\s*"([^"]+)"\s*,\s*"args"\s*:\s*(\{[\s\S]*?\})\s*\}' ) # Schema for the constrained retry when a free-form coach reply fails # validation. Tool calls and constrained decoding cannot be combined in one # request, so the tool loop runs free and this schema guards the final answer. COACH_REPLY_JSON_SCHEMA: dict[str, object] = { "type": "object", "properties": { "reply": { "type": "string", "minLength": 1, "maxLength": 2000, "description": "The coach's conversational answer to the user, plain text.", }, }, "required": ["reply"], "additionalProperties": False, } _MAX_CHAT_REPLY_CHARS = 4000 def _validate_chat_reply(content: str) -> str | None: """Return a cleaned coach reply, or None when the text fails validation. Rejects empty output and replies that are leaked JSON (tool-call spill or a raw object instead of prose), which Nemotron occasionally emits. """ trimmed = content.strip() if not trimmed: return None if trimmed.startswith("{") and trimmed.endswith("}"): return None if _PROMPT_TOOL_PATTERN.search(trimmed): return None if len(trimmed) > _MAX_CHAT_REPLY_CHARS: trimmed = trimmed[:_MAX_CHAT_REPLY_CHARS].rstrip() + "..." return trimmed class LlmClient: """Calls the remote Nemotron endpoint and runs tools locally.""" def __init__(self, config: LlmConfig | None = None): self.config = config or get_llm_config() self._openai = OpenAI( base_url=self.config.base_url, api_key=self.config.api_key, timeout=self.config.request_timeout_seconds, ) def _to_api_messages(self, messages: list[ChatMessage]) -> list[ChatCompletionMessageParam]: """Convert internal messages to OpenAI API format.""" api_messages: list[ChatCompletionMessageParam] = [] for message in messages: if message.role == "tool": api_messages.append( { "role": "tool", "content": message.content, "tool_call_id": message.tool_call_id or "", } ) continue api_messages.append( { "role": message.role, "content": message.content, } ) return api_messages def _parse_tool_arguments(self, raw_arguments: str) -> dict[str, object]: """Parse tool call arguments JSON from the model.""" if not raw_arguments.strip(): return {} parsed = json.loads(raw_arguments) if not isinstance(parsed, dict): raise ValueError("Tool arguments must be a JSON object.") return cast(dict[str, object], parsed) def _extract_prompt_based_tool_call(self, content: str) -> ToolCallRequest | None: """Fallback: parse {"tool": "...", "args": {...}} from plain text.""" match = _PROMPT_TOOL_PATTERN.search(content) if match is None: return None tool_name = match.group(1) args_json = match.group(2) arguments = self._parse_tool_arguments(args_json) return ToolCallRequest(id="prompt_tool_0", name=tool_name, arguments=arguments) def _completion_to_response( self, message_content: str | None, message, finish_reason: str | None = None, ) -> LlmResponse: """Map an OpenAI completion message to an internal LlmResponse.""" tool_calls: list[ToolCallRequest] = [] raw_tool_calls = getattr(message, "tool_calls", None) if raw_tool_calls: for index, tool_call in enumerate(raw_tool_calls): function_block = tool_call.function arguments = self._parse_tool_arguments(function_block.arguments or "{}") tool_calls.append( ToolCallRequest( id=tool_call.id or f"tool_{index}", name=function_block.name, arguments=arguments, ) ) content = message_content or "" if not tool_calls: fallback = self._extract_prompt_based_tool_call(content) if fallback is not None: tool_calls = [fallback] content = "" resolved_finish_reason = finish_reason or getattr(message, "finish_reason", None) return LlmResponse( content=content, tool_calls=tool_calls, raw_finish_reason=resolved_finish_reason, ) def complete( self, messages: list[ChatMessage], use_tools: bool = True, json_mode: bool = False, temperature: float | None = None, json_schema: dict[str, object] | None = None, ) -> LlmResponse: """Send one chat completion request to the remote LLM.""" api_messages = self._to_api_messages(messages) request_temperature = temperature if request_temperature is None: request_temperature = ( self.config.json_temperature if json_mode else self.config.tool_temperature ) kwargs: dict[str, object] = { "model": self.config.model_id, "messages": api_messages, "temperature": request_temperature, "max_tokens": self.config.max_tokens, } if use_tools: kwargs["tools"] = TOOL_SCHEMAS kwargs["tool_choice"] = "auto" if json_schema is not None: kwargs["response_format"] = _json_schema_response_format(json_schema) elif json_mode: kwargs["response_format"] = {"type": "json_object"} completion = self._openai.chat.completions.create(**kwargs) choice = completion.choices[0] message = choice.message return self._completion_to_response( message.content, message, finish_reason=choice.finish_reason, ) def _run_completion_on_api_messages( self, api_messages: list[ChatCompletionMessageParam], use_tools: bool, json_mode: bool = False, temperature: float | None = None, tools_override: list[dict[str, object]] | None = None, max_tokens: int | None = None, json_schema: dict[str, object] | None = None, ) -> LlmResponse: """Send a completion using raw OpenAI API message dicts.""" request_temperature = temperature if request_temperature is None: request_temperature = ( self.config.json_temperature if json_mode else self.config.tool_temperature ) kwargs: dict[str, object] = { "model": self.config.model_id, "messages": api_messages, "temperature": request_temperature, "max_tokens": max_tokens if max_tokens is not None else self.config.max_tokens, } if use_tools: kwargs["tools"] = tools_override if tools_override is not None else TOOL_SCHEMAS kwargs["tool_choice"] = "auto" if json_schema is not None: kwargs["response_format"] = _json_schema_response_format(json_schema) elif json_mode: kwargs["response_format"] = {"type": "json_object"} completion = self._openai.chat.completions.create(**kwargs) choice = completion.choices[0] message = choice.message return self._completion_to_response( message.content, message, finish_reason=choice.finish_reason, ) def chat( self, user_message: str, extra_messages: list[ChatMessage] | None = None, use_tools: bool = True, ) -> str: """Run a full coach conversation turn with optional tool loop.""" system_prompt = trim_prompt_to_budget( build_system_prompt(include_tools_instruction=use_tools), self.config.context_limit_tokens, ) api_messages: list[ChatCompletionMessageParam] = [ {"role": "system", "content": system_prompt}, ] if extra_messages: api_messages.extend(self._to_api_messages(extra_messages)) api_messages.append({"role": "user", "content": user_message}) food_log_requested = user_requests_food_log(user_message, extra_messages) for _iteration in range(self.config.max_tool_iterations): response = self._run_completion_on_api_messages(api_messages, use_tools=use_tools) if not response.tool_calls: if food_log_requested and not turn_had_successful_tool(api_messages, "log_food"): logger.warning( "[chat] User requested food log but model replied without log_food — " "nudging for write tool.", ) api_messages.append( { "role": "user", "content": ( "Before answering, call log_food now with meal_type and " "food_description from what I said I ate. Do not claim you " "logged anything until log_food returns success:true." ), } ) continue validated_reply = _validate_chat_reply(response.content) if validated_reply is not None: return validated_reply logger.warning( "[chat] Free-form reply failed validation (%s chars) — " "retrying with strict reply schema.", len(response.content), ) break assistant_tool_calls: list[dict[str, object]] = [] for tool_call in response.tool_calls: assistant_tool_calls.append( { "id": tool_call.id, "type": "function", "function": { "name": tool_call.name, "arguments": json.dumps(tool_call.arguments), }, } ) # Drop assistant text when tool calls are present — Nemotron often # streams a partial answer in content alongside tools, which breaks # the tool loop and can interleave with local tool trace output. api_messages.append( { "role": "assistant", "content": "" if response.tool_calls else (response.content or ""), "tool_calls": assistant_tool_calls, } ) for tool_call in response.tool_calls: result = execute_tool(tool_call.name, tool_call.arguments) if result.success: logger.info( "[chat] Tool %s OK — %s", tool_call.name, json.dumps(result.result)[:200], ) else: logger.warning( "[chat] Tool %s failed: %s (args=%s)", tool_call.name, result.error, json.dumps(tool_call.arguments), ) api_messages.append( { "role": "tool", "content": tool_result_to_json(result), "tool_call_id": tool_call.id, } ) if food_log_requested and not turn_had_successful_tool(api_messages, "log_food"): context_text = recent_user_text(user_message, extra_messages) description = extract_food_description(context_text) if description: meal_type = infer_meal_type_from_text(context_text) fallback = execute_tool( "log_food", { "meal_type": meal_type, "food_description": description, }, ) if fallback.success: logger.info( "[chat] Fallback log_food saved %s: %s", meal_type, description, ) return ( f"Got it — I've logged your {meal_type}: {description}. " "You can see it on the Nutrition tab." ) logger.warning("[chat] Fallback log_food failed: %s", fallback.error) return self._final_chat_reply_with_schema(api_messages) def _final_chat_reply_with_schema( self, api_messages: list[ChatCompletionMessageParam], ) -> str: """Force the final coach answer through the strict reply schema. Mirrors the onboarding/daily-planner approach: constrained decoding guarantees shape, then the reply field is validated before use. """ constrained_messages = list(api_messages) constrained_messages.append( { "role": "user", "content": ( "Respond to the conversation above with your final answer as JSON " 'matching {"reply": ""}. The reply must be plain ' "conversational text — no tool calls, no nested JSON." ), } ) response: LlmResponse | None = None try: response = self._run_completion_on_api_messages( constrained_messages, use_tools=False, json_mode=True, json_schema=COACH_REPLY_JSON_SCHEMA, ) except Exception as error: # Older vLLM builds reject json_schema response_format; fall back. logger.warning( "[chat] json_schema response_format failed (%s); " "falling back to json_object mode.", error, ) if response is None: response = self._run_completion_on_api_messages( constrained_messages, use_tools=False, json_mode=True, ) try: payload = self._parse_json_object(response.content, log_context="chat") reply_value = payload.get("reply") if isinstance(reply_value, str): validated_reply = _validate_chat_reply(reply_value) if validated_reply is not None: return validated_reply except ValueError as error: logger.warning("[chat] Schema-constrained reply was not valid JSON: %s", error) # Last resort: surface whatever text exists rather than an empty bubble. fallback = response.content.strip() if fallback: return fallback return "I couldn't put together a reply just now. Please try again." def generate_json( self, user_prompt: str, system_addition: str = "", retry_once: bool = True, ) -> dict[str, object]: """Request structured JSON output with validation and one retry.""" system_prompt = trim_prompt_to_budget( build_system_prompt(include_tools_instruction=False), self.config.context_limit_tokens, ) if system_addition.strip(): system_prompt = f"{system_prompt}\n\n{system_addition.strip()}" messages: list[ChatMessage] = [ ChatMessage(role="system", content=system_prompt), ChatMessage(role="user", content=user_prompt), ] last_error = "Unknown JSON parse error." attempts = 2 if retry_once else 1 for attempt in range(attempts): response = self.complete(messages, use_tools=False, json_mode=True) try: return self._parse_json_object(response.content) except ValueError as error: last_error = str(error) if attempt + 1 >= attempts: break messages.append( ChatMessage( role="user", content=( "Your previous response was not valid JSON matching the required " f"schema. Error: {last_error}. Return ONLY valid JSON." ), ) ) raise ValueError(f"Failed to parse JSON from LLM response: {last_error}") def generate_onboarding_json( self, user_prompt: str, system_prompt: str, retry_once: bool = True, json_schema: dict[str, object] | None = None, ) -> dict[str, object]: """Request onboarding JSON using a dedicated system prompt (no runtime DB context).""" from llm.system_prompt import ROLE_DEFINITION, SAFETY_GUARDRAIL full_system = trim_prompt_to_budget( f"{ROLE_DEFINITION}\n\n{SAFETY_GUARDRAIL}\n\n{system_prompt.strip()}", self.config.context_limit_tokens, ) messages: list[ChatMessage] = [ ChatMessage(role="system", content=full_system), ChatMessage(role="user", content=user_prompt), ] last_error = "Unknown JSON parse error." attempts = 2 if retry_once else 1 schema_supported = json_schema is not None for attempt in range(attempts): response: LlmResponse | None = None if schema_supported and json_schema is not None: try: response = self.complete( messages, use_tools=False, json_mode=True, json_schema=json_schema, ) except Exception as error: # Older vLLM builds reject json_schema; fall back permanently. logger.warning( "[onboarding] json_schema response_format failed (%s); " "falling back to json_object mode.", error, ) schema_supported = False if response is None: response = self.complete(messages, use_tools=False, json_mode=True) try: return self._parse_json_object(response.content, log_context="onboarding") except ValueError as error: last_error = str(error) if attempt + 1 >= attempts: break messages.append( ChatMessage( role="user", content=( "Your previous response was not valid JSON matching the required " f"schema. Error: {last_error}. Return ONLY valid JSON." ), ) ) raise ValueError(f"Failed to parse JSON from LLM response: {last_error}") def _build_daily_schedule_messages( self, full_system: str, user_prompt: str, last_error: str, attempt: int, ) -> list[ChatCompletionMessageParam]: """Start a fresh message list for one daily-plan attempt.""" messages: list[ChatCompletionMessageParam] = [ {"role": "system", "content": full_system}, {"role": "user", "content": user_prompt}, ] if attempt > 0: messages.append( { "role": "user", "content": ( "Your previous JSON plan failed to parse. " f"Error: {last_error}. " "Return ONE compact valid JSON object: 5 hydration jobs, " "3 meals, 1 exercise. Keep each context under 60 characters. " "No markdown, no commentary." ), } ) return messages def _request_daily_plan_json( self, api_messages: list[ChatCompletionMessageParam], json_schema: dict[str, object] | None, ) -> dict[str, object]: """Request the daily-plan JSON, preferring schema-constrained decoding.""" schedule_tokens = self.config.daily_schedule_max_tokens json_response: LlmResponse | None = None if json_schema is not None: try: json_response = self._run_completion_on_api_messages( api_messages, use_tools=False, json_mode=True, max_tokens=schedule_tokens, json_schema=json_schema, ) except Exception as error: # Older vLLM builds reject json_schema response_format; fall back. logger.warning( "[daily_schedule] json_schema response_format failed (%s); " "falling back to json_object mode.", error, ) json_response = None if json_response is None: json_response = self._run_completion_on_api_messages( api_messages, use_tools=False, json_mode=True, max_tokens=schedule_tokens, ) if not json_response.content.strip(): logger.warning( "[daily_schedule] json_mode returned empty content; retrying without response_format." ) api_messages.append( { "role": "user", "content": "Output ONLY the JSON object. No other text before or after.", } ) json_response = self._run_completion_on_api_messages( api_messages, use_tools=False, json_mode=False, max_tokens=schedule_tokens, ) if json_response.raw_finish_reason == "length": preview = json_response.content[:300] logger.warning( "[daily_schedule] JSON truncated at %s tokens. Preview: %s...", schedule_tokens, preview, ) raise ValueError( f"JSON response truncated (max_tokens={schedule_tokens})." ) return self._parse_json_object(json_response.content, log_context="daily_schedule") def generate_daily_schedule_json( self, user_prompt: str, system_prompt: str, json_schema: dict[str, object] | None = None, ) -> dict[str, object]: """Generate the daily plan in one schema-constrained call with retries. All planning context (profile, frameworks, meds, logs, weather) is embedded in the user prompt, so no tool loop is needed — this keeps the full 8K context window available for the JSON output. """ from llm.system_prompt import ROLE_DEFINITION, SAFETY_GUARDRAIL full_system = trim_prompt_to_budget( f"{ROLE_DEFINITION}\n\n{SAFETY_GUARDRAIL}\n\n{system_prompt.strip()}", self.config.context_limit_tokens, ) last_error = "Unknown daily schedule error." attempts = self.config.daily_schedule_max_attempts for attempt in range(attempts): api_messages = self._build_daily_schedule_messages( full_system, user_prompt, last_error, attempt, ) try: return self._request_daily_plan_json(api_messages, json_schema) except Exception as error: last_error = str(error) logger.warning( "[daily_schedule] Attempt %s/%s failed: %s", attempt + 1, attempts, last_error, ) if attempt + 1 >= attempts: break delay = self.config.daily_schedule_retry_delay_seconds logger.info( "[daily_schedule] Retrying in %ss with a fresh context...", delay, ) time.sleep(delay) raise ValueError(f"Daily schedule LLM failed after {attempts} attempts: {last_error}") def _parse_json_object( self, content: str, log_context: str = "llm", ) -> dict[str, object]: """Extract and validate a JSON object from model output.""" stripped = content.strip() if not stripped: logger.warning("[%s] Empty JSON response from model.", log_context) raise ValueError("Empty JSON response.") try: parsed = json.loads(stripped) if isinstance(parsed, dict): return cast(dict[str, object], parsed) except json.JSONDecodeError as error: logger.warning( "[%s] JSON parse error: %s. Response preview: %s", log_context, error, stripped[:400], ) match = _JSON_BLOCK_PATTERN.search(stripped) if match is None: raise ValueError("No JSON object found in response.") try: parsed = json.loads(match.group(0)) except json.JSONDecodeError as error: logger.warning( "[%s] Extracted JSON block still invalid: %s. Block preview: %s", log_context, error, match.group(0)[:400], ) raise ValueError(f"Malformed JSON in response: {error}") from error if not isinstance(parsed, dict): raise ValueError("JSON response must be an object.") return cast(dict[str, object], parsed) _client: LlmClient | None = None def get_llm_client() -> LlmClient: """Return the shared LLM client instance.""" global _client if _client is None: _client = LlmClient() return _client def reset_llm_client() -> None: """Reset the shared client (used by tests).""" global _client _client = None