| 1 | import os |
| 2 | import logging |
| 3 | import asyncio |
| 4 | from typing import List, Tuple, Dict, Optional |
| 5 | from langchain_core.messages import HumanMessage, SystemMessage |
| 6 | from langchain.chat_models import init_chat_model |
| 7 | from pydantic import BaseModel, Field |
| 8 | from langchain.output_parsers import PydanticOutputParser |
| 9 | from interfaces import Event, Scene |
| 10 | from interfaces import CharacterInScene, CharacterInEvent, CharacterInNovel |
| 11 | from tenacity import retry, stop_after_attempt |
| 12 | |
| 13 | |
| 14 | system_prompt_template_merge_characters_across_scenes_in_event = \ |
| 15 | """ |
| 16 | You are an expert script analysis and character fusion specialist. Your role is to intelligently analyze multiple script scenes, identify characters that represent the same entity across different scenes, and merge them into a unified character list with consistent identifiers. |
| 17 | |
| 18 | **TASK** |
| 19 | Process the input scenes, each containing a script and characters with their names and features. Identify and merge characters that are logically the same across scenes, even if they have different names or slight variations in description. Output a consolidated list of characters for the entire event. Each character in the list must have a unique identifier, along with the scene numbers where they appear and the name used in each scene. You also need to aggregate the static features of the same characters together. |
| 20 | |
| 21 | **INPUT** |
| 22 | A sequence of scenes. Each scene is enclosed within <SCENE_N_START> and <SCENE_N_END> tags, where N is the scene number(starting from 0). |
| 23 | Each scene includes a screnplay script and a sequence of character names. |
| 24 | The screenplay script is enclosed within <SCRIPT_START> and <SCRIPT_END> tags. |
| 25 | The sequence of character is enclosed within <CHARACTERS_START> and <CHARACTERS_END> tags. Each character in the list is enclosed within <CHARACTER_M_START> and <CHARACTER_M_END> tags, where M is the character number(starting from 0). |
| 26 | |
| 27 | Below is an example of one scene: |
| 28 | |
| 29 | <SCENE_0_START> |
| 30 | |
| 31 | <SCRIPT_START> |
| 32 | John enters the room and sees Mary. |
| 33 | John: Hi Mary, how are you? |
| 34 | Mary: I'm good, John. Thanks for asking! |
| 35 | <SCRIPT_END> |
| 36 | |
| 37 | <CHARACTERS_START> |
| 38 | |
| 39 | <CHARACTER_0_START> |
| 40 | John [visible] |
| 41 | static features: John is a tall man with short black hair and brown eyes. |
| 42 | dynamic features: Wearing a blue shirt and black pants. |
| 43 | <CHARACTER_0_END> |
| 44 | |
| 45 | <CHARACTER_1_START> |
| 46 | Mary [visible] |
| 47 | static features: Mary is a young woman with long brown hair and green eyes. |
| 48 | dynamic features: Wearing a floral dress and a denim jacket. |
| 49 | <CHARACTER_1_END> |
| 50 | |
| 51 | <CHARACTERS_END> |
| 52 | |
| 53 | <SCENE_0_END> |
| 54 | |
| 55 | |
| 56 | |
| 57 | **OUTPUT** |
| 58 | {format_instructions} |
| 59 | |
| 60 | **GUIDELINES** |
| 61 | 1. Character Fusion: Analyze contextual clues (e.g., dialogue style, role in plot, relationships, descriptions) to determine if characters from different scenes are the same person, even if names vary. |
| 62 | 2. Unique Identifier: Assign a consistent, unique ID (e.g., primary/canonical name) to each merged character. Use the most frequent or contextually appropriate name as the identifier, if possible. |
| 63 | 3. Scene Mapping: For each character, list all scenes they appear in and the exact name used in each scene. |
| 64 | 4. Completeness: Ensure all characters from all scenes are included in the final list. No duplicate, omitted, or extraneous characters. |
| 65 | 5. If a character undergoes significant changes across different scenes, it is necessary to split them into separate roles. For example, if Character A is a child in Scene 0 but an adult in Scene 1, they should be divided into two distinct characters (meaning two different actors are required to portray them). |
| 66 | 6. The language of outputs in values should be same as the input text. |
| 67 | """ |
| 68 | |
| 69 | |
| 70 | human_prompt_template_merge_characters_across_scenes_in_event = \ |
| 71 | """ |
| 72 | {scenes_sequence} |
| 73 | """ |
| 74 | |
| 75 | class MergeCharactersAcrossScenesInEventResponse(BaseModel): |
| 76 | characters: List[CharacterInEvent] = Field( |
| 77 | description="List of merged characters with their identifiers", |
| 78 | ) |
| 79 | |
| 80 | |
| 81 | |
| 82 | |
| 83 | system_prompt_template_merge_characters_to_existing_characters_in_novel = \ |
| 84 | """ |
| 85 | You are an information integration expert skilled in accurately identifying, matching, and merging character information. Your responsibility is to ensure consistency in character attributes and efficiently maintain and update the global character list. |
| 86 | |
| 87 | **TASK** |
| 88 | Merge the character list extracted from the current event (which may include new or existing characters) into the global character list. For existing characters, ensure their feature descriptions remain consistent; for new characters, add them to the global list. |
| 89 | |
| 90 | **INPUT** |
| 91 | 1. Existing Characters in the Novel: A list of characters already present in the novel, each with a unique index, identifier, and static features. The list is enclosed within <EXISTING_CHARACTERS_START> and <EXISTING_CHARACTERS_END> tags. Each character in the list is enclosed within <CHARACTER_P_START> and <CHARACTER_P_END> tags, where P is the character number(starting from 0). |
| 92 | 2. Characters in the Current Event: A list of characters identified in the current event, each with an index, identifier, active scenes, and static features. The list is enclosed within <EVENT_CHARACTERS_START> and <EVENT_CHARACTERS_END> tags. Each character in the list is enclosed within <CHARACTER_Q_START> and <CHARACTER_Q_END> tags, where Q is the character number(starting from 0). |
| 93 | |
| 94 | |
| 95 | **OUTPUT** |
| 96 | {format_instructions} |
| 97 | |
| 98 | **GUIDELINES** |
| 99 | 1. Feature Consistency: Strictly compare the features of the current event characters with those of existing characters. Some character's identifier may be the same as existing role identifier, but their features differ, such as youth and old age. You need to distinguish them as two separate characters. |
| 100 | 2. Efficient Merging: Avoid duplicate characters to ensure the list remains concise. |
| 101 | 3. Feature Update: If an existing character's features are expanded or modified based on new information from the current event, update their description accordingly. |
| 102 | """ |
| 103 | |
| 104 | human_prompt_template_merge_characters_to_existing_characters_in_novel = \ |
| 105 | """ |
| 106 | <EXISTING_CHARACTERS_START> |
| 107 | {existing_characters_in_novel} |
| 108 | <EXISTING_CHARACTERS_END> |
| 109 | |
| 110 | <EVENT_CHARACTERS_START> |
| 111 | {characters_in_event} |
| 112 | <EVENT_CHARACTERS_END> |
| 113 | """ |
| 114 | |
| 115 | |
| 116 | class CharacterForMergingToNovel(BaseModel): |
| 117 | index_in_event: int = Field( |
| 118 | description="The index of the character in the list of characters in the current event.", |
| 119 | examples=[0, 1, 2], |
| 120 | ) |
| 121 | index_in_novel: int = Field( |
| 122 | description="The index of the character in the list of existing characters in the novel. If this is a new character, set it to -1.", |
| 123 | examples=[0, 7, -1], |
| 124 | ) |
| 125 | identifier_in_novel: str = Field( |
| 126 | description="The unique identifier for the character in the novel. If this is a new character, ensure the name does not conflict with existing characters. If this is not a new character, this should match the identifier in the existing characters list.", |
| 127 | examples=["Alice", "Bob the Builder"], |
| 128 | ) |
| 129 | modified_features: str = Field( |
| 130 | description="The modified static features of the character after merging. If the character is new, this should be the full static features. If the character is existing and their features are expanded or modified, this should be filled in the complete modified features. If the character is existing and their features remain unchanged, this should be the same as the existing character's static features.", |
| 131 | ) |
| 132 | |
| 133 | class MergeCharactersToExistingCharactersInNovelResponse(BaseModel): |
| 134 | characters: List[CharacterForMergingToNovel] = Field( |
| 135 | description="List of characters in the event with their corresponding index in the existing characters in the novel. If the character is new, the index_in_novel should be -1. The number of characters in this list should be the same as the number of characters in the event.", |
| 136 | ) |
| 137 | |
| 138 | |
| 139 | |
| 140 | class GlobalInformationPlanner: |
| 141 | def __init__( |
| 142 | self, |
| 143 | api_key: str, |
| 144 | base_url: str, |
| 145 | chat_model: str, |
| 146 | ): |
| 147 | self.chat_model = init_chat_model( |
| 148 | model=chat_model, |
| 149 | model_provider="openai", |
| 150 | api_key=api_key, |
| 151 | base_url=base_url, |
| 152 | ) |
| 153 | |
| 154 | @retry( |
| 155 | stop=stop_after_attempt(3), |
| 156 | after=lambda retry_state: logging.warning(f"Retrying due to {retry_state.outcome.exception()}"), |
| 157 | ) |
| 158 | async def merge_characters_across_scenes_in_event( |
| 159 | self, |
| 160 | event_idx: int, |
| 161 | scenes: List[Scene], # Scene.characters is List[CharacterInScene] |
| 162 | ) -> List[CharacterInEvent]: |
| 163 | scenes_sequence_str = "" |
| 164 | for scene in scenes: |
| 165 | scene_str = f"<SCENE_{scene.idx}_START>\n" |
| 166 | scene_str += "<SCRIPT_START>\n" |
| 167 | scene_str += scene.script + "\n" |
| 168 | scene_str += "<SCRIPT_END>\n\n" |
| 169 | scene_str += "<CHARACTERS_START>\n" |
| 170 | for character in scene.characters: |
| 171 | scene_str += f"<CHARACTER_{character.idx}_START>\n" |
| 172 | scene_str += str(character) |
| 173 | scene_str += f"<CHARACTER_{character.idx}_END>\n" |
| 174 | scene_str += "<CHARACTERS_END>\n" |
| 175 | scene_str += f"<SCENE_{scene.idx}_END>\n" |
| 176 | scenes_sequence_str += scene_str |
| 177 | |
| 178 | parser = PydanticOutputParser(pydantic_object=MergeCharactersAcrossScenesInEventResponse) |
| 179 | |
| 180 | messages = [ |
| 181 | SystemMessage( |
| 182 | content=system_prompt_template_merge_characters_across_scenes_in_event.format( |
| 183 | format_instructions=parser.get_format_instructions(), |
| 184 | ), |
| 185 | ), |
| 186 | HumanMessage( |
| 187 | content=human_prompt_template_merge_characters_across_scenes_in_event.format( |
| 188 | scenes_sequence=scenes_sequence_str, |
| 189 | ) |
| 190 | ) |
| 191 | ] |
| 192 | |
| 193 | chain = self.chat_model | parser |
| 194 | response: MergeCharactersAcrossScenesInEventResponse = await chain.ainvoke(messages) |
| 195 | characters_in_event = response.characters |
| 196 | |
| 197 | # check the output is valid |
| 198 | flags = [{c.identifier_in_scene: False for c in s.characters} for s in scenes] |
| 199 | |
| 200 | # check if all character identifiers can be found in the scenes |
| 201 | for character in characters_in_event: |
| 202 | for scene_idx, identifier_in_scene in character.active_scenes.items(): |
| 203 | if identifier_in_scene not in [c.identifier_in_scene for c in scenes[scene_idx].characters]: |
| 204 | raise ValueError(f"Character {identifier_in_scene} not found in scene {scene_idx} of event {event_idx}") |
| 205 | else: |
| 206 | flags[scene_idx][identifier_in_scene] = True |
| 207 | |
| 208 | # check if all characters are included |
| 209 | for scene_idx, flag in enumerate(flags): |
| 210 | for identifier_in_scene, included in flag.items(): |
| 211 | if not included: |
| 212 | raise ValueError(f"Character {identifier_in_scene} in scene {scene_idx} of event {event_idx} not included in the merged characters") |
| 213 | |
| 214 | return characters_in_event |
| 215 | |
| 216 | @retry( |
| 217 | stop=stop_after_attempt(3), |
| 218 | after=lambda retry_state: logging.warning(f"Retrying due to {retry_state.outcome.exception()}"), |
| 219 | ) |
| 220 | def merge_characters_to_existing_characters_in_novel( |
| 221 | self, |
| 222 | event_idx: int, |
| 223 | existing_characters_in_novel: List[CharacterInNovel], |
| 224 | characters_in_event: List[CharacterInEvent], |
| 225 | ) -> List[CharacterInNovel]: |
| 226 | existing_characters_str = "" |
| 227 | for character in existing_characters_in_novel: |
| 228 | existing_characters_str += f"<CHARACTER_{character.index}_START>\n" |
| 229 | existing_characters_str += str(character) |
| 230 | existing_characters_str += f"<CHARACTER_{character.index}_END>\n" |
| 231 | |
| 232 | characters_in_event_str = "" |
| 233 | for character in characters_in_event: |
| 234 | characters_in_event_str += f"<CHARACTER_{character.index}_START>\n" |
| 235 | characters_in_event_str += character.identifier_in_event + "\n" |
| 236 | characters_in_event_str += "Static features: " + character.static_features + "\n" |
| 237 | characters_in_event_str += f"<CHARACTER_{character.index}_END>\n" |
| 238 | |
| 239 | parser = PydanticOutputParser(pydantic_object=MergeCharactersToExistingCharactersInNovelResponse) |
| 240 | |
| 241 | messages = [ |
| 242 | SystemMessage( |
| 243 | content=system_prompt_template_merge_characters_to_existing_characters_in_novel.format( |
| 244 | format_instructions=parser.get_format_instructions(), |
| 245 | ), |
| 246 | ), |
| 247 | HumanMessage( |
| 248 | content=human_prompt_template_merge_characters_to_existing_characters_in_novel.format( |
| 249 | existing_characters_in_novel=existing_characters_str, |
| 250 | characters_in_event=characters_in_event_str, |
| 251 | ) |
| 252 | ) |
| 253 | ] |
| 254 | |
| 255 | chain = self.chat_model | parser |
| 256 | response: MergeCharactersToExistingCharactersInNovelResponse = chain.invoke(messages) |
| 257 | |
| 258 | for character in response.characters: |
| 259 | if character.index_in_novel == -1: |
| 260 | # new character, add to existing characters |
| 261 | new_character = CharacterInNovel( |
| 262 | index=len(existing_characters_in_novel), |
| 263 | identifier_in_novel=character.identifier_in_novel, |
| 264 | static_features=character.modified_features, |
| 265 | active_events={event_idx: characters_in_event[character.index_in_event].identifier_in_event}, |
| 266 | ) |
| 267 | existing_characters_in_novel.append(new_character) |
| 268 | else: |
| 269 | existing_characters_in_novel[character.index_in_novel].static_features = character.modified_features |
| 270 | existing_characters_in_novel[character.index_in_novel].active_events.update({event_idx: characters_in_event[character.index_in_event].identifier_in_event}) |
| 271 | |
| 272 | return existing_characters_in_novel |
| 273 | |
| 274 | |
| 275 | # # TODO: 如果是长篇小说，事件太多，很容易报错，出场的角色会分不清在哪个事件里，也很容易漏，需要想办法解决 |
| 276 | # @retry( |
| 277 | # stop=stop_after_attempt(3), |
| 278 | # after=lambda retry_state: logging.warning(f"Retrying due to {retry_state.outcome.exception()}"), |
| 279 | # ) |
| 280 | # def merge_characters_across_events_in_novel( |
| 281 | # self, |
| 282 | # events: List[Event], |
| 283 | # characters_in_event: List[List[CharacterInEvent]], |
| 284 | # ) -> List[CharacterInNovelWithoutStaticFeatures]: |
| 285 | # events_sequence_str = "" |
| 286 | # for event, characters in zip(events, characters_in_event): |
| 287 | # event_str = f"<EVENT_{event.index}_START>\n\n" |
| 288 | # event_str += "<DESCRIPTION_START>\n" |
| 289 | # event_str += event.description + "\n" |
| 290 | # event_str += "<DESCRIPTION_END>\n\n" |
| 291 | # event_str += "<PROCESS_CHAIN_START>\n" |
| 292 | # for process in event.process_chain: |
| 293 | # event_str += process + "\n" |
| 294 | # event_str += "<PROCESS_CHAIN_END>\n\n" |
| 295 | # event_str += "<CHARACTERS_START>\n" |
| 296 | # for i, character in enumerate(characters): |
| 297 | # event_str += f"<CHARACTER_{i}_START>{character.identifier_in_event}<CHARACTER_{i}_END>\n" |
| 298 | # event_str += "<CHARACTERS_END>\n\n" |
| 299 | # event_str += f"<EVENT_{event.index}_END>\n\n" |
| 300 | # events_sequence_str += event_str |
| 301 | |
| 302 | # parser = PydanticOutputParser(pydantic_object=MergeCharactersAcrossEventsInNovelResponse) |
| 303 | |
| 304 | # messages = [ |
| 305 | # SystemMessage( |
| 306 | # content=system_prompt_template_merge_characters_across_events.format( |
| 307 | # format_instructions=parser.get_format_instructions(), |
| 308 | # ), |
| 309 | # ), |
| 310 | # HumanMessage( |
| 311 | # content=human_prompt_template_merge_characters_across_events.format( |
| 312 | # events_sequence=events_sequence_str, |
| 313 | # ) |
| 314 | # ) |
| 315 | # ] |
| 316 | |
| 317 | # chain = self.chat_model | parser |
| 318 | # response: MergeCharactersAcrossEventsInNovelResponse = chain.invoke(messages) |
| 319 | # characters_in_novel = response.characters |
| 320 | |
| 321 | # # check the output is valid |
| 322 | # flags = [{c.identifier_in_event: False for c in characters} for characters in characters_in_event] |
| 323 | |
| 324 | # # check if all character identifiers can be found in the events |
| 325 | # for character in characters_in_novel: |
| 326 | # for event_idx, identifier_in_event in character.active_events.items(): |
| 327 | # if identifier_in_event not in [c.identifier_in_event for c in characters_in_event[event_idx]]: |
| 328 | # raise ValueError(f"Character {identifier_in_event} not found in event {event_idx}") |
| 329 | # else: |
| 330 | # flags[event_idx][identifier_in_event] = True |
| 331 | |
| 332 | # # check if all characters are included |
| 333 | # # for event_idx, flag in enumerate(flags): |
| 334 | # # for identifier_in_event, included in flag.items(): |
| 335 | # # if not included: |
| 336 | # # raise ValueError(f"Character {identifier_in_event} in event {event_idx} not included in the merged characters") |
| 337 | |
| 338 | # return characters_in_novel |
| 339 | |
| 340 | |
| 341 | |
| 342 | # async def extract_static_feature_for_character_in_novel( |
| 343 | # self, |
| 344 | # relevant_chunks: List[str], |
| 345 | # character: CharacterInNovelWithoutStaticFeatures, |
| 346 | # ) -> str: |
| 347 | # context_fragments_str = "" |
| 348 | # for i, chunk in enumerate(relevant_chunks): |
| 349 | # context_fragments_str += f"<CONTEXT_FRAGMENT_{i}_START>\n" |
| 350 | # context_fragments_str += chunk + "\n" |
| 351 | # context_fragments_str += f"<CONTEXT_FRAGMENT_{i}_END>\n" |
| 352 | |
| 353 | # parser = None # no need to parse the output, just return the text |
| 354 | |
| 355 | # messages = [ |
| 356 | # SystemMessage( |
| 357 | # content=system_prompt_template_extract_static_feature_for_character_in_novel, |
| 358 | # ), |
| 359 | # HumanMessage( |
| 360 | # content=human_prompt_template_extract_static_feature_for_character_in_novel.format( |
| 361 | # character_name=character.identifier_in_novel, |
| 362 | # context_fragments=context_fragments_str, |
| 363 | # ) |
| 364 | # ) |
| 365 | # ] |
| 366 | |
| 367 | # base_features = await self.chat_model.ainvoke(messages) |
| 368 | # return base_features.content |
| 369 |