src/photools3/llm_session.py

@@ -22,6 +22,28 @@ from photools3.settings import settings
 logger = logging.getLogger(__name__)
 
 
+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:
     base_url: str = str(settings.llm_api_base_url)
     api_key: str = settings.llm_api_key
@@ -74,35 +96,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