Spaces:
Sleeping
Sleeping
File size: 9,434 Bytes
b5c24ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 | """FastAPI server for the LangGraph content generation agent."""
from src.agents import create_agent
import os
import json
import asyncio
import tempfile
import base64
from typing import Optional, List
from pathlib import Path
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse
from pydantic import BaseModel
from dotenv import load_dotenv
load_dotenv()
# Initialize FastAPI app
app = FastAPI(
title="Magic Hour LangGraph API",
description="Content generation API using LangChain and LangGraph",
version="1.0.0"
)
# CORS configuration
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Configure appropriately for production
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Global agent instance (per-thread isolation via thread_id)
agent = create_agent(
fal_model_name=os.getenv(
"FAL_MODEL_NAME", "google/gemini-3-flash-preview"),
temperature=float(os.getenv("LLM_TEMPERATURE", "0.7")),
max_tokens=int(os.getenv("LLM_MAX_TOKENS", "4096")),
)
class ChatMessage(BaseModel):
"""A single chat message."""
role: str # "user" or "assistant"
content: str
image_paths: Optional[List[str]] = None
class ChatRequest(BaseModel):
"""Request body for chat endpoint."""
message: str
image_paths: Optional[List[str]] = None # Paths to images to include
image_data: Optional[List[str]] = None # Base64 encoded images
thread_id: str = "default"
settings: Optional[dict] = None # mode, aspect_ratio, etc.
class ChatResponse(BaseModel):
"""Response body for chat endpoint."""
message: str
generated_content: List[str] = []
thread_id: str
def _save_base64_image(base64_data: str) -> str:
"""Save a base64 encoded image to a temp file and return the path."""
# Handle data URL format
if "," in base64_data:
base64_data = base64_data.split(",")[1]
image_bytes = base64.b64decode(base64_data)
temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
temp_file.write(image_bytes)
temp_file.close()
return temp_file.name
def _extract_response_text(messages) -> str:
"""Extract the final response text from agent messages."""
from langchain_core.messages import AIMessage
for msg in reversed(messages):
if isinstance(msg, AIMessage):
if isinstance(msg.content, str):
return msg.content
elif isinstance(msg.content, list):
# Handle multimodal responses
text_parts = [p.get("text", "") for p in msg.content if isinstance(
p, dict) and p.get("type") == "text"]
return " ".join(text_parts)
return ""
@app.get("/health")
async def health_check():
"""Health check endpoint."""
return {"status": "healthy", "service": "langgraph-agent"}
@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
"""Process a chat message and return the response."""
# Prepare image paths
image_paths = request.image_paths or []
# Save any base64 images to temp files
if request.image_data:
for b64_data in request.image_data:
try:
path = _save_base64_image(b64_data)
image_paths.append(path)
except Exception as e:
print(f"Failed to save base64 image: {e}")
# Invoke the agent
try:
result = agent.invoke(
message=request.message,
image_paths=image_paths if image_paths else None,
thread_id=request.thread_id,
settings=request.settings
)
response_text = _extract_response_text(result.get("messages", []))
generated_content = result.get("generated_content", [])
return ChatResponse(
message=response_text,
generated_content=generated_content,
thread_id=request.thread_id
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/chat/stream")
async def chat_stream(request: ChatRequest):
"""Stream the agent's response as Server-Sent Events."""
async def event_generator():
# Prepare image paths
image_paths = request.image_paths or []
# Save any base64 images to temp files
if request.image_data:
for b64_data in request.image_data:
try:
path = _save_base64_image(b64_data)
image_paths.append(path)
except Exception as e:
yield f"data: {json.dumps({'type': 'error', 'content': f'Failed to save image: {e}'})}\n\n"
try:
# Stream agent responses
for event in agent.stream(
message=request.message,
image_paths=image_paths if image_paths else None,
thread_id=request.thread_id,
settings=request.settings
):
messages = event.get("messages", [])
generated = event.get("generated_content", [])
# Process messages
for msg in messages:
if hasattr(msg, "tool_calls") and msg.tool_calls:
# Tool call event
for tool_call in msg.tool_calls:
yield f"data: {json.dumps({'type': 'tool_call', 'name': tool_call.get('name', 'unknown'), 'args': tool_call.get('args', {})})}\n\n"
elif hasattr(msg, "content"):
# Check message type
msg_type = type(msg).__name__
if msg_type == "ToolMessage":
# Tool result
content = msg.content if isinstance(
msg.content, str) else str(msg.content)
yield f"data: {json.dumps({'type': 'tool_result', 'content': content})}\n\n"
# Check for generated content paths
import re
paths = re.findall(
r'(/[^\s]+\.(?:png|jpg|jpeg|webp|mp4))', content)
for path in paths:
if os.path.exists(path):
if path.endswith(".mp4"):
yield f"data: {json.dumps({'type': 'video_complete', 'path': path})}\n\n"
else:
yield f"data: {json.dumps({'type': 'image_complete', 'path': path})}\n\n"
elif msg_type == "AIMessage":
content = msg.content if isinstance(
msg.content, str) else str(msg.content)
yield f"data: {json.dumps({'type': 'assistant', 'content': content})}\n\n"
# Allow other tasks to run
await asyncio.sleep(0.01)
# Send completion event
yield f"data: {json.dumps({'type': 'done', 'generated_content': generated})}\n\n"
except Exception as e:
yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n"
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
}
)
@app.get("/api/media")
async def get_media(path: str):
"""Serve generated media files."""
# Security: Validate path is within temp directory
temp_dir = tempfile.gettempdir()
resolved_path = Path(path).resolve()
if not str(resolved_path).startswith(temp_dir):
raise HTTPException(status_code=403, detail="Access denied")
if not resolved_path.exists():
raise HTTPException(status_code=404, detail="File not found")
# Determine media type
suffix = resolved_path.suffix.lower()
media_types = {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".webp": "image/webp",
".gif": "image/gif",
".mp4": "video/mp4",
".webm": "video/webm",
}
media_type = media_types.get(suffix, "application/octet-stream")
return FileResponse(
path=str(resolved_path),
media_type=media_type,
filename=resolved_path.name
)
@app.get("/history/{thread_id}")
async def get_history(thread_id: str):
"""Get conversation history for a thread."""
messages = agent.get_conversation_history(thread_id)
generated = agent.get_generated_content(thread_id)
# Convert messages to serializable format
history = []
for msg in messages:
msg_type = type(msg).__name__
content = msg.content if isinstance(
msg.content, str) else str(msg.content)
history.append({
"type": msg_type,
"content": content
})
return {
"thread_id": thread_id,
"messages": history,
"generated_content": generated
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(
"api_server:app",
host="0.0.0.0",
port=int(os.getenv("PORT", "8000")),
reload=True
)
|