diff --git a/src/pyob/core_utils.py b/src/pyob/core_utils.py
index 47ce417..91877ab 100644
--- a/src/pyob/core_utils.py
+++ b/src/pyob/core_utils.py
@@ -10,7 +10,7 @@
 import time
 from typing import Callable, Optional
 
-from .models import (
+from pyob.models import (
     OPENROUTER_KEY,
     get_valid_llm_response_engine,
     raw_gemini_keys,
diff --git a/src/pyob/xml_mixin.py b/src/pyob/xml_mixin.py
index fd1a462..c472551 100644
--- a/src/pyob/xml_mixin.py
+++ b/src/pyob/xml_mixin.py
@@ -275,7 +275,9 @@ def _attempt_line_by_line_match(
                     break
             if match:
                 new_code_lines = (
-                    code_lines[:i] + replace_lines + code_lines[i + len(search_lines) :]
+                    code_lines[:i]
+                    + replace_lines
+                    + code_lines[i + len(search_lines_stripped) :]
                 )
                 new_code = "\n".join(new_code_lines)
                 if not new_code.endswith("\n") and source.endswith("\n"):