使用GPT的小说创作助手Agent Python代码

Agent的主要作用是保持情节的连贯性。

import openai
import os

class GPTClient:
    def __init__(self, api_key):
        openai.api_key = api_key

    def generate_text(self, prompt, max_tokens=100):
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {"role": "system", "content": "You are a creative writing assistant."},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=max_tokens
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            print(f"Error in GPT API call: {e}")
            return "Error in generating text."

class Character:
    def __init__(self, name, personality, background):
        self.name = name
        self.personality = personality
        self.background = background
        self.arc = []

class Event:
    def __init__(self, description, characters):
        self.description = description
        self.characters = characters

class NovelWritingAgent:
    def __init__(self, gpt_client):
        self.characters = []
        self.events = []
        self.timeline = []
        self.theme = ""
        self.gpt_client = gpt_client

    def add_character(self, name, personality, background):
        character = Character(name, personality, background)
        self.characters.append(character)
        return character

    def add_event(self, description, characters):
        event = Event(description, characters)
        self.events.append(event)
        return event

    def set_theme(self, theme):
        self.theme = theme

    def generate_plot(self):
        characters_info = "; ".join([f"{c.name} ({c.personality})" for c in self.characters])
        prompt = f"Create a plot outline for a story with the theme '{self.theme}'. Characters: {characters_info}. Include {len(self.events)} major events."
        plot_outline = self.gpt_client.generate_text(prompt, max_tokens=200)
        
        # Parse the generated plot outline and update the timeline
        events = plot_outline.split('\n')
        self.timeline = [Event(event.strip(), self.characters) for event in events if event.strip()]

    def ensure_coherence(self):
        coherent_timeline = []
        for i in range(len(self.timeline)):
            coherent_timeline.append(self.timeline[i])
            if i < len(self.timeline) - 1:
                transition_event = self.create_transition_event(self.timeline[i], self.timeline[i+1])
                coherent_timeline.append(transition_event)
        self.timeline = coherent_timeline

    def create_transition_event(self, event1, event2):
        prompt = f"Create a brief transition event connecting these two events:\n1. {event1.description}\n2. {event2.description}"
        transition_description = self.gpt_client.generate_text(prompt, max_tokens=50)
        return Event(transition_description, self.characters)

    def develop_character_arcs(self):
        for character in self.characters:
            character_events = [event.description for event in self.timeline if character in event.characters]
            prompt = f"Develop a character arc for {character.name} ({character.personality}) based on these events: {', '.join(character_events)}"
            arc_description = self.gpt_client.generate_text(prompt, max_tokens=100)
            character.arc = arc_description

    def generate_story(self):
        self.generate_plot()
        self.ensure_coherence()
        self.develop_character_arcs()
        
        story_outline = f"Theme: {self.theme}\n\n"
        story_outline += "\n".join([event.description for event in self.timeline])
        
        prompt = f"Write a short story based on this outline:\n\n{story_outline}\n\nInclude character development and ensure the story flows coherently."
        final_story = self.gpt_client.generate_text(prompt, max_tokens=1000)
        return final_story

# Usage example
api_key = os.environ.get("OPENAI_API_KEY")  # Make sure to set your API key as an environment variable
gpt_client = GPTClient(api_key)
agent = NovelWritingAgent(gpt_client)

# Add characters
hero = agent.add_character("Alex", "brave and idealistic", "ordinary background")
villain = agent.add_character("Dr. Nefarious", "cunning and ambitious", "wealthy background")

# Set theme
agent.set_theme("The triumph of hope over adversity")

# Generate the story
story = agent.generate_story()
print(story)