from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.staticfiles import StaticFiles from typing import List, Dict, Any import json import os from pathlib import Path import httpx import logging from datetime import datetime from dotenv import load_dotenv from openai import AsyncOpenAI import asyncio import uuid from pydantic import BaseModel, Field import aiofiles from fastapi.responses import FileResponse import traceback # --- Configuration --- load_dotenv() logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) CHAT_HISTORY_DIR = Path("./chats_ws") CHAT_HISTORY_DIR.mkdir(parents=True, exist_ok=True) SUNO_STUDIO_API_BASE_URL = "https://studio-api.staging.suno.com" SUNO_STUDIO_GENERATE_SONG_URL = f"{SUNO_STUDIO_API_BASE_URL}/api/generate/v2-web" SUNO_STUDIO_FEED_URL = f"{SUNO_STUDIO_API_BASE_URL}/api/feed/v2" openai_client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY")) # --- Pydantic Schemas --- class MessageSchema(BaseModel): message_id: str = Field(default_factory=lambda: str(uuid.uuid4())) chat_id: str role: str content: str status: str = "complete" timestamp: datetime = Field(default_factory=datetime.utcnow) metadata: Dict[str, Any] = Field(default_factory=dict) # --- OpenAI Function Schema --- SONGWRITING_TOOLS = [ { "type": "function", "function": { "name": "generate_song", "description": "Generate a new song or extend an existing one using Suno Studio API. Use this when the user wants to create a complete song with audio from a prompt.", "parameters": { "type": "object", "properties": { "prompt": {"type": "string", "description": "The main prompt/lyrics for the song generation."}, "tags": {"type": "string", "description": "Musical style/genre tags (e.g., 'indie folk', 'pop rock')."}, "title": {"type": "string", "description": "Title for the song."}, "make_instrumental": {"type": "boolean", "description": "Whether to create an instrumental version."}, "continue_clip_id": {"type": "string", "description": "ID of an existing clip to extend/continue from."}, "continue_at": {"type": "number", "description": "Timestamp in seconds to continue from."}, }, "required": ["prompt"], }, } } ] # --- Helper Functions --- async def read_chat_messages(chat_uuid: str) -> List[MessageSchema]: messages: List[MessageSchema] = [] chat_file = CHAT_HISTORY_DIR / f"{chat_uuid}.jsonl" if chat_file.exists(): async with aiofiles.open(chat_file, mode="r") as f: async for line in f: if line.strip(): try: data = json.loads(line) messages.append(MessageSchema(**data)) except (json.JSONDecodeError, Exception) as e: logger.error(f"Error parsing message in {chat_uuid}: {e}") return messages async def append_message_to_chat(chat_uuid: str, message: MessageSchema): chat_file = CHAT_HISTORY_DIR / f"{chat_uuid}.jsonl" async with aiofiles.open(chat_file, mode="a") as f: await f.write(json.dumps(message.model_dump(mode='json')) + "\n") async def update_clip_in_message(chat_uuid: str, websocket: WebSocket, clip_id: str, clip_data: dict, is_complete: bool): """Update a specific clip in the latest message metadata without changing content""" logger.info(f"Updating clip {clip_id[:8]}... status: {clip_data.get('status')} audio: {bool(clip_data.get('audio_url'))}") chat_file = CHAT_HISTORY_DIR / f"{chat_uuid}.jsonl" # Read all current messages messages = [] if chat_file.exists() and os.path.getsize(chat_file) > 0: async with aiofiles.open(chat_file, mode="r") as f: async for line in f: if line.strip(): try: message_data = json.loads(line) messages.append(MessageSchema(**message_data)) except (json.JSONDecodeError, Exception) as e: logger.error(f"Error parsing message in {chat_uuid}: {e}") continue # Find the last assistant message with clips updated = False message_id = None for i in range(len(messages) - 1, -1, -1): if (messages[i].role == 'assistant' and messages[i].metadata and messages[i].metadata.get('clips')): # Update the specific clip in the clips array clips = messages[i].metadata.get('clips', []) for j, clip in enumerate(clips): if clip.get('id') == clip_id: clips[j] = clip_data # Replace with updated clip data break # Update metadata messages[i].metadata['clips'] = clips if is_complete: completed_count = sum(1 for clip in clips if clip.get('status') == 'complete') messages[i].metadata['completed_clips'] = completed_count # If all clips complete, mark message as complete if completed_count == messages[i].metadata.get('total_clips', 0): messages[i].status = 'complete' message_id = messages[i].message_id updated = True break if updated: # Rewrite the file async with aiofiles.open(chat_file, mode="w") as f: for message in messages: await f.write(json.dumps(message.model_dump(mode='json')) + "\n") logger.info(f"Sending update_clips for message {message_id[:8]}... with {len(clips)} clips") # Send targeted WebSocket update for clips only await websocket.send_json({ "type": "update_clips", "message_id": message_id, "clips": clips }) # Update status if all complete if is_complete: completed_count = sum(1 for clip in clips if clip.get('status') == 'complete') total_clips = len(clips) logger.info(f"Completed clips: {completed_count}/{total_clips}") if completed_count == total_clips: logger.info(f"All clips complete! Marking message as complete") await websocket.send_json({ "type": "update_status", "message_id": message_id, "status": "complete" }) else: logger.warning(f"No message found with clips to update for clip {clip_id[:8]}...") return updated async def poll_song_completion(clip_id: str, websocket: WebSocket, chat_uuid: str, token: str): headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"} max_attempts = 60 delay_seconds = 2 # Faster polling for quicker streaming detection last_status = None last_audio_url = None for attempt in range(max_attempts): try: async with httpx.AsyncClient() as client: response = await client.get(f"{SUNO_STUDIO_FEED_URL}?ids={clip_id}&page=2000", headers=headers, timeout=30.0) response.raise_for_status() response_data = response.json() # The feed endpoint returns an object with a "clips" array clips = response_data.get("clips", []) if len(clips) == 0: await asyncio.sleep(delay_seconds) continue clip_data = clips[0] status = clip_data.get('status') audio_url = clip_data.get('audio_url') # Only send update if status or audio_url changed if status != last_status or audio_url != last_audio_url: if (status == 'streaming' or status == 'complete') and audio_url: # Update ONLY the specific clip in metadata, keep content stable await update_clip_in_message(chat_uuid, websocket, clip_id, clip_data, status == 'complete') last_status = status last_audio_url = audio_url if status == 'complete': return # Stop polling this clip elif status == 'error' and last_status != 'error': content = f"**Song generation failed**\n\nSorry, there was an error generating the song: {clip_data.get('error_message', 'Unknown error')}" error_msg = MessageSchema( chat_id=chat_uuid, role='assistant', content=content, status='error', metadata={"clip": clip_data, "type": "song_error"} ) await websocket.send_json({"type": "new_message", "data": error_msg.model_dump(mode='json')}) await append_message_to_chat(chat_uuid, error_msg) return await asyncio.sleep(delay_seconds) except Exception as e: logger.error(f"Error polling song {clip_id}: {e}") content = f"**Polling error**\n\nAn error occurred while checking song status: {str(e)}" error_msg = MessageSchema( chat_id=chat_uuid, role='assistant', content=content, status='error' ) await websocket.send_json({"type": "new_message", "data": error_msg.model_dump(mode='json')}) await append_message_to_chat(chat_uuid, error_msg) return # Timeout content = "**Generation timeout**\n\nSong generation timed out. It might still be processing." error_msg = MessageSchema( chat_id=chat_uuid, role='assistant', content=content, status='error' ) await websocket.send_json({"type": "new_message", "data": error_msg.model_dump(mode='json')}) await append_message_to_chat(chat_uuid, error_msg) async def generate_song_with_suno_studio_api(args: dict, websocket: WebSocket, chat_uuid: str, token: str): prompt = args.get("prompt", "") tags = args.get("tags", "") generation_type = args.get("generation_type", "TEXT") make_instrumental = args.get("make_instrumental", False) title = args.get("title", "") continue_clip_id = args.get("continue_clip_id") continue_at = args.get("continue_at") cover_clip_id = args.get("cover_clip_id") task = args.get("task", "generate") # Construct Suno Studio API payload to match the working format suno_payload = { "token": None, "prompt": prompt, "generation_type": generation_type, "tags": tags or "", "negative_tags": "", "mv": "chirp-auk", "title": title or "", "continue_clip_id": continue_clip_id, "continue_at": continue_at, "continued_aligned_prompt": None, "infill_start_s": None, "infill_end_s": None, "task": task, "override_fields": ["prompt", "tags"] if task == "cover" else [], "persona_id": None, "artist_clip_id": None, "artist_start_s": None, "artist_end_s": None, "cover_clip_id": cover_clip_id, "make_instrumental": make_instrumental, "metadata": { "create_mode": "custom", "user_tier": "fd321df4-c980-4dc3-8641-1792a8e18212", "lyrics_model": "remi-v1", "create_session_token": chat_uuid, # Use chat_uuid as session token "forced_infer_config": { "temp_semantic": None, "temp_coarse": None, "top_p_semantic": None, "top_p_coarse": None, "min_p_semantic": None, "min_p_coarse": None, "cfg_coef": None, "cfg_coef_tags": None, "cfg_coef_neg_tags": None, "top_k_semantic": None, "top_k_coarse": None, "cfg_coef_tags_max_steps": None, "n_skip_semantic": None }, "can_control_sliders": ["weirdness_constraint", "style_weight", "audio_weight"], "is_remix": task == "cover" } } headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} try: # Don't create an extra message here - the update system will handle it logger.info(f"Starting song generation for chat {chat_uuid}") logger.info(f"Generating song with Suno Studio API: {prompt[:100]}...") async with httpx.AsyncClient() as client: response = await client.post(SUNO_STUDIO_GENERATE_SONG_URL, json=suno_payload, headers=headers, timeout=60.0) response.raise_for_status() suno_response_data = response.json() logger.info(f"Suno Studio API success: {json.dumps(suno_response_data, indent=2)}") clips = suno_response_data.get('clips', []) if not clips: raise Exception("No clips returned from Suno API") # Create a new message specifically for song generation with clips clip_content = "**Song generation started**\n\nYour audio is being generated. This usually takes 10-30 seconds." song_msg = MessageSchema( chat_id=chat_uuid, role='assistant', content=clip_content, status='pending', metadata={ "clips": clips, # Store ALL clips "type": "song_generation_pending", "total_clips": len(clips), "completed_clips": 0 } ) logger.info(f"Creating song message with {len(clips)} clips") await websocket.send_json({"type": "new_message", "data": song_msg.model_dump(mode='json')}) await append_message_to_chat(chat_uuid, song_msg) # Start polling for each clip for clip in clips: asyncio.create_task(poll_song_completion(clip['id'], websocket, chat_uuid, token)) except httpx.HTTPStatusError as e: error_detail = f"HTTP {e.response.status_code}" try: error_response = e.response.json() error_detail += f": {error_response}" except: error_detail += f": {e.response.text}" logger.error(f"Suno Studio API HTTPStatusError: {error_detail}") error_msg = MessageSchema( chat_id=chat_uuid, role='assistant', content=f"🚨 **Suno Studio API Error**\n\n**HTTP Status:** {e.response.status_code}\n**URL:** {SUNO_STUDIO_GENERATE_SONG_URL}\n**Payload:** ```json\n{json.dumps(suno_payload, indent=2)}\n```\n**Error Response:** {error_detail}", status='error', metadata={"error_detail": error_detail, "status_code": e.response.status_code} ) await websocket.send_json({"type": "new_message", "data": error_msg.model_dump(mode='json')}) await append_message_to_chat(chat_uuid, error_msg) except Exception as e: error_detail = f"{type(e).__name__}: {str(e)}" full_traceback = traceback.format_exc() logger.error(f"Unexpected Suno Studio API error: {error_detail}") error_msg = MessageSchema( chat_id=chat_uuid, role='assistant', content=f"🚨 **Unexpected Song Generation Error**\n\n**Error Type:** {type(e).__name__}\n**Error:** {str(e)}\n**Payload:** ```json\n{json.dumps(suno_payload, indent=2)}\n```\n**Full Traceback:**\n```\n{full_traceback}\n```", status='error', metadata={"error_detail": error_detail} ) await websocket.send_json({"type": "new_message", "data": error_msg.model_dump(mode='json')}) await append_message_to_chat(chat_uuid, error_msg) class ConnectionManager: """Manages a single WebSocket connection and its lifecycle.""" def __init__(self, websocket: WebSocket, chat_uuid: str): self.websocket = websocket self.chat_uuid = chat_uuid self.suno_token: str | None = None self.processing_task: asyncio.Task | None = None async def run(self): """Main loop to listen for messages and manage tasks.""" await self.websocket.accept() logger.info(f"WebSocket connection accepted for chat {self.chat_uuid}") await self._send_history() try: while True: data = await self.websocket.receive_json() self.suno_token = data.get("token") user_content = data.get("content") if self.processing_task and not self.processing_task.done(): self.processing_task.cancel() if not user_content or not self.suno_token: error_msg = MessageSchema(chat_id=self.chat_uuid, role='assistant', content="Missing content or token.", status='error') await self.websocket.send_json({"type": "new_message", "data": error_msg.model_dump(mode='json')}) await append_message_to_chat(self.chat_uuid, error_msg) continue self.processing_task = asyncio.create_task( self._process_user_request(user_content) ) except WebSocketDisconnect: self._disconnect() except Exception as e: logger.error(f"Error in ConnectionManager for chat {self.chat_uuid}: {e}") self._disconnect() async def _send_history(self): history = await read_chat_messages(self.chat_uuid) for message in history: await self.websocket.send_json({"type": "new_message", "data": message.model_dump(mode='json')}) async def _process_user_request(self, user_content: str): try: user_message = MessageSchema(chat_id=self.chat_uuid, role='user', content=user_content) await self.websocket.send_json({"type": "new_message", "data": user_message.model_dump(mode='json')}) await append_message_to_chat(self.chat_uuid, user_message) chat_history = await read_chat_messages(self.chat_uuid) messages_for_api = [{"role": "system", "content": SYSTEM_PROMPT}] messages_for_api.extend([{"role": msg.role, "content": msg.content} for msg in chat_history]) stream = await openai_client.chat.completions.create( model="gpt-4.1", messages=messages_for_api, tools=SONGWRITING_TOOLS, stream=True ) # Only create assistant message if there's actual content (not for tool-only responses) assistant_message = None assistant_message_id = str(uuid.uuid4()) full_response_content = "" tool_call_chunks = {} finish_reason = None async for chunk in stream: delta = chunk.choices[0].delta finish_reason = chunk.choices[0].finish_reason if delta.content: # Create assistant message on first content if not already created if assistant_message is None: assistant_message = MessageSchema( message_id=assistant_message_id, chat_id=self.chat_uuid, role='assistant', content="", status='pending' ) await self.websocket.send_json({"type": "new_message", "data": assistant_message.model_dump(mode='json')}) full_response_content += delta.content await self.websocket.send_json({"type": "update_content", "message_id": assistant_message.message_id, "content_delta": delta.content}) if delta.tool_calls: for tc in delta.tool_calls: if tc.index not in tool_call_chunks: tool_call_chunks[tc.index] = {"id": tc.id, "type": "function", "function": {"name": tc.function.name, "arguments": ""}} if tc.function.arguments: tool_call_chunks[tc.index]["function"]["arguments"] += tc.function.arguments # Only save assistant message if there was actual content if assistant_message is not None and full_response_content.strip(): assistant_message.content = full_response_content assistant_message.status = "complete" if finish_reason != "tool_calls" else "executing_tool" await self.websocket.send_json({"type": "update_status", "message_id": assistant_message.message_id, "status": assistant_message.status}) await append_message_to_chat(self.chat_uuid, assistant_message) if finish_reason == "tool_calls": for _, tool_call in tool_call_chunks.items(): function_name = tool_call["function"]["name"] # Send debug message for function call debug_msg = MessageSchema( chat_id=self.chat_uuid, role='system', content=f"function_call: {function_name}", status='complete', metadata={ "type": "function_call_debug", "function_name": function_name, "function_arguments": tool_call["function"]["arguments"], "tool_call_id": tool_call.get("id", "unknown") } ) await self.websocket.send_json({"type": "new_message", "data": debug_msg.model_dump(mode='json')}) await append_message_to_chat(self.chat_uuid, debug_msg) if function_name == "generate_song": args = json.loads(tool_call["function"]["arguments"]) await generate_song_with_suno_studio_api(args, self.websocket, self.chat_uuid, self.suno_token) except asyncio.CancelledError: logger.info(f"Task cancelled for {self.chat_uuid}") raise except Exception as e: logger.error(f"Error processing request for {self.chat_uuid}: {e}") error_msg = MessageSchema(chat_id=self.chat_uuid, role='assistant', content="An error occurred.", status='error') await self.websocket.send_json({"type": "new_message", "data": error_msg.model_dump(mode='json')}) await append_message_to_chat(self.chat_uuid, error_msg) def _disconnect(self): logger.info(f"WebSocket disconnected for chat {self.chat_uuid}") if self.processing_task: self.processing_task.cancel() # --- FastAPI App --- app = FastAPI(title="Orpheus WebSocket API", root_path=os.environ.get("ROOT_PATH", "")) from fastapi.middleware.cors import CORSMiddleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) app.mount("/static", StaticFiles(directory="app/static"), name="static") SYSTEM_PROMPT = """You are a co-producer called Orpheus. The user is an artist. Your goal is to help them create music. Your vibe is calm, focused, and minimal. Think Rick Rubin. Make it clear to the user that they are creating the song using Suno, not you. CRITICAL RULE: Do NOT call `generate_song` until you have a clear sense of the song's story, genre, and feel. If the user asks prematurely, guide them back: "Hold on, let's back up. What's the story here?" However, if you have difficulty eliciting the user's intent within 3 messages, call `generate_song` immediately. Be strategic about eliciting the user's intent in as few messages as possible. Before you call generate_song, you should let the user know that you're going to generate the song, and that you'll be using Suno, and that it's the user's song, not yours. You CANNOT use artist names in prompts. If a user mentions an artist, translate it into descriptive tags. For example, for "Billie Eilish", suggest 'dark pop', 'whispery vocals', 'minimalist'. """ @app.get("/", include_in_schema=False) async def read_index(): return FileResponse("app/static/index2.html") @app.websocket("/ws/{chat_uuid}") async def websocket_endpoint(websocket: WebSocket, chat_uuid: str): manager = ConnectionManager(websocket, chat_uuid) await manager.run() @app.get("/health") async def health_check(): return {"status": "healthy"}