diff --git a/main.py b/main.py new file mode 100644 index 0000000..21a9c33 --- /dev/null +++ b/main.py @@ -0,0 +1,123 @@ +import os +import asyncio +from typing import TypedDict, List, Dict, Any + +# LangGraph imports +from langgraph.graph import StateGraph, START +from langgraph.types import interrupt, Command +from langgraph.checkpoint.memory import InMemorySaver + +# LangChain LLM +from langchain_openai import ChatOpenAI + +# Console interaction +import questionary + +# ---------- 1. Define state --------------------------------- +class GameState(TypedDict): + theme: str + intro: str + options: List[str] + choice: str | None + ending: str + +# ---------- 2. LLM client ----------------------------------- +llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1", + api_key=os.getenv("JOURNAL_MCP_PAT"), + temperature=0.7, +) + +# ---------- 3. Node to generate intro and options --------- +async def generate_intro(state: GameState) -> Dict: + theme = state["theme"] + prompt = ( + f"You are a creative storyteller.\n" + f"Theme: {theme}.\n" + "Generate a short opening paragraph (2-3 sentences).\n" + "Then provide exactly three distinct actions the protagonist can take, one per line, prefixed with numbers 1) 2) 3).\n" + "Return only the intro followed by the numbered list." + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + text = response.content + # Split into intro and options + parts = text.split("\n") + intro_lines = [] + options: List[str] = [] + for line in parts: + if line.strip().startswith(tuple(str(i)+")" for i in range(1,4)): + options.append(line.strip()) + else: + intro_lines.append(line) + state["intro"] = " ".join(intro_lines).strip() + state["options"] = options + # Trigger interrupt to ask user + payload = { + "type": "choice", + "question": f"{state['intro']}\n\nWhat do you do?", + "choices": options, + } + return interrupt(payload) + +# ---------- 4. Node to process choice and generate ending ----- +async def generate_ending(state: GameState) -> Dict: + # state now contains 'choice' + theme = state["theme"] + intro = state["intro"] + choice = state.get("choice", "") + prompt = ( + f"You are a creative storyteller.\n" + f"Theme: {theme}.\n" + f"Intro: {intro}.\n" + f"The player chose: {choice}.\n" + "Write a short ending (2-3 sentences) that follows from this choice." + ) + response = await llm.ainvoke([HumanMessage(content=prompt)]) + state["ending"] = response.content.strip() + return state + +# ---------- 5. Build graph --------------------------------- +builder = StateGraph(GameState) +builder.add_node("intro", generate_intro) +builder.add_node("ending", generate_ending) +builder.set_entry_point("intro") +builder.add_edge(START, "intro") +builder.add_edge("intro", "ending") +# No edge from ending; graph ends after ending node +graph = builder.compile(checkpointer=InMemorySaver()) + +# ---------- 6. Run loop with interrupt handling --------- +async def run_game(theme: str): + thread_id = f"thread-{theme.replace(' ', '_')}" + config = {"configurable": {"thread_id": thread_id}} + # Start stream + async for event in graph.astream({"theme": theme}, config, stream_mode="messages"): + if isinstance(event, tuple): + msg, _meta = event + print(msg.content) + elif "__interrupt__" in event: + interrupt_payload = event["__interrupt__"][0].value + # Show question and choices using questionary + answer = questionary.select( + interrupt_payload["question"], + choices=[c.split(') ',1)[1] if ') ' in c else c for c in interrupt_payload["choices"]] + ).ask() + # Map back to full choice string + full_choice = next(c for c in interrupt_payload["choices"] if c.endswith(answer)) + # Resume with user answer + resume_payload = {**interrupt_payload, "choice": full_choice} + async for sub_event in graph.astream(Command(resume=resume_payload), config, stream_mode="messages"): + if isinstance(sub_event, tuple): + msg, _meta = sub_event + print(msg.content) + # After completion, print final state + final_state = await graph.aget_state(config) + print("\n--- Final Story ---") + print(final_state["intro"]) + print(f"Choice: {final_state['choice']}") + print(final_state["ending"]) + +if __name__ == "__main__": + theme = questionary.text("Enter a theme for the adventure:").ask() + asyncio.run(run_game(theme)) \ No newline at end of file