initial commit

backend/services/entity_extractor.py

@@ -0,0 +1,134 @@
+import json
+import re
+from typing import Any, Dict, List, Optional
+
+from models.bible import StoryBible
+
+from services.llm_router import llm_router
+
+
+class EntityExtractor:
+    def __init__(self):
+        self.llm_router = llm_router
+
+    def _build_extraction_prompt(
+        self, text: str, existing_bible: Optional[StoryBible]
+    ) -> str:
+        """Формирует промпт для извлечения сущностей"""
+
+        existing_chars = ""
+        existing_locs = ""
+        existing_events = ""
+
+        if existing_bible:
+            if existing_bible.characters:
+                existing_chars = "Известные персонажи: " + ", ".join(
+                    [c.name for c in existing_bible.characters]
+                )
+            if existing_bible.locations:
+                existing_locs = "Известные локации: " + ", ".join(
+                    [l.name for l in existing_bible.locations]
+                )
+            if existing_bible.events:
+                existing_events = "Известные события: " + ", ".join(
+                    [e.title for e in existing_bible.events]
+                )
+
+        prompt = f"""Ты — эксперт по анализу художественных текстов. Твоя задача — извлечь из текста новых персонажей, локации и события.
+
+ВАЖНО: Отвечай СТРОГО НА РУССКОМ ЯЗЫКЕ. Все имена, описания и названия должны быть на русском языке.
+
+{existing_chars}
+{existing_locs}
+{existing_events}
+
+ТЕКСТ ДЛЯ АНАЛИЗА:
+---
+{text}
+---
+
+ИНСТРУКЦИИ:
+1. Извлеки ТОЛЬКО НОВЫЕ сущности, которых еще нет в списке известных
+2. Для персонажей укажи: name, description (краткое), traits (список черт), appearance (если описана), role (protagonist/antagonist/supporting/minor)
+3. Для локаций укажи: name, description, atmosphere (если есть), type (city/building/nature/interior/other)
+4. Для событий укажи: title, description, importance (major/minor/background)
+5. Если обнаружены новые черты для УЖЕ ИЗВЕСТНЫХ персонажей, укажи их в updated_traits
+6. НЕ выдумывай информацию, которой нет в тексте
+7. Если новых сущностей нет, верни пустые списки
+8. Все ответы СТРОГО НА РУССКОМ ЯЗЫКЕ
+
+ОТВЕТЬ СТРОГО В ФОРМАТЕ JSON (без markdown, без пояснений, только чистый JSON):
+{{
+  "new_characters": [
+    {{"name": "...", "description": "...", "traits": ["...", "..."], "appearance": "...", "role": "..."}}
+  ],
+  "new_locations": [
+    {{"name": "...", "description": "...", "atmosphere": "...", "type": "..."}}
+  ],
+  "new_events": [
+    {{"title": "...", "description": "...", "importance": "..."}}
+  ],
+  "updated_traits": [
+    {{"character_name": "...", "new_traits": ["...", "..."]}}
+  ]
+}}"""
+
+        return prompt
+
+    def _parse_json_response(self, response_text: str) -> Dict[str, Any]:
+        """Парсит JSON из ответа модели, обрабатывая возможные markdown обертки"""
+        # Убираем возможные markdown обертки
+        cleaned = response_text.strip()
+
+        # Убираем ```json и ```
+        if cleaned.startswith("```"):
+            cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned)
+            cleaned = re.sub(r"\s*```$", "", cleaned)
+
+        # Пытаемся найти JSON в тексте
+        json_match = re.search(r"\{[\s\S]*\}", cleaned)
+        if json_match:
+            cleaned = json_match.group(0)
+
+        try:
+            return json.loads(cleaned)
+        except json.JSONDecodeError as e:
+            raise ValueError(
+                f"Failed to parse JSON from LLM response: {e}\n"
+                f"Raw response: {response_text[:500]}"
+            )
+
+    async def extract_entities(
+        self,
+        text: str,
+        existing_bible: Optional[StoryBible] = None,
+        provider: str = "openrouter",
+        model: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Извлекает сущности из текста"""
+
+        prompt = self._build_extraction_prompt(text, existing_bible)
+
+        # Используем пресет editor для точности (низкая температура)
+        generated_text = await self.llm_router.generate(
+            prompt=prompt,
+            provider=provider,
+            model=model,
+            preset="editor",
+        )
+
+        parsed = self._parse_json_response(generated_text)
+
+        # Валидация и нормализация ответа
+        result = {
+            "new_characters": parsed.get("new_characters", []),
+            "new_locations": parsed.get("new_locations", []),
+            "new_events": parsed.get("new_events", []),
+            "updated_traits": parsed.get("updated_traits", []),
+        }
+
+        return result
+
+
+# Singleton instance
+entity_extractor = EntityExtractor()