Spaces:
Sleeping
Sleeping
File size: 6,116 Bytes
b4c9cb7 172064c b4c9cb7 172064c b4c9cb7 172064c b4c9cb7 39cdf57 172064c d62d2dd 172064c d62d2dd 172064c b4c9cb7 172064c b4c9cb7 39cdf57 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 160 161 162 |
from fastapi import APIRouter, Depends
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, Field
import json
import asyncio
from src.apis.middlewares.auth_middleware import get_current_user
from typing import Annotated
from src.apis.models.user_models import User
user_dependency = Annotated[User, Depends(get_current_user)]
class GenScriptRequest(BaseModel):
video_link: str = Field(..., description="Video link")
target_word_count: int = Field(
2500, ge=2000, le=12000, description="Target word count"
)
language: str = Field(..., description="Language")
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, user: user_dependency):
"""
Generate scripts with streaming response
"""
config = RunnableConfig()
input_graph = {
"video_link": request.video_link,
"target_word_count": request.target_word_count,
"language": request.language,
}
return StreamingResponse(
message_generator(input_graph, config),
media_type="text/plain",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"Content-Type": "text/event-stream",
},
)
|