Update handler.py

handler.py

@@ -317,10 +317,11 @@ class EndpointHandler:
             s = re.sub(r"^[^\w]+|[^\w]+$", "", s)
             return s
 
-        def _extract_from_target(target: str, phrase_lower: str) -> str:
-            idx = target.lower().find(phrase_lower)
+        def _extract_from_text(text: str, phrase_lower: str) -> str:
+            """Extract phrase from text preserving original case."""
+            idx = text.lower().find(phrase_lower)
             if idx >= 0:
-                return target[idx:idx+len(phrase_lower)]
+                return text[idx:idx+len(phrase_lower)]
             return phrase_lower
 
         STOP = {
@@ -366,7 +367,23 @@ class EndpointHandler:
         
         logger.info(f"After dedup: {kept_norms}")
 
-        cleaned = [_extract_from_target(target_text, n) for n in kept_norms]
+        # Extract spans from either target or context (whichever contains them)
+        cleaned = []
+        for n in kept_norms:
+            # Try target first, then context
+            if n in target_lower:
+                extracted = _extract_from_text(target_text, n)
+                logger.info(f"  Extracted '{extracted}' from TARGET")
+                cleaned.append(extracted)
+            elif n in context_text.lower():
+                extracted = _extract_from_text(context_text, n)
+                logger.info(f"  Extracted '{extracted}' from CONTEXT")
+                cleaned.append(extracted)
+            else:
+                # Fallback - shouldn't happen given earlier validation
+                logger.warning(f"  Phrase '{n}' not found in target or context, using normalized")
+                cleaned.append(n)
+        
         logger.info(f"Final cleaned spans: {cleaned}")
 
         if not cleaned and spans:
@@ -387,6 +404,6 @@ class EndpointHandler:
                     if best:
                         break
             if best:
-                return [_extract_from_target(target_text, best)]
+                return [_extract_from_text(target_text, best)]
 
         return cleaned[:3]