135 lines
5.2 KiB
Python
135 lines
5.2 KiB
Python
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()
|