0.5 - Разбор нескольких моделей в одной строке

parse.py

@@ -31,7 +31,8 @@ SHOW_IO = os.environ.get("SHOW_IO", "1").lower() in ("1", "true", "yes")
 RECORD_DELIMITERS = ("\n", "\r\n", "\r")
 
 EXTRACT_PROMPT = """Extract brand, model, engine_code, date_range from the text.
-Return JSON object only: {"brand":"","model":"","engine_code":"","date_range":""}
+If 1 record — return JSON object: {"brand":"","model":"","engine_code":"","date_range":""}
+If 2+ records on the line — return JSON array: [{"brand":"",...},{"brand":"",...}]
 Empty fields as "". No other text."""
 
 
@@ -42,8 +43,17 @@ def prepare_records(content: str) -> list[str]:
     return [r.strip() for r in content.split(RECORD_DELIMITERS[0]) if r.strip()]
 
 
-def extract_record(client: Client, record: str) -> dict:
-    """2) Отправка промпт + 1 запись в нейронку. 3) Получение ответа."""
+def _normalize_obj(obj: dict) -> dict:
+    return {
+        "brand": obj.get("brand", ""),
+        "model": obj.get("model", ""),
+        "engine_code": obj.get("engine_code", ""),
+        "date_range": obj.get("date_range", ""),
+    }
+
+
+def extract_record(client: Client, record: str) -> list[dict]:
+    """2) Отправка промпт + 1 запись в нейронку. 3) Получение ответа (1 или более объектов)."""
     for attempt in range(1, RETRY_COUNT + 1):
         try:
             response = client.chat(
@@ -59,24 +69,21 @@ def extract_record(client: Client, record: str) -> dict:
                 result = result.split("```")[1]
                 if result.strip().lower().startswith("json"):
                     result = result.strip()[4:].strip()
-            obj = json.loads(result)
-            return {
-                "brand": obj.get("brand", ""),
-                "model": obj.get("model", ""),
-                "engine_code": obj.get("engine_code", ""),
-                "date_range": obj.get("date_range", ""),
-            }
-        except (json.JSONDecodeError, KeyError) as e:
+            parsed = json.loads(result)
+            if isinstance(parsed, list):
+                return [_normalize_obj(o) for o in parsed]
+            return [_normalize_obj(parsed)]
+        except (json.JSONDecodeError, KeyError):
             if attempt < RETRY_COUNT:
                 time.sleep(RETRY_DELAY)
             else:
-                return {"brand": "", "model": "", "engine_code": "", "date_range": record}
+                return [{"brand": "", "model": "", "engine_code": "", "date_range": record}]
         except Exception:
             if attempt < RETRY_COUNT:
                 time.sleep(RETRY_DELAY)
             else:
                 raise
-    return {"brand": "", "model": "", "engine_code": "", "date_range": record}
+    return [{"brand": "", "model": "", "engine_code": "", "date_range": record}]
 
 
 def process_file(client: Client, filepath: pathlib.Path) -> bool:
@@ -96,18 +103,18 @@ def process_file(client: Client, filepath: pathlib.Path) -> bool:
 
     for i, record in enumerate(records, 1):
         print(f"  [{i}/{total}] ", end="", flush=True)
-        obj = extract_record(client, record)
-        data.append(obj)
+        objs = extract_record(client, record)
+        data.extend(objs)
         out_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
         if SHOW_IO:
             inp = record[:60] + "…" if len(record) > 60 else record
-            out = json.dumps(obj, ensure_ascii=False)
+            out = json.dumps(objs, ensure_ascii=False)
             print(f"вход: {inp}", flush=True)
             print(f"      выход: {out}", flush=True)
         else:
             print("✓", flush=True)
     elapsed = time.perf_counter() - t0
-    print(f"  {filepath.name} — {total} записей, {elapsed:.1f}s — OK")
+    print(f"  {filepath.name} — {len(data)} объектов, {elapsed:.1f}s — OK")
     return True