diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..47dfe77 --- /dev/null +++ b/agent.py @@ -0,0 +1,136 @@ +import os +import json +import asyncio +from typing import TypedDict, List, Dict, Any + +import questionary +from langchain_openai import ChatOpenAI +from langgraph.graph import StateGraph, START, END +from langgraph.checkpoint.memory import InMemorySaver +from langgraph import interrupt, Command + +# Define the state schema according to the plan +class StoryState(TypedDict): + theme: str + scene_text: str + options: List[str] + chosen_option: str + ending_text: str + +# Configure the LLM via environment variables +openai_api_key = os.getenv("OPENAI_API_KEY") +openai_base_url = os.getenv("OPENAI_BASE_URL") +llm = ChatOpenAI( + model="gpt-4o", + temperature=0, + openai_api_key=openai_api_key, + base_url=openai_base_url, +) + +async def call_llm(prompt: str) -> str: + """Invoke the LLM and return the text content, retrying once on failure.""" + for attempt in range(2): + try: + response = await llm.invoke(prompt) + # ChatResult contains a list of messages + if hasattr(response, "messages") and response.messages: + content = response.messages[0].content + else: + # Fallback to direct attribute + content = getattr(response, "content", "") + if not content: + raise ValueError("Empty LLM response") + return content + except Exception as e: + if attempt == 1: + raise ValueError(f"LLM call failed twice: {e}") from e + await asyncio.sleep(1) + raise ValueError("LLM call failed") + +async def generate_scene(state: StoryState) -> Dict[str, Any]: + theme = state.get("theme", "Магический лес") + prompt = ( + f"Сгенерируй сцену на тему '{theme}'.\n" + "Ответ должен быть в формате JSON:\n" + "{\"scene\": \"<описание>\", \"options\": [\"<вариант1>\",\"<вариант2>\",\"<вариант3>\"]}\n" + ) + text = await call_llm(prompt) + try: + data = json.loads(text) + except Exception as e: + raise ValueError(f"Failed to parse scene JSON: {e}") + scene = data.get("scene", "") + options = data.get("options", []) + if not isinstance(options, list) or len(options) != 3: + raise ValueError("LLM must return exactly 3 options") + return {"scene_text": scene, "options": options} + +async def interrupt_choice(state: StoryState) -> Dict[str, Any]: + scene = state["scene_text"] + options = state["options"] + payload = { + "type": "choice", + "question": f"{scene} Что делаем?", + "options": options, + } + # Interrupt the graph, returning a special dictionary + return interrupt(Command(resume=payload)) + +async def generate_ending(state: StoryState) -> Dict[str, Any]: + theme = state.get("theme", "Магический лес") + scene = state.get("scene_text", "") + choice = state.get("chosen_option", "") + prompt = ( + f"На основе темы '{theme}', сцены '{scene}' и выбранного варианта '{choice}',\n" + "сгенерируй короткую концовку истории.\n" + "Ответ должен быть только текстом." + ) + ending = (await call_llm(prompt)).strip() + return {"ending_text": ending} + +# Build the graph +builder = StateGraph(StoryState) +builder.add_node("generate_scene", generate_scene) +builder.add_node("interrupt_choice", interrupt_choice) +builder.add_node("generate_ending", generate_ending) + +builder.set_entry_point("generate_scene") +builder.add_edge("generate_scene", "interrupt_choice") +builder.add_edge("interrupt_choice", "generate_ending") +builder.add_edge("generate_ending", END) + +# Persist state across interrupts +builder.add_persisted_state(saver=InMemorySaver()) + +graph = builder.compile() + +async def main() -> None: + theme = questionary.text("Введите тему истории:").ask() + if not theme: + theme = "Магический лес" + initial_state: StoryState = { + "theme": theme, + "scene_text": "", + "options": [], + "chosen_option": "", + "ending_text": "", + } + # Run the graph until the first interrupt + partial_state = await graph.ainvoke(initial_state) + if "_interrupt" in partial_state: + payload = partial_state["_interrupt"] + choice = questionary.select(payload["question"], choices=payload["options"]).ask() + partial_state["chosen_option"] = choice + thread_id = partial_state.get("thread_id") + config = {"configurable": {"thread_id": thread_id}} if thread_id else {} + final_state = await graph.ainvoke(partial_state, config=config) + else: + final_state = partial_state + # Print the final story + print("\n=== Итоговая история ===\n") + print(final_state.get("scene_text", "")) + print(f"\nВы выбрали: {final_state.get('chosen_option', '')}\n") + print(final_state.get("ending_text", "")) + +if __name__ == "__main__": + asyncio.run(main())