initial commit

backend/services/context_builder.py

@@ -0,0 +1,123 @@
+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()