124 lines
4.8 KiB
Python
124 lines
4.8 KiB
Python
import re
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from models.bible import Character, Event, Location, StoryBible
|
|
|
|
|
|
class ContextBuilder:
|
|
def extract_mentioned_entities(
|
|
self, text: str, bible: StoryBible
|
|
) -> Dict[str, List[Any]]:
|
|
"""Извлекает упоминания сущностей из текста"""
|
|
text_lower = text.lower()
|
|
|
|
mentioned_chars: List[Character] = []
|
|
mentioned_locs: List[Location] = []
|
|
mentioned_events: List[Event] = []
|
|
|
|
# Ищем упоминания персонажей
|
|
for char in bible.characters:
|
|
name_patterns = [
|
|
char.name.lower(),
|
|
char.name.lower().split()[0] if " " in char.name else char.name.lower(),
|
|
]
|
|
|
|
for pattern in name_patterns:
|
|
if pattern and len(pattern) > 2:
|
|
if re.search(r"\b" + re.escape(pattern) + r"\b", text_lower):
|
|
mentioned_chars.append(char)
|
|
break
|
|
|
|
# Ищем упоминания локаций
|
|
for loc in bible.locations:
|
|
name_patterns = [
|
|
loc.name.lower(),
|
|
loc.name.lower().split()[0] if " " in loc.name else loc.name.lower(),
|
|
]
|
|
|
|
for pattern in name_patterns:
|
|
if pattern and len(pattern) > 2:
|
|
if re.search(r"\b" + re.escape(pattern) + r"\b", text_lower):
|
|
mentioned_locs.append(loc)
|
|
break
|
|
|
|
# Ищем упоминания событий (по ключевым словам в названии)
|
|
for event in bible.events:
|
|
keywords = re.findall(r"\w{4,}", event.title.lower())
|
|
for keyword in keywords:
|
|
if keyword in text_lower:
|
|
mentioned_events.append(event)
|
|
break
|
|
|
|
return {
|
|
"characters": mentioned_chars,
|
|
"locations": mentioned_locs,
|
|
"events": mentioned_events,
|
|
}
|
|
|
|
def build_context_string(self, mentioned_entities: Dict[str, List[Any]]) -> str:
|
|
"""Строит строку контекста для промпта"""
|
|
context_parts: List[str] = []
|
|
|
|
# Добавляем информацию о персонажах
|
|
if mentioned_entities["characters"]:
|
|
chars_info: List[str] = []
|
|
for char in mentioned_entities["characters"][:5]:
|
|
traits_str = (
|
|
", ".join(char.traits[:3]) if char.traits else "нет описания"
|
|
)
|
|
chars_info.append(
|
|
f"- {char.name}: {char.description} (черты: {traits_str})"
|
|
)
|
|
|
|
context_parts.append("ПЕРСОНАЖИ В СЦЕНЕ:\n" + "\n".join(chars_info))
|
|
|
|
# Добавляем информацию о локациях
|
|
if mentioned_entities["locations"]:
|
|
locs_info: List[str] = []
|
|
for loc in mentioned_entities["locations"][:3]:
|
|
atm_str = f"Атмосфера: {loc.atmosphere}" if loc.atmosphere else ""
|
|
locs_info.append(f"- {loc.name}: {loc.description} {atm_str}".strip())
|
|
|
|
context_parts.append("ЛОКАЦИИ:\n" + "\n".join(locs_info))
|
|
|
|
# Добавляем информацию о событиях
|
|
if mentioned_entities["events"]:
|
|
events_info: List[str] = []
|
|
for event in mentioned_entities["events"][:3]:
|
|
events_info.append(f"- {event.title}: {event.description}")
|
|
|
|
context_parts.append("СОБЫТИЯ:\n" + "\n".join(events_info))
|
|
|
|
return "\n\n".join(context_parts)
|
|
|
|
def enrich_prompt(
|
|
self,
|
|
original_prompt: str,
|
|
selected_text: str,
|
|
bible: Optional[StoryBible],
|
|
) -> str:
|
|
"""Обогащает промпт контекстом из Story Bible"""
|
|
if not bible or not selected_text:
|
|
return original_prompt
|
|
|
|
mentioned = self.extract_mentioned_entities(selected_text, bible)
|
|
|
|
if not any(mentioned.values()):
|
|
return original_prompt
|
|
|
|
context_str = self.build_context_string(mentioned)
|
|
|
|
enriched_prompt = (
|
|
f"КОНТЕКСТ ПРОИЗВЕДЕНИЯ:\n{context_str}\n\n"
|
|
f"ВАЖНО: При выполнении задачи учитывай приведенный контекст. "
|
|
f"Сохраняй консистентность с описаниями персонажей, локаций и событий. "
|
|
f"Не противоречь установленным фактам.\n\n"
|
|
f"ЗАДАЧА:\n{original_prompt}"
|
|
)
|
|
|
|
return enriched_prompt
|
|
|
|
|
|
# Singleton instance
|
|
context_builder = ContextBuilder()
|