diff --git a/src/photools3/llm_session.py b/src/photools3/llm_session.py index e3b1f4a..4530efe 100644 --- a/src/photools3/llm_session.py +++ b/src/photools3/llm_session.py @@ -17,11 +17,31 @@ from photools3.models import ( LlmResponseQuality, ) from photools3.photo import Photo -from photools3.settings import Settings +from photools3.settings import settings logger = logging.getLogger(__name__) -settings = Settings() # type: ignore[reportCallIssue] + +def _clean_llm_response(txt: str): + # remove ```json ... ``` wrap + txt = re.sub(r"^```json(.*)```$", r"\1", txt, flags=re.DOTALL).strip() + return txt + + +def _correction_prompt( + original_request: LlmRequest, exception: Exception, origial_json +): + correction_prompt = ( + f"The previous response was invalid due to: {exception}. " + f"Please provide the corrected JSON based on this content: " + f"{origial_json}" + ) + original_request.messages.append( + LlmRequestMessage( + role=LlmMessageRole.user, + content=correction_prompt, + ) + ) class LlmSession: @@ -75,35 +95,22 @@ class LlmSession: for attempt in range(max_retries + 1): try: result_raw = self.client_chat_completion_create(request) - # Clean the response (strip markdown code blocks if present) - json_str = re.sub( - r"^```json(.*)```$", r"\1", result_raw, flags=re.DOTALL - ).strip() - - # Parse and validate against the class-level response_model + json_str = _clean_llm_response(result_raw) return self.response_model.model_validate_json(json_str) - except (json.JSONDecodeError, ValidationError) as e: - last_error = e + except (json.JSONDecodeError, ValidationError) as exc: + last_error = exc if attempt < max_retries: logger.info("Chatbot response is invalid; promting to correct ...") - # Prepare a correction prompt - correction_prompt = ( - f"The previous response was invalid due to: {e}. " - f"Please provide the corrected JSON based on this content: " - f"{json_str if 'json_str' in locals() else 'the original input'}" - ) - # Append correction prompt as a new user message for the next attempt - request.messages.append( - LlmRequestMessage( - role=LlmMessageRole.user, - content=correction_prompt, - ) + _correction_prompt( + request, + exc, + json_str if "json_str" in locals() else "origial json", ) continue else: raise ValueError( - f"Failed to parse or validate LLM response after {max_retries} retries: {e}" + f"Failed to parse or validate LLM response after {max_retries} retries: {exc}" ) except Exception: raise