| 1 | import os |
| 2 | import logging |
| 3 | import asyncio |
| 4 | from typing import List |
| 5 | from langchain_core.messages import HumanMessage, SystemMessage |
| 6 | from langchain_core.output_parsers import PydanticOutputParser |
| 7 | from utils.robust_json_parser import TrailingCommaTolerantPydanticOutputParser as PydanticOutputParser |
| 8 | from langchain.chat_models import init_chat_model |
| 9 | from pydantic import BaseModel, Field |
| 10 | from tenacity import retry, stop_after_attempt |
| 11 | |
| 12 | from interfaces import Event |
| 13 | |
| 14 | system_prompt_template_extract_events = \ |
| 15 | """ |
| 16 | You are a highly skilled Literary Analyst AI. Your expertise is in narrative structure, plot deconstruction, and thematic analysis. You meticulously read and interpret prose to break down a story into its fundamental sequential events. |
| 17 | |
| 18 | **TASK** |
| 19 | Extract the next event from the provided novel, following the sequence of the story and building upon the partially extracted events. |
| 20 | |
| 21 | **INPUT** |
| 22 | 1. The full text of the novel, which is enclosed within <NOVEL_TEXT_START> and <NOVEL_TEXT_END> tags |
| 23 | 2. A sequence of already-extracted events (in order), which is enclosed within <EXTRACTED_EVENTS_START> and <EXTRACTED_EVENTS_END> tags. The sequence may be empty. Each event contains multiple processes and constitutes a complete causal chain. |
| 24 | |
| 25 | Below is an example input: |
| 26 | |
| 27 | <NOVEL_TEXT_START> |
| 28 | The night was as dark as ink when the piercing alarm of the city museum suddenly shattered the silence. A thief, moving with phantom-like agility, had just pried open the display case and snatched the blue gem known as the "Heart of the Ocean" when the blaring alarm echoed through the hall. |
| 29 | ... (more novel text) ... |
| 30 | <NOVEL_TEXT_END> |
| 31 | |
| 32 | <EXTRACTED_EVENTS_START> |
| 33 | <Event 0> |
| 34 | Description: A thief who stole a gem from a museum was caught after a rooftop chase with guards, and the gem was recovered. |
| 35 | Process Chain: |
| 36 | - A thief steals a gem from a museum, triggering the alarm. Guards notice and begin the chase. |
| 37 | - The thief rushes out the museum's back door and dashes through narrow alleys, with guards closely pursuing and calling for backup. |
| 38 | - ... (more processes) ... |
| 39 | |
| 40 | <Event 1> |
| 41 | Description: ... (more description) ... |
| 42 | Process Chain: |
| 43 | - ... (more processes) ... |
| 44 | |
| 45 | <EXTRACTED_EVENTS_END> |
| 46 | |
| 47 | |
| 48 | **OUTPUT** |
| 49 | {format_instructions} |
| 50 | |
| 51 | **GUIDELINES** |
| 52 | 1. Focus on events that are critical to the plot, character development, or thematic depth. |
| 53 | 2. Ensure the event is logically distinct from previous and subsequent events. |
| 54 | 3. If the event spans multiple scenes, unify them under a single dramatic goal. For example, a chase sequence might begin in a city market, continue through back alleys, and conclude on a rooftop—all comprising a single event because they collectively achieve the dramatic purpose of "the protagonist evading capture." |
| 55 | 4. Maintain objectivity: describe events based on the text without interpretation or judgment. |
| 56 | 5. For the process field, provide a detailed, step-by-step account of the event's progression, including key actions, decisions, and turning points. Each step should be clear and concise, illustrating how the event unfolds over time. |
| 57 | Below is an example: |
| 58 | Timeframe: The following morning, after acquiring the information about the Temple. |
| 59 | Characters: Elara (protagonist) and Kaelen (her rival treasure hunter). |
| 60 | Cause: Both seek the same artifact and are determined to reach it first. |
| 61 | Process: The event begins with Elara hastily purchasing supplies in the port town (scene 1), where she spots Kaelen already hiring a crew, raising the stakes. It continues as she races to secure her own ship and captain, negotiating fiercely under time pressure (scene 2). The event culminates in a direct confrontation on the docks (scene 3), where Kaelen attempts to sabotage her vessel, leading to a brief but intense sword fight between the two rivals. |
| 62 | Outcome: Elara successfully defends her ship and sets sail, but the conflict solidifies a bitter personal rivalry with Kaelen, ensuring their race to the temple will be fraught with direct opposition and danger. |
| 63 | 6. Every detail in your event description must be directly supported by the input novel. Do not add, assume, or invent any information. |
| 64 | 7. The language of outputs in values should be same as the input text. |
| 65 | """ |
| 66 | |
| 67 | human_prompt_template_extract_next_event = \ |
| 68 | """ |
| 69 | <NOVEL_TEXT_START> |
| 70 | {novel_text} |
| 71 | <NOVEL_TEXT_END> |
| 72 | |
| 73 | <EXTRACTED_EVENTS_START> |
| 74 | {extracted_events} |
| 75 | <EXTRACTED_EVENTS_END> |
| 76 | """ |
| 77 | |
| 78 | |
| 79 | |
| 80 | class EventExtractor: |
| 81 | def __init__( |
| 82 | self, |
| 83 | api_key: str, |
| 84 | base_url: str, |
| 85 | chat_model: str, |
| 86 | ): |
| 87 | self.chat_model = init_chat_model( |
| 88 | model=chat_model, |
| 89 | model_provider="openai", |
| 90 | api_key=api_key, |
| 91 | base_url=base_url, |
| 92 | ) |
| 93 | self.parser = PydanticOutputParser(pydantic_object=Event) |
| 94 | |
| 95 | |
| 96 | # Cap on extracted events: is_last is asserted by the LLM only, so without a |
| 97 | # bound a model that never sets it would loop (and spend tokens) forever. |
| 98 | max_events = 50 |
| 99 | |
| 100 | def __call__( |
| 101 | self, |
| 102 | novel_text: str, |
| 103 | ): |
| 104 | logging.info("Extracting events from novel...") |
| 105 | |
| 106 | events = [] |
| 107 | while True: |
| 108 | if len(events) >= self.max_events: |
| 109 | raise RuntimeError( |
| 110 | f"Event extraction exceeded the maximum of {self.max_events} events " |
| 111 | "without an is_last marker; aborting to avoid unbounded LLM calls." |
| 112 | ) |
| 113 | event = self.extract_next_event(novel_text, events) |
| 114 | |
| 115 | events.append(event) |
| 116 | logging.info(f"Extracted event: \n{event}") |
| 117 | if event.is_last: |
| 118 | break |
| 119 | |
| 120 | return events |
| 121 | |
| 122 | |
| 123 | @retry( |
| 124 | stop=stop_after_attempt(3), |
| 125 | after=lambda retry_state: logging.warning(f"Retrying extract_next_event due to error: {retry_state.outcome.exception()}"), |
| 126 | ) |
| 127 | def extract_next_event( |
| 128 | self, |
| 129 | novel_text: str, |
| 130 | extracted_events: List[Event] |
| 131 | ) -> Event: |
| 132 | |
| 133 | extracted_events_str = "\n\n".join([str(e) for e in extracted_events]) |
| 134 | |
| 135 | messages = [ |
| 136 | SystemMessage( |
| 137 | content=system_prompt_template_extract_events.format(format_instructions=self.parser.get_format_instructions()), |
| 138 | ), |
| 139 | HumanMessage( |
| 140 | content=human_prompt_template_extract_next_event.format( |
| 141 | novel_text=novel_text, |
| 142 | extracted_events=extracted_events_str, |
| 143 | ) |
| 144 | ) |
| 145 | ] |
| 146 | |
| 147 | chain = self.chat_model | self.parser |
| 148 | |
| 149 | event: Event = chain.invoke(messages) |
| 150 | |
| 151 | assert event.index == len(extracted_events), f"Extracted event index {event.index} does not match the expected index {len(extracted_events)}" |
| 152 | |
| 153 | return event |
| 154 | |
| 155 | |
| 156 | |
| 157 |