Spaces:
Running
Running
| import sqlite3 | |
| import operator | |
| from typing import TypedDict, List, Dict, Any, Annotated | |
| from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage | |
| from langgraph.graph import StateGraph, END | |
| from langgraph.checkpoint.sqlite import SqliteSaver | |
| from langchain_ollama import ChatOllama | |
| class SongState(TypedDict): | |
| script: str | |
| extracted_themes: str | |
| interviewer_history: Annotated[List[BaseMessage], operator.add] | |
| user_preferences: Annotated[List[str], operator.add] | |
| generated_lyrics: str | |
| music_style_tags: str | |
| question_count: int | |
| local_llm = ChatOllama(model="llama3", temperature=0.5) # Lower temperature for more stable adherence | |
| # --- GRAPH NODES --- | |
| def analyze_script_node(state: SongState): | |
| print("\n[System]: Analyzing script tones...") | |
| prompt = f"Analyze this script and output raw musical characteristics (Mood, Theme, Tempo):\n{state['script']}" | |
| response = local_llm.invoke(prompt) | |
| return {"extracted_themes": response.content, "question_count": 0, "interviewer_history": []} | |
| def interview_user_node(state: SongState): | |
| history = state.get("interviewer_history", []) | |
| system_instruction = SystemMessage( | |
| content=( | |
| f"You are a music producer interviewing an artist based on this script analysis:\n{state['extracted_themes']}\n\n" | |
| "Ask EXACTLY ONE short, unique question to narrow down style/instruments.\n" | |
| "Do NOT repeat yourself. Look at the history and move to a new detail (e.g., pace, vocals, era).\n" | |
| "Output ONLY the question text. No conversational filler." | |
| ) | |
| ) | |
| response = local_llm.invoke([system_instruction] + history) | |
| print(f"\n[AI Producer]: {response.content}") | |
| return { | |
| "interviewer_history": [AIMessage(content=response.content)], | |
| "question_count": state["question_count"] + 1 | |
| } | |
| def human_feedback_node(state: SongState): | |
| """Pauses internally in the terminal to grab your answer naturally!""" | |
| user_reply = input("[Your Answer]: ") | |
| return { | |
| "user_preferences": [user_reply], | |
| "interviewer_history": [HumanMessage(content=user_reply)] | |
| } | |
| def lyrics_and_planner_node(state: SongState): | |
| print("\n[System]: Finalizing lyrics and music tags...") | |
| answers = "\n".join(state.get("user_preferences", [])) | |
| prompt = f"Write structured song lyrics ([Verse], [Chorus], [Bridge]) based on:\n{state['script']}\nStyle adjustments:\n{answers}" | |
| lyrics_out = local_llm.invoke(prompt) | |
| tag_prompt = f"Convert these preferences into a brief comma-separated list of musical style tags (e.g., 'lo-fi, 120bpm, synth'):\n{answers}" | |
| tags_out = local_llm.invoke(tag_prompt) | |
| return {"generated_lyrics": lyrics_out.content, "music_style_tags": tags_out.content} | |
| # --- CONTROL ROUTER --- | |
| def interview_router(state: SongState): | |
| if int(state.get("question_count", 0)) >= 3: | |
| return "compile_tracks" | |
| return "ask_more" | |
| # --- COMPILING THE WORKFLOW --- | |
| conn = sqlite3.connect("song_memory.db", check_same_thread=False) | |
| memory = SqliteSaver(conn) | |
| workflow = StateGraph(SongState) | |
| workflow.add_node("analyzer", analyze_script_node) | |
| workflow.add_node("interviewer", interview_user_node) | |
| workflow.add_node("human_input", human_feedback_node) # New node | |
| workflow.add_node("composer", lyrics_and_planner_node) | |
| workflow.set_entry_point("analyzer") | |
| workflow.add_edge("analyzer", "interviewer") | |
| # Route to human input first, then check the counter loop | |
| workflow.add_edge("interviewer", "human_input") | |
| workflow.add_conditional_edges( | |
| "human_input", | |
| interview_router, | |
| {"ask_more": "interviewer", "compile_tracks": "composer"} | |
| ) | |
| workflow.add_edge("composer", END) | |
| langgraph_app = workflow.compile(checkpointer=memory) |