Spaces:
Sleeping
Sleeping
File size: 6,857 Bytes
172064c |
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 |
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from langchain_core.messages import AIMessageChunk
from langchain_core.runnables import RunnableConfig
from src.agents.agent_transcript.flow import script_writer_agent
from src.utils.logger import logger
from pydantic import BaseModel
import json
import asyncio
class GenScriptRequest(BaseModel):
video_link: str
target_word_count: int = 50000 # Default 2500 words
router = APIRouter()
async def message_generator(
input_graph: dict,
config: RunnableConfig,
):
try:
last_output_state = None
try:
async for event in script_writer_agent.astream(
input=input_graph, stream_mode=["messages", "values"], config=config
):
try:
event_type, event_message = event
logger.info(f"Event type: {event_type}")
if event_type == "messages":
message, metadata = event_message
if isinstance(message, AIMessageChunk):
# Stream AI message chunks
node = metadata.get("node")
chunk_data = {
"type": "message_chunk",
"content": message.content,
"metadata": metadata,
"node_step": node,
}
logger.info(f"Chunk data: {chunk_data}")
yield f"data: {json.dumps(chunk_data)}\n\n"
elif event_type == "values":
# Stream state updates
state_data = {"type": "state_update", "state": event_message}
last_output_state = event_message
# Handle specific data extractions
if "transcript" in event_message and event_message["transcript"]:
transcript_data = {
"type": "transcript_extracted",
"transcript": event_message["transcript"][:500] + "..." if len(event_message["transcript"]) > 500 else event_message["transcript"],
"full_length": len(event_message["transcript"])
}
yield f"data: {json.dumps(transcript_data)}\n\n"
if "comment" in event_message and event_message["comment"]:
comment_data = {
"type": "comment_extracted",
"comment": event_message["comment"][:500] + "..." if len(event_message["comment"]) > 500 else event_message["comment"],
"full_length": len(event_message["comment"])
}
yield f"data: {json.dumps(comment_data)}\n\n"
if "script_count" in event_message:
script_count_data = {
"type": "script_count_calculated",
"script_count": event_message["script_count"],
"target_word_count": event_message.get("target_word_count", 8000)
}
yield f"data: {json.dumps(script_count_data)}\n\n"
# Handle individual script updates
if "script_writer_response" in event_message and "current_script_index" in event_message:
current_scripts = event_message["script_writer_response"]
current_index = event_message["current_script_index"]
script_count = event_message.get("script_count", 10)
if current_scripts:
individual_script_data = {
"type": "individual_script",
"script_index": current_index,
"script_content": current_scripts[-1] if current_scripts else "",
"progress": f"{current_index}/{script_count}",
"scripts": current_scripts
}
yield f"data: {json.dumps(individual_script_data)}\n\n"
yield f"data: {json.dumps(state_data, default=str)}\n\n"
except Exception as e:
logger.error(f"Error processing event: {e}")
error_data = {"type": "error", "message": str(e)}
yield f"data: {json.dumps(error_data)}\n\n"
except Exception as e:
logger.error(f"Error in streaming: {e}")
error_data = {"type": "error", "message": str(e)}
yield f"data: {json.dumps(error_data)}\n\n"
# Send final result
if last_output_state:
final_data = {
"type": "final_result",
"scripts": last_output_state.get("script_writer_response", []),
"total_scripts": len(
last_output_state.get("script_writer_response", [])
),
}
yield f"data: {json.dumps(final_data, default=str)}\n\n"
except Exception as e:
logger.error(f"Fatal error in message_generator: {e}")
yield f"data: {json.dumps({'type': 'fatal_error', 'message': str(e)})}\n\n"
@router.post("/gen-script")
async def gen_script(request: GenScriptRequest):
"""
Generate scripts with streaming response
"""
config = RunnableConfig()
input_graph = {
"video_link": request.video_link,
"target_word_count": request.target_word_count
}
return StreamingResponse(
message_generator(input_graph, config),
media_type="text/plain",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"Content-Type": "text/event-stream",
},
)
@router.post("/gen-script-sync")
def gen_script_sync(request: GenScriptRequest):
"""
Generate scripts with synchronous response (non-streaming)
"""
response = script_writer_agent.invoke({
"video_link": request.video_link,
"target_word_count": request.target_word_count
})
return {
"scripts": response.get("script_writer_response", []),
"total_scripts": len(response.get("script_writer_response", [])),
"full_response": response,
}
|