chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+22
@@ -0,0 +1,22 @@
|
||||
# Python-generated files
|
||||
__pycache__/
|
||||
*.py[oc]
|
||||
build/
|
||||
dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
|
||||
# Virtual environments
|
||||
.venv
|
||||
|
||||
# working directory
|
||||
.working_dir/
|
||||
.vimax/
|
||||
.test/
|
||||
# Node dependencies
|
||||
node_modules/
|
||||
|
||||
# Local agent secrets
|
||||
configs/*.local.json
|
||||
configs/*.local.yaml
|
||||
!configs/agent.local.yaml
|
||||
@@ -0,0 +1,6 @@
|
||||
We provide QR codes for joining the HKUDS discussion groups on WeChat and Feishu.
|
||||
|
||||
You can join by scanning the QR codes below:
|
||||
|
||||
<img src="https://github.com/HKUDS/.github/blob/main/profile/QR.png" alt="WeChat QR Code" width="400"/>
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`HKUDS/ViMax`
|
||||
- 原始仓库：https://github.com/HKUDS/ViMax
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
+451
@@ -0,0 +1,451 @@
|
||||
<div align="center">
|
||||
<img src="./assets/vimax.png">
|
||||
<br>
|
||||
<br>
|
||||
<h1 align="center">ViMax: Agentic Video Generation</h1>
|
||||
|
||||
<div align="center">
|
||||
</div>
|
||||
|
||||
<p align="center">
|
||||
<img src="https://img.shields.io/badge/🐍Python-3.12-00d9ff?style=for-the-badge&logo=python&logoColor=white&labelColor=1a1a2e">
|
||||
<a href="https://github.com/astral-sh/uv"><img src="https://img.shields.io/badge/⚡uv-Ready-ff6b6b?style=for-the-badge&logo=python&logoColor=white&labelColor=1a1a2e"></a>
|
||||
<img src="https://img.shields.io/badge/License-MIT-4ecdc4?style=for-the-badge&logo=opensourceinitiative&logoColor=white" alt="MIT License">
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="./Communication.md"><img src="https://img.shields.io/badge/💬Feishu-Group-07c160?style=for-the-badge&logoColor=white&labelColor=1a1a2e"></a>
|
||||
<a href="./Communication.md"><img src="https://img.shields.io/badge/WeChat-Group-07c160?style=for-the-badge&logo=wechat&logoColor=white&labelColor=1a1a2e"></a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href='https://www.youtube.com/@AI-Creator-is-here'><img src='https://img.shields.io/badge/YouTube-ff0000?style=for-the-badge&logo=youtube&logoColor=white&labelColor=1a1a2e' /></a>
|
||||
<a href='https://arxiv.org/abs/2606.07649'><img src='https://img.shields.io/badge/arXiv-2606.07649-b31b1b?style=for-the-badge&logo=arxiv&logoColor=white&labelColor=1a1a2e' /></a>
|
||||
</p>
|
||||
|
||||
</div>
|
||||
<div align="center">
|
||||
|
||||
<p>
|
||||
<a href="readme.md"><img src="https://img.shields.io/badge/English-1a1a2e?style=for-the-badge"></a>
|
||||
<a href="README_ZH.md"><img src="https://img.shields.io/badge/中文版-1a1a2e?style=for-the-badge"></a>
|
||||
</p>
|
||||
<a href="#quick-start" style="text-decoration: none;">
|
||||
<img src="https://img.shields.io/badge/Quick%20Start-Get%20Started%20Now-FFC107?style=for-the-badge&logo=rocket&logoColor=white&labelColor=1a1a2e">
|
||||
</a>
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
<div align="center">
|
||||
|
||||
|
||||
|
||||
https://github.com/user-attachments/assets/5bad46b2-8276-4e1d-9480-3522640744b2
|
||||
|
||||
|
||||
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
|
||||
## 📑 目录
|
||||
|
||||
- [💡 核心特性](#核心特性)
|
||||
- [🔮 演示示例](#演示示例)
|
||||
- [🏗️ 系统架构](#️-系统架构)
|
||||
- [🚀 快速开始](#快速开始)
|
||||
|
||||
---
|
||||
## 💡核心特性
|
||||
|
||||
<br/>
|
||||
|
||||
<table align="center" width="100%" style="border: none; table-layout: fixed;">
|
||||
<tr>
|
||||
<td width="25%" align="center" style="vertical-align: top; padding: 20px;">
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<h3 style="margin: 0; padding: 0;">🌟 <strong>创意到视频</strong></h3>
|
||||
</div>
|
||||
|
||||
<div align="center" style="margin: 15px 0;">
|
||||
<img src="https://img.shields.io/badge/创意-生成-ff6b6b?style=for-the-badge&logo=algorithm&logoColor=white" alt="算法徽章" />
|
||||
</div>
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>从灵感到银幕</strong></p>
|
||||
</div>
|
||||
|
||||
<div style="height: 60px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center">通过智能多智能体工作流，将<strong>原始创意</strong>自动转化为完整视频故事，涵盖<strong>叙事构建、角色设计与视频制作</strong>全流程。
|
||||
</p>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
</td>
|
||||
<td width="25%" align="center" style="vertical-align: top; padding: 20px;">
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<h3 style="margin: 0; padding: 0;">🎨 <strong>小说到视频</strong></h3>
|
||||
</div>
|
||||
|
||||
<div align="center" style="margin: 15px 0;">
|
||||
<img src="https://img.shields.io/badge/小说-改编-4ecdc4?style=for-the-badge&logo=book&logoColor=white" alt="前端徽章" />
|
||||
</div>
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>智能文学改编引擎</strong></p>
|
||||
</div>
|
||||
|
||||
<div style="height: 60px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center">将<strong>完整小说</strong>智能压缩并转化为<strong>分集视频内容</strong>，实现角色追踪、叙事压缩与逐场景视觉化改编。</p>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
</td>
|
||||
<td width="25%" align="center" style="vertical-align: top; padding: 20px;">
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<h3 style="margin: 0; padding: 0;">⚙️ <strong>剧本到视频</strong></h3>
|
||||
</div>
|
||||
|
||||
<div align="center" style="margin: 15px 0;">
|
||||
<img src="https://img.shields.io/badge/剧本-改编-9b59b6?style=for-the-badge&logo=server&logoColor=white" alt="后端徽章" />
|
||||
</div>
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>无限剧本视频创作</strong></p>
|
||||
</div>
|
||||
|
||||
<div style="height: 60px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center">自由创作<strong>任意剧本</strong>——从个人故事到史诗冒险，全面掌控视觉叙事的每个细节。</p>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
</td>
|
||||
<td width="25%" align="center" style="vertical-align: top; padding: 20px;">
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<h3 style="margin: 0; padding: 0;">🤳 <strong>智能客串</strong></h3>
|
||||
</div>
|
||||
|
||||
<div align="center" style="margin: 15px 0;">
|
||||
<img src="https://img.shields.io/badge/互动-生成-FFC107?style=for-the-badge&logo=server&logoColor=white" alt="后端徽章" />
|
||||
</div>
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>用你的照片生成视频</strong></p>
|
||||
</div>
|
||||
|
||||
<div style="height: 60px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>创建属于你的客串视频</strong>，将自己融入无限创意剧本、电影级镜头与互动剧情中，成为故事中的明星角色。</p>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<br/>
|
||||
|
||||
|
||||
---
|
||||
<table>
|
||||
<tr>
|
||||
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/c2fb27b0-218c-4976-b3d6-2abf8ea06be7" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/bfa566a8-688d-4d53-a9e2-6cedeb4a399d" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/49f61134-4f78-4285-9a9e-bb5e3e0c4abf" controls width="100%"></video>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/a950f449-a15c-449b-a1b8-c393951aa9be" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/bb3ff0fd-9433-4806-886a-3f77b61d06ec" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/2624a3f0-9f66-4fa4-b527-45c0ea0353fc" controls width="100%"></video>
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/5dbb80f7-aff0-4211-940c-a898f91fb80c" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/cc0b0bcd-e7db-4839-950b-0b03949637bd" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/85919b59-80f0-461a-af7e-a93d3fb412fc" controls width="100%"></video>
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
</table>
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
### 🎯 **端到端视频创作引擎**
|
||||
|
||||
**面临的挑战**：
|
||||
|
||||
- 🌅 **参考图像**：获取、整理并精准对齐能准确表达角色、物体、位置与环境的参考帧，耗时费力。
|
||||
|
||||
- 🫠 **一致性校验**：即使提供了正确的角色、位置、环境参考图与提示词，图像生成器有时仍会产出不可用图像。
|
||||
|
||||
- 📄 **剧本生成**：专业高质量视频需要高信息密度与结构化设计。
|
||||
|
||||
- 📝 **分镜设计**：将故事转化为视觉叙事，需要摄影、构图与视觉叙事的专业知识，而大多数创作者并不具备。
|
||||
|
||||
- 🎬 **镜头设计**：在复杂场景中保持叙事连贯性的同时，设计合理的镜头角度、转场与节奏。
|
||||
|
||||
- 🎨 **风格一致性**：在长视频中确保数百个镜头的角色外观、环境与艺术风格保持一致。
|
||||
|
||||
- ⏱️ **制作效率**：传统视频制作依赖多个专业人员与冗长流程，阻碍了独立创作者与快速原型开发。
|
||||
|
||||
- 🎥 **AI视频扩展性**：AI生成视频通常仅几秒，而分钟级甚至小时级的高质量长视频需要复杂的跨场景连续性与多分镜协同处理能力。
|
||||
|
||||
**ViMAX**：通过自动化从叙事输入到最终视频输出的完整流程，彻底消除上述制作瓶颈。
|
||||
|
||||
---
|
||||
|
||||
|
||||
### 🔥 **为什么选择 ViMax？**
|
||||
|
||||
| 🧠 **一键生成** | 🚀 **完全创作自由** | 🔊 **音画同步** | 🎨 **专业品质** | 🤩 **互动视频**
|
||||
|:---:|:---:|:---:|:---:|:---:|
|
||||
| 一句话生成完整视频 | 任何叙事皆可成真 | 音画完美融合 | 电影级输出 | 生成你的专属客串视频
|
||||
| 无需技术细节——只需描述你的创意，ViMax 自动完成剧本生成、分镜设计、镜头规划、参考管理与一致性验证 | 创意无边界——无论是预告片、短篇故事、小说章节还是原创概念，ViMax 都能智能构建叙事并设计镜头语言，将任何想法变为现实 | 无缝融合角色语音与音效，打造沉浸式视听体验 | 自动质量控制确保角色一致性、场景构图合理、每帧画面均达专业水准 | 上传你的照片即可在自己的故事中互动出演——ViMax 智能将你作为角色融入视频，保持外观一致并实现自然交互
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
|
||||
## 🏗️ 系统架构
|
||||
|
||||
### 📊 **系统概览**
|
||||
|
||||
**ViMax** 是一个多智能体视频生成框架，支持自动化多镜头视频生成，并确保角色与场景的一致性。系统能将你的创意无缝转化为对应视频，让你专注于讲故事，而非技术实现。
|
||||
|
||||
🎯 **技术能力**：
|
||||
|
||||
🧬 **智能长剧本生成**
|
||||
基于 RAG 的长剧本引擎，可智能分析小说级长文本，并自动切分为多场景剧本格式，精准保留关键情节与角色对话。
|
||||
|
||||
🪄 **表现力分镜设计**
|
||||
基于用户需求与目标受众，运用电影语言生成富有表现力的镜头级分镜，为后续视频生成奠定叙事节奏。
|
||||
|
||||
🔮 **多机位拍摄模拟**
|
||||
模拟多机位拍摄，提供沉浸式观看体验，同时确保同一场景内角色位置与背景的一致性。
|
||||
|
||||
🧸 **智能参考图选择**
|
||||
智能选取当前视频首帧所需的参考图（包括前序时间线中的分镜），确保视频越长，多角色与环境元素越准确。
|
||||
|
||||
⚙️ **自动化图像生成**
|
||||
基于所选参考图与前序时间线的视觉逻辑，自动生成图像生成器提示词，合理安排角色与环境的空间交互位置。
|
||||
|
||||
✅ **图像生成一致性校验**
|
||||
并行生成多张图像，并通过 MLLM/VLM 选择最一致的图像作为首帧，模拟人类创作者的工作流程。
|
||||
|
||||
⚡ **高效并行镜头生成**
|
||||
对同一机位拍摄的连续镜头进行并行处理，极大提升视频生产效率。
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
### 🤖 <strong>多智能体视频生成流水线</strong>
|
||||
|
||||
<div align="center">
|
||||
<table align="center" width="100%" style="border: none; border-collapse: collapse;">
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px; color: white; font-weight: bold;">
|
||||
🧠 <strong>输入层</strong><br/>
|
||||
📝 创意/剧本/小说 • 💭 自然语言提示 • 🖼️ 参考图像 • 🎨 风格指令 • 🧩 配置参数
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="20"></td></tr>
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 15px; background: linear-gradient(135deg, #ff6b6b 0%, #ee5a24 100%); border-radius: 12px; color: white; font-weight: bold;">
|
||||
🧭 <strong>中央调度</strong><br/>
|
||||
智能体调度 • 阶段切换 • 资源管理 • 重试/降级逻辑
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<td align="center" style="padding: 12px; background: linear-gradient(135deg, #3742fa 0%, #2f3542 100%); border-radius: 10px; color: white; width: 50%;">
|
||||
🧾 <strong>剧本理解</strong><br/>
|
||||
<small>角色/环境提取 • 场景边界识别 • 风格意图解析</small>
|
||||
</td>
|
||||
<td width="10"></td>
|
||||
<td align="center" style="padding: 12px; background: linear-gradient(135deg, #8c7ae6 0%, #9c88ff 100%); border-radius: 10px; color: white; width: 50%;">
|
||||
🎥 <strong>场景与镜头规划</strong><br/>
|
||||
<small>分镜步骤 • 镜头列表 • 关键帧与节奏点</small>
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 15px; background: linear-gradient(135deg, #00d2d3 0%, #54a0ff 100%); border-radius: 12px; color: white; font-weight: bold;">
|
||||
🧪 <strong>视觉资产规划</strong><br/>
|
||||
参考图选择 • 外观/风格引导 • 提示词条件化
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<!-- 左：资产索引 -->
|
||||
<td align="center" style="padding: 12px; background: linear-gradient(135deg, #e056fd 0%, #f368e0 100%); border-radius: 10px; color: white; width: 50%;">
|
||||
🗂️ <strong>资产索引</strong><br/>
|
||||
<small>帧/参考图目录 • 嵌入向量 • 复用检索</small>
|
||||
</td>
|
||||
<td width="10"></td>
|
||||
<!-- 右：一致性与连续性 -->
|
||||
<td align="center" style="padding: 12px; background: linear-gradient(135deg, #ffa726 0%, #ff7043 100%); border-radius: 10px; color: white; width: 50%;">
|
||||
♻️ <strong>一致性与连续性</strong><br/>
|
||||
<small>角色/环境追踪 • 参考匹配 • 时序连贯性</small>
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 15px; background: linear-gradient(135deg, #26de81 0%, #20bf6b 100%); border-radius: 12px; color: white; font-weight: bold;">
|
||||
✂️ <strong>视觉合成与组装</strong><br/>
|
||||
图像生成 • 最佳帧选择 • 首尾帧→视频 • 剪辑与时间线合成
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 20px; background: linear-gradient(135deg, #045de9 0%, #09c6f9 100%); border-radius: 15px; color: white; font-weight: bold;">
|
||||
🚀 <strong>输出层</strong><br/>
|
||||
🖼️ 帧图像 • 🎞️ 片段与最终视频 • 📜 日志 • 📦 工作目录产物
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
## 🚀Quick Start
|
||||
|
||||
### 🖥️ **Environment**
|
||||
|
||||
```
|
||||
OS: Linux, Windows
|
||||
```
|
||||
|
||||
### 📥 **Clone and Install**
|
||||
We use uv to manage the environment. For uv installation, please refer to the https://docs.astral.sh/uv/getting-started/installation/.
|
||||
```bash
|
||||
git clone https://github.com/HKUDS/ViMax.git
|
||||
cd ViMax
|
||||
uv sync
|
||||
```
|
||||
|
||||
|
||||
### 🧠 **Agent TUI / Agents Loop**
|
||||
ViMax also provides a minimal TUI for interactive agent-based video creation. Configure `configs/agent.local.yaml` for the LLM, image, and video providers, then start the TUI from the ViMax root directory.
|
||||
```bash
|
||||
vimax tui
|
||||
```
|
||||
|
||||
Start a new session or resume an existing one:
|
||||
```bash
|
||||
vimax tui new
|
||||
vimax tui resume
|
||||
vimax tui resume <session_id>
|
||||
```
|
||||
|
||||
### 🎯 **Usage**
|
||||
main_idea2video.py is used to convert your ideas into videos.
|
||||
You need to configure the model and API key information in the configs/idea2video.yaml file, including three parts—the chat model, the image generator, and the video generator, as shown below
|
||||
```yaml
|
||||
chat_model:
|
||||
init_args:
|
||||
model: google/gemini-2.5-flash-lite-preview-09-2025
|
||||
model_provider: openai
|
||||
api_key: <YOUR_API_KEY>
|
||||
base_url: https://openrouter.ai/api/v1
|
||||
|
||||
image_generator:
|
||||
class_path: tools.ImageGeneratorNanobananaGoogleAPI
|
||||
init_args:
|
||||
api_key: <YOUR_API_KEY>
|
||||
|
||||
video_generator:
|
||||
class_path: tools.VideoGeneratorVeoGoogleAPI
|
||||
init_args:
|
||||
api_key: <YOUR_API_KEY>
|
||||
|
||||
working_dir: .working_dir/idea2video
|
||||
```
|
||||
|
||||
Then, provide a simple yet thoughtful idea and the corresponding creative requirements in main_idea2video.py.
|
||||
```bash
|
||||
idea = \
|
||||
"""
|
||||
If a cat and a dog are best friends, what would happen when they meet a new cat?
|
||||
"""
|
||||
user_requirement = \
|
||||
"""
|
||||
For children, do not exceed 3 scenes.
|
||||
"""
|
||||
style = "Cartoon"
|
||||
```
|
||||
|
||||
main_script2video.py generates a video based on a specific script.
|
||||
You similarly need to set up the API configuration in configs/script2video.yaml file. Then, provide a scene script and the corresponding creative requirements in main_script2video.py, as shown below.
|
||||
```python
|
||||
script = \
|
||||
"""
|
||||
EXT. SCHOOL GYM - DAY
|
||||
A group of students are practicing basketball in the gym. The gym is large and open, with a basketball hoop at one end and a large crowd of spectators at the other end. John (18, male, tall, athletic) is the star player, and he is practicing his dribble and shot. Jane (17, female, short, athletic) is the assistant coach, and she is helping John with his practice. The other students are watching the practice and cheering for John.
|
||||
John: (dribbling the ball) I'm going to score a basket!
|
||||
Jane: (smiling) Good job, John!
|
||||
John: (shooting the ball) Yes!
|
||||
...
|
||||
"""
|
||||
user_requirement = \
|
||||
"""
|
||||
Fast-paced with no more than 20 shots.
|
||||
"""
|
||||
style = "Animate Style"
|
||||
```
|
||||
|
||||
|
||||
---
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = ["AgentLoop", "SessionIndex", "ToolRegistry", "build_runtime"]
|
||||
|
||||
|
||||
def build_runtime(*args, **kwargs):
|
||||
from .loop import build_runtime as _build_runtime
|
||||
|
||||
return _build_runtime(*args, **kwargs)
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
if name == "AgentLoop":
|
||||
from .loop import AgentLoop
|
||||
|
||||
return AgentLoop
|
||||
if name == "SessionIndex":
|
||||
from .session_index import SessionIndex
|
||||
|
||||
return SessionIndex
|
||||
if name == "ToolRegistry":
|
||||
from .tools import ToolRegistry
|
||||
|
||||
return ToolRegistry
|
||||
raise AttributeError(name)
|
||||
@@ -0,0 +1,134 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
DEFAULT_LLM_MODEL = "gpt-5.5"
|
||||
DEFAULT_LLM_MODEL_PROVIDER = "openai"
|
||||
DEFAULT_LLM_BASE_URL = "https://yunwu.ai/v1"
|
||||
DEFAULT_IMAGE_MODEL = "gemini-3.1-flash-image-preview"
|
||||
DEFAULT_IMAGE_BASE_URL = "https://yunwu.ai"
|
||||
DEFAULT_VIDEO_MODEL = "veo3.1-fast"
|
||||
DEFAULT_VIDEO_BASE_URL = "https://openrouter.ai/api/v1"
|
||||
DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small"
|
||||
DEFAULT_EMBEDDING_MODEL_PROVIDER = "openai"
|
||||
DEFAULT_RERANKER_MODEL = "BAAI/bge-reranker-v2-m3"
|
||||
|
||||
|
||||
@lru_cache(maxsize=4)
|
||||
def load_agent_config(workspace_root: str | Path = ".") -> dict[str, Any]:
|
||||
path = Path(workspace_root).resolve() / "configs" / "agent.local.yaml"
|
||||
if not path.exists():
|
||||
return {}
|
||||
try:
|
||||
payload = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
except yaml.YAMLError as exc:
|
||||
raise RuntimeError(f"Invalid configs/agent.local.yaml: {exc}") from exc
|
||||
if not isinstance(payload, dict):
|
||||
raise RuntimeError("configs/agent.local.yaml must be a YAML mapping")
|
||||
return payload
|
||||
|
||||
|
||||
def config_value(section: str, key: str, env_names: list[str], default: str = "", workspace_root: str | Path = ".") -> str:
|
||||
for env_name in env_names:
|
||||
value = os.environ.get(env_name)
|
||||
if value:
|
||||
return value
|
||||
section_payload = load_agent_config(workspace_root).get(section, {})
|
||||
if isinstance(section_payload, dict):
|
||||
value = section_payload.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return default
|
||||
|
||||
|
||||
def llm_model(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("llm", "model", ["VIMAX_LLM_MODEL"], DEFAULT_LLM_MODEL, workspace_root)
|
||||
|
||||
|
||||
def llm_model_provider(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("llm", "model_provider", ["VIMAX_LLM_MODEL_PROVIDER"], DEFAULT_LLM_MODEL_PROVIDER, workspace_root)
|
||||
|
||||
|
||||
def llm_base_url(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("llm", "base_url", ["VIMAX_LLM_BASE_URL"], DEFAULT_LLM_BASE_URL, workspace_root)
|
||||
|
||||
|
||||
def llm_api_key(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("llm", "api_key", ["VIMAX_LLM_API_KEY", "VIMAX_API_KEY"], "", workspace_root)
|
||||
|
||||
|
||||
def image_model(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("image", "model", ["VIMAX_IMAGE_MODEL"], DEFAULT_IMAGE_MODEL, workspace_root)
|
||||
|
||||
|
||||
def image_base_url(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("image", "base_url", ["VIMAX_IMAGE_BASE_URL"], DEFAULT_IMAGE_BASE_URL, workspace_root)
|
||||
|
||||
|
||||
def image_api_key(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("image", "api_key", ["VIMAX_IMAGE_API_KEY", "VIMAX_LLM_API_KEY", "VIMAX_API_KEY"], llm_api_key(workspace_root), workspace_root)
|
||||
|
||||
|
||||
|
||||
def embedding_model(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("embedding", "model", ["VIMAX_EMBEDDING_MODEL"], DEFAULT_EMBEDDING_MODEL, workspace_root)
|
||||
|
||||
|
||||
def embedding_model_provider(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("embedding", "model_provider", ["VIMAX_EMBEDDING_MODEL_PROVIDER"], DEFAULT_EMBEDDING_MODEL_PROVIDER, workspace_root)
|
||||
|
||||
|
||||
def embedding_base_url(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("embedding", "base_url", ["VIMAX_EMBEDDING_BASE_URL"], "", workspace_root)
|
||||
|
||||
|
||||
def embedding_api_key(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("embedding", "api_key", ["VIMAX_EMBEDDING_API_KEY"], "", workspace_root)
|
||||
|
||||
|
||||
def reranker_model(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("reranker", "model", ["VIMAX_RERANKER_MODEL"], DEFAULT_RERANKER_MODEL, workspace_root)
|
||||
|
||||
|
||||
def reranker_base_url(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("reranker", "base_url", ["VIMAX_RERANKER_BASE_URL"], "", workspace_root)
|
||||
|
||||
|
||||
def reranker_api_key(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("reranker", "api_key", ["VIMAX_RERANKER_API_KEY"], "", workspace_root)
|
||||
|
||||
|
||||
def video_model(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("video", "model", ["VIMAX_VIDEO_MODEL"], DEFAULT_VIDEO_MODEL, workspace_root)
|
||||
|
||||
|
||||
def video_base_url(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("video", "base_url", ["VIMAX_VIDEO_BASE_URL"], DEFAULT_VIDEO_BASE_URL, workspace_root)
|
||||
|
||||
|
||||
def video_api_key(workspace_root: str | Path = ".") -> str:
|
||||
return config_value("video", "api_key", ["VIMAX_VIDEO_API_KEY", "VIMAX_LLM_API_KEY", "VIMAX_API_KEY"], llm_api_key(workspace_root), workspace_root)
|
||||
|
||||
|
||||
def api_provider_from_base_url(base_url: str) -> str:
|
||||
normalized = base_url.strip().lower()
|
||||
if "openrouter.ai" in normalized:
|
||||
return "openrouter"
|
||||
if "yunwu.ai" in normalized:
|
||||
return "yunwu"
|
||||
return ""
|
||||
|
||||
|
||||
def video_provider(workspace_root: str | Path = ".") -> str:
|
||||
"""Infer the video API relay/provider from video.base_url.
|
||||
|
||||
This is not a model provider setting. OpenRouter/Yunwu are transport/API
|
||||
gateways here, so users should configure base_url and let the adapter pick
|
||||
the matching implementation.
|
||||
"""
|
||||
return api_provider_from_base_url(video_base_url(workspace_root))
|
||||
@@ -0,0 +1,254 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
|
||||
SUMMARY_SECTIONS = [
|
||||
"Reference Context Only",
|
||||
"Active Task",
|
||||
"Completed Actions",
|
||||
"Important Files",
|
||||
"Decisions",
|
||||
"Errors & Risks",
|
||||
"Remaining Work",
|
||||
"Critical Context",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CompactionResult:
|
||||
summary: str
|
||||
preserved_messages: list[dict[str, Any]]
|
||||
compacted_message_count: int
|
||||
estimated_tokens_before: int
|
||||
estimated_tokens_after: int
|
||||
reason: str
|
||||
mode: str
|
||||
created_at: str = field(default_factory=lambda: datetime.now().isoformat(timespec="seconds"))
|
||||
|
||||
|
||||
class ContextCompactor:
|
||||
def __init__(
|
||||
self,
|
||||
llm: Any | None = None,
|
||||
*,
|
||||
token_threshold: int | None = None,
|
||||
buffer_tokens: int | None = None,
|
||||
preserve_last_n: int | None = None,
|
||||
max_messages: int | None = None,
|
||||
summary_max_chars: int | None = None,
|
||||
) -> None:
|
||||
self.llm = llm
|
||||
configured_threshold = token_threshold if token_threshold is not None else _default_token_threshold()
|
||||
self.token_threshold = _env_int("VIMAX_AUTO_COMPACT_TOKEN_THRESHOLD", configured_threshold)
|
||||
self.buffer_tokens = _env_int("VIMAX_AUTO_COMPACT_BUFFER_TOKENS", buffer_tokens if buffer_tokens is not None else 20000)
|
||||
self.preserve_last_n = _env_int("VIMAX_COMPACT_PRESERVE_LAST_N", preserve_last_n if preserve_last_n is not None else 6)
|
||||
self.max_messages = _env_int("VIMAX_COMPACT_MAX_MESSAGES", max_messages if max_messages is not None else 48)
|
||||
self.summary_max_chars = _env_int("VIMAX_COMPACT_SUMMARY_MAX_CHARS", summary_max_chars if summary_max_chars is not None else 6000)
|
||||
|
||||
def compact_target_tokens(self) -> int:
|
||||
if self.token_threshold <= 0:
|
||||
return 0
|
||||
return max(0, self.token_threshold - max(0, self.buffer_tokens))
|
||||
|
||||
def estimate_message_tokens(self, message: dict[str, Any]) -> int:
|
||||
role = str(message.get("role", "user") or "user")
|
||||
content = str(message.get("content", "") or "")
|
||||
metadata = {key: value for key, value in message.items() if key not in {"role", "content"}}
|
||||
word_count = len(re.findall(r"\w+", content))
|
||||
line_count = content.count("\n") + 1 if content else 0
|
||||
punctuation_count = len(re.findall(r"[^\w\s]", content))
|
||||
role_overhead = {"system": 18, "user": 12, "assistant": 14, "tool": 16}.get(role, 12)
|
||||
metadata_bonus = min(300, len(json.dumps(metadata, ensure_ascii=False, default=str)) // 6) if metadata else 0
|
||||
tool_bonus = 80 if "tool_calls" in message or role == "tool" else 0
|
||||
return max(role_overhead, role_overhead + len(content) // 4 + word_count // 2 + line_count * 2 + punctuation_count // 4 + metadata_bonus + tool_bonus)
|
||||
|
||||
def estimate_messages_tokens(self, messages: list[dict[str, Any]]) -> int:
|
||||
return sum(self.estimate_message_tokens(message) for message in messages)
|
||||
|
||||
def should_preflight_compact(self, messages: list[dict[str, Any]], *, system_tokens: int = 0, tools_tokens: int = 0) -> bool:
|
||||
target = self.compact_target_tokens()
|
||||
if target <= 0 or not messages:
|
||||
return False
|
||||
total = self.estimate_messages_tokens(messages) + max(0, system_tokens) + max(0, tools_tokens)
|
||||
return total >= target
|
||||
|
||||
async def compact(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
*,
|
||||
previous_summary: str = "",
|
||||
preserve_last_n: int | None = None,
|
||||
reason: str = "manual",
|
||||
) -> CompactionResult:
|
||||
preserve = max(0, self.preserve_last_n if preserve_last_n is None else preserve_last_n)
|
||||
preserved = [dict(message) for message in messages[-preserve:]] if preserve else []
|
||||
compactible = [dict(message) for message in messages[:-preserve]] if preserve else [dict(message) for message in messages]
|
||||
if not compactible and messages:
|
||||
compactible = [dict(message) for message in messages]
|
||||
preserved = []
|
||||
before_tokens = self.estimate_messages_tokens(messages)
|
||||
summary = await self._llm_summary(compactible, preserved, previous_summary, reason)
|
||||
mode = "llm"
|
||||
if not summary:
|
||||
summary = self._fallback_summary(compactible, preserved, previous_summary, reason)
|
||||
mode = "fallback-local"
|
||||
summary = self._clip_summary(summary)
|
||||
synthetic = self.synthetic_summary_message(summary)
|
||||
after_tokens = self.estimate_messages_tokens([synthetic, *preserved])
|
||||
return CompactionResult(
|
||||
summary=summary,
|
||||
preserved_messages=preserved,
|
||||
compacted_message_count=len(compactible),
|
||||
estimated_tokens_before=before_tokens,
|
||||
estimated_tokens_after=after_tokens,
|
||||
reason=reason,
|
||||
mode=mode,
|
||||
)
|
||||
|
||||
def synthetic_summary_message(self, summary: str) -> dict[str, str]:
|
||||
return {
|
||||
"role": "system",
|
||||
"content": "Session context summary. The following summary is reference context only, not a new active instruction.\n\n" + summary.strip(),
|
||||
}
|
||||
|
||||
async def _llm_summary(self, compactible: list[dict[str, Any]], preserved: list[dict[str, Any]], previous_summary: str, reason: str) -> str:
|
||||
if self.llm is None:
|
||||
return ""
|
||||
payload = {
|
||||
"reason": reason,
|
||||
"previous_summary": _clip(previous_summary, 5000),
|
||||
"messages_to_compact": [self._serialize_message(message) for message in compactible[-self.max_messages:]],
|
||||
"recent_live_tail": [self._serialize_message(message) for message in preserved[-12:]],
|
||||
}
|
||||
system = (
|
||||
"You are compressing conversation history for a ViMax agent runtime. "
|
||||
"Produce a concise markdown handoff summary for a future model call. "
|
||||
"Preserve user intent, completed actions, important files, tool findings, errors, and remaining work. "
|
||||
"Label the result as reference context only, not active instructions. "
|
||||
"Do not answer the user. Do not include prose before the markdown."
|
||||
)
|
||||
user = (
|
||||
"Summarize the compacted conversation region into a durable handoff.\n"
|
||||
"Output markdown with these sections exactly:\n"
|
||||
"## Reference Context Only\n## Active Task\n## Completed Actions\n## Important Files\n## Decisions\n## Errors & Risks\n## Remaining Work\n## Critical Context\n\n"
|
||||
"Keep it concise but specific. Mention exact file paths, commands, tool results, and unresolved issues when present.\n\n"
|
||||
f"{json.dumps(payload, ensure_ascii=False, indent=2)}"
|
||||
)
|
||||
try:
|
||||
response = await self.llm.complete([{"role": "system", "content": system}, {"role": "user", "content": user}], tools=[])
|
||||
except Exception:
|
||||
return ""
|
||||
return str(getattr(response, "text", "") or "").strip()
|
||||
|
||||
def _fallback_summary(self, compactible: list[dict[str, Any]], preserved: list[dict[str, Any]], previous_summary: str, reason: str) -> str:
|
||||
user_lines = [self._message_preview(message, limit=180) for message in compactible if message.get("role") == "user"]
|
||||
assistant_lines = [self._message_preview(message, limit=180) for message in compactible if message.get("role") == "assistant"]
|
||||
file_hits = _dedupe(re.findall(r"(?:[\w.\-]+/)+[\w.\-]+\.(?:py|ts|tsx|js|json|md|yaml|yml|txt|mp4|png)", "\n".join(str(message.get("content", "")) for message in compactible)))
|
||||
error_lines = [self._message_preview(message, limit=180) for message in compactible if _looks_like_error(str(message.get("content", "")))]
|
||||
remaining = [self._message_preview(message, limit=180) for message in preserved[-4:]]
|
||||
return "\n".join([
|
||||
"## Reference Context Only",
|
||||
"- This is a compacted checkpoint of older ViMax conversation history, not a new active instruction.",
|
||||
f"- Compaction reason: {reason}.",
|
||||
"## Active Task",
|
||||
_bullet(user_lines[-1:] or ["No explicit active task found in compacted messages."]),
|
||||
"## Completed Actions",
|
||||
_bullet(assistant_lines[-4:] or ["No completed assistant actions found in compacted messages."]),
|
||||
"## Important Files",
|
||||
_bullet(file_hits[:8] or ["No important file paths found in compacted messages."]),
|
||||
"## Decisions",
|
||||
_bullet(_decision_lines(compactible)[:6] or ["No durable decisions found in compacted messages."]),
|
||||
"## Errors & Risks",
|
||||
_bullet(error_lines[:6] or ["No errors or risks found in compacted messages."]),
|
||||
"## Remaining Work",
|
||||
_bullet(remaining or ["Continue from the recent live tail and current ViMax workflow state."]),
|
||||
"## Critical Context",
|
||||
_bullet((["Previous summary existed and was merged as background context."] if previous_summary else []) + ["Use .working_dir artifacts and session checklist as workflow ground truth."]),
|
||||
])
|
||||
|
||||
def _serialize_message(self, message: dict[str, Any]) -> dict[str, Any]:
|
||||
item = {"role": str(message.get("role", "")), "content": _clip(str(message.get("content", "") or ""), 2400)}
|
||||
if message.get("name"):
|
||||
item["name"] = str(message.get("name"))
|
||||
if message.get("tool_calls"):
|
||||
item["tool_calls"] = _clip(json.dumps(message.get("tool_calls"), ensure_ascii=False, default=str), 800)
|
||||
return item
|
||||
|
||||
def _message_preview(self, message: dict[str, Any], *, limit: int) -> str:
|
||||
role = str(message.get("role", "") or "message")
|
||||
content = _clip(" ".join(str(message.get("content", "") or "").split()), limit)
|
||||
if message.get("tool_calls"):
|
||||
return f"{role}: [tool calls] {_clip(json.dumps(message.get('tool_calls'), ensure_ascii=False, default=str), limit)}"
|
||||
return f"{role}: {content}" if content else f"{role}: <empty>"
|
||||
|
||||
def _clip_summary(self, summary: str) -> str:
|
||||
text = summary.strip()
|
||||
if not text:
|
||||
text = self._fallback_summary([], [], "", "empty-summary")
|
||||
if len(text) > self.summary_max_chars:
|
||||
text = text[: max(0, self.summary_max_chars - 3)].rstrip() + "..."
|
||||
return text
|
||||
|
||||
|
||||
def _default_token_threshold() -> int:
|
||||
context_window = _env_int("VIMAX_CONTEXT_WINDOW_TOKENS", 200000)
|
||||
ratio = _env_float("VIMAX_AUTO_COMPACT_RATIO", 0.90)
|
||||
ratio = min(1.0, max(0.0, ratio))
|
||||
return int(context_window * ratio)
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
try:
|
||||
return int(os.environ.get(name, str(default)))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _env_float(name: str, default: float) -> float:
|
||||
try:
|
||||
return float(os.environ.get(name, str(default)))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _clip(text: str, limit: int) -> str:
|
||||
compact = " ".join(str(text or "").split())
|
||||
if len(compact) <= limit:
|
||||
return compact
|
||||
return compact[: max(0, limit - 3)].rstrip() + "..."
|
||||
|
||||
|
||||
def _bullet(items: list[str]) -> str:
|
||||
return "\n".join(f"- {item}" for item in items if str(item).strip())
|
||||
|
||||
|
||||
def _dedupe(items: list[str]) -> list[str]:
|
||||
seen: list[str] = []
|
||||
for item in items:
|
||||
normalized = " ".join(str(item).split())
|
||||
if normalized and normalized not in seen:
|
||||
seen.append(normalized)
|
||||
return seen
|
||||
|
||||
|
||||
def _looks_like_error(text: str) -> bool:
|
||||
lowered = text.lower()
|
||||
return any(token in lowered for token in ("error", "failed", "failure", "timeout", "not found", "blocked", "permission"))
|
||||
|
||||
|
||||
def _decision_lines(messages: list[dict[str, Any]]) -> list[str]:
|
||||
tokens = ("decision", "decided", "prefer", "keep ", "switch ", "use ", "preserve ", "avoid ")
|
||||
rows: list[str] = []
|
||||
for message in messages:
|
||||
content = str(message.get("content", "") or "")
|
||||
for raw in content.splitlines():
|
||||
line = raw.strip(" -")
|
||||
if line and any(token in line.lower() for token in tokens):
|
||||
rows.append(_clip(line, 180))
|
||||
return _dedupe(rows)
|
||||
@@ -0,0 +1,136 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from openai import APIConnectionError, APITimeoutError, AsyncOpenAI
|
||||
|
||||
from .config import llm_api_key, llm_base_url, llm_model
|
||||
from .models import ToolCall
|
||||
|
||||
|
||||
LLM_MAX_ATTEMPTS = 3
|
||||
LLM_RETRY_BACKOFF_SECONDS = (1.0, 4.0)
|
||||
LLM_REQUEST_TIMEOUT_SECONDS = 300.0
|
||||
|
||||
|
||||
def _is_retryable_llm_error(exc: BaseException) -> bool:
|
||||
status = getattr(exc, "status_code", None)
|
||||
if status is not None:
|
||||
try:
|
||||
status = int(status)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
return status == 429 or status >= 500
|
||||
return isinstance(exc, (APIConnectionError, APITimeoutError))
|
||||
|
||||
|
||||
class LLMResponseShapeError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AssistantMessage:
|
||||
text: str = ""
|
||||
tool_calls: list[ToolCall] = field(default_factory=list)
|
||||
raw_message: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class OpenAICompatibleLLM:
|
||||
def __init__(self, model: str | None = None, base_url: str | None = None, api_key: str | None = None) -> None:
|
||||
self.model = model or llm_model()
|
||||
self.base_url = base_url or llm_base_url()
|
||||
self.api_key = api_key or llm_api_key()
|
||||
if not self.api_key:
|
||||
raise RuntimeError("VIMAX_LLM_API_KEY is required for the agent LLM client")
|
||||
self.client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url, timeout=LLM_REQUEST_TIMEOUT_SECONDS)
|
||||
|
||||
async def complete(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> AssistantMessage:
|
||||
shape_attempts = [
|
||||
{"tools": tools or None, "tool_choice": "auto" if tools else None},
|
||||
{"tools": tools or None, "tool_choice": "auto" if tools else None},
|
||||
]
|
||||
if tools:
|
||||
shape_attempts.append({"tools": None, "tool_choice": None})
|
||||
|
||||
last_shape_error: Exception | None = None
|
||||
for attempt in shape_attempts:
|
||||
try:
|
||||
response = await self._create_completion_with_retries(messages, attempt["tools"], attempt["tool_choice"])
|
||||
return _assistant_message_from_response(response)
|
||||
except LLMResponseShapeError as exc:
|
||||
last_shape_error = exc
|
||||
continue
|
||||
assert last_shape_error is not None
|
||||
raise last_shape_error
|
||||
|
||||
async def _create_completion_with_retries(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, tool_choice: str | None) -> Any:
|
||||
for attempt in range(LLM_MAX_ATTEMPTS):
|
||||
try:
|
||||
return await self._create_completion(messages, tools, tool_choice)
|
||||
except Exception as exc:
|
||||
if isinstance(exc, LLMResponseShapeError) or attempt == LLM_MAX_ATTEMPTS - 1 or not _is_retryable_llm_error(exc):
|
||||
raise
|
||||
delay = LLM_RETRY_BACKOFF_SECONDS[min(attempt, len(LLM_RETRY_BACKOFF_SECONDS) - 1)]
|
||||
logging.warning("LLM call failed (%s); retrying in %.1fs (attempt %d/%d)", exc, delay, attempt + 1, LLM_MAX_ATTEMPTS)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
async def _create_completion(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, tool_choice: str | None) -> Any:
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
}
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
if tool_choice:
|
||||
kwargs["tool_choice"] = tool_choice
|
||||
return await self.client.chat.completions.create(**kwargs)
|
||||
|
||||
|
||||
def _assistant_message_from_response(response: Any) -> AssistantMessage:
|
||||
message = _extract_message(response)
|
||||
text = _message_value(message, "content") or ""
|
||||
calls: list[ToolCall] = []
|
||||
for call in _message_value(message, "tool_calls") or []:
|
||||
function = _message_value(call, "function") or {}
|
||||
try:
|
||||
arguments = json.loads(_message_value(function, "arguments") or "{}")
|
||||
except json.JSONDecodeError:
|
||||
arguments = {}
|
||||
calls.append(ToolCall(id=_message_value(call, "id") or f"tool-{uuid4().hex[:12]}", name=_message_value(function, "name"), arguments=arguments))
|
||||
return AssistantMessage(text=text, tool_calls=calls, raw_message=_dump_message(message))
|
||||
|
||||
|
||||
def _extract_message(response: Any) -> Any:
|
||||
if isinstance(response, str):
|
||||
try:
|
||||
response = json.loads(response)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise LLMResponseShapeError(f"LLM provider returned a string instead of a chat completion object: {response[:300]}") from exc
|
||||
choices = _message_value(response, "choices")
|
||||
if not choices:
|
||||
raise LLMResponseShapeError(f"LLM provider response missing choices: {str(response)[:500]}")
|
||||
first_choice = choices[0]
|
||||
message = _message_value(first_choice, "message")
|
||||
if message is None:
|
||||
raise LLMResponseShapeError(f"LLM provider response missing choice.message: {str(response)[:500]}")
|
||||
return message
|
||||
|
||||
|
||||
def _message_value(obj: Any, key: str) -> Any:
|
||||
if isinstance(obj, dict):
|
||||
return obj.get(key)
|
||||
return getattr(obj, key, None)
|
||||
|
||||
|
||||
def _dump_message(message: Any) -> dict[str, Any]:
|
||||
if isinstance(message, dict):
|
||||
return message
|
||||
if hasattr(message, "model_dump"):
|
||||
return message.model_dump()
|
||||
return {"content": str(message)}
|
||||
@@ -0,0 +1,171 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
from .context_compactor import ContextCompactor, CompactionResult
|
||||
from .llm import OpenAICompatibleLLM
|
||||
from .models import ToolCall, ToolResult, TurnControl
|
||||
from .prompts import PromptBuilder
|
||||
from .session_index import SessionIndex
|
||||
from .tool_executor import ToolExecutor
|
||||
from .tools import ToolRegistry, build_builtin_registry
|
||||
|
||||
MAX_TOOL_PASSES = 50
|
||||
|
||||
|
||||
class AgentLoop:
|
||||
def __init__(self, session_index: SessionIndex, prompt_builder: PromptBuilder, tool_registry: ToolRegistry, tool_executor: ToolExecutor, llm: Any, context_compactor: ContextCompactor | None = None) -> None:
|
||||
self.session_index = session_index
|
||||
self.prompt_builder = prompt_builder
|
||||
self.tool_registry = tool_registry
|
||||
self.tool_executor = tool_executor
|
||||
self.llm = llm
|
||||
self.context_compactor = context_compactor or ContextCompactor(llm)
|
||||
self.history: list[dict[str, Any]] = []
|
||||
|
||||
async def compact_history(self, *, reason: str = "manual") -> str:
|
||||
if not self.history:
|
||||
return "No conversation history to compact."
|
||||
session = self.session_index.active() or self.session_index.create()
|
||||
result = await self.context_compactor.compact(
|
||||
self.history,
|
||||
previous_summary=str(session.get("compacted_summary", "") or ""),
|
||||
reason=reason,
|
||||
)
|
||||
self.history = [self.context_compactor.synthetic_summary_message(result.summary), *result.preserved_messages]
|
||||
self.session_index.update_compaction(session["session_id"], _compaction_record(result))
|
||||
return f"Compacted context {result.estimated_tokens_before} -> {result.estimated_tokens_after} ({result.mode})."
|
||||
|
||||
async def stream_events(self, user_input: str) -> AsyncIterator[dict[str, Any]]:
|
||||
control = TurnControl()
|
||||
yield {"type": "turn", "turn_id": control.turn_id, "turn": {"id": control.turn_id}}
|
||||
tool_schemas = self.tool_registry.list_function_tools()
|
||||
parts = self.prompt_builder.build_parts(user_input)
|
||||
system = "\n\n".join(f"## {part.title}\n{part.body}" for part in parts if part.id != "request.user")
|
||||
if self.context_compactor.should_preflight_compact(
|
||||
[*self.history, {"role": "user", "content": user_input}],
|
||||
system_tokens=_prompt_tokens(parts),
|
||||
tools_tokens=_tool_schema_tokens(tool_schemas),
|
||||
):
|
||||
yield {"type": "status", "turn_id": control.turn_id, "phase": "compact", "message": "Compacting context before sampling"}
|
||||
await self.compact_history(reason="token-pressure")
|
||||
parts = self.prompt_builder.build_parts(user_input)
|
||||
system = "\n\n".join(f"## {part.title}\n{part.body}" for part in parts if part.id != "request.user")
|
||||
yield {"type": "prompt_trace", "turn_id": control.turn_id, "prompt_trace": self.prompt_builder.trace(parts)}
|
||||
runtime_messages: list[dict[str, Any]] = [{"role": "system", "content": system}, *self.history, {"role": "user", "content": user_input}]
|
||||
assistant_turns: list[dict[str, Any]] = []
|
||||
tool_rounds: list[dict[str, Any]] = []
|
||||
transitions: list[dict[str, str]] = []
|
||||
all_tool_results: list[ToolResult] = []
|
||||
final_text = ""
|
||||
status = "completed"
|
||||
tool_round = 0
|
||||
|
||||
while True:
|
||||
yield {"type": "status", "turn_id": control.turn_id, "phase": "sampling_assistant", "message": "Sampling assistant"}
|
||||
try:
|
||||
assistant = await self.llm.complete(runtime_messages, tools=tool_schemas)
|
||||
except Exception as exc:
|
||||
status = "failed"
|
||||
final_text = f"Agent LLM request failed: {exc}"
|
||||
transitions.append(_transition("sampling_assistant", "finalizing_answer", "llm_sampling_failed"))
|
||||
yield {"type": "error", "turn_id": control.turn_id, "message": final_text, "metadata": {"error_type": "llm_sampling_failed"}}
|
||||
break
|
||||
assistant_turns.append({"phase": "initial" if tool_round == 0 else f"followup_{tool_round}", "text": assistant.text, "tool_calls": [call.as_dict() for call in assistant.tool_calls]})
|
||||
if not assistant.tool_calls:
|
||||
transitions.append(_transition("sampling_assistant", "finalizing_answer", "assistant_finished_without_tools"))
|
||||
final_text = assistant.text
|
||||
if final_text:
|
||||
yield {"type": "token", "turn_id": control.turn_id, "delta": final_text}
|
||||
break
|
||||
transitions.append(_transition("sampling_assistant", "executing_tools", "assistant_requested_tools"))
|
||||
if tool_round >= MAX_TOOL_PASSES:
|
||||
status = "halted"
|
||||
final_text = "Tool loop halted after max tool passes."
|
||||
transitions.append(_transition("executing_tools", "finalizing_answer", "max_tool_passes_reached"))
|
||||
yield {"type": "error", "turn_id": control.turn_id, "message": final_text, "metadata": {"max_tool_passes": MAX_TOOL_PASSES}}
|
||||
break
|
||||
tool_round += 1
|
||||
yield {"type": "status", "turn_id": control.turn_id, "phase": "executing_tools", "message": f"Running tools (round {tool_round})"}
|
||||
runtime_messages.append({"role": "assistant", "content": assistant.text or "", "tool_calls": [_openai_tool_call(call) for call in assistant.tool_calls]})
|
||||
round_results: list[ToolResult] = []
|
||||
|
||||
for call in assistant.tool_calls:
|
||||
yield {"type": "tool_start", "turn_id": control.turn_id, "tool": call.as_dict()}
|
||||
progress_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
|
||||
|
||||
def on_progress(event: dict[str, Any]) -> None:
|
||||
progress_queue.put_nowait(event)
|
||||
|
||||
task = asyncio.create_task(self.tool_executor.execute(call, control, progress_callback=on_progress))
|
||||
while not task.done():
|
||||
try:
|
||||
yield await asyncio.wait_for(progress_queue.get(), timeout=0.1)
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
while not progress_queue.empty():
|
||||
yield progress_queue.get_nowait()
|
||||
record = await task
|
||||
result = record.result
|
||||
round_results.append(result)
|
||||
all_tool_results.append(result)
|
||||
yield {"type": "tool_result", "turn_id": control.turn_id, "tool_result": result.as_dict()}
|
||||
runtime_messages.append({"role": "tool", "tool_call_id": call.id, "name": result.name, "content": json.dumps(result.as_dict(), ensure_ascii=False)})
|
||||
tool_rounds.append({"tool_round": tool_round, "requested_tools": [call.as_dict() for call in assistant.tool_calls], "tool_results": [result.as_dict() for result in round_results]})
|
||||
transitions.append(_transition("executing_tools", "post_tool_decision", "tool_round_completed"))
|
||||
transitions.append(_transition("post_tool_decision", "sampling_assistant", "runtime_continuation_after_tools"))
|
||||
|
||||
self.history.extend([{"role": "user", "content": user_input}, {"role": "assistant", "content": final_text}])
|
||||
turn_record = {"turn_id": control.turn_id, "status": status, "raw_user_input": user_input, "assistant_turns": assistant_turns, "tool_rounds": tool_rounds, "transitions": transitions, "final_assistant_text": final_text, "created_at": datetime.now().isoformat(timespec="seconds")}
|
||||
final_session = self.session_index.active() or self.session_index.create()
|
||||
self.session_index.append_turn_record(final_session["session_id"], turn_record)
|
||||
yield {"type": "done", "turn_id": control.turn_id, "assistant": final_text, "tool_results": [result.as_dict() for result in all_tool_results]}
|
||||
yield {"type": "session", "turn_id": control.turn_id, "session": self.session_index.snapshot()}
|
||||
|
||||
|
||||
def _compaction_record(result: CompactionResult) -> dict[str, Any]:
|
||||
return {
|
||||
"summary": result.summary,
|
||||
"preserved_message_count": len(result.preserved_messages),
|
||||
"compacted_message_count": result.compacted_message_count,
|
||||
"estimated_tokens_before": result.estimated_tokens_before,
|
||||
"estimated_tokens_after": result.estimated_tokens_after,
|
||||
"reason": result.reason,
|
||||
"mode": result.mode,
|
||||
"created_at": result.created_at,
|
||||
}
|
||||
|
||||
|
||||
def _prompt_tokens(parts: list[Any]) -> int:
|
||||
return sum(max(1, len(str(getattr(part, "body", ""))) // 4) for part in parts)
|
||||
|
||||
|
||||
def _tool_schema_tokens(tool_schemas: list[dict[str, Any]]) -> int:
|
||||
try:
|
||||
return max(0, len(json.dumps(tool_schemas, ensure_ascii=False, default=str)) // 4)
|
||||
except TypeError:
|
||||
return max(0, len(str(tool_schemas)) // 4)
|
||||
|
||||
|
||||
def _transition(src: str, dst: str, reason: str) -> dict[str, str]:
|
||||
return {"from": src, "to": dst, "reason": reason}
|
||||
|
||||
|
||||
def _openai_tool_call(call: ToolCall) -> dict[str, Any]:
|
||||
return {"id": call.id, "type": "function", "function": {"name": call.name, "arguments": json.dumps(call.arguments, ensure_ascii=False)}}
|
||||
|
||||
|
||||
def build_runtime(workspace_root: str | Path = ".", llm: Any | None = None, adapter_specs: list[Any] | None = None) -> AgentLoop:
|
||||
from .vimax_adapters import build_vimax_adapter_specs
|
||||
root = Path(workspace_root).resolve()
|
||||
session_index = SessionIndex(root)
|
||||
specs = adapter_specs if adapter_specs is not None else build_vimax_adapter_specs(root, session_index)
|
||||
registry = build_builtin_registry(root, session_index, specs)
|
||||
executor = ToolExecutor(registry, session_index)
|
||||
prompt_builder = PromptBuilder(root / "prompts", session_index, registry)
|
||||
resolved_llm = llm or OpenAICompatibleLLM()
|
||||
return AgentLoop(session_index, prompt_builder, registry, executor, resolved_llm, ContextCompactor(resolved_llm))
|
||||
@@ -0,0 +1,60 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from threading import Event
|
||||
from time import time
|
||||
from typing import Any, Literal
|
||||
from uuid import uuid4
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ToolCall:
|
||||
name: str
|
||||
arguments: dict[str, Any] = field(default_factory=dict)
|
||||
id: str = field(default_factory=lambda: f"tool-{uuid4().hex[:12]}")
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
return {"id": self.id, "name": self.name, "arguments": self.arguments}
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ToolResult:
|
||||
name: str
|
||||
ok: bool
|
||||
content: str
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
return {"name": self.name, "ok": self.ok, "content": self.content, "metadata": dict(self.metadata)}
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class TurnControl:
|
||||
turn_id: str = field(default_factory=lambda: f"turn-{uuid4().hex[:12]}")
|
||||
cancel_event: Event = field(default_factory=Event)
|
||||
cancel_reason: str = ""
|
||||
|
||||
def cancel(self, reason: str = "") -> None:
|
||||
self.cancel_reason = reason.strip()
|
||||
self.cancel_event.set()
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SessionRecord:
|
||||
session_id: str
|
||||
working_dir: str
|
||||
idea: str = ""
|
||||
user_requirement: str = ""
|
||||
style: str = ""
|
||||
stage: str = "created"
|
||||
summary: str = ""
|
||||
stale: dict[str, bool] = field(default_factory=dict)
|
||||
created_at: str = ""
|
||||
updated_at: str = ""
|
||||
|
||||
|
||||
StreamEventType = Literal["turn", "status", "token", "tool_start", "tool_progress", "tool_result", "terminal", "done", "session", "error", "prompt_trace"]
|
||||
|
||||
|
||||
def now_ts() -> float:
|
||||
return time()
|
||||
@@ -0,0 +1,105 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PromptPart:
|
||||
id: str
|
||||
title: str
|
||||
body: str
|
||||
zone: str
|
||||
category: str
|
||||
cacheable: bool = False
|
||||
|
||||
|
||||
class PromptBuilder:
|
||||
def __init__(self, prompt_dir: str | Path, session_index: Any, tool_registry: Any) -> None:
|
||||
self.prompt_dir = Path(prompt_dir)
|
||||
self.session_index = session_index
|
||||
self.tool_registry = tool_registry
|
||||
|
||||
def build_parts(self, user_input: str) -> list[PromptPart]:
|
||||
return [
|
||||
PromptPart("agent.core", "Agent", self._read_prompt("agent.md"), "stable", "agent", True),
|
||||
PromptPart("workflow.core", "Workflow", self._read_prompt("workflow.md"), "stable", "workflow", True),
|
||||
PromptPart("tool.manifest", "Tools", self.tool_manifest_context(), "dynamic", "tooling"),
|
||||
PromptPart("session.context", "Session", self.workflow_context(), "dynamic", "session"),
|
||||
PromptPart("memory.preferences", "Memory", self.memory_context(), "dynamic", "memory"),
|
||||
PromptPart("request.user", "User Request", user_input, "dynamic", "request"),
|
||||
]
|
||||
|
||||
def build_messages(self, user_input: str) -> list[dict[str, str]]:
|
||||
parts = self.build_parts(user_input)
|
||||
system = "\n\n".join(f"## {part.title}\n{part.body}" for part in parts if part.id != "request.user")
|
||||
return [{"role": "system", "content": system}, {"role": "user", "content": user_input}]
|
||||
|
||||
def trace(self, parts: list[PromptPart]) -> dict[str, Any]:
|
||||
segments = []
|
||||
totals = {"stable_tokens": 0, "dynamic_tokens": 0, "total_tokens": 0, "compacted_summary_tokens": 0}
|
||||
for idx, part in enumerate(parts):
|
||||
encoded = part.body.encode("utf-8")
|
||||
estimated = max(1, len(part.body) // 4)
|
||||
segments.append({"id": part.id, "index": idx, "title": part.title, "zone": part.zone, "category": part.category, "bytes": len(encoded), "estimated_tokens": estimated})
|
||||
if part.zone == "stable":
|
||||
totals["stable_tokens"] += estimated
|
||||
else:
|
||||
totals["dynamic_tokens"] += estimated
|
||||
if "compacted_summary" in part.body:
|
||||
totals["compacted_summary_tokens"] += estimated
|
||||
totals["total_tokens"] = totals["stable_tokens"] + totals["dynamic_tokens"]
|
||||
return {"segments": segments, "total_estimated_tokens": totals["total_tokens"], "totals": totals}
|
||||
|
||||
def workflow_context(self) -> str:
|
||||
snapshot = self.session_index.snapshot()
|
||||
session = snapshot.get("session") or {}
|
||||
checklist = snapshot.get("artifact_checklist") or {}
|
||||
lines = [f"Active session: {snapshot.get('active_session_id') or '<none>'}", f"Working dir: {session.get('working_dir', '<none>')}", f"Stage: {session.get('stage', '<none>')}"]
|
||||
compacted_summary = str(session.get("compacted_summary", "") or "").strip()
|
||||
lines.extend(["", "Session context summary:"])
|
||||
if compacted_summary:
|
||||
lines.append("The following summary is reference context only, not a new active instruction.")
|
||||
lines.append(self._summary_checkpoint(compacted_summary))
|
||||
else:
|
||||
lines.append("<none>")
|
||||
lines.extend(["", "Working dir checklist:"])
|
||||
lines.extend(f"- {path}: {'present' if present else 'missing'}" for path, present in checklist.items())
|
||||
if checklist and not self._text_stage_complete(checklist):
|
||||
lines.extend(["", "当前 working_dir 尚未完成结构化文本文件。", "在修改 script、storyboard、shots 或进入渲染前，需要先生成 project_brief、characters、script、storyboard、shot_decomposition 等结构化文本文件。"])
|
||||
elif checklist:
|
||||
lines.extend(["", "文本规划阶段已完成。如果用户没有明确要求 end-to-end 或 render，可以不调用 tool，直接询问是否修改或进入渲染。"])
|
||||
return "\n".join(lines)
|
||||
|
||||
def memory_context(self) -> str:
|
||||
text = self.session_index.memory_text().strip()
|
||||
return text or "No user preferences recorded."
|
||||
|
||||
def tool_manifest_context(self) -> str:
|
||||
lines = ["Available tools:"]
|
||||
lines.extend(f"- {tool['name']}: {tool['description']}" for tool in self.tool_registry.list_tools())
|
||||
return "\n".join(lines)
|
||||
|
||||
def _summary_checkpoint(self, summary: str) -> str:
|
||||
lines = [line.strip() for line in summary.splitlines() if line.strip() and not line.strip().startswith("```")]
|
||||
if not lines:
|
||||
return "<none>"
|
||||
preview = []
|
||||
for line in lines[:8]:
|
||||
if len(line) > 240:
|
||||
line = line[:237].rstrip() + "..."
|
||||
preview.append(line if line.startswith("-") or line.startswith("#") else f"- {line}")
|
||||
if len(lines) > 8:
|
||||
preview.append(f"- <trimmed +{len(lines) - 8} lines>")
|
||||
return "\n".join(preview)
|
||||
|
||||
def _read_prompt(self, name: str) -> str:
|
||||
path = self.prompt_dir / name
|
||||
return path.read_text(encoding="utf-8") if path.exists() else ""
|
||||
|
||||
def _text_stage_complete(self, checklist: dict[str, bool]) -> bool:
|
||||
idea_mode_complete = bool(checklist.get("idea2video/story.txt") and checklist.get("idea2video/characters.json") and checklist.get("idea2video/script.json") and checklist.get("idea2video/scene_*/storyboard.json") and checklist.get("idea2video/scene_*/shots/*/shot_description.json") and checklist.get("idea2video/scene_*/camera_tree.json"))
|
||||
script_mode_complete = bool(checklist.get("script2video/script.txt") and checklist.get("script2video/characters.json") and checklist.get("script2video/storyboard.json") and checklist.get("script2video/shots/*/shot_description.json") and checklist.get("script2video/camera_tree.json"))
|
||||
novel_mode_complete = bool(checklist.get("novel2video/novel/novel_compressed.txt") and checklist.get("novel2video/events/event_*.json") and checklist.get("novel2video/relevant_chunks/event_*") and checklist.get("novel2video/scenes/event_*/scene_*.json") and checklist.get("novel2video/global_information/characters/novel_level/*.json"))
|
||||
return idea_mode_complete or script_mode_complete or novel_mode_complete
|
||||
@@ -0,0 +1,336 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime
|
||||
from functools import wraps
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover - non-POSIX platforms
|
||||
fcntl = None
|
||||
|
||||
|
||||
STALE_KEYS = ["story", "characters", "script", "storyboard", "shot_descriptions", "camera_tree", "frames", "clips", "final_video"]
|
||||
|
||||
|
||||
def _synchronized(method):
|
||||
"""Hold the index file lock across a read-modify-write cycle.
|
||||
|
||||
Every mutator loads the whole sessions file, edits it, and saves it back;
|
||||
without a lock, two concurrent writers (threads or processes) silently
|
||||
drop each other's updates.
|
||||
"""
|
||||
|
||||
@wraps(method)
|
||||
def wrapper(self, *args, **kwargs):
|
||||
with self._locked():
|
||||
return method(self, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class SessionIndex:
|
||||
def __init__(self, workspace_root: str | Path) -> None:
|
||||
self.workspace_root = Path(workspace_root).resolve()
|
||||
self.vimax_dir = self.workspace_root / ".vimax"
|
||||
self.sessions_path = self.vimax_dir / "sessions.json"
|
||||
self.memory_path = self.vimax_dir / "memory.md"
|
||||
self.logs_dir = self.vimax_dir / "logs"
|
||||
self.working_root = self.workspace_root / ".working_dir"
|
||||
self.vimax_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.logs_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.working_root.mkdir(parents=True, exist_ok=True)
|
||||
if not self.memory_path.exists():
|
||||
self.memory_path.write_text("# User Preferences\n", encoding="utf-8")
|
||||
if not self.sessions_path.exists():
|
||||
self.save({"active_session_id": "", "sessions": {}})
|
||||
|
||||
@contextmanager
|
||||
def _locked(self):
|
||||
if fcntl is None:
|
||||
yield
|
||||
return
|
||||
lock_path = self.vimax_dir / "sessions.lock"
|
||||
with open(lock_path, "a+", encoding="utf-8") as handle:
|
||||
fcntl.flock(handle, fcntl.LOCK_EX)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(handle, fcntl.LOCK_UN)
|
||||
|
||||
def load(self) -> dict[str, Any]:
|
||||
try:
|
||||
return json.loads(self.sessions_path.read_text(encoding="utf-8"))
|
||||
except FileNotFoundError:
|
||||
return {"active_session_id": "", "sessions": {}}
|
||||
except json.JSONDecodeError:
|
||||
# A corrupt file usually means a crash mid-write. Returning empty
|
||||
# state is fine for this call, but the next save() would overwrite
|
||||
# the file and destroy every session — keep the evidence first.
|
||||
backup = self.sessions_path.with_name(f"sessions.json.corrupt-{datetime.now().strftime('%Y%m%d-%H%M%S-%f')}")
|
||||
try:
|
||||
os.replace(self.sessions_path, backup)
|
||||
logging.error("sessions.json was corrupt; preserved it at %s and starting with empty state", backup)
|
||||
except OSError:
|
||||
logging.error("sessions.json is corrupt and could not be backed up; starting with empty state")
|
||||
return {"active_session_id": "", "sessions": {}}
|
||||
|
||||
def save(self, data: dict[str, Any]) -> None:
|
||||
tmp_path = self.sessions_path.with_name("sessions.json.tmp")
|
||||
tmp_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
os.replace(tmp_path, self.sessions_path)
|
||||
|
||||
def active(self) -> dict[str, Any] | None:
|
||||
data = self.load()
|
||||
session_id = str(data.get("active_session_id", ""))
|
||||
if not session_id:
|
||||
return None
|
||||
record = data.get("sessions", {}).get(session_id)
|
||||
return self._with_session_defaults(record) if isinstance(record, dict) else None
|
||||
|
||||
def get(self, session_id: str) -> dict[str, Any] | None:
|
||||
normalized = self._normalize_session_id(session_id)
|
||||
record = self.load().get("sessions", {}).get(normalized)
|
||||
return self._with_session_defaults(record) if isinstance(record, dict) else None
|
||||
|
||||
@_synchronized
|
||||
def create(self, idea: str = "", user_requirement: str = "", style: str = "", session_id: str | None = None) -> dict[str, Any]:
|
||||
data = self.load()
|
||||
sessions = data.setdefault("sessions", {})
|
||||
final_id = self._normalize_session_id(session_id) if session_id else self._new_session_id(idea or user_requirement or "vimax", sessions)
|
||||
if final_id in sessions:
|
||||
final_id = self._dedupe_session_id(final_id, sessions)
|
||||
now = datetime.now().isoformat(timespec="seconds")
|
||||
working_dir = self._working_dir_for_id(final_id)
|
||||
(working_dir / "idea2video").mkdir(parents=True, exist_ok=True)
|
||||
(working_dir / "script2video").mkdir(parents=True, exist_ok=True)
|
||||
record = {
|
||||
"session_id": final_id,
|
||||
"working_dir": str(working_dir.relative_to(self.workspace_root)),
|
||||
"idea": idea,
|
||||
"user_requirement": user_requirement,
|
||||
"style": style,
|
||||
"stage": "created",
|
||||
"summary": "",
|
||||
"stale": {key: False for key in STALE_KEYS},
|
||||
"recent_turn_records": [],
|
||||
"compacted_summary": "",
|
||||
"compacted_turns": 0,
|
||||
"compaction_snapshots": [],
|
||||
"last_compaction_reason": "",
|
||||
"last_compaction_at": "",
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
}
|
||||
sessions[final_id] = record
|
||||
data["active_session_id"] = final_id
|
||||
self.save(data)
|
||||
return record
|
||||
|
||||
def get_or_create_active(self, idea: str = "", user_requirement: str = "", style: str = "") -> dict[str, Any]:
|
||||
active = self.active()
|
||||
if active is not None:
|
||||
return active
|
||||
return self.create(idea=idea, user_requirement=user_requirement, style=style)
|
||||
|
||||
@_synchronized
|
||||
def set_active(self, session_id: str) -> dict[str, Any]:
|
||||
normalized = self._normalize_session_id(session_id)
|
||||
data = self.load()
|
||||
if normalized not in data.get("sessions", {}):
|
||||
raise KeyError(f"Unknown session_id: {session_id}")
|
||||
data["active_session_id"] = normalized
|
||||
self.save(data)
|
||||
return dict(data["sessions"][normalized])
|
||||
|
||||
@_synchronized
|
||||
def update_stage(self, session_id: str, stage: str, summary: str = "") -> None:
|
||||
data = self.load()
|
||||
record = data.get("sessions", {}).get(session_id)
|
||||
if not isinstance(record, dict):
|
||||
raise KeyError(f"Unknown session_id: {session_id}")
|
||||
record["stage"] = stage
|
||||
if summary:
|
||||
record["summary"] = summary
|
||||
record["updated_at"] = datetime.now().isoformat(timespec="seconds")
|
||||
self.save(data)
|
||||
|
||||
@_synchronized
|
||||
def mark_stale(self, session_id: str, keys: list[str]) -> None:
|
||||
data = self.load()
|
||||
record = data.get("sessions", {}).get(session_id)
|
||||
if not isinstance(record, dict):
|
||||
raise KeyError(f"Unknown session_id: {session_id}")
|
||||
stale = record.setdefault("stale", {key: False for key in STALE_KEYS})
|
||||
for key in keys:
|
||||
stale[key] = True
|
||||
record["updated_at"] = datetime.now().isoformat(timespec="seconds")
|
||||
self.save(data)
|
||||
|
||||
@_synchronized
|
||||
def update_compaction(self, session_id: str, result: dict[str, Any]) -> None:
|
||||
data = self.load()
|
||||
session = data.get("sessions", {}).get(session_id)
|
||||
if not isinstance(session, dict):
|
||||
raise KeyError(f"Unknown session_id: {session_id}")
|
||||
summary = str(result.get("summary", "") or "")
|
||||
compacted_count = int(result.get("compacted_message_count", 0) or 0)
|
||||
snapshot = {
|
||||
"level": len(session.get("compaction_snapshots", []) or []) + 1,
|
||||
"reason": str(result.get("reason", "manual") or "manual"),
|
||||
"mode": str(result.get("mode", "unknown") or "unknown"),
|
||||
"summary": summary,
|
||||
"preserved_messages": int(result.get("preserved_message_count", 0) or 0),
|
||||
"compacted_message_count": compacted_count,
|
||||
"estimated_tokens_before": int(result.get("estimated_tokens_before", 0) or 0),
|
||||
"estimated_tokens_after": int(result.get("estimated_tokens_after", 0) or 0),
|
||||
"created_at": str(result.get("created_at", "") or datetime.now().isoformat(timespec="seconds")),
|
||||
}
|
||||
session["compacted_summary"] = summary
|
||||
session["compacted_turns"] = int(session.get("compacted_turns", 0) or 0) + max(1, compacted_count // 2)
|
||||
snapshots = list(session.get("compaction_snapshots", []) or [])
|
||||
snapshots.append(snapshot)
|
||||
session["compaction_snapshots"] = snapshots[-8:]
|
||||
session["last_compaction_reason"] = snapshot["reason"]
|
||||
session["last_compaction_at"] = snapshot["created_at"]
|
||||
session["updated_at"] = datetime.now().isoformat(timespec="seconds")
|
||||
self.save(data)
|
||||
self.append_log("loop_history", {"session_id": session_id, "event": "context_compacted", "compaction": snapshot})
|
||||
|
||||
def compacted_summary(self, session_id: str | None = None) -> str:
|
||||
record = self.get(session_id) if session_id else self.active()
|
||||
return str((record or {}).get("compacted_summary", "") or "")
|
||||
|
||||
@_synchronized
|
||||
def append_turn_record(self, session_id: str, record: dict[str, Any]) -> None:
|
||||
data = self.load()
|
||||
session = data.get("sessions", {}).get(session_id)
|
||||
if isinstance(session, dict):
|
||||
recent = session.setdefault("recent_turn_records", [])
|
||||
recent.append({
|
||||
"turn_id": record.get("turn_id", ""),
|
||||
"status": record.get("status", ""),
|
||||
"tool_round_count": len(record.get("tool_rounds", [])),
|
||||
"final_preview": str(record.get("final_assistant_text", ""))[:240],
|
||||
"created_at": record.get("created_at", ""),
|
||||
})
|
||||
session["recent_turn_records"] = recent[-6:]
|
||||
session["updated_at"] = datetime.now().isoformat(timespec="seconds")
|
||||
self.save(data)
|
||||
self.append_log("loop_history", {"session_id": session_id, **record})
|
||||
|
||||
def working_dir(self, session_id: str | None = None) -> Path:
|
||||
record = self.get(session_id) if session_id else self.active()
|
||||
if record is None:
|
||||
record = self.create()
|
||||
path = (self.workspace_root / str(record["working_dir"])).resolve()
|
||||
if path != self.working_root and self.working_root not in path.parents:
|
||||
raise ValueError(f"Session working_dir escapes .working_dir: {record.get('working_dir')}")
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
def artifact_checklist(self, session_id: str | None = None) -> dict[str, bool]:
|
||||
root = self.working_dir(session_id)
|
||||
idea_dir = root / "idea2video"
|
||||
idea_scene_dirs = sorted(path for path in idea_dir.glob("scene_*") if path.is_dir()) if idea_dir.exists() else []
|
||||
idea_scene_storyboards = [path / "storyboard.json" for path in idea_scene_dirs]
|
||||
idea_scene_camera_trees = [path / "camera_tree.json" for path in idea_scene_dirs]
|
||||
idea_scene_shot_desc_groups = [list((scene / "shots").glob("*/shot_description.json")) for scene in idea_scene_dirs]
|
||||
idea_scene_selector_outputs = [output for scene in idea_scene_dirs for output in (scene / "shots").glob("*/*_selector_output.json")]
|
||||
|
||||
script_shots = root / "script2video" / "shots"
|
||||
script_shot_descs = list(script_shots.glob("*/shot_description.json")) if script_shots.exists() else []
|
||||
script_selector_outputs = list(script_shots.glob("*/*_selector_output.json")) if script_shots.exists() else []
|
||||
|
||||
novel_dir = root / "novel2video"
|
||||
novel_events = list((novel_dir / "events").glob("event_*.json")) if novel_dir.exists() else []
|
||||
novel_relevant_chunks = [path for path in (novel_dir / "relevant_chunks").glob("event_*/*") if path.is_file()] if novel_dir.exists() else []
|
||||
novel_scenes = list((novel_dir / "scenes").glob("event_*/scene_*.json")) if novel_dir.exists() else []
|
||||
novel_event_chars = list((novel_dir / "global_information" / "characters" / "event_level").glob("event_*_characters.json")) if novel_dir.exists() else []
|
||||
novel_level_chars = list((novel_dir / "global_information" / "characters" / "novel_level").glob("novel_characters_after_event_*.json")) if novel_dir.exists() else []
|
||||
return {
|
||||
"idea2video/story.txt": (idea_dir / "story.txt").exists(),
|
||||
"idea2video/characters.json": (idea_dir / "characters.json").exists(),
|
||||
"idea2video/script.json": (idea_dir / "script.json").exists(),
|
||||
"idea2video/scene_*/storyboard.json": bool(idea_scene_storyboards) and all(path.exists() for path in idea_scene_storyboards),
|
||||
"idea2video/scene_*/camera_tree.json": bool(idea_scene_camera_trees) and all(path.exists() for path in idea_scene_camera_trees),
|
||||
"idea2video/scene_*/shots/*/shot_description.json": bool(idea_scene_shot_desc_groups) and all(idea_scene_shot_desc_groups),
|
||||
"idea2video/scene_*/shots/*/*_selector_output.json": bool(idea_scene_selector_outputs),
|
||||
"idea2video/final_video.mp4": (idea_dir / "final_video.mp4").exists(),
|
||||
"script2video/script.txt": (root / "script2video" / "script.txt").exists(),
|
||||
"script2video/characters.json": (root / "script2video" / "characters.json").exists(),
|
||||
"script2video/storyboard.json": (root / "script2video" / "storyboard.json").exists(),
|
||||
"script2video/shots/*/shot_description.json": bool(script_shot_descs),
|
||||
"script2video/camera_tree.json": (root / "script2video" / "camera_tree.json").exists(),
|
||||
"script2video/shots/*/*_selector_output.json": bool(script_selector_outputs),
|
||||
"script2video/final_video.mp4": (root / "script2video" / "final_video.mp4").exists(),
|
||||
"novel2video/novel/novel.txt": (novel_dir / "novel" / "novel.txt").exists(),
|
||||
"novel2video/novel/novel_compressed.txt": (novel_dir / "novel" / "novel_compressed.txt").exists(),
|
||||
"novel2video/events/event_*.json": bool(novel_events),
|
||||
"novel2video/relevant_chunks/event_*": bool(novel_relevant_chunks),
|
||||
"novel2video/scenes/event_*/scene_*.json": bool(novel_scenes),
|
||||
"novel2video/global_information/characters/event_level/*.json": bool(novel_event_chars),
|
||||
"novel2video/global_information/characters/novel_level/*.json": bool(novel_level_chars),
|
||||
}
|
||||
|
||||
def memory_text(self) -> str:
|
||||
return self.memory_path.read_text(encoding="utf-8") if self.memory_path.exists() else ""
|
||||
|
||||
def write_memory(self, text: str) -> None:
|
||||
self.memory_path.write_text(text, encoding="utf-8")
|
||||
|
||||
def append_log(self, name: str, payload: dict[str, Any]) -> None:
|
||||
event = {"timestamp": datetime.now().isoformat(timespec="seconds"), **payload}
|
||||
path = self.logs_dir / f"{name}.jsonl"
|
||||
with path.open("a", encoding="utf-8") as f:
|
||||
f.write(json.dumps(event, ensure_ascii=False, default=str) + "\n")
|
||||
|
||||
def snapshot(self) -> dict[str, Any]:
|
||||
active = self.active()
|
||||
if active is None:
|
||||
return {"active_session_id": "", "session": None}
|
||||
return {"active_session_id": active["session_id"], "session": active, "artifact_checklist": self.artifact_checklist(active["session_id"])}
|
||||
|
||||
def _with_session_defaults(self, record: dict[str, Any]) -> dict[str, Any]:
|
||||
item = dict(record)
|
||||
item.setdefault("compacted_summary", "")
|
||||
item.setdefault("compacted_turns", 0)
|
||||
item.setdefault("compaction_snapshots", [])
|
||||
item.setdefault("last_compaction_reason", "")
|
||||
item.setdefault("last_compaction_at", "")
|
||||
item.setdefault("recent_turn_records", [])
|
||||
return item
|
||||
|
||||
def _new_session_id(self, source: str, sessions: dict[str, Any]) -> str:
|
||||
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
slug = (re.sub(r"[^a-zA-Z0-9]+", "-", source.lower()).strip("-")[:32].strip("-") or "vimax")
|
||||
return self._dedupe_session_id(f"{stamp}-{slug}", sessions)
|
||||
|
||||
def _dedupe_session_id(self, base: str, sessions: dict[str, Any]) -> str:
|
||||
candidate = base
|
||||
counter = 2
|
||||
while candidate in sessions:
|
||||
candidate = f"{base}-{counter}"
|
||||
counter += 1
|
||||
return candidate
|
||||
|
||||
def _normalize_session_id(self, session_id: str | None) -> str:
|
||||
raw = str(session_id or "").strip()
|
||||
if not raw:
|
||||
raise ValueError("session_id cannot be empty")
|
||||
normalized = re.sub(r"[^a-zA-Z0-9]+", "-", raw).strip("-")[:96]
|
||||
if not normalized:
|
||||
raise ValueError(f"Invalid session_id: {session_id}")
|
||||
return normalized
|
||||
|
||||
def _working_dir_for_id(self, session_id: str) -> Path:
|
||||
path = (self.working_root / session_id).resolve()
|
||||
if path != self.working_root and self.working_root not in path.parents:
|
||||
raise ValueError(f"Session path escapes .working_dir: {session_id}")
|
||||
return path
|
||||
@@ -0,0 +1,46 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from time import time
|
||||
from typing import Any, Callable
|
||||
|
||||
from .models import ToolCall, ToolResult, TurnControl
|
||||
from .tools import ToolRegistry, ToolRuntimeContext
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ToolExecutionRecord:
|
||||
requested_name: str
|
||||
canonical_name: str
|
||||
arguments_before: dict[str, Any]
|
||||
arguments_after: dict[str, Any]
|
||||
result: ToolResult
|
||||
started_at: float
|
||||
finished_at: float
|
||||
telemetry: dict[str, Any]
|
||||
|
||||
|
||||
class ToolExecutor:
|
||||
def __init__(self, registry: ToolRegistry, session_index: Any) -> None:
|
||||
self.registry = registry
|
||||
self.session_index = session_index
|
||||
|
||||
async def execute(self, call: ToolCall, control: TurnControl, progress_callback: Callable[[dict[str, Any]], None] | None = None) -> ToolExecutionRecord:
|
||||
requested_name = call.name
|
||||
canonical_name = self.registry.resolve_name(call.name)
|
||||
before = deepcopy(call.arguments)
|
||||
started_at = time()
|
||||
validated, validation_error = self.registry.validate_arguments(canonical_name, call.arguments)
|
||||
arguments = validated if validated is not None else call.arguments
|
||||
runtime = ToolRuntimeContext(requested_name=requested_name, canonical_name=canonical_name, turn_id=control.turn_id, cancel_event=control.cancel_event, progress_callback=progress_callback, metadata={"cancel_reason": control.cancel_reason})
|
||||
if validation_error:
|
||||
result = ToolResult(canonical_name, False, validation_error, {"validation_error": True})
|
||||
elif control.cancel_event.is_set():
|
||||
result = ToolResult(canonical_name, False, control.cancel_reason or "Tool execution cancelled", {"cancelled": True})
|
||||
else:
|
||||
result = await self.registry.execute(canonical_name, arguments, runtime=runtime)
|
||||
finished_at = time()
|
||||
telemetry = {"duration_ms": int((finished_at - started_at) * 1000), "requested_name": requested_name, "canonical_name": canonical_name, "result_ok": result.ok}
|
||||
self.session_index.append_log("tool_calls", {"turn_id": control.turn_id, "tool": canonical_name, "arguments_preview": str(before)[:500], "ok": result.ok, "content_preview": result.content[:500], **telemetry})
|
||||
return ToolExecutionRecord(requested_name, canonical_name, before, deepcopy(arguments), result, started_at, finished_at, telemetry)
|
||||
@@ -0,0 +1,404 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import glob
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from threading import Event
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
from .models import ToolCall, ToolResult
|
||||
|
||||
ToolHandler = Callable[..., Awaitable[ToolResult] | ToolResult]
|
||||
ProgressCallback = Callable[[dict[str, Any]], None]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ToolArgumentSchema:
|
||||
type: type | tuple[type, ...]
|
||||
required: bool = False
|
||||
default: Any = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ToolSpec:
|
||||
name: str
|
||||
description: str
|
||||
handler: ToolHandler
|
||||
aliases: tuple[str, ...] = ()
|
||||
permission_mode: str = "workspace-write"
|
||||
schema: dict[str, ToolArgumentSchema] | None = None
|
||||
json_schema: dict[str, Any] | None = None
|
||||
concurrency_safe: bool = False
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ToolRuntimeContext:
|
||||
requested_name: str
|
||||
canonical_name: str
|
||||
turn_id: str = ""
|
||||
cancel_event: Event | None = None
|
||||
progress_callback: ProgressCallback | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def emit_progress(self, message: str, *, stage: str = "running", metadata: dict[str, Any] | None = None) -> None:
|
||||
if self.progress_callback is None:
|
||||
return
|
||||
payload: dict[str, Any] = {
|
||||
"type": "tool_progress",
|
||||
"tool": {"requested_name": self.requested_name, "name": self.canonical_name},
|
||||
"progress": {"stage": stage, "message": message, "metadata": metadata or {}},
|
||||
}
|
||||
if self.turn_id:
|
||||
payload["turn_id"] = self.turn_id
|
||||
self.progress_callback(payload)
|
||||
|
||||
def emit_terminal(self, line: str, *, stream: str = "stdout") -> None:
|
||||
if self.progress_callback is None:
|
||||
return
|
||||
if not line:
|
||||
return
|
||||
payload: dict[str, Any] = {"type": "terminal", "stream": stream, "line": line}
|
||||
if self.turn_id:
|
||||
payload["turn_id"] = self.turn_id
|
||||
self.progress_callback(payload)
|
||||
|
||||
def is_cancelled(self) -> bool:
|
||||
return self.cancel_event.is_set() if self.cancel_event is not None else False
|
||||
|
||||
def raise_if_cancelled(self, default_reason: str = "Tool execution cancelled") -> None:
|
||||
if self.is_cancelled():
|
||||
raise RuntimeError(str(self.metadata.get("cancel_reason") or default_reason))
|
||||
|
||||
|
||||
class ToolRegistry:
|
||||
def __init__(self, specs: list[ToolSpec] | None = None) -> None:
|
||||
self._specs: dict[str, ToolSpec] = {}
|
||||
self._aliases: dict[str, str] = {}
|
||||
for spec in specs or []:
|
||||
self.register(spec)
|
||||
|
||||
def register(self, spec: ToolSpec) -> None:
|
||||
self._specs[spec.name] = spec
|
||||
for alias in spec.aliases:
|
||||
self._aliases[alias] = spec.name
|
||||
|
||||
def list_tools(self) -> list[dict[str, str]]:
|
||||
return sorted([{"name": spec.name, "description": spec.description, "permission_mode": spec.permission_mode} for spec in self._specs.values()], key=lambda item: item["name"])
|
||||
|
||||
def list_function_tools(self) -> list[dict[str, Any]]:
|
||||
tools = []
|
||||
for spec in sorted(self._specs.values(), key=lambda item: item.name):
|
||||
parameters = spec.json_schema or _argument_schema_to_json_schema(spec.schema or {})
|
||||
tools.append({"type": "function", "function": {"name": spec.name, "description": spec.description, "parameters": parameters}})
|
||||
return tools
|
||||
|
||||
def get_spec(self, name: str) -> ToolSpec | None:
|
||||
return self._specs.get(self.resolve_name(name))
|
||||
|
||||
def resolve_name(self, name: str) -> str:
|
||||
normalized = name.strip()
|
||||
return self._aliases.get(normalized, normalized)
|
||||
|
||||
def validate_arguments(self, name: str, arguments: dict[str, Any]) -> tuple[dict[str, Any] | None, str | None]:
|
||||
spec = self.get_spec(name)
|
||||
if spec is None:
|
||||
return None, f"Unknown tool: {name}"
|
||||
schema = spec.schema or {}
|
||||
normalized = dict(arguments or {})
|
||||
for field_name, field_spec in schema.items():
|
||||
if field_name not in normalized:
|
||||
if field_spec.required and field_spec.default is None:
|
||||
return None, f"Missing required argument '{field_name}' for {spec.name}"
|
||||
if field_spec.default is not None:
|
||||
normalized[field_name] = field_spec.default
|
||||
continue
|
||||
value = normalized[field_name]
|
||||
expected = field_spec.type
|
||||
if expected is bool and isinstance(value, str) and value.lower() in {"true", "false"}:
|
||||
normalized[field_name] = value.lower() == "true"
|
||||
continue
|
||||
if expected is int and isinstance(value, str):
|
||||
try:
|
||||
normalized[field_name] = int(value)
|
||||
continue
|
||||
except ValueError:
|
||||
return None, f"Argument '{field_name}' for {spec.name} must be an integer"
|
||||
if not isinstance(normalized[field_name], expected):
|
||||
expected_name = ", ".join(t.__name__ for t in expected) if isinstance(expected, tuple) else expected.__name__
|
||||
return None, f"Argument '{field_name}' for {spec.name} must be {expected_name}"
|
||||
return normalized, None
|
||||
|
||||
def is_concurrency_safe(self, name: str) -> bool:
|
||||
spec = self.get_spec(name)
|
||||
return bool(spec and spec.concurrency_safe)
|
||||
|
||||
def partition_calls(self, calls: list[ToolCall]) -> list[list[ToolCall]]:
|
||||
batches: list[list[ToolCall]] = []
|
||||
for call in calls:
|
||||
if self.is_concurrency_safe(call.name) and batches and all(self.is_concurrency_safe(item.name) for item in batches[-1]):
|
||||
batches[-1].append(call)
|
||||
else:
|
||||
batches.append([call])
|
||||
return batches
|
||||
|
||||
async def execute(self, name: str, arguments: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
|
||||
canonical = self.resolve_name(name)
|
||||
spec = self._specs.get(canonical)
|
||||
if spec is None:
|
||||
return ToolResult(name=name, ok=False, content=f"Unknown tool: {name}", metadata={"error_type": "unknown_tool"})
|
||||
handler = spec.handler
|
||||
try:
|
||||
params = inspect.signature(handler).parameters
|
||||
result = handler(arguments, runtime) if runtime is not None and len(params) >= 2 else handler(arguments)
|
||||
if inspect.isawaitable(result):
|
||||
return await result
|
||||
return result
|
||||
except Exception as exc:
|
||||
return ToolResult(name=canonical, ok=False, content=str(exc), metadata={"error_type": "exception"})
|
||||
|
||||
|
||||
def _argument_schema_to_json_schema(schema: dict[str, ToolArgumentSchema]) -> dict[str, Any]:
|
||||
properties: dict[str, Any] = {}
|
||||
required: list[str] = []
|
||||
for field_name, field_spec in schema.items():
|
||||
field_schema = _type_to_json_schema(field_spec.type)
|
||||
if field_spec.default is not None:
|
||||
field_schema["default"] = field_spec.default
|
||||
properties[field_name] = field_schema
|
||||
if field_spec.required and field_spec.default is None:
|
||||
required.append(field_name)
|
||||
payload: dict[str, Any] = {"type": "object", "properties": properties, "additionalProperties": False}
|
||||
if required:
|
||||
payload["required"] = required
|
||||
return payload
|
||||
|
||||
|
||||
def _type_to_json_schema(tp: type | tuple[type, ...]) -> dict[str, Any]:
|
||||
if isinstance(tp, tuple):
|
||||
return {"anyOf": [_type_to_json_schema(item) for item in tp]}
|
||||
return {str: {"type": "string"}, int: {"type": "integer"}, bool: {"type": "boolean"}, dict: {"type": "object", "additionalProperties": True}, list: {"type": "array", "items": {}}}.get(tp, {"type": "string"})
|
||||
|
||||
|
||||
def build_builtin_registry(workspace_root: str | Path, session_index: Any, adapter_specs: list[ToolSpec] | None = None) -> ToolRegistry:
|
||||
root = Path(workspace_root).resolve()
|
||||
|
||||
def safe_path(raw: Any) -> Path:
|
||||
path = (root / str(raw)).resolve()
|
||||
if root not in path.parents and path != root:
|
||||
raise ValueError(f"Path escapes workspace: {raw}")
|
||||
return path
|
||||
|
||||
def _legacy_virtual_read(raw_path: Any, *, as_json: bool) -> ToolResult | None:
|
||||
"""Compatibility for paths older prompts/models may hallucinate.
|
||||
|
||||
The authoritative session state is .vimax/sessions.json and logs are
|
||||
.vimax/logs/*.jsonl, but some model turns ask for per-session files like
|
||||
.working_dir/<session>/session.json or .vimax/logs/<session>.log.
|
||||
"""
|
||||
path = safe_path(raw_path)
|
||||
try:
|
||||
rel = path.relative_to(root)
|
||||
except ValueError:
|
||||
return None
|
||||
parts = rel.parts
|
||||
if len(parts) == 3 and parts[0] == ".working_dir" and parts[2] == "session.json":
|
||||
session_id = parts[1]
|
||||
record = session_index.get(session_id)
|
||||
if record is None:
|
||||
return None
|
||||
payload = {
|
||||
"session": record,
|
||||
"artifact_checklist": session_index.artifact_checklist(session_id),
|
||||
"source": ".vimax/sessions.json",
|
||||
"virtual_path": rel.as_posix(),
|
||||
}
|
||||
content = json.dumps(payload, ensure_ascii=False, indent=2)
|
||||
return ToolResult("read_json" if as_json else "read_file", True, content, {"virtual_path": True, "source": ".vimax/sessions.json"})
|
||||
if len(parts) == 3 and parts[0] == ".vimax" and parts[1] == "logs" and parts[2].endswith(".log"):
|
||||
session_id = parts[2][:-4]
|
||||
rows: list[dict[str, Any]] = []
|
||||
for log_name in ("loop_history", "tool_calls", "revisions"):
|
||||
log_path = session_index.logs_dir / f"{log_name}.jsonl"
|
||||
if not log_path.exists():
|
||||
continue
|
||||
for line in log_path.read_text(encoding="utf-8", errors="replace").splitlines():
|
||||
if session_id not in line:
|
||||
continue
|
||||
try:
|
||||
item = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
item = {"raw": line}
|
||||
item["_log"] = log_name
|
||||
rows.append(item)
|
||||
payload = {
|
||||
"session_id": session_id,
|
||||
"source": ".vimax/logs/*.jsonl",
|
||||
"virtual_path": rel.as_posix(),
|
||||
"records": rows,
|
||||
}
|
||||
content = json.dumps(payload, ensure_ascii=False, indent=2)
|
||||
return ToolResult("read_json" if as_json else "read_file", True, content, {"virtual_path": True, "source": ".vimax/logs/*.jsonl", "record_count": len(rows)})
|
||||
return None
|
||||
|
||||
def read_file(args: dict[str, Any]) -> ToolResult:
|
||||
path = safe_path(args["path"])
|
||||
if not path.exists():
|
||||
virtual = _legacy_virtual_read(args["path"], as_json=False)
|
||||
if virtual is not None:
|
||||
return virtual
|
||||
return ToolResult("read_file", False, f"File not found: {path}")
|
||||
return ToolResult("read_file", True, path.read_text(encoding="utf-8"))
|
||||
|
||||
def read_json(args: dict[str, Any]) -> ToolResult:
|
||||
path = safe_path(args["path"])
|
||||
if not path.exists():
|
||||
virtual = _legacy_virtual_read(args["path"], as_json=True)
|
||||
if virtual is not None:
|
||||
return virtual
|
||||
return ToolResult("read_json", False, f"File not found: {path}")
|
||||
try:
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
except json.JSONDecodeError as exc:
|
||||
return ToolResult("read_json", False, f"Invalid JSON: {exc}", {"error_type": "invalid_json"})
|
||||
return ToolResult("read_json", True, json.dumps(payload, ensure_ascii=False, indent=2))
|
||||
|
||||
def write_json(args: dict[str, Any]) -> ToolResult:
|
||||
path = safe_path(args["path"])
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(args["data"], ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return ToolResult("write_json", True, f"Wrote JSON {path.relative_to(root)}")
|
||||
|
||||
def list_files(args: dict[str, Any]) -> ToolResult:
|
||||
path = safe_path(args.get("path", "."))
|
||||
if not path.exists():
|
||||
return ToolResult("list_files", False, f"Path not found: {path}")
|
||||
rows = [str(item.relative_to(root)) for item in sorted(path.iterdir())]
|
||||
return ToolResult("list_files", True, "\n".join(rows) or "No entries")
|
||||
|
||||
def glob_files(args: dict[str, Any]) -> ToolResult:
|
||||
pattern = str(args["pattern"])
|
||||
matches = [str(Path(item).resolve().relative_to(root)) for item in glob.glob(str(root / pattern), recursive=True)]
|
||||
return ToolResult("glob_files", True, "\n".join(matches[:200]) or "No matches")
|
||||
|
||||
def search_text(args: dict[str, Any]) -> ToolResult:
|
||||
needle = str(args["query"])
|
||||
base = safe_path(args.get("path", "."))
|
||||
rows: list[str] = []
|
||||
paths = base.rglob("*") if base.is_dir() else [base]
|
||||
for path in paths:
|
||||
if not path.is_file():
|
||||
continue
|
||||
try:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
except UnicodeDecodeError:
|
||||
continue
|
||||
for idx, line in enumerate(text.splitlines(), start=1):
|
||||
if needle in line:
|
||||
rows.append(f"{path.relative_to(root)}:{idx}: {line}")
|
||||
if len(rows) >= int(args.get("max_results", 100)):
|
||||
return ToolResult("search_text", True, "\n".join(rows))
|
||||
return ToolResult("search_text", True, "\n".join(rows) or "No matches")
|
||||
|
||||
def memory_read(args: dict[str, Any]) -> ToolResult:
|
||||
return ToolResult("memory_read", True, session_index.memory_text())
|
||||
|
||||
def memory_write(args: dict[str, Any]) -> ToolResult:
|
||||
session_index.write_memory(str(args["content"]))
|
||||
return ToolResult("memory_write", True, "Updated .vimax/memory.md")
|
||||
|
||||
def todo_path() -> Path:
|
||||
return root / ".vimax" / "todo.json"
|
||||
|
||||
def todo_read(args: dict[str, Any]) -> ToolResult:
|
||||
path = todo_path()
|
||||
if not path.exists():
|
||||
return ToolResult("todo_read", True, json.dumps({"items": []}, ensure_ascii=False, indent=2), {"items": []})
|
||||
try:
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
except json.JSONDecodeError as exc:
|
||||
return ToolResult("todo_read", False, f"Invalid todo JSON: {exc}", {"error_type": "invalid_json"})
|
||||
items = payload.get("items")
|
||||
if not isinstance(items, list):
|
||||
return ToolResult("todo_read", False, "Invalid todo JSON: expected an items array", {"error_type": "invalid_todo"})
|
||||
return ToolResult("todo_read", True, json.dumps({"items": items}, ensure_ascii=False, indent=2), {"items": items})
|
||||
|
||||
def todo_write(args: dict[str, Any]) -> ToolResult:
|
||||
items = args.get("items")
|
||||
if not isinstance(items, list):
|
||||
return ToolResult("todo_write", False, "items must be an array", {"error_type": "invalid_arguments"})
|
||||
normalized: list[dict[str, Any]] = []
|
||||
for index, item in enumerate(items):
|
||||
if not isinstance(item, dict):
|
||||
return ToolResult("todo_write", False, f"items[{index}] must be an object", {"error_type": "invalid_arguments", "index": index})
|
||||
content = str(item.get("content", "")).strip()
|
||||
if not content:
|
||||
return ToolResult("todo_write", False, f"items[{index}].content is required", {"error_type": "invalid_arguments", "index": index})
|
||||
status = str(item.get("status", "pending")).strip() or "pending"
|
||||
if status not in {"pending", "in_progress", "completed"}:
|
||||
return ToolResult("todo_write", False, f"items[{index}].status must be pending, in_progress, or completed", {"error_type": "invalid_arguments", "index": index})
|
||||
normalized.append({"content": content, "status": status})
|
||||
path = todo_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps({"items": normalized}, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return ToolResult("todo_write", True, f"Updated .vimax/todo.json with {len(normalized)} item(s)", {"items": normalized, "item_count": len(normalized)})
|
||||
|
||||
async def sleep_tool(args: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
|
||||
seconds = float(args.get("seconds", 0))
|
||||
if seconds < 0 or seconds > 300:
|
||||
return ToolResult("sleep", False, "seconds must be between 0 and 300")
|
||||
if runtime:
|
||||
runtime.emit_progress(f"Sleeping for {seconds:g}s", stage="running")
|
||||
await asyncio.sleep(seconds)
|
||||
return ToolResult("sleep", True, f"Slept for {seconds:g}s")
|
||||
|
||||
async def run_shell(args: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
|
||||
if os.environ.get("VIMAX_ENABLE_RUN_SHELL") != "1":
|
||||
return ToolResult("run_shell", False, "run_shell is disabled by default. Set VIMAX_ENABLE_RUN_SHELL=1 to enable bounded shell commands.", {"error_type": "disabled"})
|
||||
command = str(args["command"]).strip()
|
||||
timeout_seconds = min(max(int(args.get("timeout_seconds", 30)), 1), 120)
|
||||
output_limit = min(max(int(args.get("output_limit", 20000)), 1000), 50000)
|
||||
denied_tokens = ["rm ", "rm -", "sudo", "chmod", "chown", "mkfs", "dd ", ":(){", "curl ", "wget ", "ssh ", "printenv", "env", "export"]
|
||||
lowered = command.lower()
|
||||
if any(token in lowered for token in denied_tokens):
|
||||
return ToolResult("run_shell", False, "Command rejected by run_shell policy.", {"error_type": "command_rejected"})
|
||||
if runtime:
|
||||
runtime.emit_progress("Starting shell command", stage="starting", metadata={"command": command, "timeout_seconds": timeout_seconds})
|
||||
proc = await asyncio.create_subprocess_shell(command, cwd=root, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout_seconds)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await proc.communicate()
|
||||
return ToolResult("run_shell", False, f"Command timed out after {timeout_seconds}s", {"error_type": "timeout", "timeout_seconds": timeout_seconds})
|
||||
content = ""
|
||||
if stdout:
|
||||
content += stdout.decode(errors="replace")
|
||||
if stderr:
|
||||
content += stderr.decode(errors="replace")
|
||||
truncated = len(content) > output_limit
|
||||
if truncated:
|
||||
content = content[:output_limit] + "\n...[truncated]"
|
||||
return ToolResult("run_shell", proc.returncode == 0, content, {"returncode": proc.returncode, "truncated": truncated})
|
||||
|
||||
specs = [
|
||||
ToolSpec("read_file", "Read a UTF-8 text file inside the workspace. Also resolves virtual legacy session paths like .vimax/logs/<session>.log.", read_file, schema={"path": ToolArgumentSchema(str, True)}, concurrency_safe=True),
|
||||
ToolSpec("read_json", "Read and parse a JSON file inside the workspace. Also resolves virtual legacy session paths like .working_dir/<session>/session.json.", read_json, schema={"path": ToolArgumentSchema(str, True)}, concurrency_safe=True),
|
||||
ToolSpec("write_json", "Write formatted JSON inside the workspace.", write_json, schema={"path": ToolArgumentSchema(str, True), "data": ToolArgumentSchema((dict, list), True)}),
|
||||
ToolSpec("list_files", "List direct children of a workspace path.", list_files, schema={"path": ToolArgumentSchema(str, False, ".")}, concurrency_safe=True),
|
||||
ToolSpec("glob_files", "Find workspace files with a glob pattern.", glob_files, schema={"pattern": ToolArgumentSchema(str, True)}, concurrency_safe=True),
|
||||
ToolSpec("search_text", "Search text in workspace files.", search_text, schema={"query": ToolArgumentSchema(str, True), "path": ToolArgumentSchema(str, False, "."), "max_results": ToolArgumentSchema(int, False, 100)}, concurrency_safe=True),
|
||||
ToolSpec("memory_read", "Read .vimax/memory.md user preferences.", memory_read, schema={}, concurrency_safe=True),
|
||||
ToolSpec("memory_write", "Replace .vimax/memory.md with user preference notes only.", memory_write, schema={"content": ToolArgumentSchema(str, True)}),
|
||||
ToolSpec("todo_read", "Read short-term todo items from .vimax/todo.json. This is not a task or team system.", todo_read, schema={}, concurrency_safe=True),
|
||||
ToolSpec("todo_write", "Replace short-term todo items in .vimax/todo.json. Items require content and may use pending, in_progress, or completed status.", todo_write, schema={"items": ToolArgumentSchema(list, True)}),
|
||||
ToolSpec("sleep", "Wait for a bounded number of seconds.", sleep_tool, schema={"seconds": ToolArgumentSchema(int, False, 0)}, concurrency_safe=True),
|
||||
ToolSpec("run_shell", "Run a bounded shell command in the workspace. Disabled unless VIMAX_ENABLE_RUN_SHELL=1; rejects dangerous commands, enforces timeout, and truncates output.", run_shell, schema={"command": ToolArgumentSchema(str, True), "timeout_seconds": ToolArgumentSchema(int, False, 30), "output_limit": ToolArgumentSchema(int, False, 20000)}),
|
||||
]
|
||||
for spec in adapter_specs or []:
|
||||
specs.append(spec)
|
||||
return ToolRegistry(specs)
|
||||
@@ -0,0 +1,810 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
from contextlib import contextmanager, redirect_stderr, redirect_stdout
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from langchain.chat_models import init_chat_model
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
from tenacity import RetryError
|
||||
|
||||
from interfaces import CharacterInScene
|
||||
from agents.event_extractor import EventExtractor
|
||||
from agents.global_information_planner import GlobalInformationPlanner
|
||||
from agents.novel_compressor import NovelCompressor
|
||||
from agents.scene_extractor import SceneExtractor
|
||||
from pipelines.novel2movie_pipeline import Novel2MoviePipeline
|
||||
from pipelines.idea2video_pipeline import Idea2VideoPipeline
|
||||
from pipelines.script2video_pipeline import Script2VideoPipeline
|
||||
from tools.image_generator_nanobanana_yunwu_api import ImageGeneratorNanobananaYunwuAPI
|
||||
from tools.reranker_bge_silicon_api import RerankerBgeSiliconapi
|
||||
from tools.video_generator_openrouter_api import VideoGeneratorOpenRouterAPI
|
||||
from tools.video_generator_veo_yunwu_api import VideoGeneratorVeoYunwuAPI
|
||||
|
||||
from .config import embedding_api_key, embedding_base_url, embedding_model, embedding_model_provider, image_api_key, image_base_url, image_model, llm_api_key, llm_base_url, llm_model, llm_model_provider, reranker_api_key, reranker_base_url, reranker_model, video_api_key, video_base_url, video_model, video_provider
|
||||
from .models import ToolResult
|
||||
from .tools import ToolArgumentSchema, ToolRuntimeContext, ToolSpec
|
||||
|
||||
|
||||
class _UnavailableGenerator:
|
||||
async def generate_single_image(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise RuntimeError("Image generator is not available in narrative planning mode")
|
||||
|
||||
async def generate_single_video(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise RuntimeError("Video generator is not available in narrative planning mode")
|
||||
|
||||
|
||||
def build_vimax_adapter_specs(workspace_root: str | Path, session_index: Any) -> list[ToolSpec]:
|
||||
adapter = ViMaxAdapters(Path(workspace_root), session_index)
|
||||
return [
|
||||
ToolSpec(
|
||||
name="vimax_narrative_planning",
|
||||
description=(
|
||||
"Create or revise ViMax structured text artifacts for the active session. "
|
||||
"Idea mode writes story, characters, script, and scene-level storyboard/shot_decomposition/camera_tree under idea2video/scene_<idx>/. "
|
||||
"Script mode writes characters, storyboard, shot_decomposition, and camera_tree under script2video/. "
|
||||
"For a new video idea or new script, omit session_id or pass the new idea/script; the adapter will create a new session instead of reusing mismatched artifacts. If idea/script/revision_target are omitted and the active session has an idea, continue that session and fill missing structured text artifacts. "
|
||||
"It does not generate keyframes, video clips, or final video. Call this before revising storyboard/shots when those artifacts do not exist."
|
||||
),
|
||||
handler=adapter.vimax_narrative_planning,
|
||||
schema={
|
||||
"session_id": ToolArgumentSchema(str, required=False, default=""),
|
||||
"idea": ToolArgumentSchema(str, required=False, default=""),
|
||||
"script": ToolArgumentSchema(str, required=False, default=""),
|
||||
"user_requirement": ToolArgumentSchema(str, required=False, default=""),
|
||||
"style": ToolArgumentSchema(str, required=False, default=""),
|
||||
"revision_target": ToolArgumentSchema(str, required=False, default=""),
|
||||
"revision_instruction": ToolArgumentSchema(str, required=False, default=""),
|
||||
},
|
||||
),
|
||||
ToolSpec(
|
||||
name="vimax_novel_planning",
|
||||
description=(
|
||||
"Create ViMax structured text artifacts from a novel or novel excerpt. "
|
||||
"This writes novel2video/novel, events, relevant_chunks, scenes, and global_information text artifacts. "
|
||||
"Use this when the user provides long prose, a novel excerpt, or asks for novel-to-video planning. "
|
||||
"It does not generate character portraits, scene videos, or final video."
|
||||
),
|
||||
handler=adapter.vimax_novel_planning,
|
||||
schema={
|
||||
"session_id": ToolArgumentSchema(str, required=False, default=""),
|
||||
"novel_text": ToolArgumentSchema(str, required=True),
|
||||
"user_requirement": ToolArgumentSchema(str, required=False, default=""),
|
||||
"style": ToolArgumentSchema(str, required=False, default=""),
|
||||
},
|
||||
),
|
||||
ToolSpec(
|
||||
name="vimax_render_video",
|
||||
description=(
|
||||
"Render keyframes, video clips, and final video for the active ViMax session. "
|
||||
"This checks that structured text artifacts exist before rendering and reports missing dependencies instead of pretending render started."
|
||||
),
|
||||
handler=adapter.vimax_render_video,
|
||||
schema={
|
||||
"session_id": ToolArgumentSchema(str, required=False, default=""),
|
||||
"mode": ToolArgumentSchema(str, required=False, default="foreground"),
|
||||
"force": ToolArgumentSchema(bool, required=False, default=False),
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ViMaxAdapters:
|
||||
def __init__(self, workspace_root: Path, session_index: Any) -> None:
|
||||
self.workspace_root = workspace_root.resolve()
|
||||
self.session_index = session_index
|
||||
|
||||
async def vimax_narrative_planning(self, args: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
|
||||
idea = str(args.get("idea", "") or "").strip()
|
||||
script = str(args.get("script", "") or "").strip()
|
||||
user_requirement = str(args.get("user_requirement", "") or "").strip()
|
||||
requested_style = str(args.get("style", "") or "").strip()
|
||||
style = requested_style
|
||||
session = self._resolve_session(str(args.get("session_id", "") or ""), idea=idea, script=script, user_requirement=user_requirement, style=requested_style)
|
||||
session_id = session["session_id"]
|
||||
working_dir = self.session_index.working_dir(session_id)
|
||||
idea_dir = working_dir / "idea2video"
|
||||
script_dir = working_dir / "script2video"
|
||||
idea_dir.mkdir(parents=True, exist_ok=True)
|
||||
script_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if not idea and not script:
|
||||
revision_target = str(args.get("revision_target") or "").strip()
|
||||
if revision_target:
|
||||
return await self._revise_narrative_artifact(session_id, working_dir, revision_target, str(args.get("revision_instruction") or "").strip(), runtime)
|
||||
session_idea = str(session.get("idea") or "").strip()
|
||||
if session_idea:
|
||||
idea = session_idea
|
||||
user_requirement = user_requirement or str(session.get("user_requirement") or "").strip()
|
||||
style = requested_style or str(session.get("style") or "").strip() or "Cinematic, coherent, 16:9"
|
||||
else:
|
||||
return ToolResult("vimax_narrative_planning", False, "Provide `idea`, `script`, a revision target, or an active session with an existing idea for narrative planning.", {"error_type": "missing_input", "session_id": session_id})
|
||||
|
||||
style = style or str(session.get("style") or "").strip() or "Cinematic, coherent, 16:9"
|
||||
self._update_session_metadata(session_id, idea="", user_requirement="", style=style)
|
||||
|
||||
try:
|
||||
self.session_index.update_stage(session_id, "narrative_planning", "Generating structured text artifacts")
|
||||
if runtime:
|
||||
runtime.emit_progress("Starting narrative planning", stage="starting", metadata={"session_id": session_id})
|
||||
await asyncio.sleep(0)
|
||||
generated_before = self.session_index.artifact_checklist(session_id)
|
||||
if runtime:
|
||||
runtime.emit_progress("Initializing bounded chat model", stage="initializing_llm", metadata={"session_id": session_id, "timeout_seconds": _llm_request_timeout_seconds(), "max_tokens": _narrative_max_tokens()})
|
||||
await asyncio.sleep(0)
|
||||
chat_model = _build_chat_model()
|
||||
if runtime:
|
||||
runtime.emit_progress("Bounded chat model initialized", stage="chat_model_ready", metadata={"session_id": session_id})
|
||||
await asyncio.sleep(0)
|
||||
dummy = _UnavailableGenerator()
|
||||
# Do not globally redirect stdout/stderr while the JSONL CLI is streaming events.
|
||||
# The adapter exposes pipeline progress through explicit tool_progress events instead.
|
||||
if idea:
|
||||
idea_pipeline = Idea2VideoPipeline(chat_model=chat_model, image_generator=dummy, video_generator=dummy, working_dir=str(idea_dir))
|
||||
if runtime:
|
||||
runtime.emit_progress("Idea pipeline initialized", stage="idea_pipeline_ready", metadata={"session_id": session_id})
|
||||
await asyncio.sleep(0)
|
||||
story = await _run_planning_step(
|
||||
"Developing story from user idea",
|
||||
"develop_story",
|
||||
idea_pipeline.develop_story(idea=idea, user_requirement=user_requirement, quiet=True),
|
||||
runtime,
|
||||
{"session_id": session_id},
|
||||
)
|
||||
characters = await _run_planning_step(
|
||||
"Extracting characters from story",
|
||||
"extract_characters",
|
||||
idea_pipeline.extract_characters(story=story, quiet=True),
|
||||
runtime,
|
||||
{"session_id": session_id},
|
||||
)
|
||||
scene_scripts = await _run_planning_step(
|
||||
"Writing scene scripts from story",
|
||||
"write_script",
|
||||
idea_pipeline.write_script_based_on_story(story=story, user_requirement=user_requirement, quiet=True),
|
||||
runtime,
|
||||
{"session_id": session_id},
|
||||
)
|
||||
for idx, scene_script in enumerate(scene_scripts if isinstance(scene_scripts, list) else [scene_scripts]):
|
||||
scene_dir = idea_dir / f"scene_{idx}"
|
||||
scene_text = scene_script if isinstance(scene_script, str) else json.dumps(scene_script, ensure_ascii=False, indent=2)
|
||||
script_pipeline = Script2VideoPipeline(chat_model=chat_model, image_generator=dummy, video_generator=dummy, working_dir=str(scene_dir))
|
||||
await _run_planning_step(
|
||||
f"Planning scene {idx} storyboard and shots",
|
||||
"plan_scene",
|
||||
script_pipeline.plan_text_artifacts(script=scene_text, user_requirement=user_requirement, style=style, characters=characters, progress=_pipeline_progress(runtime, session_id, scene_index=idx), quiet=True),
|
||||
runtime,
|
||||
{"session_id": session_id, "scene_index": idx},
|
||||
)
|
||||
else:
|
||||
(script_dir / "script.txt").write_text(script, encoding="utf-8")
|
||||
script_pipeline = Script2VideoPipeline(chat_model=chat_model, image_generator=dummy, video_generator=dummy, working_dir=str(script_dir))
|
||||
if runtime:
|
||||
runtime.emit_progress("Script pipeline initialized", stage="script_pipeline_ready", metadata={"session_id": session_id})
|
||||
await asyncio.sleep(0)
|
||||
await _run_planning_step(
|
||||
"Planning storyboard and shots from provided script",
|
||||
"plan_script",
|
||||
script_pipeline.plan_text_artifacts(script=script, user_requirement=user_requirement, style=style, progress=_pipeline_progress(runtime, session_id), quiet=True),
|
||||
runtime,
|
||||
{"session_id": session_id},
|
||||
)
|
||||
except Exception as exc:
|
||||
self.session_index.update_stage(session_id, "error", f"Narrative planning failed: {exc}")
|
||||
checklist = self.session_index.artifact_checklist(session_id)
|
||||
payload = {
|
||||
"session_id": session_id,
|
||||
"working_dir": str(working_dir.relative_to(self.workspace_root)),
|
||||
"error_type": "recoverable_planning_step_failed",
|
||||
"retryable": True,
|
||||
"error": str(exc),
|
||||
"present": [path for path, present in checklist.items() if present],
|
||||
"missing": [path for path, present in checklist.items() if not present],
|
||||
}
|
||||
if runtime:
|
||||
runtime.emit_progress("Narrative planning failed; partial artifacts were kept", stage="planning_failed", metadata=payload)
|
||||
return ToolResult("vimax_narrative_planning", False, f"Narrative planning failed: {exc}", payload)
|
||||
|
||||
checklist = self.session_index.artifact_checklist(session_id)
|
||||
generated = [path for path, present in checklist.items() if present and not generated_before.get(path)]
|
||||
reused = [path for path, present in checklist.items() if present and generated_before.get(path)]
|
||||
ready_for_render = _ready_for_render(checklist)
|
||||
self.session_index.update_stage(session_id, "narrative_planned", "Structured text planning complete" if ready_for_render else "Structured text planning partially complete")
|
||||
if runtime:
|
||||
runtime.emit_progress("Narrative planning complete", stage="completed", metadata={"ready_for_render": ready_for_render})
|
||||
payload = {
|
||||
"session_id": session_id,
|
||||
"working_dir": str(working_dir.relative_to(self.workspace_root)),
|
||||
"generated": generated,
|
||||
"reused": reused,
|
||||
"missing": [path for path, present in checklist.items() if not present],
|
||||
"ready_for_render": ready_for_render,
|
||||
}
|
||||
return ToolResult("vimax_narrative_planning", True, json.dumps(payload, ensure_ascii=False, indent=2), payload)
|
||||
|
||||
async def _revise_narrative_artifact(self, session_id: str, working_dir: Path, revision_target: str, revision_instruction: str, runtime: ToolRuntimeContext | None = None) -> ToolResult:
|
||||
if not revision_instruction:
|
||||
self.session_index.update_stage(session_id, "error", "Revision failed: missing revision_instruction")
|
||||
return ToolResult("vimax_narrative_planning", False, "revision_instruction is required when revision_target is provided.", {"error_type": "missing_revision_instruction", "session_id": session_id, "revision_target": revision_target})
|
||||
try:
|
||||
target_path = _resolve_artifact_path(working_dir, revision_target)
|
||||
except ValueError as exc:
|
||||
self.session_index.update_stage(session_id, "error", f"Revision failed: {exc}")
|
||||
return ToolResult("vimax_narrative_planning", False, str(exc), {"error_type": "invalid_revision_target", "session_id": session_id, "revision_target": revision_target})
|
||||
if not target_path.exists():
|
||||
self.session_index.update_stage(session_id, "error", f"Revision failed: target does not exist: {revision_target}")
|
||||
return ToolResult("vimax_narrative_planning", False, f"Revision target does not exist: {revision_target}", {"error_type": "dependency_missing", "session_id": session_id, "revision_target": revision_target})
|
||||
try:
|
||||
self.session_index.update_stage(session_id, "narrative_planning", "Revising structured text artifact")
|
||||
if runtime:
|
||||
runtime.emit_progress("Revising structured text artifact", stage="revising", metadata={"session_id": session_id, "revision_target": revision_target})
|
||||
chat_model = _build_chat_model()
|
||||
before = target_path.read_text(encoding="utf-8")
|
||||
revised = await _revise_artifact_with_llm(chat_model, target_path.relative_to(working_dir).as_posix(), before, revision_instruction)
|
||||
if target_path.suffix == ".json":
|
||||
try:
|
||||
revised_payload = json.loads(revised)
|
||||
except json.JSONDecodeError as exc:
|
||||
self.session_index.update_stage(session_id, "error", f"Revision failed: invalid JSON output: {exc}")
|
||||
return ToolResult("vimax_narrative_planning", False, f"Revision output was not valid JSON: {exc}", {"error_type": "invalid_revision_json", "session_id": session_id, "revision_target": revision_target})
|
||||
revised = json.dumps(revised_payload, ensure_ascii=False, indent=2)
|
||||
target_path.write_text(revised, encoding="utf-8")
|
||||
except Exception as exc:
|
||||
self.session_index.update_stage(session_id, "error", f"Revision failed: {exc}")
|
||||
raise
|
||||
|
||||
stale = _stale_keys_for_revision(target_path.relative_to(working_dir).as_posix())
|
||||
if stale:
|
||||
self.session_index.mark_stale(session_id, stale)
|
||||
self.session_index.append_log("revisions", {"session_id": session_id, "target": target_path.relative_to(working_dir).as_posix(), "instruction": revision_instruction, "stale": stale, "before_preview": before[:500], "after_preview": revised[:500]})
|
||||
checklist = self.session_index.artifact_checklist(session_id)
|
||||
ready_for_render = _ready_for_render(checklist)
|
||||
self.session_index.update_stage(session_id, "narrative_planned" if ready_for_render else "narrative_planning", "Revised structured text artifact")
|
||||
payload = {
|
||||
"session_id": session_id,
|
||||
"working_dir": str(working_dir.relative_to(self.workspace_root)),
|
||||
"generated": [],
|
||||
"reused": [path for path, present in checklist.items() if present],
|
||||
"revised": [target_path.relative_to(working_dir).as_posix()],
|
||||
"missing": [path for path, present in checklist.items() if not present],
|
||||
"stale": stale,
|
||||
"ready_for_render": ready_for_render,
|
||||
"revision_target": target_path.relative_to(working_dir).as_posix(),
|
||||
}
|
||||
return ToolResult("vimax_narrative_planning", True, json.dumps(payload, ensure_ascii=False, indent=2), payload)
|
||||
|
||||
async def vimax_novel_planning(self, args: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
|
||||
novel_text = str(args.get("novel_text", "") or "").strip()
|
||||
user_requirement = str(args.get("user_requirement", "") or "").strip()
|
||||
style = str(args.get("style", "") or "").strip() or "Cinematic, coherent, 16:9"
|
||||
if not novel_text:
|
||||
return ToolResult("vimax_novel_planning", False, "novel_text is required for novel planning.", {"error_type": "missing_input"})
|
||||
|
||||
session_id_arg = str(args.get("session_id", "") or "").strip()
|
||||
session = self.session_index.create(idea=novel_text, user_requirement=user_requirement, style=style, session_id=session_id_arg or None)
|
||||
session_id = session["session_id"]
|
||||
working_dir = self.session_index.working_dir(session_id)
|
||||
novel_dir = working_dir / "novel2video"
|
||||
novel_dir.mkdir(parents=True, exist_ok=True)
|
||||
generated_before = self.session_index.artifact_checklist(session_id)
|
||||
|
||||
try:
|
||||
self.session_index.update_stage(session_id, "novel_planning", "Generating novel structured text artifacts")
|
||||
if runtime:
|
||||
runtime.emit_progress("Starting novel planning", stage="starting", metadata={"session_id": session_id})
|
||||
await asyncio.sleep(0)
|
||||
pipeline = _build_novel_pipeline(novel_dir)
|
||||
await _run_planning_step(
|
||||
"Planning novel structured text artifacts",
|
||||
"novel_plan_text_artifacts",
|
||||
pipeline.plan_text_artifacts(
|
||||
novel_text=novel_text,
|
||||
user_requirement=user_requirement,
|
||||
style=style,
|
||||
progress=_pipeline_progress(runtime, session_id),
|
||||
quiet=True,
|
||||
),
|
||||
runtime,
|
||||
{"session_id": session_id},
|
||||
)
|
||||
except Exception as exc:
|
||||
self.session_index.update_stage(session_id, "error", f"Novel planning failed: {exc}")
|
||||
return ToolResult("vimax_novel_planning", False, str(exc), {"error_type": "exception", "session_id": session_id})
|
||||
|
||||
checklist = self.session_index.artifact_checklist(session_id)
|
||||
generated = [path for path, present in checklist.items() if path.startswith("novel2video/") and present and not generated_before.get(path)]
|
||||
reused = [path for path, present in checklist.items() if path.startswith("novel2video/") and present and generated_before.get(path)]
|
||||
missing = [path for path, present in checklist.items() if path.startswith("novel2video/") and not present]
|
||||
ready = _novel_text_ready(checklist)
|
||||
self.session_index.update_stage(session_id, "novel_planned" if ready else "novel_planning", "Novel structured text planning complete" if ready else "Novel structured text planning partially complete")
|
||||
if runtime:
|
||||
runtime.emit_progress("Novel planning complete", stage="completed", metadata={"session_id": session_id, "ready_for_scene_render": False})
|
||||
payload = {
|
||||
"session_id": session_id,
|
||||
"working_dir": str(working_dir.relative_to(self.workspace_root)),
|
||||
"generated": generated,
|
||||
"reused": reused,
|
||||
"missing": missing,
|
||||
"ready_for_scene_render": False,
|
||||
}
|
||||
return ToolResult("vimax_novel_planning", True, json.dumps(payload, ensure_ascii=False, indent=2), payload)
|
||||
|
||||
async def vimax_render_video(self, args: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
|
||||
session_id = str(args.get("session_id", "") or "").strip()
|
||||
session = self.session_index.get(session_id) if session_id else self.session_index.active()
|
||||
if session is None:
|
||||
return ToolResult("vimax_render_video", False, "No active session to render.", {"error_type": "missing_session"})
|
||||
session_id = session["session_id"]
|
||||
checklist = self.session_index.artifact_checklist(session_id)
|
||||
missing = _missing_render_dependencies(checklist)
|
||||
working_dir = self.session_index.working_dir(session_id)
|
||||
if missing:
|
||||
payload = {"error_type": "dependency_missing", "missing": missing, "session_id": session_id}
|
||||
_write_render_status(working_dir, status="dependency_missing", payload=payload)
|
||||
return ToolResult("vimax_render_video", False, f"Dependency missing: {', '.join(missing)}", payload)
|
||||
|
||||
self.session_index.update_stage(session_id, "rendering", "Rendering video artifacts")
|
||||
_write_render_status(working_dir, status="rendering", payload={"session_id": session_id, "render_started": True, "render_completed": False})
|
||||
try:
|
||||
chat_model = _build_chat_model()
|
||||
image_generator = _build_image_generator()
|
||||
video_generator = _build_video_generator()
|
||||
if runtime:
|
||||
runtime.emit_progress("Starting video render", stage="rendering", metadata={"session_id": session_id})
|
||||
if _idea_mode_ready(checklist):
|
||||
idea_pipeline = Idea2VideoPipeline(chat_model=chat_model, image_generator=image_generator, video_generator=video_generator, working_dir=str(working_dir / "idea2video"))
|
||||
with _suppress_pipeline_output():
|
||||
final_video = await idea_pipeline(idea=str(session.get("idea", "")), user_requirement=str(session.get("user_requirement", "")), style=str(session.get("style", "")), quiet=True)
|
||||
self.session_index.update_stage(session_id, "rendered", "Final video rendered")
|
||||
payload = {"session_id": session_id, "render_mode": "idea2video", "render_started": True, "render_completed": True, "final_video_path": str(Path(final_video).relative_to(self.workspace_root)), "missing": []}
|
||||
_write_render_status(working_dir, status="rendered", payload=payload)
|
||||
return ToolResult("vimax_render_video", True, json.dumps(payload, ensure_ascii=False, indent=2), payload)
|
||||
if _script_mode_ready(checklist):
|
||||
script_dir = working_dir / "script2video"
|
||||
script_text = _load_script_text(working_dir)
|
||||
characters = _load_characters(script_dir / "characters.json")
|
||||
pipeline = Script2VideoPipeline(chat_model=chat_model, image_generator=image_generator, video_generator=video_generator, working_dir=str(script_dir))
|
||||
with _suppress_pipeline_output():
|
||||
final_video = await pipeline(script=script_text, user_requirement=str(session.get("user_requirement", "")), style=str(session.get("style", "")), characters=characters, quiet=True, progress=_pipeline_progress(runtime, session_id))
|
||||
self.session_index.update_stage(session_id, "rendered", "Final video rendered")
|
||||
payload = {"session_id": session_id, "render_mode": "script2video", "render_started": True, "render_completed": True, "final_video_path": str(Path(final_video).relative_to(self.workspace_root)), "missing": []}
|
||||
_write_render_status(working_dir, status="rendered", payload=payload)
|
||||
return ToolResult("vimax_render_video", True, json.dumps(payload, ensure_ascii=False, indent=2), payload)
|
||||
if _novel_mode_ready(checklist):
|
||||
novel_dir = working_dir / "novel2video"
|
||||
pipeline = _build_novel_render_pipeline(novel_dir, chat_model, image_generator, video_generator)
|
||||
with _suppress_pipeline_output():
|
||||
render_result = await pipeline.render_video_artifacts(style=str(session.get("style", "")), user_requirement=str(session.get("user_requirement", "")), quiet=True, progress=_pipeline_progress(runtime, session_id))
|
||||
scene_videos_dir = Path(render_result["scene_videos_dir"])
|
||||
self.session_index.update_stage(session_id, "novel_scene_rendered", "Novel scene videos rendered")
|
||||
payload = {
|
||||
"session_id": session_id,
|
||||
"render_mode": "novel2video",
|
||||
"render_started": True,
|
||||
"render_completed": True,
|
||||
"scene_render_completed": True,
|
||||
"final_video_path": None,
|
||||
"scene_videos_dir": str(scene_videos_dir.relative_to(self.workspace_root)),
|
||||
"scene_video_dirs": [str(Path(path).relative_to(self.workspace_root)) for path in render_result.get("scene_video_dirs", [])],
|
||||
"scene_count": render_result.get("scene_count", 0),
|
||||
"missing": [],
|
||||
}
|
||||
_write_render_status(working_dir, status="rendered", payload=payload)
|
||||
return ToolResult("vimax_render_video", True, json.dumps(payload, ensure_ascii=False, indent=2), payload)
|
||||
except Exception as exc:
|
||||
unwrapped = _unwrap_retry_error(exc)
|
||||
error_text = _sanitize_error_text(str(unwrapped))
|
||||
wrapped_error_text = _sanitize_error_text(str(exc))
|
||||
self.session_index.update_stage(session_id, "error", f"Render failed: {error_text}")
|
||||
checklist = self.session_index.artifact_checklist(session_id)
|
||||
payload = {
|
||||
"error_type": "render_failed",
|
||||
"retryable": _is_retryable_render_error(unwrapped),
|
||||
"session_id": session_id,
|
||||
"error": error_text,
|
||||
"wrapped_error": wrapped_error_text,
|
||||
"present": [path for path, present in checklist.items() if present],
|
||||
"missing": [path for path, present in checklist.items() if not present],
|
||||
}
|
||||
_write_render_status(working_dir, status="error", payload=payload)
|
||||
if runtime:
|
||||
runtime.emit_progress("Render failed; partial artifacts were kept", stage="render_failed", metadata=payload)
|
||||
return ToolResult("vimax_render_video", False, f"Render failed: {error_text}", payload)
|
||||
payload = {"error_type": "dependency_missing", "session_id": session_id}
|
||||
_write_render_status(working_dir, status="dependency_missing", payload=payload)
|
||||
return ToolResult("vimax_render_video", False, "No render mode matched current session.", payload)
|
||||
|
||||
def _resolve_session(self, session_id: str, *, idea: str, script: str, user_requirement: str, style: str) -> dict[str, Any]:
|
||||
requested_source = idea or script
|
||||
if session_id:
|
||||
session = self.session_index.get(session_id)
|
||||
if session is None:
|
||||
session = self.session_index.create(idea=requested_source, user_requirement=user_requirement, style=style, session_id=session_id)
|
||||
elif requested_source and _is_new_source_for_session(session, requested_source):
|
||||
session = self.session_index.create(idea=requested_source, user_requirement=user_requirement, style=style)
|
||||
else:
|
||||
self.session_index.set_active(session_id)
|
||||
else:
|
||||
if requested_source:
|
||||
session = self.session_index.create(idea=requested_source, user_requirement=user_requirement, style=style)
|
||||
else:
|
||||
session = self.session_index.active() or self.session_index.create(idea=requested_source, user_requirement=user_requirement, style=style)
|
||||
self._update_session_metadata(session["session_id"], idea=requested_source, user_requirement=user_requirement, style=style)
|
||||
return self.session_index.get(session["session_id"]) or session
|
||||
|
||||
def _update_session_metadata(self, session_id: str, *, idea: str, user_requirement: str, style: str) -> None:
|
||||
data = self.session_index.load()
|
||||
record = data.get("sessions", {}).get(session_id)
|
||||
if not isinstance(record, dict):
|
||||
return
|
||||
if idea and not record.get("idea"):
|
||||
record["idea"] = idea
|
||||
if user_requirement:
|
||||
record["user_requirement"] = user_requirement
|
||||
if style:
|
||||
record["style"] = style
|
||||
self.session_index.save(data)
|
||||
|
||||
|
||||
class _DiscardStream:
|
||||
def write(self, text: str) -> int:
|
||||
return len(text)
|
||||
|
||||
def flush(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
_PIPELINE_OUTPUT_SINK = _DiscardStream()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _suppress_pipeline_output():
|
||||
previous_disable_level = logging.root.manager.disable
|
||||
logging.disable(logging.WARNING)
|
||||
try:
|
||||
with redirect_stdout(_PIPELINE_OUTPUT_SINK), redirect_stderr(_PIPELINE_OUTPUT_SINK):
|
||||
yield
|
||||
finally:
|
||||
logging.disable(previous_disable_level)
|
||||
|
||||
|
||||
def _narrative_step_timeout_seconds() -> float:
|
||||
raw = os.environ.get("VIMAX_NARRATIVE_STEP_TIMEOUT_SECONDS", "900")
|
||||
try:
|
||||
return max(0.0, float(raw))
|
||||
except ValueError:
|
||||
return 900.0
|
||||
|
||||
|
||||
async def _run_planning_step(
|
||||
message: str,
|
||||
stage: str,
|
||||
awaitable: Any,
|
||||
runtime: ToolRuntimeContext | None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
timeout_seconds = _narrative_step_timeout_seconds()
|
||||
event_metadata = dict(metadata or {})
|
||||
event_metadata["timeout_seconds"] = timeout_seconds
|
||||
if runtime:
|
||||
runtime.emit_progress(message, stage=stage, metadata=event_metadata)
|
||||
await asyncio.sleep(0)
|
||||
try:
|
||||
with _suppress_pipeline_output():
|
||||
if timeout_seconds <= 0:
|
||||
return await awaitable
|
||||
return await asyncio.wait_for(awaitable, timeout=timeout_seconds)
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise RuntimeError(f"{message} timed out after {timeout_seconds:g}s") from exc
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"{message} failed: {exc}") from exc
|
||||
|
||||
|
||||
def _is_new_source_for_session(session: dict[str, Any], requested_source: str) -> bool:
|
||||
current = str(session.get("idea") or "").strip()
|
||||
requested = requested_source.strip()
|
||||
if not current or not requested:
|
||||
return False
|
||||
return current != requested
|
||||
|
||||
|
||||
def _llm_request_timeout_seconds() -> float:
|
||||
raw = os.environ.get("VIMAX_LLM_REQUEST_TIMEOUT_SECONDS", "300")
|
||||
try:
|
||||
return max(1.0, float(raw))
|
||||
except ValueError:
|
||||
return 300.0
|
||||
|
||||
|
||||
def _narrative_max_tokens() -> int:
|
||||
raw = os.environ.get("VIMAX_NARRATIVE_MAX_TOKENS", "4096")
|
||||
try:
|
||||
return max(256, int(raw))
|
||||
except ValueError:
|
||||
return 4096
|
||||
|
||||
|
||||
def _pipeline_progress(runtime: ToolRuntimeContext | None, session_id: str, *, scene_index: int | None = None):
|
||||
if runtime is None:
|
||||
return None
|
||||
|
||||
def emit(stage: str, message: str, metadata: dict[str, Any] | None = None) -> None:
|
||||
payload = dict(metadata or {})
|
||||
payload["session_id"] = session_id
|
||||
if scene_index is not None:
|
||||
payload["scene_index"] = scene_index
|
||||
runtime.emit_progress(message, stage=stage, metadata=payload)
|
||||
|
||||
return emit
|
||||
|
||||
|
||||
def _build_chat_model() -> Any:
|
||||
api_key = llm_api_key()
|
||||
if not api_key:
|
||||
raise RuntimeError("VIMAX_LLM_API_KEY or configs/agent.local.yaml llm.api_key is required for narrative planning")
|
||||
return init_chat_model(
|
||||
model=llm_model(),
|
||||
model_provider=llm_model_provider(),
|
||||
api_key=api_key,
|
||||
base_url=llm_base_url(),
|
||||
timeout=_llm_request_timeout_seconds(),
|
||||
max_retries=0,
|
||||
max_completion_tokens=_narrative_max_tokens(),
|
||||
)
|
||||
|
||||
|
||||
def _build_image_generator() -> ImageGeneratorNanobananaYunwuAPI:
|
||||
api_key = image_api_key()
|
||||
if not api_key:
|
||||
raise RuntimeError("VIMAX_IMAGE_API_KEY, VIMAX_LLM_API_KEY, or configs/agent.local.yaml image/llm api_key is required for image generation")
|
||||
return ImageGeneratorNanobananaYunwuAPI(api_key=api_key, model=image_model(), base_url=image_base_url())
|
||||
|
||||
|
||||
def _build_video_generator() -> VideoGeneratorVeoYunwuAPI | VideoGeneratorOpenRouterAPI:
|
||||
api_key = video_api_key()
|
||||
if not api_key:
|
||||
raise RuntimeError("VIMAX_VIDEO_API_KEY, VIMAX_LLM_API_KEY, or configs/agent.local.yaml video/llm api_key is required for video generation")
|
||||
model = video_model()
|
||||
base_url = video_base_url()
|
||||
provider = video_provider().strip().lower()
|
||||
if provider == "openrouter":
|
||||
return VideoGeneratorOpenRouterAPI(api_key=api_key, model=model, base_url=base_url)
|
||||
if provider == "yunwu":
|
||||
return VideoGeneratorVeoYunwuAPI(api_key=api_key, t2v_model=model, ff2v_model=model, base_url=base_url)
|
||||
raise RuntimeError(f"Unsupported video base_url for automatic provider matching: {base_url}")
|
||||
|
||||
|
||||
class _IdentityRewriter:
|
||||
async def __call__(self, prompt: str) -> str:
|
||||
return prompt
|
||||
|
||||
|
||||
def _build_embedding_model() -> Any:
|
||||
api_key = embedding_api_key()
|
||||
base_url = embedding_base_url()
|
||||
provider = embedding_model_provider().strip().lower()
|
||||
if not api_key or not base_url:
|
||||
raise RuntimeError("VIMAX_EMBEDDING_API_KEY or configs/agent.local.yaml embedding api_key/base_url is required for novel planning")
|
||||
if provider != "openai":
|
||||
raise RuntimeError(f"Unsupported embedding model_provider: {provider}")
|
||||
return OpenAIEmbeddings(model=embedding_model(), api_key=api_key, base_url=base_url)
|
||||
|
||||
|
||||
def _build_reranker() -> RerankerBgeSiliconapi:
|
||||
api_key = reranker_api_key()
|
||||
base_url = reranker_base_url()
|
||||
if not api_key or not base_url:
|
||||
raise RuntimeError("VIMAX_RERANKER_API_KEY or configs/agent.local.yaml reranker api_key/base_url is required for novel planning")
|
||||
return RerankerBgeSiliconapi(api_key=api_key, base_url=base_url, model=reranker_model())
|
||||
|
||||
|
||||
def _build_novel_pipeline(working_dir: Path) -> Novel2MoviePipeline:
|
||||
api_key = llm_api_key()
|
||||
if not api_key:
|
||||
raise RuntimeError("VIMAX_LLM_API_KEY or configs/agent.local.yaml llm.api_key is required for novel planning")
|
||||
base_url = llm_base_url()
|
||||
model = llm_model()
|
||||
dummy = _UnavailableGenerator()
|
||||
return Novel2MoviePipeline(
|
||||
novel_compressor=NovelCompressor(api_key=api_key, base_url=base_url, chat_model=model),
|
||||
event_extractor=EventExtractor(api_key=api_key, base_url=base_url, chat_model=model),
|
||||
embeddings=_build_embedding_model(),
|
||||
rerank_model=_build_reranker(),
|
||||
scene_extractor=SceneExtractor(api_key=api_key, base_url=base_url, chat_model=model),
|
||||
global_information_planner=GlobalInformationPlanner(api_key=api_key, base_url=base_url, chat_model=model),
|
||||
image_generator=dummy,
|
||||
rewriter=_IdentityRewriter(),
|
||||
script2video_pipeline=dummy,
|
||||
working_dir=str(working_dir),
|
||||
)
|
||||
|
||||
|
||||
def _build_novel_render_pipeline(working_dir: Path, chat_model: Any, image_generator: Any, video_generator: Any) -> Novel2MoviePipeline:
|
||||
api_key = llm_api_key()
|
||||
if not api_key:
|
||||
raise RuntimeError("VIMAX_LLM_API_KEY or configs/agent.local.yaml llm.api_key is required for novel rendering")
|
||||
base_url = llm_base_url()
|
||||
model = llm_model()
|
||||
script_pipeline = Script2VideoPipeline(chat_model=chat_model, image_generator=image_generator, video_generator=video_generator, working_dir=str(working_dir / "videos"))
|
||||
return Novel2MoviePipeline(
|
||||
novel_compressor=NovelCompressor(api_key=api_key, base_url=base_url, chat_model=model),
|
||||
event_extractor=EventExtractor(api_key=api_key, base_url=base_url, chat_model=model),
|
||||
embeddings=_build_embedding_model(),
|
||||
rerank_model=_build_reranker(),
|
||||
scene_extractor=SceneExtractor(api_key=api_key, base_url=base_url, chat_model=model),
|
||||
global_information_planner=GlobalInformationPlanner(api_key=api_key, base_url=base_url, chat_model=model),
|
||||
image_generator=image_generator,
|
||||
rewriter=_IdentityRewriter(),
|
||||
script2video_pipeline=script_pipeline,
|
||||
working_dir=str(working_dir),
|
||||
)
|
||||
|
||||
|
||||
def _unwrap_retry_error(exc: Exception) -> Exception:
|
||||
if isinstance(exc, RetryError):
|
||||
try:
|
||||
return exc.last_attempt.exception() or exc
|
||||
except Exception:
|
||||
return exc
|
||||
return exc
|
||||
|
||||
|
||||
def _is_retryable_render_error(exc: Exception) -> bool:
|
||||
text = str(exc).lower()
|
||||
if isinstance(exc, AttributeError):
|
||||
return False
|
||||
if "http 403" in text or "key limit exceeded" in text or "quota" in text:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _sanitize_error_text(text: str) -> str:
|
||||
sanitized = text
|
||||
for marker in ("workspaces/default/keys/",):
|
||||
if marker in sanitized:
|
||||
prefix, rest = sanitized.split(marker, 1)
|
||||
key_id = []
|
||||
for char in rest:
|
||||
if char.isalnum() or char in "-_":
|
||||
key_id.append(char)
|
||||
continue
|
||||
break
|
||||
sanitized = prefix + marker + "<redacted>" + rest[len(key_id):]
|
||||
if "sk-" in sanitized:
|
||||
prefix, rest = sanitized.split("sk-", 1)
|
||||
token = []
|
||||
for char in rest:
|
||||
if char.isalnum() or char in "-_":
|
||||
token.append(char)
|
||||
continue
|
||||
break
|
||||
sanitized = prefix + "sk-<redacted>" + rest[len(token):]
|
||||
return sanitized
|
||||
|
||||
|
||||
def _write_render_status(working_dir: Path, *, status: str, payload: dict[str, Any]) -> None:
|
||||
working_dir.mkdir(parents=True, exist_ok=True)
|
||||
event = {
|
||||
"timestamp": datetime.now().isoformat(timespec="seconds"),
|
||||
"status": status,
|
||||
**payload,
|
||||
}
|
||||
(working_dir / "render_status.json").write_text(json.dumps(event, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
with (working_dir / "render_events.jsonl").open("a", encoding="utf-8") as handle:
|
||||
handle.write(json.dumps(event, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
def _write_characters_if_missing(path: Path, characters: list[CharacterInScene]) -> None:
|
||||
if path.exists():
|
||||
return
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps([character.model_dump() for character in characters], ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
def _load_characters(path: Path) -> list[CharacterInScene]:
|
||||
return [CharacterInScene.model_validate(item) for item in json.loads(path.read_text(encoding="utf-8"))]
|
||||
|
||||
|
||||
def _load_script_text(working_dir: Path) -> str:
|
||||
script_text = working_dir / "script2video" / "script.txt"
|
||||
if script_text.exists():
|
||||
return script_text.read_text(encoding="utf-8")
|
||||
idea_script = working_dir / "idea2video" / "script.json"
|
||||
if idea_script.exists():
|
||||
payload = json.loads(idea_script.read_text(encoding="utf-8"))
|
||||
return json.dumps(payload, ensure_ascii=False, indent=2) if not isinstance(payload, str) else payload
|
||||
story = working_dir / "idea2video" / "story.txt"
|
||||
if story.exists():
|
||||
return story.read_text(encoding="utf-8")
|
||||
return ""
|
||||
|
||||
|
||||
def _resolve_artifact_path(working_dir: Path, revision_target: str) -> Path:
|
||||
rel = Path(revision_target)
|
||||
if rel.is_absolute():
|
||||
raise ValueError(f"revision_target must be relative to session working_dir: {revision_target}")
|
||||
path = (working_dir / rel).resolve()
|
||||
if path != working_dir and working_dir not in path.parents:
|
||||
raise ValueError(f"revision_target escapes session working_dir: {revision_target}")
|
||||
return path
|
||||
|
||||
|
||||
async def _revise_artifact_with_llm(chat_model: Any, target: str, current_text: str, instruction: str) -> str:
|
||||
prompt = (
|
||||
"Revise this ViMax structured artifact exactly as requested. "
|
||||
"Return only the complete replacement file content, with no Markdown fences or explanation. "
|
||||
"If the file is JSON, preserve valid JSON and the existing schema shape.\n\n"
|
||||
f"Target: {target}\n"
|
||||
f"Revision instruction: {instruction}\n\n"
|
||||
"Current file content:\n"
|
||||
f"{current_text}"
|
||||
)
|
||||
if hasattr(chat_model, "ainvoke"):
|
||||
response = await chat_model.ainvoke(prompt)
|
||||
elif hasattr(chat_model, "invoke"):
|
||||
response = chat_model.invoke(prompt)
|
||||
else:
|
||||
raise RuntimeError("chat_model does not support invoke/ainvoke for revision mode")
|
||||
content = getattr(response, "content", response)
|
||||
if isinstance(content, list):
|
||||
content = "".join(str(item.get("text", item)) if isinstance(item, dict) else str(item) for item in content)
|
||||
return _strip_markdown_fences(str(content).strip())
|
||||
|
||||
|
||||
def _strip_markdown_fences(text: str) -> str:
|
||||
if not text.startswith("```"):
|
||||
return text
|
||||
lines = text.splitlines()
|
||||
if lines and lines[0].startswith("```"):
|
||||
lines = lines[1:]
|
||||
if lines and lines[-1].strip() == "```":
|
||||
lines = lines[:-1]
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
|
||||
def _stale_keys_for_revision(target: str) -> list[str]:
|
||||
if "storyboard.json" in target:
|
||||
return ["shot_descriptions", "camera_tree", "frames", "clips", "final_video"]
|
||||
if "shot_description.json" in target:
|
||||
return ["frames", "clips", "final_video"]
|
||||
if "camera_tree.json" in target:
|
||||
return ["frames", "clips", "final_video"]
|
||||
if target.endswith("script.json") or target.endswith("story.txt"):
|
||||
return ["storyboard", "shot_descriptions", "camera_tree", "frames", "clips", "final_video"]
|
||||
if target.endswith("characters.json"):
|
||||
return ["storyboard", "shot_descriptions", "frames", "clips", "final_video"]
|
||||
return ["frames", "clips", "final_video"]
|
||||
|
||||
|
||||
def _ready_for_render(checklist: dict[str, bool]) -> bool:
|
||||
return _idea_mode_ready(checklist) or _script_mode_ready(checklist) or _novel_mode_ready(checklist)
|
||||
|
||||
|
||||
def _missing_render_dependencies(checklist: dict[str, bool]) -> list[str]:
|
||||
if _ready_for_render(checklist):
|
||||
return []
|
||||
idea_required = ["idea2video/story.txt", "idea2video/characters.json", "idea2video/script.json", "idea2video/scene_*/storyboard.json", "idea2video/scene_*/shots/*/shot_description.json", "idea2video/scene_*/camera_tree.json"]
|
||||
script_required = ["script2video/script.txt", "script2video/characters.json", "script2video/storyboard.json", "script2video/shots/*/shot_description.json", "script2video/camera_tree.json"]
|
||||
novel_required = ["novel2video/novel/novel_compressed.txt", "novel2video/events/event_*.json", "novel2video/relevant_chunks/event_*", "novel2video/scenes/event_*/scene_*.json", "novel2video/global_information/characters/event_level/*.json", "novel2video/global_information/characters/novel_level/*.json"]
|
||||
return [f"idea mode: {path}" for path in idea_required if not checklist.get(path)] + [f"script mode: {path}" for path in script_required if not checklist.get(path)] + [f"novel mode: {path}" for path in novel_required if not checklist.get(path)]
|
||||
|
||||
|
||||
def _idea_mode_ready(checklist: dict[str, bool]) -> bool:
|
||||
return bool(checklist.get("idea2video/story.txt") and checklist.get("idea2video/characters.json") and checklist.get("idea2video/script.json") and checklist.get("idea2video/scene_*/storyboard.json") and checklist.get("idea2video/scene_*/shots/*/shot_description.json") and checklist.get("idea2video/scene_*/camera_tree.json"))
|
||||
|
||||
|
||||
def _novel_text_ready(checklist: dict[str, bool]) -> bool:
|
||||
return _novel_mode_ready(checklist)
|
||||
|
||||
|
||||
def _novel_mode_ready(checklist: dict[str, bool]) -> bool:
|
||||
return bool(checklist.get("novel2video/novel/novel_compressed.txt") and checklist.get("novel2video/events/event_*.json") and checklist.get("novel2video/relevant_chunks/event_*") and checklist.get("novel2video/scenes/event_*/scene_*.json") and checklist.get("novel2video/global_information/characters/event_level/*.json") and checklist.get("novel2video/global_information/characters/novel_level/*.json"))
|
||||
|
||||
|
||||
def _script_mode_ready(checklist: dict[str, bool]) -> bool:
|
||||
return bool(checklist.get("script2video/script.txt") and checklist.get("script2video/characters.json") and checklist.get("script2video/storyboard.json") and checklist.get("script2video/shots/*/shot_description.json") and checklist.get("script2video/camera_tree.json"))
|
||||
@@ -0,0 +1,15 @@
|
||||
from .screenwriter import Screenwriter
|
||||
from .storyboard_artist import StoryboardArtist
|
||||
from .camera_image_generator import CameraImageGenerator
|
||||
from .character_extractor import CharacterExtractor
|
||||
from .character_portraits_generator import CharacterPortraitsGenerator
|
||||
from .reference_image_selector import ReferenceImageSelector
|
||||
|
||||
__all__ = [
|
||||
"Screenwriter",
|
||||
"StoryboardArtist",
|
||||
"CameraImageGenerator",
|
||||
"CharacterExtractor",
|
||||
"CharacterPortraitsGenerator",
|
||||
"ReferenceImageSelector",
|
||||
]
|
||||
@@ -0,0 +1,147 @@
|
||||
import logging
|
||||
from typing import List, Tuple
|
||||
from pydantic import BaseModel, Field
|
||||
from tenacity import retry, stop_after_attempt
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain.chat_models import init_chat_model
|
||||
from utils.image import image_path_to_b64
|
||||
|
||||
|
||||
|
||||
system_prompt_template_select_most_consistent_image = \
|
||||
"""
|
||||
[Role]
|
||||
You are a professional visual assessment expert. Your expertise includes identifying Character Consistency and Spatial Consistency between candidate image and reference image, and assessing semantic consistency between candidate image and text description.
|
||||
|
||||
[Task]
|
||||
Based on the reference image provided by the user, the text description of the target image, and several candidate images, evaluate which candidate image performs best in the following aspects:
|
||||
- Character Consistency: Whether the character features (a. gender, b.ethnicity, c.age, d.facial features, e.body shape, f.outlook, g. hairstyle) in the candidate image align with those of the character in the reference image.
|
||||
- Spatial Consistency: Whether the relative positions between characters (e.g. Character A is on the left, character B is on the right, scene layout, perspective, and other spatial relationships) in the candidate image are consistent with those in the reference image.
|
||||
- Description Accuracy: Whether the candidate image accurately reflects the content described in the text (Note: The text description describes the target image we want, which is not an editing instruction).
|
||||
|
||||
[Input]
|
||||
The user will provide the following content:
|
||||
- Reference images: These include images of characters or other perspectives, each along with a brief text description. For example, "Reference Image 0: A young girl with long brown hair wearing a red dress." then follow the corresponding image. The index starts from 0.
|
||||
- Candidate images: The candidate images to be evaluated. For example, "Generated Image 0", then follow a generated image. The index starts from 0.
|
||||
- Text description for target image: This describes what the generated image should contain. It is enclosed <TARGET_DESCRIPTION_START> and <TARGET_DESCRIPTION_END> tags.
|
||||
|
||||
[Output]
|
||||
{format_instructions}
|
||||
|
||||
[Guidelines]
|
||||
- Prioritize Character Consistency: Ensure that the characters in the generated image are highly consistent with those in the reference image in terms of visual features (e.g., a. gender b.ethnicity, c.age, d.facial features, e.body shape, f.outlook, g. hairstyle etc.).
|
||||
- Focus on Spatial Consistency: Verify whether the relative positions of characters, object arrangements, and perspectives align logically with the reference image (e.g., if Character A is on the left and Character B is on the right in the reference image, the generated image should not reverse this).
|
||||
- Strictly Compare with Text Description: The generated image must adhere to key elements in the text description (e.g., actions, scenes, objects, etc.), while disregarding parts related to editing instructions (as the input description reflects the expected outcome rather than directives).
|
||||
- If multiple images partially meet the criteria, select the one with the highest overall consistency; if none are ideal, choose the relatively best option and explain its shortcomings.
|
||||
- Ensure the key elements described in the text are present in the selected image.
|
||||
- Avoid subjective preferences; base all analysis on objective comparisons.
|
||||
- Prioritize images without white borders, black edges, or any additional framing.
|
||||
"""
|
||||
|
||||
human_prompt_template_select_most_consistent_image = \
|
||||
"""
|
||||
<TARGET_DESCRIPTION_START>
|
||||
{target_description}
|
||||
<TARGET_DESCRIPTION_END>
|
||||
"""
|
||||
|
||||
|
||||
class BestImageResponse(BaseModel):
|
||||
best_image_index: int = Field(
|
||||
...,
|
||||
description="The index of the best image."
|
||||
)
|
||||
reason: str = Field(
|
||||
...,
|
||||
description="The reason why the image is the best."
|
||||
)
|
||||
|
||||
|
||||
class BestImageSelector:
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
chat_model: str,
|
||||
):
|
||||
|
||||
self.chat_model = init_chat_model(
|
||||
model=chat_model,
|
||||
model_provider="openai",
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
after=lambda retry_state: logging.warning(f"Retrying best image selection due to {retry_state.outcome.exception()}"),
|
||||
)
|
||||
async def __call__(
|
||||
self,
|
||||
reference_image_path_and_text_pairs: List[Tuple[str, str]],
|
||||
target_description: str,
|
||||
candidate_image_paths: List[str],
|
||||
) -> str:
|
||||
"""
|
||||
Args:
|
||||
ref_image_path_and_text_pairs:
|
||||
A list of tuples containing reference image paths and their descriptions.
|
||||
|
||||
target_description:
|
||||
The description of the target image.
|
||||
|
||||
candidate_image_paths:
|
||||
A list of paths to the candidate images to be evaluated.
|
||||
"""
|
||||
|
||||
if not candidate_image_paths:
|
||||
logging.warning("No candidate images provided; skipping best image selection")
|
||||
raise ValueError("No candidate images to select from")
|
||||
|
||||
logging.info(f"Selecting the best image from candidates: {candidate_image_paths}")
|
||||
|
||||
human_content = []
|
||||
for idx, (ref_image_path, text) in enumerate(reference_image_path_and_text_pairs):
|
||||
human_content.append({
|
||||
"type": "text",
|
||||
"text": f"Reference Image {idx}: {text}"
|
||||
})
|
||||
human_content.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_path_to_b64(ref_image_path, mime=True)}
|
||||
})
|
||||
|
||||
for idx, candidate_image_path in enumerate(candidate_image_paths):
|
||||
human_content.append({
|
||||
"type": "text",
|
||||
"text": f"Candidate Image {idx}"
|
||||
})
|
||||
human_content.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_path_to_b64(candidate_image_path, mime=True)}
|
||||
})
|
||||
human_content.append({
|
||||
"type": "text",
|
||||
"text": human_prompt_template_select_most_consistent_image.format(target_description=target_description)
|
||||
})
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=BestImageResponse)
|
||||
|
||||
messages = [
|
||||
SystemMessage(content=system_prompt_template_select_most_consistent_image.format(format_instructions=parser.get_format_instructions())),
|
||||
HumanMessage(content=human_content)
|
||||
]
|
||||
|
||||
chain = self.chat_model | parser
|
||||
|
||||
response = await chain.ainvoke(messages)
|
||||
idx = response.best_image_index
|
||||
if not isinstance(idx, int) or idx < 0 or idx >= len(candidate_image_paths):
|
||||
logging.warning(f"Received invalid best_image_index={idx}; defaulting to 0")
|
||||
idx = 0
|
||||
best_image_path = candidate_image_paths[idx]
|
||||
logging.info(f"Best image selected: {best_image_path}")
|
||||
logging.info(f"Selection reason: {response.reason}")
|
||||
return best_image_path
|
||||
@@ -0,0 +1,272 @@
|
||||
import os
|
||||
import logging
|
||||
import cv2
|
||||
from typing import List, Tuple, Union, Optional
|
||||
from pydantic import BaseModel, Field
|
||||
from tenacity import retry, stop_after_attempt
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from scenedetect import open_video, SceneManager, split_video_ffmpeg
|
||||
from scenedetect.detectors import ContentDetector
|
||||
|
||||
from interfaces import ShotDescription, ShotBriefDescription, Camera, ImageOutput, VideoOutput
|
||||
|
||||
|
||||
from moviepy import VideoFileClip
|
||||
from PIL import Image
|
||||
|
||||
|
||||
system_prompt_template_select_reference_camera = \
|
||||
"""
|
||||
[Role]
|
||||
You are a professional video editing expert specializing in multi-camera shot analysis and scene structure modeling. You have deep knowledge of cinematic language, enabling you to understand shot sizes (e.g., wide shot, medium shot, close-up) and content inclusion relationships. You can infer hierarchical structures between camera positions based on corresponding shot descriptions.
|
||||
|
||||
[Task]
|
||||
Your task is to analyze the input camera position data to construct a "camera position tree". This tree structure represents a relationship where a parent camera's content encompasses that of a child camera. Specifically, you need to identify the parent camera for each camera position (if one exists) and determine the dependent shot indices (i.e., the specific shots within the parent camera's footage that contain the child camera's content). If a camera position has no parent, output None.
|
||||
|
||||
[Input]
|
||||
The input is a sequence of cameras. The sequence will be enclosed within <CAMERA_SEQ> and </CAMERA_SEQ>.
|
||||
Each camera contains a sequence of shots filmed by the camera, which will be enclosed within <CAMERA_N> and </CAMERA_N>, where N is the index of the camera.
|
||||
|
||||
Below is an example of the input format:
|
||||
|
||||
<CAMERA_SEQ>
|
||||
<CAMERA_0>
|
||||
Shot 0: Medium shot of the street. Alice and Bob are walking towards each other.
|
||||
Shot 2: Medium shot of the street. Alice and Bob hug each other.
|
||||
</CAMERA_0>
|
||||
<CAMERA_1>
|
||||
Shot 1: Close-up of the Alice's face. Her expression shifts from surprise to delight as she recognizes Bob.
|
||||
</CAMERA_1>
|
||||
</CAMERA_SEQ>
|
||||
|
||||
|
||||
[Output]
|
||||
{format_instructions}
|
||||
|
||||
[Guidelines]
|
||||
- The language of all output values (not include keys) should be consistent with the language of the input.
|
||||
- Content Inclusion Check: The parent camera should as fully as possible contain the child camera's content in certain shots (e.g., a parent medium two-shot encompasses a child over-the-shoulder reverse shot). Analyze shot descriptions by comparing keywords (e.g., characters, actions, setting) to ensure the parent shot's field of view covers the child shot's.
|
||||
- Transition Smoothness Priority: Larger shot size as parent camera is preferred, such as Wide Shot -> Medium Shot or Medium Shot -> Close-up. The shot sizes of adjacent parent and child nodes should be as similar as possible. A direct transition from a long shot to a close-up is not allowed unless absolutely necessary.
|
||||
- Temporal Proximity: Each camera is described by its corresponding first shot, and the parent camera is located based on the description of the first shot. The shot index of the parent camera should be as close as possible to the first shot index of the child camera.
|
||||
- Logical Consistency: The camera tree should be acyclic, avoid circular dependencies. If a camera is contained by multiple potential parents, select the best match (based on shot size and content). If there is no suitable parent camera, output None.
|
||||
- When a broader perspective is not available, choose the shot with the largest overlapping field of view as the parent (the one with the most information overlap), or a shot can also serve as the parent of a reverse shot. When two cameras can be the parent of each other, choose the one with the smaller index as the parent of the camera with the larger index.
|
||||
- Only one camera can exist without a parent.
|
||||
- When describing the elements lost in a shot, carefully compare the details between the parent shot and the child shot. For example, the parent shot is a medium shot of Character A and Character B facing each other (both in profile to the camera), while the child shot is a close-up of Character A (with Character A facing the camera directly). In this case, the child shot lacks the frontal view information of Character A.
|
||||
- The first camera must be the root of the camera tree.
|
||||
"""
|
||||
|
||||
|
||||
human_prompt_template_select_reference_camera = \
|
||||
"""
|
||||
<CAMERA_SEQ>
|
||||
{camera_seq_str}
|
||||
</CAMERA_SEQ>
|
||||
"""
|
||||
|
||||
|
||||
class CameraParentItem(BaseModel):
|
||||
parent_cam_idx: Optional[int] = Field(
|
||||
default=None,
|
||||
description="The index of the parent camera. Set to None if the camera has no parent (e.g., for a root camera).",
|
||||
examples=[0, 1, None],
|
||||
)
|
||||
parent_shot_idx: Optional[int] = Field(
|
||||
default=None,
|
||||
description="The index of the dependent shot. Set to None if the camera has no parent (e.g., for a root camera).",
|
||||
examples=[0, 3, None],
|
||||
)
|
||||
reason: str = Field(
|
||||
description="The reason for the selection of the parent camera. If the camera has no parent, it should explain why it's a root camera.",
|
||||
examples=[
|
||||
"The parent shot's field of view covers the child shot's field of view (from medium shot to close-up)",
|
||||
"The parent shot and the child shot have a shot/reverse shot relationship.",
|
||||
"CAMERA_0 (Shot 0) establishes the entire scene and contains all characters and the setting. It is the root camera." # 补充 LLM 实际输出的例子
|
||||
],
|
||||
)
|
||||
is_parent_fully_covers_child: Optional[bool] = Field(
|
||||
default=None,
|
||||
description="Whether the parent camera fully covers the child camera's content. Set to None if the camera has no parent.",
|
||||
examples=[True, False, None],
|
||||
)
|
||||
missing_info: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The missing elements in the child shot that are not covered by the parent shot. If the parent shot fully covers the child shot, set this to None.",
|
||||
examples=[
|
||||
"The frontal view of Alice.",
|
||||
None,
|
||||
],
|
||||
)
|
||||
|
||||
class CameraTreeResponse(BaseModel):
|
||||
camera_parent_items: List[Optional[CameraParentItem]] = Field(
|
||||
description="The parent camera items for each camera. If a camera has no parent, set this to None. The length of the list should be the same as the number of cameras.",
|
||||
)
|
||||
|
||||
|
||||
|
||||
class CameraImageGenerator:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chat_model,
|
||||
image_generator,
|
||||
video_generator,
|
||||
):
|
||||
self.chat_model = chat_model
|
||||
self.image_generator = image_generator
|
||||
self.video_generator = video_generator
|
||||
|
||||
|
||||
async def construct_camera_tree(
|
||||
self,
|
||||
cameras: List[Camera],
|
||||
shot_descs: List[Union[ShotDescription, ShotBriefDescription]],
|
||||
) -> List[Camera]:
|
||||
parser = PydanticOutputParser(pydantic_object=CameraTreeResponse)
|
||||
shot_desc_by_idx = {shot.idx: shot for shot in shot_descs}
|
||||
|
||||
camera_seq_str = "<CAMERA_SEQ>\n"
|
||||
for cam in cameras:
|
||||
camera_seq_str += f"<CAMERA_{cam.idx}>\n"
|
||||
for shot_idx in cam.active_shot_idxs:
|
||||
shot_desc = shot_desc_by_idx.get(shot_idx)
|
||||
if shot_desc is None:
|
||||
raise ValueError(f"Camera {cam.idx} references missing shot {shot_idx}")
|
||||
camera_seq_str += f"Shot {shot_idx}: {shot_desc.visual_desc}\n"
|
||||
camera_seq_str += f"</CAMERA_{cam.idx}>\n"
|
||||
camera_seq_str += "</CAMERA_SEQ>"
|
||||
|
||||
messages = [
|
||||
SystemMessage(content=system_prompt_template_select_reference_camera.format(format_instructions=parser.get_format_instructions())),
|
||||
HumanMessage(content=human_prompt_template_select_reference_camera.format(camera_seq_str=camera_seq_str)),
|
||||
]
|
||||
|
||||
chain = self.chat_model | parser
|
||||
response: CameraTreeResponse = await chain.ainvoke(messages)
|
||||
parent_items = response.camera_parent_items
|
||||
if len(parent_items) != len(cameras):
|
||||
raise ValueError(f"Camera tree response length mismatch: expected {len(cameras)}, got {len(parent_items)}")
|
||||
|
||||
valid_camera_idxs = {cam.idx for cam in cameras}
|
||||
valid_shot_idxs = set(shot_desc_by_idx)
|
||||
parent_by_camera = {}
|
||||
for cam, parent_cam_item in zip(cameras, parent_items):
|
||||
parent_cam_idx = parent_cam_item.parent_cam_idx if parent_cam_item is not None else None
|
||||
parent_shot_idx = parent_cam_item.parent_shot_idx if parent_cam_item is not None else None
|
||||
if parent_cam_idx is not None and parent_cam_idx not in valid_camera_idxs:
|
||||
raise ValueError(f"Camera {cam.idx} has invalid parent camera {parent_cam_idx}")
|
||||
if parent_cam_idx == cam.idx:
|
||||
raise ValueError(f"Camera {cam.idx} cannot be its own parent")
|
||||
if parent_shot_idx is not None and parent_shot_idx not in valid_shot_idxs:
|
||||
raise ValueError(f"Camera {cam.idx} has invalid parent shot {parent_shot_idx}")
|
||||
parent_by_camera[cam.idx] = parent_cam_idx
|
||||
|
||||
for cam in cameras:
|
||||
seen = set()
|
||||
current = cam.idx
|
||||
while parent_by_camera.get(current) is not None:
|
||||
current = parent_by_camera[current]
|
||||
if current in seen:
|
||||
raise ValueError(f"Camera tree contains a cycle involving camera {cam.idx}")
|
||||
seen.add(current)
|
||||
|
||||
for cam, parent_cam_item in zip(cameras, parent_items):
|
||||
cam.parent_cam_idx = parent_cam_item.parent_cam_idx if parent_cam_item is not None else None
|
||||
cam.parent_shot_idx = parent_cam_item.parent_shot_idx if parent_cam_item is not None else None
|
||||
cam.reason = parent_cam_item.reason if parent_cam_item is not None else None
|
||||
cam.is_parent_fully_covers_child = parent_cam_item.is_parent_fully_covers_child if parent_cam_item is not None else None
|
||||
cam.missing_info = parent_cam_item.missing_info if parent_cam_item is not None else None
|
||||
return cameras
|
||||
|
||||
|
||||
async def generate_transition_video(
|
||||
self,
|
||||
first_shot_visual_desc: str,
|
||||
second_shot_visual_desc: str,
|
||||
first_shot_ff_path: str,
|
||||
progress=None,
|
||||
) -> VideoOutput:
|
||||
|
||||
prompt = f"Two shots. The transition between the shots is a cut to. The style of the two shots should be consistent."
|
||||
prompt += f"\nThe first shot description: {first_shot_visual_desc}."
|
||||
prompt += f"\nThe second shot description: {second_shot_visual_desc}."
|
||||
reference_image_paths = [first_shot_ff_path]
|
||||
video_output = await self.video_generator.generate_single_video(
|
||||
prompt=prompt,
|
||||
reference_image_paths=reference_image_paths,
|
||||
progress=progress,
|
||||
)
|
||||
return video_output
|
||||
|
||||
|
||||
def get_new_camera_image(
|
||||
self,
|
||||
transition_video_path: str,
|
||||
) -> ImageOutput:
|
||||
video = open_video(transition_video_path)
|
||||
scene_manager = SceneManager()
|
||||
scene_manager.add_detector(ContentDetector())
|
||||
scene_manager.detect_scenes(video, show_progress=False)
|
||||
scene_list = scene_manager.get_scene_list()
|
||||
output_dir = os.path.join(os.path.dirname(transition_video_path), "cache")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
split_video_ffmpeg(transition_video_path, scene_list, output_dir, show_progress=True)
|
||||
|
||||
|
||||
video_name = os.path.basename(transition_video_path).split('.')[0]
|
||||
second_video_path = os.path.join(output_dir, f"{video_name}-Scene-002.mp4")
|
||||
if os.path.exists(second_video_path):
|
||||
# use first frame of second shot as new camera image
|
||||
clip = VideoFileClip(second_video_path)
|
||||
ff = clip.get_frame(0)
|
||||
ff = Image.fromarray(ff.astype('uint8'), 'RGB')
|
||||
return ImageOutput(fmt="pil", ext="png", data=ff)
|
||||
else:
|
||||
# use last frame of transition video to instead
|
||||
clip = VideoFileClip(transition_video_path)
|
||||
lf_time = clip.duration - (1 / clip.fps)
|
||||
lf_time = max(0, lf_time)
|
||||
lf = clip.get_frame(lf_time)
|
||||
lf = Image.fromarray(lf.astype('uint8'), 'RGB')
|
||||
return ImageOutput(fmt="pil", ext="png", data=lf)
|
||||
|
||||
|
||||
async def generate_first_frame(
|
||||
self,
|
||||
shot_desc: ShotDescription,
|
||||
character_portrait_path_and_text_pairs: List[Tuple[str, str]],
|
||||
) -> ImageOutput:
|
||||
prompt = ""
|
||||
reference_image_paths = []
|
||||
for i,(path, text )in enumerate(character_portrait_path_and_text_pairs):
|
||||
prompt += f"Image {i}: {text}\n"
|
||||
reference_image_paths.append(path)
|
||||
prompt += f"Generate an image based on the following description: {shot_desc.ff_desc}."
|
||||
image_output = await self.image_generator.generate_single_image(
|
||||
prompt=prompt,
|
||||
reference_image_paths=reference_image_paths,
|
||||
size="1600x900",
|
||||
)
|
||||
return image_output
|
||||
|
||||
|
||||
|
||||
def _validate_camera_tree(cameras: List[Camera]) -> None:
|
||||
"""Reject parent assignments that would deadlock frame generation."""
|
||||
by_idx = {cam.idx: cam for cam in cameras}
|
||||
for cam in cameras:
|
||||
if cam.parent_cam_idx is None:
|
||||
continue
|
||||
if cam.parent_cam_idx == cam.idx:
|
||||
raise ValueError(f"Camera {cam.idx} lists itself as its parent.")
|
||||
if cam.parent_cam_idx not in by_idx:
|
||||
raise ValueError(f"Camera {cam.idx} references unknown parent camera {cam.parent_cam_idx}.")
|
||||
for cam in cameras:
|
||||
seen = set()
|
||||
current = cam
|
||||
while current.parent_cam_idx is not None:
|
||||
if current.idx in seen:
|
||||
raise ValueError(f"Cycle detected in camera parent graph involving camera {current.idx}.")
|
||||
seen.add(current.idx)
|
||||
current = by_idx[current.parent_cam_idx]
|
||||
@@ -0,0 +1,89 @@
|
||||
import logging
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain.chat_models.base import BaseChatModel
|
||||
from langchain.chat_models import init_chat_model
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List
|
||||
from tenacity import retry, stop_after_attempt
|
||||
from interfaces import CharacterInScene
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
|
||||
from utils.retry import after_func
|
||||
|
||||
|
||||
system_prompt_template_extract_characters = \
|
||||
"""
|
||||
[Role]
|
||||
You are a top-tier movie script analysis expert.
|
||||
|
||||
[Task]
|
||||
Your task is to analyze the provided script and extract all relevant character information.
|
||||
|
||||
[Input]
|
||||
You will receive a script enclosed within <SCRIPT> and </SCRIPT>.
|
||||
|
||||
Below is a simple example of the input:
|
||||
|
||||
<SCRIPT>
|
||||
A young woman sits alone at a table, staring out the window. She takes a sip of her coffee and sighs. The liquid is no longer warm, just a bitter reminder of the time that has passed. Outside, the world moves in a blur of hurried footsteps and distant car horns, but inside the quiet café, time feels thick and heavy.
|
||||
Her finger traces the rim of the ceramic mug, following the imperfect circle over and over. The decision she had to make was supposed to be simple—a mere checkbox on the form of her life. Yesor No. Stayor Go. Yet, it had rooted itself in her chest, a tangled knot of fear and longing.
|
||||
</SCRIPT>
|
||||
|
||||
[Output]
|
||||
{format_instructions}
|
||||
|
||||
|
||||
[Guidelines]
|
||||
- Ensure that the language of all output values(not include keys) matches that used in the script.
|
||||
- Group all names referring to the same entity under one character. Select the most appropriate name as the character's identifier. If the person is a real famous person, the real person's name should be retained (e.g., Elon Musk, Bill Gates)
|
||||
- If the character's name is not mentioned, you can use reasonable pronouns to refer to them, including using their occupation or notable physical traits. For example, "the young woman" or "the barista".
|
||||
- For background characters in the script, you do not need to consider them as individual characters.
|
||||
- If a character's traits are not described or only partially outlined in the script, you need to design plausible features based on the context to make their characteristics more complete and detailed, ensuring they are vivid and evocative.
|
||||
- In static features, you need to describe the character's physical appearance, physique, and other relatively unchanging features. In dynamic features, you need to describe the character's attire, accessories, key items they carry, and other easily changeable features.
|
||||
- Don't include any information about the character's personality, role, or relationships with others in either static or dynamic features.
|
||||
- When designing character features, within reasonable limits, different character appearances should be made more distinct from each other.
|
||||
- The description of characters should be detailed, avoiding the use of abstract terms. Instead, employ descriptions that can be visualized—such as specific clothing colors and concrete physical traits (e.g., large eyes, a high nose bridge).
|
||||
"""
|
||||
|
||||
human_prompt_template_extract_characters = \
|
||||
"""
|
||||
<SCRIPT>
|
||||
{script}
|
||||
</SCRIPT>
|
||||
"""
|
||||
|
||||
|
||||
class ExtractCharactersResponse(BaseModel):
|
||||
characters: List[CharacterInScene] = Field(
|
||||
..., description="A list of characters extracted from the script."
|
||||
)
|
||||
|
||||
|
||||
|
||||
class CharacterExtractor:
|
||||
def __init__(
|
||||
self,
|
||||
chat_model,
|
||||
):
|
||||
self.chat_model = chat_model
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
after=after_func,
|
||||
)
|
||||
async def extract_characters(self, script: str) -> List[CharacterInScene]:
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=ExtractCharactersResponse)
|
||||
|
||||
messages = [
|
||||
SystemMessage(content=system_prompt_template_extract_characters.format(format_instructions=parser.get_format_instructions())),
|
||||
HumanMessage(content=human_prompt_template_extract_characters.format(script=script)),
|
||||
]
|
||||
|
||||
chain = self.chat_model | parser
|
||||
|
||||
response: ExtractCharactersResponse = await chain.ainvoke(messages)
|
||||
|
||||
return response.characters
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
import logging
|
||||
import os
|
||||
import asyncio
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain.chat_models.base import BaseChatModel
|
||||
from langchain.chat_models import init_chat_model
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Dict
|
||||
from interfaces import CharacterInScene, ImageOutput
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
|
||||
|
||||
|
||||
prompt_template_front = \
|
||||
"""
|
||||
Generate a full-body, front-view portrait of character {identifier} based on the following description, with a pure white background. Use a wide 16:9 landscape canvas, not a vertical portrait canvas. The character should be centered in the image, occupying the middle of the wide frame with enough horizontal empty space. Gazing straight ahead. Standing with arms relaxed at sides. Natural expression.
|
||||
Features: {features}
|
||||
Style: {style}
|
||||
"""
|
||||
|
||||
prompt_template_side = \
|
||||
"""
|
||||
Generate a full-body, side-view portrait of character {identifier} based on the provided front-view portrait, with a pure white background. Use a wide 16:9 landscape canvas, not a vertical portrait canvas. The character should be centered in the image, occupying the middle of the wide frame with enough horizontal empty space. Facing left. Standing with arms relaxed at sides.
|
||||
"""
|
||||
|
||||
prompt_template_back = \
|
||||
"""
|
||||
Generate a full-body, back-view portrait of character {identifier} based on the provided front-view portrait, with a pure white background. Use a wide 16:9 landscape canvas, not a vertical portrait canvas. The character should be centered in the image, occupying the middle of the wide frame with enough horizontal empty space. No facial features should be visible.
|
||||
"""
|
||||
|
||||
|
||||
class CharacterPortraitsGenerator:
|
||||
def __init__(
|
||||
self,
|
||||
image_generator,
|
||||
):
|
||||
self.image_generator = image_generator
|
||||
|
||||
|
||||
async def generate_front_portrait(
|
||||
self,
|
||||
character: CharacterInScene,
|
||||
style: str,
|
||||
) -> ImageOutput:
|
||||
features = "(static) " + character.static_features + "; (dynamic) " + character.dynamic_features
|
||||
prompt = prompt_template_front.format(
|
||||
identifier=character.identifier_in_scene,
|
||||
features=features,
|
||||
style=style,
|
||||
)
|
||||
image_output = await self.image_generator.generate_single_image(
|
||||
prompt=prompt,
|
||||
# size="512x512",
|
||||
)
|
||||
return image_output
|
||||
|
||||
async def generate_side_portrait(
|
||||
self,
|
||||
character: CharacterInScene,
|
||||
front_image_path: str,
|
||||
) -> ImageOutput:
|
||||
prompt = prompt_template_side.format(
|
||||
identifier=character.identifier_in_scene,
|
||||
)
|
||||
image_output = await self.image_generator.generate_single_image(
|
||||
prompt=prompt,
|
||||
reference_image_paths=[front_image_path],
|
||||
# size="1024x1024",
|
||||
)
|
||||
return image_output
|
||||
|
||||
|
||||
async def generate_back_portrait(
|
||||
self,
|
||||
character: CharacterInScene,
|
||||
front_image_path: str,
|
||||
) -> ImageOutput:
|
||||
prompt = prompt_template_back.format(
|
||||
identifier=character.identifier_in_scene,
|
||||
)
|
||||
image_output = await self.image_generator.generate_single_image(
|
||||
prompt=prompt,
|
||||
reference_image_paths=[front_image_path],
|
||||
# size="512x512",
|
||||
)
|
||||
return image_output
|
||||
@@ -0,0 +1,155 @@
|
||||
import os
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain.chat_models import init_chat_model
|
||||
from pydantic import BaseModel, Field
|
||||
from tenacity import retry, stop_after_attempt
|
||||
|
||||
from interfaces import Event
|
||||
|
||||
system_prompt_template_extract_events = \
|
||||
"""
|
||||
You are a highly skilled Literary Analyst AI. Your expertise is in narrative structure, plot deconstruction, and thematic analysis. You meticulously read and interpret prose to break down a story into its fundamental sequential events.
|
||||
|
||||
**TASK**
|
||||
Extract the next event from the provided novel, following the sequence of the story and building upon the partially extracted events.
|
||||
|
||||
**INPUT**
|
||||
1. The full text of the novel, which is enclosed within <NOVEL_TEXT_START> and <NOVEL_TEXT_END> tags
|
||||
2. A sequence of already-extracted events (in order), which is enclosed within <EXTRACTED_EVENTS_START> and <EXTRACTED_EVENTS_END> tags. The sequence may be empty. Each event contains multiple processes and constitutes a complete causal chain.
|
||||
|
||||
Below is an example input:
|
||||
|
||||
<NOVEL_TEXT_START>
|
||||
The night was as dark as ink when the piercing alarm of the city museum suddenly shattered the silence. A thief, moving with phantom-like agility, had just pried open the display case and snatched the blue gem known as the "Heart of the Ocean" when the blaring alarm echoed through the hall.
|
||||
... (more novel text) ...
|
||||
<NOVEL_TEXT_END>
|
||||
|
||||
<EXTRACTED_EVENTS_START>
|
||||
<Event 0>
|
||||
Description: A thief who stole a gem from a museum was caught after a rooftop chase with guards, and the gem was recovered.
|
||||
Process Chain:
|
||||
- A thief steals a gem from a museum, triggering the alarm. Guards notice and begin the chase.
|
||||
- The thief rushes out the museum's back door and dashes through narrow alleys, with guards closely pursuing and calling for backup.
|
||||
- ... (more processes) ...
|
||||
|
||||
<Event 1>
|
||||
Description: ... (more description) ...
|
||||
Process Chain:
|
||||
- ... (more processes) ...
|
||||
|
||||
<EXTRACTED_EVENTS_END>
|
||||
|
||||
|
||||
**OUTPUT**
|
||||
{format_instructions}
|
||||
|
||||
**GUIDELINES**
|
||||
1. Focus on events that are critical to the plot, character development, or thematic depth.
|
||||
2. Ensure the event is logically distinct from previous and subsequent events.
|
||||
3. If the event spans multiple scenes, unify them under a single dramatic goal. For example, a chase sequence might begin in a city market, continue through back alleys, and conclude on a rooftop—all comprising a single event because they collectively achieve the dramatic purpose of "the protagonist evading capture."
|
||||
4. Maintain objectivity: describe events based on the text without interpretation or judgment.
|
||||
5. For the process field, provide a detailed, step-by-step account of the event's progression, including key actions, decisions, and turning points. Each step should be clear and concise, illustrating how the event unfolds over time.
|
||||
Below is an example:
|
||||
Timeframe: The following morning, after acquiring the information about the Temple.
|
||||
Characters: Elara (protagonist) and Kaelen (her rival treasure hunter).
|
||||
Cause: Both seek the same artifact and are determined to reach it first.
|
||||
Process: The event begins with Elara hastily purchasing supplies in the port town (scene 1), where she spots Kaelen already hiring a crew, raising the stakes. It continues as she races to secure her own ship and captain, negotiating fiercely under time pressure (scene 2). The event culminates in a direct confrontation on the docks (scene 3), where Kaelen attempts to sabotage her vessel, leading to a brief but intense sword fight between the two rivals.
|
||||
Outcome: Elara successfully defends her ship and sets sail, but the conflict solidifies a bitter personal rivalry with Kaelen, ensuring their race to the temple will be fraught with direct opposition and danger.
|
||||
6. Every detail in your event description must be directly supported by the input novel. Do not add, assume, or invent any information.
|
||||
7. The language of outputs in values should be same as the input text.
|
||||
"""
|
||||
|
||||
human_prompt_template_extract_next_event = \
|
||||
"""
|
||||
<NOVEL_TEXT_START>
|
||||
{novel_text}
|
||||
<NOVEL_TEXT_END>
|
||||
|
||||
<EXTRACTED_EVENTS_START>
|
||||
{extracted_events}
|
||||
<EXTRACTED_EVENTS_END>
|
||||
"""
|
||||
|
||||
|
||||
|
||||
class EventExtractor:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
chat_model: str,
|
||||
):
|
||||
self.chat_model = init_chat_model(
|
||||
model=chat_model,
|
||||
model_provider="openai",
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
)
|
||||
self.parser = PydanticOutputParser(pydantic_object=Event)
|
||||
|
||||
|
||||
# Cap on extracted events: is_last is asserted by the LLM only, so without a
|
||||
# bound a model that never sets it would loop (and spend tokens) forever.
|
||||
max_events = 50
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
novel_text: str,
|
||||
):
|
||||
logging.info("Extracting events from novel...")
|
||||
|
||||
events = []
|
||||
while True:
|
||||
if len(events) >= self.max_events:
|
||||
raise RuntimeError(
|
||||
f"Event extraction exceeded the maximum of {self.max_events} events "
|
||||
"without an is_last marker; aborting to avoid unbounded LLM calls."
|
||||
)
|
||||
event = self.extract_next_event(novel_text, events)
|
||||
|
||||
events.append(event)
|
||||
logging.info(f"Extracted event: \n{event}")
|
||||
if event.is_last:
|
||||
break
|
||||
|
||||
return events
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
after=lambda retry_state: logging.warning(f"Retrying extract_next_event due to error: {retry_state.outcome.exception()}"),
|
||||
)
|
||||
def extract_next_event(
|
||||
self,
|
||||
novel_text: str,
|
||||
extracted_events: List[Event]
|
||||
) -> Event:
|
||||
|
||||
extracted_events_str = "\n\n".join([str(e) for e in extracted_events])
|
||||
|
||||
messages = [
|
||||
SystemMessage(
|
||||
content=system_prompt_template_extract_events.format(format_instructions=self.parser.get_format_instructions()),
|
||||
),
|
||||
HumanMessage(
|
||||
content=human_prompt_template_extract_next_event.format(
|
||||
novel_text=novel_text,
|
||||
extracted_events=extracted_events_str,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
chain = self.chat_model | self.parser
|
||||
|
||||
event: Event = chain.invoke(messages)
|
||||
|
||||
assert event.index == len(extracted_events), f"Extracted event index {event.index} does not match the expected index {len(extracted_events)}"
|
||||
|
||||
return event
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
import os
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Tuple, Dict, Optional
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain.chat_models import init_chat_model
|
||||
from pydantic import BaseModel, Field
|
||||
from langchain.output_parsers import PydanticOutputParser
|
||||
from interfaces import Event, Scene
|
||||
from interfaces import CharacterInScene, CharacterInEvent, CharacterInNovel
|
||||
from tenacity import retry, stop_after_attempt
|
||||
|
||||
|
||||
system_prompt_template_merge_characters_across_scenes_in_event = \
|
||||
"""
|
||||
You are an expert script analysis and character fusion specialist. Your role is to intelligently analyze multiple script scenes, identify characters that represent the same entity across different scenes, and merge them into a unified character list with consistent identifiers.
|
||||
|
||||
**TASK**
|
||||
Process the input scenes, each containing a script and characters with their names and features. Identify and merge characters that are logically the same across scenes, even if they have different names or slight variations in description. Output a consolidated list of characters for the entire event. Each character in the list must have a unique identifier, along with the scene numbers where they appear and the name used in each scene. You also need to aggregate the static features of the same characters together.
|
||||
|
||||
**INPUT**
|
||||
A sequence of scenes. Each scene is enclosed within <SCENE_N_START> and <SCENE_N_END> tags, where N is the scene number(starting from 0).
|
||||
Each scene includes a screnplay script and a sequence of character names.
|
||||
The screenplay script is enclosed within <SCRIPT_START> and <SCRIPT_END> tags.
|
||||
The sequence of character is enclosed within <CHARACTERS_START> and <CHARACTERS_END> tags. Each character in the list is enclosed within <CHARACTER_M_START> and <CHARACTER_M_END> tags, where M is the character number(starting from 0).
|
||||
|
||||
Below is an example of one scene:
|
||||
|
||||
<SCENE_0_START>
|
||||
|
||||
<SCRIPT_START>
|
||||
John enters the room and sees Mary.
|
||||
John: Hi Mary, how are you?
|
||||
Mary: I'm good, John. Thanks for asking!
|
||||
<SCRIPT_END>
|
||||
|
||||
<CHARACTERS_START>
|
||||
|
||||
<CHARACTER_0_START>
|
||||
John [visible]
|
||||
static features: John is a tall man with short black hair and brown eyes.
|
||||
dynamic features: Wearing a blue shirt and black pants.
|
||||
<CHARACTER_0_END>
|
||||
|
||||
<CHARACTER_1_START>
|
||||
Mary [visible]
|
||||
static features: Mary is a young woman with long brown hair and green eyes.
|
||||
dynamic features: Wearing a floral dress and a denim jacket.
|
||||
<CHARACTER_1_END>
|
||||
|
||||
<CHARACTERS_END>
|
||||
|
||||
<SCENE_0_END>
|
||||
|
||||
|
||||
|
||||
**OUTPUT**
|
||||
{format_instructions}
|
||||
|
||||
**GUIDELINES**
|
||||
1. Character Fusion: Analyze contextual clues (e.g., dialogue style, role in plot, relationships, descriptions) to determine if characters from different scenes are the same person, even if names vary.
|
||||
2. Unique Identifier: Assign a consistent, unique ID (e.g., primary/canonical name) to each merged character. Use the most frequent or contextually appropriate name as the identifier, if possible.
|
||||
3. Scene Mapping: For each character, list all scenes they appear in and the exact name used in each scene.
|
||||
4. Completeness: Ensure all characters from all scenes are included in the final list. No duplicate, omitted, or extraneous characters.
|
||||
5. If a character undergoes significant changes across different scenes, it is necessary to split them into separate roles. For example, if Character A is a child in Scene 0 but an adult in Scene 1, they should be divided into two distinct characters (meaning two different actors are required to portray them).
|
||||
6. The language of outputs in values should be same as the input text.
|
||||
"""
|
||||
|
||||
|
||||
human_prompt_template_merge_characters_across_scenes_in_event = \
|
||||
"""
|
||||
{scenes_sequence}
|
||||
"""
|
||||
|
||||
class MergeCharactersAcrossScenesInEventResponse(BaseModel):
|
||||
characters: List[CharacterInEvent] = Field(
|
||||
description="List of merged characters with their identifiers",
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
system_prompt_template_merge_characters_to_existing_characters_in_novel = \
|
||||
"""
|
||||
You are an information integration expert skilled in accurately identifying, matching, and merging character information. Your responsibility is to ensure consistency in character attributes and efficiently maintain and update the global character list.
|
||||
|
||||
**TASK**
|
||||
Merge the character list extracted from the current event (which may include new or existing characters) into the global character list. For existing characters, ensure their feature descriptions remain consistent; for new characters, add them to the global list.
|
||||
|
||||
**INPUT**
|
||||
1. Existing Characters in the Novel: A list of characters already present in the novel, each with a unique index, identifier, and static features. The list is enclosed within <EXISTING_CHARACTERS_START> and <EXISTING_CHARACTERS_END> tags. Each character in the list is enclosed within <CHARACTER_P_START> and <CHARACTER_P_END> tags, where P is the character number(starting from 0).
|
||||
2. Characters in the Current Event: A list of characters identified in the current event, each with an index, identifier, active scenes, and static features. The list is enclosed within <EVENT_CHARACTERS_START> and <EVENT_CHARACTERS_END> tags. Each character in the list is enclosed within <CHARACTER_Q_START> and <CHARACTER_Q_END> tags, where Q is the character number(starting from 0).
|
||||
|
||||
|
||||
**OUTPUT**
|
||||
{format_instructions}
|
||||
|
||||
**GUIDELINES**
|
||||
1. Feature Consistency: Strictly compare the features of the current event characters with those of existing characters. Some character's identifier may be the same as existing role identifier, but their features differ, such as youth and old age. You need to distinguish them as two separate characters.
|
||||
2. Efficient Merging: Avoid duplicate characters to ensure the list remains concise.
|
||||
3. Feature Update: If an existing character's features are expanded or modified based on new information from the current event, update their description accordingly.
|
||||
"""
|
||||
|
||||
human_prompt_template_merge_characters_to_existing_characters_in_novel = \
|
||||
"""
|
||||
<EXISTING_CHARACTERS_START>
|
||||
{existing_characters_in_novel}
|
||||
<EXISTING_CHARACTERS_END>
|
||||
|
||||
<EVENT_CHARACTERS_START>
|
||||
{characters_in_event}
|
||||
<EVENT_CHARACTERS_END>
|
||||
"""
|
||||
|
||||
|
||||
class CharacterForMergingToNovel(BaseModel):
|
||||
index_in_event: int = Field(
|
||||
description="The index of the character in the list of characters in the current event.",
|
||||
examples=[0, 1, 2],
|
||||
)
|
||||
index_in_novel: int = Field(
|
||||
description="The index of the character in the list of existing characters in the novel. If this is a new character, set it to -1.",
|
||||
examples=[0, 7, -1],
|
||||
)
|
||||
identifier_in_novel: str = Field(
|
||||
description="The unique identifier for the character in the novel. If this is a new character, ensure the name does not conflict with existing characters. If this is not a new character, this should match the identifier in the existing characters list.",
|
||||
examples=["Alice", "Bob the Builder"],
|
||||
)
|
||||
modified_features: str = Field(
|
||||
description="The modified static features of the character after merging. If the character is new, this should be the full static features. If the character is existing and their features are expanded or modified, this should be filled in the complete modified features. If the character is existing and their features remain unchanged, this should be the same as the existing character's static features.",
|
||||
)
|
||||
|
||||
class MergeCharactersToExistingCharactersInNovelResponse(BaseModel):
|
||||
characters: List[CharacterForMergingToNovel] = Field(
|
||||
description="List of characters in the event with their corresponding index in the existing characters in the novel. If the character is new, the index_in_novel should be -1. The number of characters in this list should be the same as the number of characters in the event.",
|
||||
)
|
||||
|
||||
|
||||
|
||||
class GlobalInformationPlanner:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
chat_model: str,
|
||||
):
|
||||
self.chat_model = init_chat_model(
|
||||
model=chat_model,
|
||||
model_provider="openai",
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
)
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
after=lambda retry_state: logging.warning(f"Retrying due to {retry_state.outcome.exception()}"),
|
||||
)
|
||||
async def merge_characters_across_scenes_in_event(
|
||||
self,
|
||||
event_idx: int,
|
||||
scenes: List[Scene], # Scene.characters is List[CharacterInScene]
|
||||
) -> List[CharacterInEvent]:
|
||||
scenes_sequence_str = ""
|
||||
for scene in scenes:
|
||||
scene_str = f"<SCENE_{scene.idx}_START>\n"
|
||||
scene_str += "<SCRIPT_START>\n"
|
||||
scene_str += scene.script + "\n"
|
||||
scene_str += "<SCRIPT_END>\n\n"
|
||||
scene_str += "<CHARACTERS_START>\n"
|
||||
for character in scene.characters:
|
||||
scene_str += f"<CHARACTER_{character.idx}_START>\n"
|
||||
scene_str += str(character)
|
||||
scene_str += f"<CHARACTER_{character.idx}_END>\n"
|
||||
scene_str += "<CHARACTERS_END>\n"
|
||||
scene_str += f"<SCENE_{scene.idx}_END>\n"
|
||||
scenes_sequence_str += scene_str
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=MergeCharactersAcrossScenesInEventResponse)
|
||||
|
||||
messages = [
|
||||
SystemMessage(
|
||||
content=system_prompt_template_merge_characters_across_scenes_in_event.format(
|
||||
format_instructions=parser.get_format_instructions(),
|
||||
),
|
||||
),
|
||||
HumanMessage(
|
||||
content=human_prompt_template_merge_characters_across_scenes_in_event.format(
|
||||
scenes_sequence=scenes_sequence_str,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
chain = self.chat_model | parser
|
||||
response: MergeCharactersAcrossScenesInEventResponse = await chain.ainvoke(messages)
|
||||
characters_in_event = response.characters
|
||||
|
||||
# check the output is valid
|
||||
flags = [{c.identifier_in_scene: False for c in s.characters} for s in scenes]
|
||||
|
||||
# check if all character identifiers can be found in the scenes
|
||||
for character in characters_in_event:
|
||||
for scene_idx, identifier_in_scene in character.active_scenes.items():
|
||||
if identifier_in_scene not in [c.identifier_in_scene for c in scenes[scene_idx].characters]:
|
||||
raise ValueError(f"Character {identifier_in_scene} not found in scene {scene_idx} of event {event_idx}")
|
||||
else:
|
||||
flags[scene_idx][identifier_in_scene] = True
|
||||
|
||||
# check if all characters are included
|
||||
for scene_idx, flag in enumerate(flags):
|
||||
for identifier_in_scene, included in flag.items():
|
||||
if not included:
|
||||
raise ValueError(f"Character {identifier_in_scene} in scene {scene_idx} of event {event_idx} not included in the merged characters")
|
||||
|
||||
return characters_in_event
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
after=lambda retry_state: logging.warning(f"Retrying due to {retry_state.outcome.exception()}"),
|
||||
)
|
||||
def merge_characters_to_existing_characters_in_novel(
|
||||
self,
|
||||
event_idx: int,
|
||||
existing_characters_in_novel: List[CharacterInNovel],
|
||||
characters_in_event: List[CharacterInEvent],
|
||||
) -> List[CharacterInNovel]:
|
||||
existing_characters_str = ""
|
||||
for character in existing_characters_in_novel:
|
||||
existing_characters_str += f"<CHARACTER_{character.index}_START>\n"
|
||||
existing_characters_str += str(character)
|
||||
existing_characters_str += f"<CHARACTER_{character.index}_END>\n"
|
||||
|
||||
characters_in_event_str = ""
|
||||
for character in characters_in_event:
|
||||
characters_in_event_str += f"<CHARACTER_{character.index}_START>\n"
|
||||
characters_in_event_str += character.identifier_in_event + "\n"
|
||||
characters_in_event_str += "Static features: " + character.static_features + "\n"
|
||||
characters_in_event_str += f"<CHARACTER_{character.index}_END>\n"
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=MergeCharactersToExistingCharactersInNovelResponse)
|
||||
|
||||
messages = [
|
||||
SystemMessage(
|
||||
content=system_prompt_template_merge_characters_to_existing_characters_in_novel.format(
|
||||
format_instructions=parser.get_format_instructions(),
|
||||
),
|
||||
),
|
||||
HumanMessage(
|
||||
content=human_prompt_template_merge_characters_to_existing_characters_in_novel.format(
|
||||
existing_characters_in_novel=existing_characters_str,
|
||||
characters_in_event=characters_in_event_str,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
chain = self.chat_model | parser
|
||||
response: MergeCharactersToExistingCharactersInNovelResponse = chain.invoke(messages)
|
||||
|
||||
for character in response.characters:
|
||||
if character.index_in_novel == -1:
|
||||
# new character, add to existing characters
|
||||
new_character = CharacterInNovel(
|
||||
index=len(existing_characters_in_novel),
|
||||
identifier_in_novel=character.identifier_in_novel,
|
||||
static_features=character.modified_features,
|
||||
active_events={event_idx: characters_in_event[character.index_in_event].identifier_in_event},
|
||||
)
|
||||
existing_characters_in_novel.append(new_character)
|
||||
else:
|
||||
existing_characters_in_novel[character.index_in_novel].static_features = character.modified_features
|
||||
existing_characters_in_novel[character.index_in_novel].active_events.update({event_idx: characters_in_event[character.index_in_event].identifier_in_event})
|
||||
|
||||
return existing_characters_in_novel
|
||||
|
||||
|
||||
# # TODO: 如果是长篇小说，事件太多，很容易报错，出场的角色会分不清在哪个事件里，也很容易漏，需要想办法解决
|
||||
# @retry(
|
||||
# stop=stop_after_attempt(3),
|
||||
# after=lambda retry_state: logging.warning(f"Retrying due to {retry_state.outcome.exception()}"),
|
||||
# )
|
||||
# def merge_characters_across_events_in_novel(
|
||||
# self,
|
||||
# events: List[Event],
|
||||
# characters_in_event: List[List[CharacterInEvent]],
|
||||
# ) -> List[CharacterInNovelWithoutStaticFeatures]:
|
||||
# events_sequence_str = ""
|
||||
# for event, characters in zip(events, characters_in_event):
|
||||
# event_str = f"<EVENT_{event.index}_START>\n\n"
|
||||
# event_str += "<DESCRIPTION_START>\n"
|
||||
# event_str += event.description + "\n"
|
||||
# event_str += "<DESCRIPTION_END>\n\n"
|
||||
# event_str += "<PROCESS_CHAIN_START>\n"
|
||||
# for process in event.process_chain:
|
||||
# event_str += process + "\n"
|
||||
# event_str += "<PROCESS_CHAIN_END>\n\n"
|
||||
# event_str += "<CHARACTERS_START>\n"
|
||||
# for i, character in enumerate(characters):
|
||||
# event_str += f"<CHARACTER_{i}_START>{character.identifier_in_event}<CHARACTER_{i}_END>\n"
|
||||
# event_str += "<CHARACTERS_END>\n\n"
|
||||
# event_str += f"<EVENT_{event.index}_END>\n\n"
|
||||
# events_sequence_str += event_str
|
||||
|
||||
# parser = PydanticOutputParser(pydantic_object=MergeCharactersAcrossEventsInNovelResponse)
|
||||
|
||||
# messages = [
|
||||
# SystemMessage(
|
||||
# content=system_prompt_template_merge_characters_across_events.format(
|
||||
# format_instructions=parser.get_format_instructions(),
|
||||
# ),
|
||||
# ),
|
||||
# HumanMessage(
|
||||
# content=human_prompt_template_merge_characters_across_events.format(
|
||||
# events_sequence=events_sequence_str,
|
||||
# )
|
||||
# )
|
||||
# ]
|
||||
|
||||
# chain = self.chat_model | parser
|
||||
# response: MergeCharactersAcrossEventsInNovelResponse = chain.invoke(messages)
|
||||
# characters_in_novel = response.characters
|
||||
|
||||
# # check the output is valid
|
||||
# flags = [{c.identifier_in_event: False for c in characters} for characters in characters_in_event]
|
||||
|
||||
# # check if all character identifiers can be found in the events
|
||||
# for character in characters_in_novel:
|
||||
# for event_idx, identifier_in_event in character.active_events.items():
|
||||
# if identifier_in_event not in [c.identifier_in_event for c in characters_in_event[event_idx]]:
|
||||
# raise ValueError(f"Character {identifier_in_event} not found in event {event_idx}")
|
||||
# else:
|
||||
# flags[event_idx][identifier_in_event] = True
|
||||
|
||||
# # check if all characters are included
|
||||
# # for event_idx, flag in enumerate(flags):
|
||||
# # for identifier_in_event, included in flag.items():
|
||||
# # if not included:
|
||||
# # raise ValueError(f"Character {identifier_in_event} in event {event_idx} not included in the merged characters")
|
||||
|
||||
# return characters_in_novel
|
||||
|
||||
|
||||
|
||||
# async def extract_static_feature_for_character_in_novel(
|
||||
# self,
|
||||
# relevant_chunks: List[str],
|
||||
# character: CharacterInNovelWithoutStaticFeatures,
|
||||
# ) -> str:
|
||||
# context_fragments_str = ""
|
||||
# for i, chunk in enumerate(relevant_chunks):
|
||||
# context_fragments_str += f"<CONTEXT_FRAGMENT_{i}_START>\n"
|
||||
# context_fragments_str += chunk + "\n"
|
||||
# context_fragments_str += f"<CONTEXT_FRAGMENT_{i}_END>\n"
|
||||
|
||||
# parser = None # no need to parse the output, just return the text
|
||||
|
||||
# messages = [
|
||||
# SystemMessage(
|
||||
# content=system_prompt_template_extract_static_feature_for_character_in_novel,
|
||||
# ),
|
||||
# HumanMessage(
|
||||
# content=human_prompt_template_extract_static_feature_for_character_in_novel.format(
|
||||
# character_name=character.identifier_in_novel,
|
||||
# context_fragments=context_fragments_str,
|
||||
# )
|
||||
# )
|
||||
# ]
|
||||
|
||||
# base_features = await self.chat_model.ainvoke(messages)
|
||||
# return base_features.content
|
||||
@@ -0,0 +1,171 @@
|
||||
import os
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Tuple
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain.chat_models import init_chat_model
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
|
||||
|
||||
|
||||
system_prompt_template_compress_novel_chunk = \
|
||||
"""
|
||||
You are an expert text compression assistant specialized in literary content. Your goal is to condense novels or story excerpts while preserving core narrative elements, key details, character development, and plot coherence.
|
||||
|
||||
|
||||
**TASK**
|
||||
Compress the provided input text to reduce its length significantly, eliminating redundancies, overly descriptive passages, and minor details—but without losing essential story arcs, dialogue, or emotional impact. Aim for clarity and readability in the compressed output.
|
||||
|
||||
|
||||
**INPUT**
|
||||
A segment of a novel (possibly truncated due to context length constraints). It is enclosed within <NOVEL_CHUNK_START> and <NOVEL_CHUNK_END> tags.
|
||||
|
||||
|
||||
**OUTPUT**
|
||||
A compressed version of the input text, retaining the core narrative, critical events, and character interactions.
|
||||
|
||||
**GUIDELINES**
|
||||
1. Fidelity to the Plot: Absolutely preserve all major plot points, twists, revelations, and the sequence of key events. Do not omit crucial story elements.
|
||||
2. Character Consistency: Maintain character actions, decisions, and development. Important dialogue that reveals plot or character can be condensed or paraphrased but its meaning must be kept intact.
|
||||
3. Streamline Description: Reduce lengthy descriptions of settings, characters, or objects to their most essential and evocative elements. Capture the mood and critical details without the elaborate prose.
|
||||
4. Condense Internal Monologue: Paraphrase characters' extended internal thoughts and reflections, focusing on the key realizations or decisions they lead to.
|
||||
5. Simplify Language: Use more direct and concise language. Combine sentences, eliminate redundant adverbs and adjectives, and avoid repetitive phrasing.
|
||||
6. Cohesion and Flow: Ensure the compressed text is smooth, readable, and maintains a logical narrative flow. It should not feel like a fragmented list of events.
|
||||
7. Discard any non-narrative text (e.g., "Please follow my account!", "Background setting:...", personal opinions).
|
||||
8. Produce a seamless paragraph (or paragraphs if necessary) without markers (e.g., "Chapter 1") or section breaks.
|
||||
9. The language of output should be consistent with the original text.
|
||||
"""
|
||||
|
||||
human_prompt_template_compress_novel_chunk = \
|
||||
"""
|
||||
<NOVEL_CHUNK_START>
|
||||
{novel_chunk}
|
||||
<NOVEL_CHUNK_END>
|
||||
"""
|
||||
|
||||
|
||||
system_prompt_template_aggregate = \
|
||||
"""
|
||||
You are a professional text processing assistant specializing in the aggregation and refinement of segmented text chunks. Your expertise lies in seamlessly merging sequential text fragments while intelligently handling overlapping or duplicated content expressed in different ways.
|
||||
|
||||
**TASK**
|
||||
Aggregate the provided text chunks into a coherent and continuous short story. Carefully identify and resolve overlaps where the end of one chunk and the beginning of the next chunk contain semantically similar content but with different expressions. Remove redundant repetitions while preserving the original meaning, style, and flow of the text. Ensure all non-overlapping content remains unchanged and intact.
|
||||
|
||||
|
||||
**INPUT**
|
||||
A sequence of text chunks (ordered from first to last), where each chunk may have an overlapping segment with the next chunk. The overlapping segments might vary in wording but convey similar meaning. Each chunk is enclosed within <CHUNK_N_START> and <CHUNK_N_END> tags, where N is the chunk index starting from 0.
|
||||
|
||||
**OUTPUT**
|
||||
A single, consolidated text of the short story without unnatural repetitions or disruptions. The output should maintain the original narrative structure, tone, and details, with smooth transitions between originally adjacent chunks.
|
||||
|
||||
**GUIDELINES**
|
||||
1. Analyze the input chunks sequentially. For each adjacent pair (e.g., Chunk N and Chunk N+1), compare the end of Chunk N and the beginning of Chunk N+1 to detect overlapping content.
|
||||
2. If the overlapping segments are semantically equivalent but phrased differently, merge them by retaining the most natural or contextually appropriate version (prioritize the version from the later chunk if both are equally valid, but avoid introducing inconsistency).
|
||||
3. If the overlapping segments are not perfectly equivalent (e.g., one contains additional details), integrate the meaningful information without duplication, ensuring no loss of content.
|
||||
4. Preserve all non-overlapping text exactly as it appears in the original chunks. Do not modify, paraphrase, or omit any unique content.
|
||||
5. Ensure the merged text is fluent and coherent, without abrupt jumps or redundant phrases.
|
||||
6. If no overlap is detected between two chunks, concatenate them directly without changes.
|
||||
7. Do not invent new content or alter the original narrative beyond handling the overlaps.
|
||||
8. The language of output should be consistent with the original text.
|
||||
"""
|
||||
|
||||
human_prompt_template_aggregate = \
|
||||
"""
|
||||
{chunks}
|
||||
"""
|
||||
|
||||
|
||||
|
||||
|
||||
class NovelCompressor:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
chat_model: str,
|
||||
chunk_size: int = 65536,
|
||||
chunk_overlap: int = 8192,
|
||||
):
|
||||
self.chat_model = init_chat_model(
|
||||
model=chat_model,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model_provider="openai",
|
||||
)
|
||||
|
||||
self.splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
)
|
||||
|
||||
|
||||
def split(
|
||||
self,
|
||||
novel_text: str,
|
||||
):
|
||||
novel_chunks = self.splitter.split_text(novel_text)
|
||||
return novel_chunks
|
||||
|
||||
|
||||
async def compress(
|
||||
self,
|
||||
index_chunk_pairs: List[Tuple[int, str]],
|
||||
max_concurrent_tasks: int = 5,
|
||||
) -> str:
|
||||
sem = asyncio.Semaphore(max_concurrent_tasks)
|
||||
|
||||
tasks = [
|
||||
self.compress_single_novel_chunk(sem, index, novel_chunk)
|
||||
for index, novel_chunk in index_chunk_pairs
|
||||
]
|
||||
compressed_novel_chunks = await asyncio.gather(*tasks)
|
||||
return compressed_novel_chunks
|
||||
|
||||
|
||||
async def compress_single_novel_chunk(
|
||||
self,
|
||||
semaphore: asyncio.Semaphore,
|
||||
index,
|
||||
novel_chunk: str,
|
||||
) -> str:
|
||||
async with semaphore:
|
||||
logging.info(f"Compressing novel chunk {index}")
|
||||
messages = [
|
||||
SystemMessage(
|
||||
content=system_prompt_template_compress_novel_chunk
|
||||
),
|
||||
HumanMessage(
|
||||
content=human_prompt_template_compress_novel_chunk.format(
|
||||
novel_chunk=novel_chunk
|
||||
)
|
||||
),
|
||||
]
|
||||
response = await self.chat_model.ainvoke(messages)
|
||||
compressed_novel_chunk = response.content
|
||||
logging.info(f"Compressed novel chunk {index}")
|
||||
return index, compressed_novel_chunk
|
||||
|
||||
|
||||
def aggregate(
|
||||
self,
|
||||
compressed_novel_chunks: List[str],
|
||||
):
|
||||
chunks_str = "\n".join([
|
||||
f"<CHUNK_{i}_START>\n{chunk}\n<CHUNK_{i}_END>"
|
||||
for i, chunk in enumerate(compressed_novel_chunks)
|
||||
])
|
||||
|
||||
messages = [
|
||||
SystemMessage(
|
||||
content=system_prompt_template_aggregate
|
||||
),
|
||||
HumanMessage(
|
||||
content=human_prompt_template_aggregate.format(
|
||||
chunks=chunks_str
|
||||
)
|
||||
),
|
||||
]
|
||||
response = self.chat_model.invoke(messages)
|
||||
aggregated_novel = response.content
|
||||
return aggregated_novel
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
import logging
|
||||
from typing import List, Tuple
|
||||
from tenacity import retry, stop_after_attempt
|
||||
from pydantic import BaseModel, Field
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain.chat_models import init_chat_model
|
||||
from utils.image import image_path_to_b64
|
||||
|
||||
from utils.retry import after_func
|
||||
|
||||
system_prompt_template_select_reference_images_only_text = \
|
||||
"""
|
||||
[Role]
|
||||
You are a professional visual creation assistant skilled in multimodal image analysis and reasoning.
|
||||
|
||||
[Task]
|
||||
Your core task is to intelligently select the most suitable reference images from a provided set of reference image descriptions (including multiple character reference images and existing scene images from prior frames) based on the user's text description (describing the target frame), ensuring that the subsequently generated image meets the following key consistencies:
|
||||
- Character Consistency: The appearance (e.g. gender, ethnicity, age, facial features, hairstyle, body shape), clothing, expression, posture, etc., of the generated character should highly match the reference image descriptions.
|
||||
- Environmental Consistency: The scene of the generated image (e.g., background, lighting, atmosphere, layout) should remain coherent with the existing image descriptions from prior frames.
|
||||
- Style Consistency: The visual style of the generated image (e.g., realistic, cartoon, film-like, color tone) should harmonize with the reference image descriptions.
|
||||
|
||||
[Input]
|
||||
You will receive a text description of the target frame, along with a sequence of reference image descriptions.
|
||||
- The text description of the target frame is enclosed within <FRAME_DESC> and </FRAME_DESC>.
|
||||
- The sequence of reference image descriptions is enclosed within <SEQ_DESC> and </SEQ_DESC>. Each description is prefixed with its index, starting from 0.
|
||||
|
||||
Below is an example of the input format:
|
||||
<FRAME_DESC>
|
||||
[Camera 1] Shot from Alice's over-the-shoulder perspective. Alice is on the side closer to the camera, with only her shoulder appearing in the lower left corner of the frame. Bob is on the side farther from the camera, positioned slightly right of center in the frame. Bob's expression shifts from surprise to delight as he recognizes Alice.
|
||||
</FRAME_DESC>
|
||||
|
||||
<SEQ_DESC>
|
||||
Image 0: A front-view portrait of Alice.
|
||||
Image 1: A front-view portrait of Bob.
|
||||
Image 2: [Camera 0] Medium shot of the supermarket aisle. Alice and Bob are shown in profile facing the right side of the frame. Bob is on the right side of the frame, and Alice is on the left side. Alice, looking down and pushing a shopping cart, follows closely behind Bob and accidentally bumps into his heel.
|
||||
Image 3: [Camera 1] Shot from Alice's over-the-shoulder perspective. Alice is on the side closer to the camera, with only her shoulder appearing in the lower left corner of the frame. Bob is on the side farther from the camera, positioned slightly right of center in the frame. Bob quickly turns around, and his expression shifts from neutral to surprised.
|
||||
Image 4: [Camera 2] Shot from Bob's over-the-shoulder perspective. Bob is on the side closer to the camera, with only his shoulder appearing in the lower right corner of the frame. Alice is on the side farther from the camera, positioned slightly left of center in the frame. Alice looks down, then up as she prepares to apologize. Upon realizing it's someone familiar, her expression shifts to one of surprise.
|
||||
</SEQ_DESC>
|
||||
|
||||
|
||||
[Output]
|
||||
You need to select up to 8 of the most relevant reference images based on the user's description and put the corresponding indices in the ref_image_indices field of the output. At the same time, you should generate a text prompt that describes the image to be created, specifying which elements in the generated image should reference which image description (and which elements within it).
|
||||
|
||||
{format_instructions}
|
||||
|
||||
|
||||
[Guidelines]
|
||||
- Ensure that the language of all output values (not include keys) matches that used in the frame description.
|
||||
- The reference image descriptions may depict the same character from different angles, in different outfits, or in different scenes. Identify the description closest to the version described by the user
|
||||
- Prioritize image descriptions with similar compositions, i.e., shots taken by the same camera.
|
||||
- The images from prior frames are arranged in chronological order. Give higher priority to more recent images (those closer to the end of the sequence).
|
||||
- Choose reference image descriptions that are as concise as possible and avoid including duplicate information. For example, if Image 3 depicts the facial features of Bob from the front, and Image 1 also depicts Bob's facial features from the front-view portrait, then Image 1 is redundant and should not be selected.
|
||||
- When a new character appears in the frame description, prioritize selecting their portrait image description (if available) to ensure accurate depiction of their appearance. Pay attention to whether the character is facing the camera from the front, side, or back. Choose the most suitable view as the reference image for the character.
|
||||
- For character portraits, you can only select at most one image from multiple views (front, side, back). Choose the most appropriate one based on the frame description. For example, when depicting a character from the side, choose the side view of the character.
|
||||
- Select at most **8** optimal reference image descriptions.
|
||||
"""
|
||||
|
||||
|
||||
system_prompt_template_select_reference_images_multimodal = \
|
||||
"""
|
||||
[Role]
|
||||
You are a professional visual creation assistant skilled in multimodal image analysis and reasoning.
|
||||
|
||||
[Task]
|
||||
Your core task is to intelligently select the most suitable reference images from a provided reference image library (including multiple character reference images and existing scene images from prior frames) based on the user's text description (describing the target frame), ensuring that the subsequently generated image meets the following key consistencies:
|
||||
- Character Consistency: The appearance (e.g. gender, ethnicity, age, facial features, hairstyle, body shape), clothing, expression, posture, etc., of the generated character should highly match the reference images.
|
||||
- Environmental Consistency: The scene of the generated image (e.g., background, lighting, atmosphere, layout) should remain coherent with the existing images from prior frames.
|
||||
- Style Consistency: The visual style of the generated image (e.g., realistic, cartoon, film-like, color tone) should harmonize with the reference images and existing images.
|
||||
|
||||
[Input]
|
||||
You will receive a text description of the target frame, along with a sequence of reference images.
|
||||
- The text description of the target frame is enclosed within <FRAME_DESC> and </FRAME_DESC>.
|
||||
- The sequence of reference images is enclosed within <SEQ_IMAGES> and </SEQ_IMAGES>. Each reference image is provided with a text description. The reference images are indexed starting from 0.
|
||||
|
||||
Below is an example of the input format:
|
||||
<FRAME_DESC>
|
||||
[Camera 1] Shot from Alice's over-the-shoulder perspective. <Alice> is on the side closer to the camera, with only her shoulder appearing in the lower left corner of the frame. <Bob> is on the side farther from the camera, positioned slightly right of center in the frame. <Bob>'s expression shifts from surprise to delight as he recognizes <Alice>.
|
||||
</FRAME_DESC>
|
||||
|
||||
<SEQ_IMAGES>
|
||||
Image 0: A front-view portrait of Alice.
|
||||
[Image 0 here]
|
||||
Image 1: A front-view portrait of Bob.
|
||||
[Image 1 here]
|
||||
Image 2: [Camera 0] Medium shot of the supermarket aisle. Alice and Bob are shown in profile facing the right side of the frame. Bob is on the right side of the frame, and Alice is on the left side. Alice, looking down and pushing a shopping cart, follows closely behind Bob and accidentally bumps into his heel.
|
||||
[Image 2 here]
|
||||
Image 3: [Camera 1] Shot from Alice's over-the-shoulder perspective. Alice is on the side closer to the camera, with only her shoulder appearing in the lower left corner of the frame. Bob is on the side farther from the camera, positioned slightly right of center in the frame. Bob is back to the camera.
|
||||
[Image 3 here]
|
||||
Image 4: [Camera 2] Shot from Bob's over-the-shoulder perspective. Bob is on the side closer to the camera, with only his shoulder appearing in the lower right corner of the frame. Alice is on the side farther from the camera, positioned slightly left of center in the frame. Alice looks down, then up as she prepares to apologize. Upon realizing it's someone familiar, her expression shifts to one of surprise.
|
||||
</SEQ_IMAGES>
|
||||
|
||||
[Output]
|
||||
You need to select the most relevant reference images based on the user's description and put the corresponding indices in the `ref_image_indices` field of the output. At the same time, you should generate a text prompt that describes the image to be created, specifying which elements in the generated image should reference which image (and which elements within it).
|
||||
|
||||
{format_instructions}
|
||||
|
||||
|
||||
[Guidelines]
|
||||
- Ensure that the language of all output values (not include keys) matches that used in the frame description.
|
||||
- The reference image descriptions may depict the same character from different angles, in different outfits, or in different scenes. Identify the description closest to the version described by the user
|
||||
- Prioritize image descriptions with similar compositions, i.e., shots taken by the same camera.
|
||||
- The images from prior frames are arranged in chronological order. Give higher priority to more recent images (those closer to the end of the sequence).
|
||||
- Choose reference image descriptions that are as concise as possible and avoid including duplicate information. For example, if Image 3 depicts the facial features of Bob from the front, and Image 1 also depicts Bob's facial features from the front-view portrait, then Image 1 is redundant and should not be selected.
|
||||
- For character portraits, you can only select at most one image from multiple views (front, side, back). Choose the most appropriate one based on the frame description. For example, when depicting a character from the side, choose the side view of the character.
|
||||
- Select at most **8** optimal reference image descriptions.
|
||||
- The text guiding image editing should be as concise as possible.
|
||||
"""
|
||||
|
||||
|
||||
human_prompt_template_select_reference_images = \
|
||||
"""
|
||||
<FRAME_DESC>
|
||||
{frame_description}
|
||||
</FRAME_DESC>
|
||||
"""
|
||||
|
||||
|
||||
|
||||
|
||||
class RefImageIndicesAndTextPrompt(BaseModel):
|
||||
ref_image_indices: List[int] = Field(
|
||||
description="Indices of reference images selected from the provided images. For example, [0, 2, 5] means selecting the first, third, and sixth images. The indices should be 0-based.",
|
||||
examples=[
|
||||
[1, 3]
|
||||
]
|
||||
)
|
||||
text_prompt: str = Field(
|
||||
description="Text description to guide the image generation. You need to describe the image to be generated, specifying which elements in the generated image should reference which image (and which elements within it). For example, 'Create an image following the given description: \nThe man is standing in the landscape. The man should reference Image 0. The landscape should reference Image 1.' Here, the index of the reference image should refer to its position in the ref_image_indices list, not the sequence number in the provided image list. Refer to the reference image must be in the format of Image N. Do not use any other word except Image.",
|
||||
examples=[
|
||||
"Create an image based on the following guidance: \n Make modifications based on Image 1: Bob's body turns to face the camera, while all other elements remain unchanged. Bob's appearance should refer to Image 0.",
|
||||
"Create an image following the given description: \nThe man is standing in the landscape. The man should reference Image 0. The landscape should reference Image 1."
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
|
||||
class ReferenceImageSelector:
|
||||
def __init__(
|
||||
self,
|
||||
chat_model,
|
||||
):
|
||||
|
||||
self.chat_model = chat_model
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
after=after_func,
|
||||
)
|
||||
async def select_reference_images_and_generate_prompt(
|
||||
self,
|
||||
available_image_path_and_text_pairs: List[Tuple[str, str]],
|
||||
frame_description: str,
|
||||
):
|
||||
filtered_image_path_and_text_pairs = available_image_path_and_text_pairs
|
||||
|
||||
# 1. filter images using text-only model
|
||||
if len(available_image_path_and_text_pairs) >= 8:
|
||||
human_content = []
|
||||
for idx, (_, text) in enumerate(available_image_path_and_text_pairs):
|
||||
human_content.append({
|
||||
"type": "text",
|
||||
"text": f"Image {idx}: {text}"
|
||||
})
|
||||
human_content.append({
|
||||
"type": "text",
|
||||
"text": human_prompt_template_select_reference_images.format(frame_description=frame_description)
|
||||
})
|
||||
parser = PydanticOutputParser(pydantic_object=RefImageIndicesAndTextPrompt)
|
||||
|
||||
messages = [
|
||||
SystemMessage(content=system_prompt_template_select_reference_images_only_text.format(format_instructions=parser.get_format_instructions())),
|
||||
HumanMessage(content=human_content)
|
||||
]
|
||||
|
||||
chain = self.chat_model | parser
|
||||
|
||||
try:
|
||||
ref = await chain.ainvoke(messages)
|
||||
filtered_image_path_and_text_pairs = select_pairs_by_indices(available_image_path_and_text_pairs, ref.ref_image_indices)
|
||||
logging.info(f"Filtered image idx:{ref.ref_image_indices}")
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error get image prompt: \n{e}")
|
||||
raise e
|
||||
|
||||
# 2. filter images using multimodal model
|
||||
human_content = []
|
||||
for idx, (image_path, text) in enumerate(filtered_image_path_and_text_pairs):
|
||||
human_content.append({
|
||||
"type": "text",
|
||||
"text": f"Image {idx}: {text}"
|
||||
})
|
||||
human_content.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_path_to_b64(image_path)}
|
||||
})
|
||||
human_content.append({
|
||||
"type": "text",
|
||||
"text": human_prompt_template_select_reference_images.format(frame_description=frame_description)
|
||||
})
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=RefImageIndicesAndTextPrompt)
|
||||
|
||||
messages = [
|
||||
SystemMessage(content=system_prompt_template_select_reference_images_multimodal.format(format_instructions=parser.get_format_instructions())),
|
||||
HumanMessage(content=human_content)
|
||||
]
|
||||
|
||||
chain = self.chat_model | parser
|
||||
|
||||
try:
|
||||
response = await chain.ainvoke(messages)
|
||||
reference_image_path_and_text_pairs = select_pairs_by_indices(filtered_image_path_and_text_pairs, response.ref_image_indices)
|
||||
return {
|
||||
"reference_image_path_and_text_pairs": reference_image_path_and_text_pairs,
|
||||
"text_prompt": response.text_prompt,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error get image prompt: \n{e}")
|
||||
raise e
|
||||
|
||||
|
||||
|
||||
|
||||
def select_pairs_by_indices(pairs, indices):
|
||||
"""Index into pairs with LLM-emitted indices, rejecting out-of-range values.
|
||||
|
||||
Negative indices would silently select the wrong image via Python indexing.
|
||||
"""
|
||||
invalid = [i for i in indices if i < 0 or i >= len(pairs)]
|
||||
if invalid:
|
||||
raise ValueError(f"ref_image_indices out of range: {invalid} (have {len(pairs)} images)")
|
||||
return [pairs[i] for i in indices]
|
||||
@@ -0,0 +1,109 @@
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from interfaces import Event, Scene
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain.chat_models import init_chat_model
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Literal, Tuple, Dict
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from tenacity import retry, stop_after_attempt
|
||||
import logging
|
||||
|
||||
system_prompt_template_get_next_scene = \
|
||||
"""
|
||||
You are an expert scriptwriter specializing in adapting literary works into structured screenplay scenes. Your task is to analyze event descriptions from novels and transform them into compelling screenplay scenes, leveraging relevant context while ignoring extraneous information.
|
||||
|
||||
**TASK**
|
||||
Generate the next scene for a screenplay adaptation based on the provided input. Each scene must include:
|
||||
- Environment: slugline and detailed description
|
||||
- Characters: List of characters appearing in the scene, with their static features (e.g., facial features, body shape), dynamic features (e.g., clothing, accessories), and visibility status
|
||||
- Script: Character actions and dialogues in standard screenplay format
|
||||
|
||||
**INPUT**
|
||||
- Event Description: A clear, concise summary of the event to adapt. The event description is enclosed within <EVENT_DESCRIPTION_START> and <EVENT_DESCRIPTION_END> tags.
|
||||
- Context Fragments: Multiple excerpts retrieved from the novel via RAG. These may contain irrelevant passages. Ignore any content not directly related to the event. The sequence of context fragments is enclosed within <CONTEXT_FRAGMENTS_START> and <CONTEXT_FRAGMENTS_END> tags. Each fragment in the sequence is enclosed within its own <FRAGMENT_N_START> and <FRAGMENT_N_END> tags, with N being the fragment number.
|
||||
- Previous Scenes (if any): Already adapted scenes for context (may be empty). The sequence of previous scenes is enclosed within <PREVIOUS_SCENES_START> and <PREVIOUS_SCENES_END> tags. Each scene is enclosed within its own <SCENE_N_START> and <SCENE_N_END> tags, with N being the scene number.
|
||||
|
||||
**OUTPUT**
|
||||
{format_instructions}
|
||||
|
||||
**GUIDELINES**
|
||||
1. Extract scenes based on the provided context fragments. Strive to preserve the original meaning and dialogue without making arbitrary alterations. When adapting, ensure that every line of dialogue has a corresponding or derivative basis in the original text.
|
||||
2. Focus on Relevance: Use only context fragments that directly align with the event description. Disregard any unrelated paragraphs.
|
||||
3. Dialogues and Actions: Convert descriptive prose into actionable lines and dialogues. Invent minimal necessary dialogue if implied but not explicit in the context.
|
||||
4. Conciseness: Keep descriptions brief and visual. Avoid prose-like explanations.
|
||||
5. Format Consistency: Ensure industry-standard screenplay structure.
|
||||
6. Implicit Inference: If context fragments lack exact details, infer logically from the event description or broader narrative context.
|
||||
7. No Extraneous Content: Do not include scenes, characters, or dialogues unrelated to the core event.
|
||||
8. The character must be an individual, not a group of individuals (such as a crowd of onlookers or a rescue team).
|
||||
9. When the location or time changes, a new scene should be created. The total number of scenes should not more than 5!!!
|
||||
10. The language of outputs in values should be same as the input.
|
||||
"""
|
||||
|
||||
|
||||
human_prompt_template_get_next_scene = \
|
||||
"""
|
||||
<EVENT_DESCRIPTION_START>
|
||||
{event_description}
|
||||
<EVENT_DESCRIPTION_END>
|
||||
|
||||
<CONTEXT_FRAGMENTS_START>
|
||||
{context_fragments}
|
||||
<CONTEXT_FRAGMENTS_END>
|
||||
|
||||
<PREVIOUS_SCENES_START>
|
||||
{previous_scenes}
|
||||
<PREVIOUS_SCENES_END>
|
||||
"""
|
||||
|
||||
|
||||
|
||||
|
||||
class SceneExtractor:
|
||||
def __init__(
|
||||
self,
|
||||
api_key,
|
||||
base_url,
|
||||
chat_model,
|
||||
):
|
||||
self.chat_model = init_chat_model(
|
||||
model=chat_model,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model_provider="openai",
|
||||
)
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(5),
|
||||
after=lambda retry_state: logging.warning(f"Retrying SceneExtractor.get_next_scene due to error: {retry_state.outcome.exception()}"),
|
||||
)
|
||||
async def get_next_scene(
|
||||
self,
|
||||
relevant_chunks: List[str],
|
||||
event: Event,
|
||||
previous_scenes: List[Scene]
|
||||
) -> Scene:
|
||||
|
||||
context_fragments_str = "\n".join([f"<FRAGMENT_{i}_START>\n{chunk}\n<FRAGMENT_{i}_END>" for i, chunk in enumerate(relevant_chunks)])
|
||||
|
||||
previous_scenes_str = "\n".join([f"<SCENE_{i}_START>\n{scene}\n<SCENE_{i}_END>" for i, scene in enumerate(previous_scenes)])
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=Scene)
|
||||
|
||||
messages = [
|
||||
SystemMessage(
|
||||
content=system_prompt_template_get_next_scene.format(
|
||||
format_instructions=parser.get_format_instructions(),
|
||||
),
|
||||
),
|
||||
HumanMessage(
|
||||
content=human_prompt_template_get_next_scene.format(
|
||||
event_description=str(event),
|
||||
context_fragments=context_fragments_str,
|
||||
previous_scenes=previous_scenes_str,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
chain = self.chat_model | parser
|
||||
scene = await chain.ainvoke(messages)
|
||||
return scene
|
||||
@@ -0,0 +1,167 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain.chat_models import init_chat_model
|
||||
from pydantic import BaseModel, Field
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
|
||||
from utils.retry import after_func
|
||||
|
||||
|
||||
|
||||
system_prompt_template_develop_story = \
|
||||
"""
|
||||
[Role]
|
||||
You are a seasoned creative story generation expert. You possess the following core skills:
|
||||
- Idea Expansion and Conceptualization: The ability to expand a vague idea, a one-line inspiration, or a concept into a fleshed-out, logically coherent story world.
|
||||
- Story Structure Design: Mastery of classic narrative models like the three-act structure, the hero's journey, etc., enabling you to construct engaging story arcs with a beginning, middle, and end, tailored to the story's genre.
|
||||
- Character Development: Expertise in creating three-dimensional characters with motivations, flaws, and growth arcs, and designing complex relationships between them.
|
||||
- Scene Depiction and Pacing: The skill to vividly depict various settings and precisely control the narrative rhythm, allocating detail appropriately based on the required number of scenes.
|
||||
- Audience Adaptation: The ability to adjust the language style, thematic depth, and content suitability based on the target audience (e.g., children, teenagers, adults).
|
||||
- Screenplay-Oriented Thinking: When the story is intended for short film or movie adaptation, you can naturally incorporate visual elements (e.g., scene atmosphere, key actions, dialogue) into the narrative, making the story more cinematic and filmable.
|
||||
|
||||
[Task]
|
||||
Your core task is to generate a complete, engaging story that conforms to the specified requirements, based on the user's provided "Idea" and "Requirements."
|
||||
|
||||
[Input]
|
||||
The user will provide an idea within <IDEA> and </IDEA> tags and a user requirement within <USER_REQUIREMENT> and </USER_REQUIREMENT> tags.
|
||||
- Idea: This is the core seed of the story. It could be a sentence, a concept, a setting, or a scene. For example,
|
||||
- "A programmer discovers his shadow has a consciousness of its own.",
|
||||
- "What if memories could be deleted and backed up like files?",
|
||||
- "A locked-room murder mystery occurring on a space station."
|
||||
- User Requirement (Optional): Optional constraints or guidelines the user may specify. For example,
|
||||
- Target Audience: e.g., Children (7-12), Young Adults, Adults, All Ages.
|
||||
- Story Type/Genre: e.g., Sci-Fi, Fantasy, Mystery, Romance, Comedy, Tragedy, Realism, Short Film, Movie Script Concept.
|
||||
- Length: e.g., 5 key scenes, a tight story suitable for a 10-minute short film.
|
||||
- Other: e.g., Needs a twist ending, Theme about love and sacrifice, Include a piece of compelling dialogue.
|
||||
|
||||
[Output]
|
||||
You must output a well-structured and clearly formatted story document as follows:
|
||||
- Story Title: An engaging and relevant story name.
|
||||
- Target Audience & Genre: Start by explicitly restating: "This story is targeted at [User-Specified Audience], in the [User-Specified Genre] genre."
|
||||
- Story Outline/Summary: Provide a one-paragraph (100-200 words) summary of the entire story, covering the core plot, central conflict, and outcome.
|
||||
Main Characters Introduction: Briefly introduce the core characters, including their names, key traits, and motivations.
|
||||
- Full Story Narrative:
|
||||
- If the number of scenes is unspecified, narrate the story naturally in paragraphs following the "Introduction - Development - Climax - Conclusion" structure.
|
||||
- If a specific number of scenes (e.g., N scenes) is specified, clearly divide the story into N scenes, giving each a subheading (e.g., Scene One: Code at Midnight). The description for each scene should be relatively balanced, including atmosphere, character actions, and dialogue, all working together to advance the plot.
|
||||
- The narrative should be vivid and detailed, matching the specified genre and target audience.
|
||||
- The output should begin directly with the story, without any extra words.
|
||||
|
||||
[Guidelines]
|
||||
- The language of output should be same as the input.
|
||||
- Idea-Centric: Keep the user's core idea as the foundation; do not deviate from its essence. If the user's idea is vague, you can use creativity to make reasonable expansions.
|
||||
- Logical Consistency: Ensure that event progression and character actions within the story have logical motives and internal consistency, avoiding abrupt or contradictory plots.
|
||||
- Show, Don't Tell: Reveal characters' personalities and emotions through their actions, dialogues, and details, rather than stating them flatly. For example, use "He clenched - his fist, nails digging deep into his palm" instead of "He was very angry."
|
||||
- Originality & Compliance: Generate original content based on the user's idea, avoiding direct plagiarism of well-known existing works. The generated content must be positive, healthy, and comply with general content safety policies.
|
||||
"""
|
||||
|
||||
human_prompt_template_develop_story = \
|
||||
"""
|
||||
<IDEA>
|
||||
{idea}
|
||||
</IDEA>
|
||||
|
||||
<USER_REQUIREMENT>
|
||||
{user_requirement}
|
||||
</USER_REQUIREMENT>
|
||||
"""
|
||||
|
||||
|
||||
|
||||
system_prompt_template_write_script_based_on_story = \
|
||||
"""
|
||||
[Role]
|
||||
You are a professional AI script adaptation assistant skilled in adapting stories into scripts. You possess the following skills:
|
||||
- Story Analysis Skills: Ability to deeply understand the story content, identify key plot points, character arcs, and themes.
|
||||
- Scene Segmentation Skills: Ability to break down the story into logical scene units based on continuity of time and location.
|
||||
- Script Writing Skills: Familiarity with script formats (e.g., for short films or movies), capable of crafting vivid dialogue, action descriptions, and stage directions.
|
||||
- Adaptive Adjustment Skills: Ability to adjust the script's style, language, and content based on user requirements (e.g., target audience, story genre, number of scenes).
|
||||
- Creative Enhancement Skills: Ability to appropriately add dramatic elements to enhance the script's appeal while remaining faithful to the original story.
|
||||
|
||||
[Task]
|
||||
Your task is to adapt the user's input story, along with optional requirements, into a script divided by scenes. The output should be a list of scripts, each representing a complete script for one scene. Each scene must be a continuous dramatic action unit occurring at the same time and location.
|
||||
|
||||
[Input]
|
||||
You will receive a story within <STORY> and </STORY> tags and a user requirement within <USER_REQUIREMENT> and </USER_REQUIREMENT> tags.
|
||||
- Story: A complete or partial narrative text, which may contain one or more scenes. The story will provide plot, characters, dialogues, and background descriptions.
|
||||
- User Requirement (Optional): A user requirement, which may be empty. The user requirement may include:
|
||||
- Target audience (e.g., children, teenagers, adults).
|
||||
- Script genre (e.g., micro-film, moive, short drama).
|
||||
- Desired number of scenes (e.g., "divide into 3 scenes").
|
||||
- Other specific instructions (e.g., emphasize dialogue or action).
|
||||
|
||||
[Output]
|
||||
{format_instructions}
|
||||
|
||||
[Guidelines]
|
||||
- The language of output in values should be same as the input story.
|
||||
- Scene Division Principles: Each scene must be based on the same time and location. Start a new scene when the time or location changes. If the user specifies the number of scenes, try to match the requirement. Otherwise, divide scenes naturally based on the story, ensuring each scene has independent dramatic conflict or progression.
|
||||
- Script Formatting Standards: Use standard script formatting: Scene headings in full caps or bold, character names centered or capitalized, dialogue indented, and action descriptions in parentheses.
|
||||
- Coherence and Fluidity: Ensure natural transitions between scenes and overall story flow. Avoid abrupt plot jumps.
|
||||
- Visual Enhancement Principles: All descriptions must be "filmable". Use concrete actions instead of abstract emotions (e.g., "He turns away to avoid eye contact" instead of "He feels ashamed"). Decribe rich environmental details include lighting, props, weather, etc., to enhance the atmosphere. Visualize character performances such as express internal states through facial expressions, gestures, and movements (e.g., "She bites her lip, her hands trembling" to imply nervousness).
|
||||
- Consistency: Ensure dialogue and actions align with the original story's intent, without deviating from the core plot.
|
||||
"""
|
||||
|
||||
|
||||
human_prompt_template_write_script_based_on_story = \
|
||||
"""
|
||||
<STORY>
|
||||
{story}
|
||||
</STORY>
|
||||
|
||||
<USER_REQUIREMENT>
|
||||
{user_requirement}
|
||||
</USER_REQUIREMENT>
|
||||
"""
|
||||
|
||||
|
||||
class Screenwriter:
|
||||
def __init__(
|
||||
self,
|
||||
chat_model: str,
|
||||
):
|
||||
self.chat_model = chat_model
|
||||
|
||||
async def develop_story(
|
||||
self,
|
||||
idea: str,
|
||||
user_requirement: Optional[str] = None,
|
||||
) -> str:
|
||||
messages = [
|
||||
("system", system_prompt_template_develop_story),
|
||||
("human", human_prompt_template_develop_story.format(idea=idea, user_requirement=user_requirement)),
|
||||
]
|
||||
response = await self.chat_model.ainvoke(messages)
|
||||
story = response.content
|
||||
return story
|
||||
|
||||
|
||||
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, max=30), after=after_func)
|
||||
async def write_script_based_on_story(
|
||||
self,
|
||||
story: str,
|
||||
user_requirement: Optional[str] = None,
|
||||
) -> List[str]:
|
||||
|
||||
|
||||
class WriteScriptBasedOnStoryResponse(BaseModel):
|
||||
script: List[str] = Field(
|
||||
...,
|
||||
description="The script based on the story. Each element is a scene "
|
||||
)
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=WriteScriptBasedOnStoryResponse)
|
||||
format_instructions = parser.get_format_instructions()
|
||||
|
||||
messages = [
|
||||
("system", system_prompt_template_write_script_based_on_story.format(format_instructions=format_instructions)),
|
||||
("human", human_prompt_template_write_script_based_on_story.format(story=story, user_requirement=user_requirement)),
|
||||
]
|
||||
response = await self.chat_model.ainvoke(messages)
|
||||
response = parser.parse(response.content)
|
||||
script = response.script
|
||||
return script
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
import logging
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain.chat_models import init_chat_model
|
||||
from pydantic import BaseModel, Field
|
||||
from tenacity import retry, stop_after_attempt
|
||||
|
||||
|
||||
system_prompt_template_script_enhancer = \
|
||||
"""
|
||||
[Role]
|
||||
You are a senior screenplay polishing and continuity expert.
|
||||
|
||||
[Task]
|
||||
Enhance a planned narrative script by adding specific, concrete sensory details, tightening continuity, clarifying scene transitions, and keeping terminology consistent (character names, locations, objects). Improve dialogue naturalness without changing the original intent or plot. Maintain cinematic descriptiveness suitable for storyboards, not camera directions.
|
||||
|
||||
[Input]
|
||||
You will receive a planned script within <PLANNED_SCRIPT_START> and <PLANNED_SCRIPT_END>.
|
||||
|
||||
[Output]
|
||||
{format_instructions}
|
||||
|
||||
[Guidelines]
|
||||
1. Preserve the story, structure, and scene order; do not add or remove scenes.
|
||||
2. Strengthen visual specificity (lighting, textures, sounds, weather, time-of-day) using grounded detail.
|
||||
3. Ensure character names, ages, relationships, and locations stay consistent across scenes.
|
||||
5. Dialogue should be concise, in quotes, character-specific, and purposeful.
|
||||
6. Avoid camera jargon (e.g., cut to, close-up) and voiceover formatting.
|
||||
7. No metaphors.
|
||||
8. Repetition for Precision
|
||||
Re‑state important objects/actors often (vehicle name, seat position, or character role) to remove ambiguity. Accuracy takes precedence over rhythm — redundancy is acceptable.
|
||||
9. Character Features for Dialogue
|
||||
For each character in the conversation, repeat the core voice description (e.g., male, early 50s, South African–North American accent) using the same prompt each time.
|
||||
10. Preserve the original narration symbols if exists (eg. Narration: "Everything is looking good").
|
||||
|
||||
Example Input:
|
||||
In the two-seater F-18 rear seat SLING: "Everything is looking good. All systems are green, Elon. We’re ready for takeoff."
|
||||
In the two-seater F-18 front seat Elon Musk: "Understood, Sling. Let’s get this show on the road."
|
||||
In the two‑seater F‑18 rear seat SLING: "Roger that. Strap in tight, boss. It’s gonna be a smooth ride."
|
||||
In the two‑seater F‑18 front seat ELON MUSK: "Smooth is good. Let’s keep it that way."
|
||||
|
||||
Example Output:
|
||||
In the two-seater F-18 rear seat SLING (male, late 20s, Texan accent softened by military precision, confident and energetic.): "Everything is looking good. All systems are green, Elon. We’re ready for takeoff."
|
||||
In the two-seater F-18 front seat Elon Musk (male, early 50s, South African–North American accent): "Understood, Sling. Let’s get this show on the road."
|
||||
In the two‑seater F‑18 rear seat SLING (male, late 20s, Texan accent softened by military precision, confident and energetic.): "Roger that. Strap in tight, boss. It’s gonna be a smooth ride."
|
||||
In the two‑seater F‑18 front seat ELON MUSK (male, early 50s, South African–North American accent): "Smooth is good. Let’s keep it that way."
|
||||
10. Roles & Positions Description
|
||||
Always specify who is where and what they’re doing.
|
||||
Example Input: “In the cockpit front seat of the two‑seat F‑18, the pilot checks his controls.”
|
||||
Example Output: “In the cockpit front seat of the two‑seat F‑18, Elon Musk checks his controls.”
|
||||
Avoid shorthand (“the pilot”) unless you’ve already identified them in that exact position.
|
||||
|
||||
Warnings
|
||||
No camera directions. No metaphors. Do not change the plot.
|
||||
"""
|
||||
|
||||
|
||||
human_prompt_template_script_enhancer = \
|
||||
"""
|
||||
<PLANNED_SCRIPT_START>
|
||||
{planned_script}
|
||||
<PLANNED_SCRIPT_END>
|
||||
"""
|
||||
|
||||
|
||||
class EnhancedScriptResponse(BaseModel):
|
||||
enhanced_script: str = Field(
|
||||
...,
|
||||
description="A refined script version with clearer continuity, stronger concrete detail, and improved dialogue while preserving the original story and scene order."
|
||||
)
|
||||
|
||||
|
||||
class ScriptEnhancer:
|
||||
def __init__(
|
||||
self,
|
||||
chat_model: str,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model_provider: str = "openai",
|
||||
):
|
||||
self.chat_model = init_chat_model(
|
||||
model=chat_model,
|
||||
model_provider=model_provider,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
after=lambda retry_state: logging.warning(f"Retrying enhance_script due to error: {retry_state.outcome.exception()}"),
|
||||
)
|
||||
async def enhance_script(
|
||||
self,
|
||||
planned_script: str,
|
||||
) -> EnhancedScriptResponse:
|
||||
"""
|
||||
Enhance a planned script with more concrete detail and continuity polish.
|
||||
"""
|
||||
parser = PydanticOutputParser(pydantic_object=EnhancedScriptResponse)
|
||||
prompt_template = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", system_prompt_template_script_enhancer),
|
||||
("human", human_prompt_template_script_enhancer),
|
||||
]
|
||||
)
|
||||
chain = prompt_template | self.chat_model | parser
|
||||
|
||||
try:
|
||||
logging.info("Enhancing planned script...")
|
||||
response: EnhancedScriptResponse = await chain.ainvoke(
|
||||
{
|
||||
"format_instructions": parser.get_format_instructions(),
|
||||
"planned_script": planned_script,
|
||||
}
|
||||
)
|
||||
logging.info("Script enhancement completed.")
|
||||
return response.enhanced_script
|
||||
except Exception as e:
|
||||
logging.error(f"Error enhancing script: \n{e}")
|
||||
raise e
|
||||
|
||||
|
||||
@@ -0,0 +1,433 @@
|
||||
import logging
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain.chat_models import init_chat_model
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Literal
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
|
||||
from utils.retry import after_func
|
||||
|
||||
|
||||
narrative_script_prompt_template = \
|
||||
"""
|
||||
You are a world-class creative writing and screenplay development expert with extensive experience in story structure, character development, and narrative pacing.
|
||||
|
||||
**Task**
|
||||
Your task is to transform a basic story idea into a comprehensive, engaging script with rich narrative detail, compelling character arcs, and cinematic storytelling elements.
|
||||
|
||||
**Input**
|
||||
You will receive a basic story idea or concept enclosed within <BASIC_IDEA_START> and <BASIC_IDEA_END>.
|
||||
|
||||
Below is a simple example of the input:
|
||||
|
||||
<BASIC_IDEA_START>
|
||||
A person discovers they can time travel but every time they change something, they lose a memory.
|
||||
<BASIC_IDEA_END>
|
||||
|
||||
**Output**
|
||||
{format_instructions}
|
||||
|
||||
**Guidelines**
|
||||
No metaphors allowed!!! (eg. A gust of wind rustled through it, a ghostly touch. ; an F1 car that looks less like a vehicle and more like a fighter jet stripped of its wings)
|
||||
|
||||
1. **Story Structure**: Develop a clear three-act structure with proper setup, confrontation, and resolution. Include compelling plot points, rising action, climax, develop the content according to the plot timeline, maintain a clear main plotline, and maintain coherent narrative connections. Keep the plot moving forward. Avoid summarizing events and characters, and use dialogue between key characters appropriately.
|
||||
|
||||
2. **Character Development**: Create well-rounded characters with clear motivations, flaws, and character arcs. Ensure protagonists have relatable goals and face meaningful obstacles.
|
||||
|
||||
3. **Visual Storytelling**: Write with cinematic language that emphasizes visual elements, actions, and atmospheric details rather than exposition-heavy dialogue.
|
||||
|
||||
4. **Emotional Depth**: Incorporate emotional beats, internal conflicts, and character relationships that resonate with audiences.
|
||||
|
||||
5. **Pacing and Tension**: Build suspense and maintain engagement through proper scene transitions, conflict escalation, and strategic revelation of information.
|
||||
|
||||
6. **Genre Consistency**: Maintain appropriate tone, style, and conventions for the story's genre while adding unique creative elements.
|
||||
|
||||
7. **Dialogue Quality**: When you writing some dialogue, you should use the:" " symbols (eg. Peter says: "Everything is looking good. All systems are green, Elon. We’re ready for takeoff."). Do not use voiceover format. Create natural, character-specific dialogue that advances plot and reveals personality without being overly expository.
|
||||
|
||||
8. **Thematic Elements**: Weave in meaningful themes and subtext that give the story depth and universal appeal.
|
||||
|
||||
9. **Conflict and Stakes**: Establish clear external and internal conflicts with high stakes that matter to both characters and audience.
|
||||
|
||||
10. **Satisfying Resolution**: Ensure all major plot threads are resolved and character arcs reach meaningful conclusions.
|
||||
|
||||
11. **Each dialogue should not too short or too long, (eg."Everything is looking good. All systems are green, Elon. We’re ready for takeoff." )
|
||||
|
||||
|
||||
**Warnings**
|
||||
|
||||
Don't write any camera movement in the script (eg. cut to), you should write the script by using storyboard description, not camera view.
|
||||
No metaphors allowed!!! (eg. A gust of wind rustled through it, a ghostly touch. ; an F1 car that looks less like a vehicle and more like a fighter jet stripped of its wings)
|
||||
|
||||
|
||||
**Examples of narrative scripts**
|
||||
|
||||
The starry sky is vast, the Milky Way glittering.
|
||||
On the beach, there's a fire, a portable dining table and chairs (three balloons tied to one corner, swaying in the wind), an SUV, and a camping tent. Next to the tent is an astronomical telescope. A man (Liu Peiqiang, 35, reserved) operates the telescope, while a little boy (Liu Qi, 4, Liu Peiqiang's son) observes under his father's guidance.
|
||||
Liu Peiqiang (somewhat excitedly) Quick, quick, quick... Look, it's Jupiter... the largest planet in the solar system.
|
||||
Adjusting the telescope's eyepiece's focus and position, Jupiter gradually comes into focus. Liu Qi: Dad, there's an eye on Jupiter.
|
||||
Liu Peiqiang: That's not an eye, it's a massive storm on Jupiter's surface. Liu Qi: Why...?
|
||||
Liu Peiqiang: (touching the boy's head, pointing to the balloons on the table) Jupiter is just a giant balloon, 90% hydrogen. Liu Qi: What is hydrogen?
|
||||
An old man (Han Ziang, 59, Liu Peiqiang's father-in-law and Liu Qi's grandfather) walked out of the tent and stood silently beside Liu Peiqiang and his son.
|
||||
Liu Peiqiang: Hydrogen... Hydrogen is the fuel for Dad's big rocket. The campfire flickered, and Han Ziang turned to look at Liu Peiqiang. Liu Qi: Why? Liu Peiqiang smiled and patted his son's head.
|
||||
Liu Peiqiang (O.S.): When the day comes when you can see Jupiter without a telescope, Dad will be back.
|
||||
|
||||
|
||||
|
||||
**Scriptwriting Guidelines End**
|
||||
|
||||
|
||||
"""
|
||||
|
||||
motion_script_prompt_template = \
|
||||
"""
|
||||
You are a top-tier action and motion-sequence script designer with deep visual expertise in conveying speed, force, choreography, and technical precision. Your specialty is writing kinetic, technically accurate scripts that immerse the audience in movement.
|
||||
|
||||
**Task**
|
||||
Transform a basic idea into a motion-driven script that emphasizes precise action description, clear spatial orientation, and unambiguous, technically accurate details.
|
||||
|
||||
**Input**
|
||||
You will receive a basic idea enclosed within <BASIC_IDEA_START> and <BASIC_IDEA_END>.
|
||||
|
||||
**Output**
|
||||
{format_instructions}
|
||||
|
||||
**Global Rules**
|
||||
No metaphors allowed. Less conversation
|
||||
|
||||
**Motion Style Guidelines**
|
||||
1. Technical Explicitness: Prefer precise nouns and qualifiers over poetic language. Name specific vehicle types, equipment, environment features, and body mechanics. If vehicles are implied, specify make/class if reasonable. If combat, specify stance, guard, strike type, target, and contact result.
|
||||
2. Kinetic Clarity: Make trajectories, vectors, speed/acceleration sensations, and force outcomes explicit. Describe distances and orientations when helpful (e.g., left/right, fore/aft).
|
||||
3. Spatial Cohesion: Maintain a consistent mental map of positions. Keep continuity of who/what is where. When positions change, describe how and by what path.
|
||||
4. Sequenced Action Beats: Write step-by-step beats that can be storyboarded. Each beat should be actionable and unambiguous.
|
||||
5. Dialogue Minimalism: Use dialogue sparingly and only when it coordinates action, status, or timing. Use :"dialogue" quotes for spoken lines.
|
||||
6. Keep the script length similar to the following examples.
|
||||
7. If the user does not specify, only one character can appear at most.
|
||||
8. Less character's actions close-ups, more exterior shots
|
||||
9. Don't describe the character's physical state (e.g. jowls and the loose skin around its neck to press back).
|
||||
|
||||
**Examples of motion & speed immersion fighter scripts** (should be accurate, technical, and explicit, Technical Explicitness: Consistently repeats “two seats F‑18” in each stage direction. Prioritizes precision in identifying the aircraft type and location (front seat / rear seat). Reads almost like a technical report or aviation manual, ensuring no ambiguity.)
|
||||
The immense gray flight deck of a nuclear aircraft carrier cuts through a deep blue ocean. The horizon is a clean, sharp line. Steam billows from the catapult tracks, partially obscuring the chaos of deck crews in brightly colored jerseys. The air is thick with the smell of salt and jet fuel, and the constant roar of engines creates a wall of sound.
|
||||
|
||||
An F-18, is positioned on the steam-powered catapult. Its twin engines blast waves of heat that distort the air behind it. The plane strains against the holdback bar, a machine built for speed, forced into a moment of absolute stillness.
|
||||
|
||||
Epic cinematic style with dramatic wide shots, dynamic camera movements, rich color grading, and theatrical lighting reminiscent of major Hollywood productions. Camera gradually moves forward to pilot Elon Musk (50s, sharp eyes and unwavering focus) sits in the cockpit of a F-18. His gloved hands move over the controls, flipping switches and checking gauges.
|
||||
|
||||
In the F-18 cockpit Elon Musk: "Understood, Sling. Let’s get this show on the road."
|
||||
|
||||
In the F-18 cockpit Elon Musk's left hand push on the F-18 throttle, his right grips the control stick.
|
||||
|
||||
A side view. The Shooter drops to one knee, pointing down the deck. The world seems to hold its breath. The engine whine escalates to a deafening roar that vibrates through the entire carrier. The F-18's twin vertical stabilizers shudder with contained power.
|
||||
|
||||
First-person POV from inside the cockpit of F18. With a violent jolt, the catapult fires. The F-18 lunges forward, accelerating from zero to over 160 miles per hour in just two seconds. The deck becomes a blur of motion. Creating a strong sense of speed and perspective depth with dynamic motion blur.
|
||||
|
||||
A side camera view. Then, with a surge of raw power from the afterburners igniting. The F-18 climbs, asserting its dominance over gravity. The landing gear retracts into the fuselage with a solid thud. Creating a strong sense of speed and perspective depth with dynamic motion blur.
|
||||
|
||||
Elon Musk levels the F-18 wings, the sun glinting off his visor as he scans the empty sky ahead.
|
||||
|
||||
The F-18, a sleek instrument of combat, roars to life as it pushes, slicing through the air with an elegant grace. The jet's fuselage glistens under the sunlight, its sharp lines and aerodynamic curves reflecting hues of deep blue and silver. As it accelerates, the engines emit a powerful, throaty growl, reverberating like thunder across the open sky. Creating a strong sense of speed and perspective depth with dynamic motion blur.
|
||||
|
||||
**Examples of motion & speed immersion F1 racing scripts**
|
||||
Epic cinematic style with dramatic wide shots, dynamic camera movements, rich color grading, and theatrical lighting reminiscent of major Hollywood productions. In the black and gold Formula One cockpit, Camera gradually moves forward to F1 driver Elon Musk (playing the driver, a man in his 40s, with a steely gaze and utter concentration) buckling his harness, his helmet visor which reflects the fluttering checkered flags and a blur of cheering spectators in the stands. He drives a sleek black and gold F1 car.
|
||||
|
||||
The starting lights on the track go out, and First-person POV from inside the cockpit of a black and gold F1 car which starts and speeding through the Arena. You grip the wheel — full throttle. The engine roars, gear shifts snap. The blur of the cheering spectators in the stands flashes on your left. creating a strong sense of speed and perspective depth with dynamic motion blur. are engaged in a frenetic, no-holds-barred race. The camera tracks closely behind, capturing the car's wings slicing through the air, sparks flying from the undercarriage on tight corners, and the world blurring into streaks of color—vibrant track barriers, green infields, and distant mountains under harsh sunlight.
|
||||
|
||||
The camera closely tracks the side with dynamic chasing shots., hugging the ground to capture Elon Musk's sleek black and gold F1 car slicing through the air, its APX tail wing flexing under the wind, sparks erupting from the chassis like fireworks as it powers through tight turns and begins overtaking rivals—dodging a pursuing Formula One car , nearly clipping in a heart-pounding near-miss. Cutting to another close-up on Elon Musk, his gloved hands gripping the F1 steering wheel tightly, while the background track barriers streak by in accelerated motion. Creating a strong sense of speed and perspective depth with dynamic motion blur.
|
||||
|
||||
An aerial view for a wide chase perspective, showing Elon Musk's APX Formula One car boldly overtaking another rival in a daring maneuver, debris scattering across the asphalt as it pulls ahead, the pulsating to a crescendo amidst the intensified roar of engines, whistling wind, and the stronger surge of acceleration that makes the entire frame vibrate with raw power. Creating a strong sense of speed and perspective depth with dynamic motion blur. are engaged in a frenetic, no-holds-barred race.
|
||||
|
||||
A front-mounted chase shot follows, emphasizing the APX tail wing's metallic sheen as the black and gold F1 car banks into a hairpin turn, other Formula One rivals closing in from both sides in a tense three-way battle, the movement acceleration pushing the limits as Elon Musk's black and gold F1 car breaks free, leaving F1 competitors in a cloud of dust.
|
||||
|
||||
The camera jolts into a raw handheld shot as Elon Musk’s APX black and gold F1 car rockets down a blistering straightaway, creating a strong sense of speed and perspective depth with dynamic motion blur, are engaged in a frenetic, no-holds-barred race. Rivals' red-white Formula one car closing in tight on both flanks. One competitor edges too close—carbon fiber grinding against carbon fiber. Sparks erupt in a spray of gold as Elon Musk wrenches the wheel, but the rival's red-white Formula one car fishtails, spinning out of control before slamming violently into the barrier. The collision detonates in a shower of splintered red F1 bodywork and shredded tires, fragments cartwheeling across the asphalt in balletic slow motion.
|
||||
|
||||
Wide aerial shots capture the chaos as smoke and dust mushroom upward, the track swallowed in a haze of flame-orange light. Then—an explosive cut back to full speed—Elon Musk’s sleek black and gold F1 APX car bursts through the choking smoke cloud, unbroken, streaking down the straight. Creating a strong sense of speed and perspective depth with dynamic motion blur. are engaged in a frenetic, no-holds-barred race.
|
||||
|
||||
Another extreme close-up zooms in on F1 driver Elon Musk's visor, the lens focus pronouncing the reflection of the track rushing by, capturing the intensity of his focus amid the chaos. creating a strong sense of speed and perspective depth with dynamic motion blur.
|
||||
|
||||
The sequence escalates with a low-angle chase shot from behind, creating a strong sense of speed and perspective depth with dynamic motion blur. Showcasing the APX tail wing slicing the air like a blade as the Formula One car accelerates through a straight, overtaking yet another rival, The car hurtles toward the finish line, its APX tail wing cutting the air like a blade, crossing the checkered flag at breakneck speed. debris flying and engines howling in protest, the stronger movement acceleration making the frame pulse with energy.
|
||||
|
||||
**Warnings**
|
||||
- Do not use metaphors.
|
||||
|
||||
"""
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
montage_script_prompt_template = \
|
||||
"""
|
||||
You are a top-tier montage script designer with deep expertise in compressing time, juxtaposing images, and shaping emotional arcs through shot selection and rhythm. Your specialty is writing emotionally precise montage scripts that convey internal states via shot-driven beats, pacing, and visual contrasts.
|
||||
|
||||
Task
|
||||
Transform a basic idea into an emotion-driven montage script that emphasizes internal experience through visual sequencing, clear emotional expression per shot/beat, and unambiguous psychological details.
|
||||
|
||||
Input
|
||||
You will receive a basic idea enclosed within <BASIC_IDEA_START> and <BASIC_IDEA_END>.
|
||||
|
||||
Output
|
||||
{format_instructions}
|
||||
|
||||
**Global Rules**
|
||||
No metaphors allowed.
|
||||
Keep dialogue minimal.
|
||||
Use pure paragraph.
|
||||
Convey meaning primarily through shot progression, rhythm, and visual juxtaposition.
|
||||
Montage Style Guidelines
|
||||
Use plain sentence/paragraph
|
||||
For each secene, you should write multiple shots to enhance montage effect.
|
||||
Total no less than 500 words, each paragraph no more than 50 words.
|
||||
Escalation or Resolution: Build an emotional arc across beats. Show explicit changes in emotional state and the cause for each change.
|
||||
Sound Design Minimalism: Use sparse, precise notes for sound/music that influence emotion (tempo rise, percussive cuts, breath presence). Avoid lyrical description.
|
||||
Dialogue Minimalism: Include dialogue only if it marks a clear emotional shift. Use :"dialogue" quotes.
|
||||
Visual Clarity Over Action: Limit complex external action. Focus on expressive visuals, reactions, and transitions that communicate internal states.
|
||||
No extraneous physical traits. Only describe details that influence or reveal emotion.
|
||||
**Warnings**
|
||||
Do not use metaphors.
|
||||
Avoid poetic language; prefer precise, observable details.
|
||||
|
||||
**Examples of scripts**
|
||||
Morning light across a small practice room. A girl (Lisa) around seven lifts a violin from its case. Bow slips on the first note.
|
||||
|
||||
|
||||
She (Lisa) winces, then tries again. Shoulders ease. Relief. Quiet room, a single chair creak.
|
||||
|
||||
|
||||
She (Lisa) rests her cheek on the chinrest. The string hum stabilizes.
|
||||
|
||||
|
||||
A small smile shows on Lisa.
|
||||
|
||||
|
||||
Front hall. School shoes near a folded music stand.
|
||||
|
||||
|
||||
She (Lisa) struggles with the latch. The stand clicks open. Light metal tap on tile.
|
||||
|
||||
|
||||
Afternoon window. She (Lisa) traces notes with a finger. Her mother taps a rhythm on the table.
|
||||
|
||||
|
||||
She (Lisa) frowns, then raises her elbow. Concentration holds. The bow settles. Shared stillness. Page flip, steady breath.
|
||||
|
||||
|
||||
Bathroom. She (Lisa) wipes rosin dust off the instrument, coughing once.
|
||||
|
||||
|
||||
Bedroom floor. Sheet music spread. She (Lisa) circles three notes with a red pencil.
|
||||
|
||||
|
||||
She (Lisa) plays them alone, slow, then again faster. Frustration dips, control returns. Pencil tap stops.
|
||||
|
||||
|
||||
Kitchen doorway. A metronome ticks beside a bowl of fruit. She (Lisa) dials it down two clicks. Shoulders drop. She follows the pulse, bow hand steadier with each measure.
|
||||
|
||||
|
||||
Living room. A TV murmurs. She (Lisa) crosses, lowers the volume, returns to her stand. Boundary set without words. The room holds for practice.
|
||||
|
||||
|
||||
Front steps. Case open to the sun. A neighbor waves. She (Lisa) shields the strings with her palm, smiles, and closes the lid. Protection learned.
|
||||
|
||||
|
||||
Music store aisle. Shoulder rests in a row. She (Lisa) tries one that squeaks, then another that fits. Jaw unclenches. She nods, decision made.
|
||||
|
||||
|
||||
Rain on the window. She (Lisa) misses a shift three times. Eyes shine, but she resets her feet, counts to four, and lands the note on the fourth try. Relief, not triumph. Bow lifts, still.
|
||||
|
||||
|
||||
Mirror practice. Thin tape marks on the fingerboard. She (Lisa) glances once, places a finger true, then plays without looking. Confidence grows around the guide.
|
||||
|
||||
|
||||
School hallway before recital. Cold hands under a dryer. She (Lisa) shakes out wrists. Fear thins to focus. She walks toward the stage door, steps even.
|
||||
|
||||
|
||||
Curtain edge. Small tremor at the frog. She (Lisa) loosens grip, breathes, and steps into light.
|
||||
|
||||
|
||||
Two clean phrases. One fuzzy entrance. She (Lisa) holds tempo, corrects on the next measure. Recovery without apology.
|
||||
|
||||
|
||||
Exit corridor. Water bottle cap clicks. She (Lisa) writes in a pocket notebook: “Entrance softer, elbow high.” Emotion measured by action.
|
||||
|
||||
|
||||
Saturday morning. An online tutorial freezes mid-vibrato. She (Lisa) mimics the motion without sound. Adds bow. Wobble uneven. She smiles anyway. Incremental progress accepted.
|
||||
|
||||
|
||||
Park bench. Practice mute on the bridge. Joggers pass without looking. She (Lisa) finishes a scale, closes her eyes a moment, then starts the etude. Privacy inside noise.
|
||||
|
||||
|
||||
Bedroom desk. A planner open. She (Lisa) blocks out “scales + shifts” for fifteen minutes daily. A small star beside Sunday. Plan replaces hope.
|
||||
|
||||
|
||||
Evening soreness. A red mark under her jaw. She (Lisa) folds a soft cloth over the rest, tries again. Mark fades. Comfort adjusted, practice continues.
|
||||
|
||||
|
||||
String snap. Sharp, quick. She (Lisa) flinches, then opens a spare packet, threads, winds, tunes slow. Disruption handled. Bow returns to the string.
|
||||
|
||||
|
||||
Phone buzz. A friend’s invitation lights the screen. She (Lisa) looks once, turns it face down, and plays the piece end to end. Reward after task.
|
||||
|
||||
|
||||
Audition day. Waiting chairs in a line. She (Lisa) air-bows the first phrase, eyes closed. Shoulders stay low. Name called. She stands smoothly.
|
||||
|
||||
|
||||
Small studio. Two judges, still faces. She (Lisa) tunes, begins. First note centered, breath even. A slip in the middle; tempo holds. The last note rings.
|
||||
|
||||
|
||||
Street outside. She (Lisa) exhales into cool air, checks her watch, and walks home. No jump, no slump. Next step implied.
|
||||
|
||||
|
||||
Kitchen table. Acceptance email on a tablet. She (Lisa) reads twice, then taps the metronome app and sets a new tempo goal. Celebration nested inside routine.
|
||||
|
||||
|
||||
Summer afternoon. Open window, distant mower. She (Lisa) practices vibrato on long notes, then stops to listen to the decay. Ear sharpens.
|
||||
|
||||
|
||||
Library corner. She (Lisa) copies fingerings in neat pencil on a fresh sheet. The messy draft slides into recycling. Order replaces clutter.
|
||||
|
||||
|
||||
Community center stage. A quartet rehearsal. She (Lisa) watches the leader’s breath, lifts with it, and enters together. Listening added to playing.
|
||||
|
||||
|
||||
Night lamp. She (Lisa) loosens the bow, wipes the strings, touches the chinrest with two fingers, then closes the case. Habit completes the day. Quiet returns.
|
||||
"""
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
human_prompt_template_script_planner = \
|
||||
"""
|
||||
<BASIC_IDEA_START>
|
||||
{basic_idea}
|
||||
<BASIC_IDEA_END>
|
||||
"""
|
||||
|
||||
|
||||
class IntentRouterResponse(BaseModel):
|
||||
intent: Literal["narrative", "motion", "montage"] = Field(
|
||||
..., description="Routing decision: 'narrative' for characters multi-conversation focus, 'motion' for action/kinetic focus, 'montage' for emotional montage focus"
|
||||
)
|
||||
rationale: Optional[str] = Field(
|
||||
default=None, description="Brief reason for the classification"
|
||||
)
|
||||
|
||||
|
||||
class PlannedScriptResponse(BaseModel):
|
||||
planned_script: str = Field(
|
||||
...,
|
||||
description="The full planned script with rich narrative detail, character development, dialogue, and cinematic descriptions. Should be significantly more detailed and engaging than the original basic idea."
|
||||
)
|
||||
|
||||
|
||||
|
||||
class ScriptPlanner:
|
||||
def __init__(
|
||||
self,
|
||||
chat_model: str,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model_provider: str = "openai",
|
||||
):
|
||||
self.chat_model = init_chat_model(
|
||||
model=chat_model,
|
||||
model_provider=model_provider,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, max=30), after=after_func)
|
||||
def plan_script(
|
||||
self,
|
||||
basic_idea: str,
|
||||
) -> PlannedScriptResponse:
|
||||
"""
|
||||
Plan a comprehensive script from a basic idea.
|
||||
|
||||
Args:
|
||||
basic_idea: A simple story concept or idea to be expanded
|
||||
|
||||
Returns:
|
||||
PlannedScriptResponse: A comprehensive script with structure, characters, and narrative detail
|
||||
"""
|
||||
# 1) Route intent to select the appropriate template
|
||||
router_parser = PydanticOutputParser(pydantic_object=IntentRouterResponse)
|
||||
router_prompt_template = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
(
|
||||
'system',
|
||||
"""
|
||||
You are an intent router for script planning. Classify the user's basic idea into one of following intents:
|
||||
|
||||
- narrative: The idea centers on character, plot, themes, dialogue, or broad storytelling beats.
|
||||
- motion: The idea centers on action, speed, vehicles, combat, choreography, sports, or any kinetic sequence where precise, technical motion description is primary.
|
||||
- montage: The idea centers on a series of shots that convey an emotional arc through imagery, pacing, and juxtaposition.
|
||||
|
||||
Respond using the required JSON format only
|
||||
{format_instructions}
|
||||
"""
|
||||
),
|
||||
('human', human_prompt_template_script_planner),
|
||||
]
|
||||
)
|
||||
router_chain = router_prompt_template | self.chat_model | router_parser
|
||||
|
||||
routing = router_chain.invoke(
|
||||
{
|
||||
"format_instructions": router_parser.get_format_instructions(),
|
||||
"basic_idea": basic_idea,
|
||||
}
|
||||
)
|
||||
chosen_intent = routing.intent if isinstance(routing, IntentRouterResponse) else "narrative"
|
||||
logging.info(f"[ScriptPlanner] Intent routed to: {chosen_intent}")
|
||||
|
||||
# 2) Build the planning chain with the selected template
|
||||
planning_parser = PydanticOutputParser(pydantic_object=PlannedScriptResponse)
|
||||
|
||||
# Template selection with graceful fallbacks
|
||||
def get_system_template(intent: str) -> str:
|
||||
try:
|
||||
if intent == "narrative":
|
||||
return narrative_script_prompt_template
|
||||
if intent == "motion":
|
||||
return motion_script_prompt_template
|
||||
if intent == "montage":
|
||||
return montage_script_prompt_template
|
||||
except NameError:
|
||||
# Fallbacks if specific templates not defined in scope
|
||||
pass
|
||||
# Default fallback
|
||||
return narrative_script_prompt_template
|
||||
|
||||
system_template = get_system_template(chosen_intent)
|
||||
|
||||
planning_prompt_template = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
('system', system_template),
|
||||
('human', human_prompt_template_script_planner),
|
||||
]
|
||||
)
|
||||
planning_chain = planning_prompt_template | self.chat_model | planning_parser
|
||||
|
||||
try:
|
||||
logging.info(f"Planning script from basic idea: {basic_idea[:100]}...")
|
||||
response = planning_chain.invoke(
|
||||
{
|
||||
"format_instructions": planning_parser.get_format_instructions(),
|
||||
"basic_idea": basic_idea,
|
||||
}
|
||||
)
|
||||
logging.info("Script planning completed.")
|
||||
return response
|
||||
except Exception as e:
|
||||
logging.error(f"Error planning script: \n{e}")
|
||||
raise e
|
||||
|
||||
|
||||
@@ -0,0 +1,275 @@
|
||||
from typing import List, Optional, Literal
|
||||
import asyncio
|
||||
from pydantic import BaseModel, Field
|
||||
from tenacity import retry, stop_after_attempt
|
||||
|
||||
from langchain.chat_models.base import BaseChatModel
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from interfaces import CharacterInScene, ShotDescription, ShotBriefDescription
|
||||
|
||||
from utils.retry import after_func
|
||||
|
||||
|
||||
|
||||
system_prompt_template_design_storyboard = \
|
||||
"""
|
||||
[Role]
|
||||
You are a professional storyboard artist with the following core skills:
|
||||
- Script Analysis: Ability to quickly interpret a script's text, identifying the setting, character actions, dialogue, emotions, and narrative pacing.
|
||||
- Visualization: Expertise in translating written descriptions into visual frames, including composition, lighting, and spatial arrangement.
|
||||
- Storyboarding: Proficiency in cinematic language, such as shot types (e.g., close-up, medium shot, wide shot), camera angles (e.g., high angle, eye-level), camera movements (e.g., zoom, pan), and transitions.
|
||||
- Narrative Continuity: Ability to ensure the storyboard sequence is logically smooth, highlights key plot points, and maintains emotional consistency.
|
||||
- Technical Knowledge: Understanding of basic storyboard formats and industry standards, such as using numbered shots and concise descriptions.
|
||||
|
||||
[Task]
|
||||
Your task is to design a complete storyboard based on a user-provided script (which contains only one scene). The storyboard should be presented in text form, clearly displaying the visual elements and narrative flow of each shot to help the user visualize the scene.
|
||||
|
||||
[Input]
|
||||
The user will provide the following input.
|
||||
- Script:A complete scene script containing dialogue, action descriptions, and scene settings. The script focuses on only one scene; there is no need to handle multiple scene transitions. The script input is enclosed within <SCRIPT> and </SCRIPT>.
|
||||
- Characters List: A list describing basic information for each character, such as name, personality traits, appearance (if relevant). The character list is enclosed within <CHARACTERS> and </CHARACTERS>.
|
||||
- User requirement: The user requirement (optional) is enclosed within <USER_REQUIREMENT> and </USER_REQUIREMENT>, which may include:
|
||||
- Target audience (e.g., children, teenagers, adults).
|
||||
- Storyboard style (e.g., realistic, cartoon, abstract).
|
||||
- Desired number of shots (e.g., "not more than 10 shots").
|
||||
- Other specific instructions (e.g., emphasize the characters' actions).
|
||||
|
||||
[Output]
|
||||
{format_instructions}
|
||||
|
||||
[Guidelines]
|
||||
- Ensure all output values (except keys) match the language used in the script.
|
||||
- Each shot must have a clear narrative purpose—such as establishing the setting, showing character relationships, or highlighting reactions.
|
||||
- Use cinematic language deliberately: close-ups for emotion, wide shots for context, and varied angles to direct audience attention.
|
||||
- When designing a new shot, first consider whether it can be filmed using an existing camera position. Introduce a new one only if the shot size, angle, and focus differ significantly. If the camera undergoes significant movement, it cannot be used thereafter.
|
||||
- Keep character names in visual descriptions and speaker fields consistent with the character list. In visual descriptions, enclose names in angle brackets (e.g., <Alice>), but not in dialogue or speaker fields.
|
||||
- When describing visual elements, it is necessary to indicate the position of the element within the frame. For example, Character A is on the left side of the frame, facing toward the right, with a table in front of him. The table is positioned slightly to the left of the center of the frame. Ensure that invisible elements are not included. For instance, do not describe someone behind a closed door if they cannot be seen.
|
||||
- Avoid unsafe content (violence, discrimination, etc.) in visual descriptions. Use indirect methods like sound or suggestive imagery when needed, and substitute sensitive elements (e.g., ketchup for blood).
|
||||
- Assign at most one dialogue line per character per shot. Each line of dialogue should correspond to a shot.
|
||||
- Each shot requires an independent description without reference to each other.
|
||||
- When the shot focuses on a character, describe which specific body part the focus is on.
|
||||
- When describing a character, it is necessary to indicate the direction they are facing.
|
||||
"""
|
||||
|
||||
|
||||
human_prompt_template_design_storyboard = \
|
||||
"""
|
||||
<SCRIPT>
|
||||
{script_str}
|
||||
</SCRIPT>
|
||||
|
||||
<CHARACTERS>
|
||||
{characters_str}
|
||||
</CHARACTERS>
|
||||
|
||||
<USER_REQUIREMENT>
|
||||
{user_requirement_str}
|
||||
</USER_REQUIREMENT>
|
||||
"""
|
||||
|
||||
|
||||
|
||||
system_prompt_template_decompose_visual_description = \
|
||||
"""
|
||||
[Role]
|
||||
You are a professional visual text analyst, proficient in cinematic language and shot narration. Your expertise lies in deconstructing a comprehensive shot description accurately into three core components: the static first frame, the static last frame, and the dynamic motion that connects them.
|
||||
|
||||
[Task]
|
||||
Your task is to dissect and rewrite a user-provided visual text description of a shot strictly and insightfully into three distinct parts:
|
||||
- First Frame Description: Describe the static image at the very beginning of the shot. Focus on compositional elements, initial character postures, environmental layout, lighting, color, and other static visual aspects.
|
||||
- Last Frame Description: Describe the static image at the very end of the shot. Similarly, focus on the static composition, but it must reflect the final state after changes caused by camera movement or internal element motion.
|
||||
- Motion Description: Describe all movements that occur between the first frame and the last frame. This includes camera movement (e.g., static, push-in, pull-out, pan, track, follow, tilt, etc.) and movement of elements within the shot (e.g., character movement, object displacement, changes in lighting, etc.). This is the most dynamic part of the entire description. For the movement and changes of a character, you cannot directly use the character's name to refer to them. Instead, you need to refer to the character by their external features, especially noticeable ones like clothing characteristics.
|
||||
|
||||
[Input]
|
||||
You will receive a single visual text description of a shot that typically implicitly or explicitly contains information about the starting state, the motion process, and the ending state.
|
||||
Additionally, you will receive a sequence of potential characters, each containing an identifier and a feature.
|
||||
- The description is enclosed within <VISUAL_DESC> and </VISUAL_DESC>.
|
||||
- The character list is enclosed within <CHARACTERS> and </CHARACTERS>.
|
||||
|
||||
|
||||
[Output]
|
||||
{format_instructions}
|
||||
|
||||
[Guidelines]
|
||||
- Ensure all output values (except keys) match the language used in the script.
|
||||
- Ensure the first and last frame descriptions are pure "snapshots," containing no ongoing actions (e.g., "He is about to stand up" is unacceptable; it should be "He is sitting on the chair, leaning slightly forward").
|
||||
- In the motion description, you must clearly distinguish between camera movement and on-screen movement. Use professional cinematic terminology (e.g., dolly shot, pan, zoom, etc.) as precisely as possible to describe camera movement.
|
||||
- In the motion description, you cannot directly use character names to refer to characters; instead, you should use the characters' visible characteristics to refer to them. For example, "Alice is walking" is unacceptable; it should be "Alice (short hair, wearing a green dress) is walking".
|
||||
- The last frame description must be logically consistent with the first frame description and the motion description. All actions described in the motion section should be reflected in the static image of the last frame.
|
||||
- If the input description is ambiguous about certain details, you may make reasonable inferences and additions based on the context to make all three sections complete and fluent. However, core elements must strictly adhere to the input text.
|
||||
- Use accurate, concise, and professional descriptive language. Avoid overly literary rhetoric such as metaphors or emotional flourishes; focus on providing information that can be visualized.
|
||||
- Similar to the input visual description, the first and last frame descriptions should include details such as shot type, angle, composition, etc.
|
||||
- Below are the three types of variation within a shot (not between two shots):
|
||||
(1) 'large' cases typically involve the exaggerated transition shots which means a significant change in the composition and focus, such as smoothly changing from a wide shot to a close-up. It is usually accompanied by significant camera movement (e.g., drone perspective shots across the city).
|
||||
(2) 'medium' cases often involve the introduction of new characters and a character turns from the back to face the front (facing the camera).
|
||||
(3) 'small' cases usually involve minor changes, such as expression changes, movement and pose changes of existing characters(e.g., walking, sitting down, standing up), moderate camera movements(e.g., pan, tilt, track).
|
||||
- When describing a character, it is necessary to indicate the direction they are facing.
|
||||
- The first shot must establish the overall scene environment, using the widest possible shot.
|
||||
- Use as few camera positions as possible.
|
||||
"""
|
||||
|
||||
|
||||
human_prompt_template_decompose_visual_description = \
|
||||
"""
|
||||
<VISUAL_DESC>
|
||||
{visual_desc}
|
||||
</VISUAL_DESC>
|
||||
|
||||
<CHARACTERS>
|
||||
{characters_str}
|
||||
</CHARACTERS>
|
||||
"""
|
||||
|
||||
|
||||
class VisDescDecompositionResponse(BaseModel):
|
||||
ff_desc: str = Field(
|
||||
description="A detailed description of the first frame of the shot, capturing the initial visual elements and composition.",
|
||||
# examples=[
|
||||
# "Medium shot of a supermarket aisle at eye level. Bob(a tall man wearing a blue shirt and jeans) is positioned on the right side of the frame, captured in profile and facing right, while Alice(a young woman with short hair, wearing a green dress) is on the left, shown pushing a shopping cart with her gaze lowered toward the ground. They are arranged in a front-to-back spatial relationship. Shelves line both sides of the frame, and cool-toned fluorescent lighting from above washes over the scene. The vibrant colors of product packaging contrast with the metallic gray of the shopping cart, all contained within a stable, horizontally balanced composition.",
|
||||
# "Extreme long shot. Aerial view from hundreds of meters above the ground. The boundless golden desert resembles undulating frozen waves, occupying the vast majority of the frame. At the very center of the image, a tiny, solitary explorer appears only as a faint dark speck, dragging a long, lonely trail of footprints behind him, stretching all the way to the edge of the frame.",
|
||||
# "Medium shot at eye level angle. Designer A(with a beard, wearing a white suit) leans forward passionately, speaking emphatically. Product Manager B(with a beard, wearing a white T-shirt) sits with crossed arms, looking skeptical. Between them, Development Engineer C(brown hair, wearing a blue T-shirt) appears anxious, glancing between the two. Project Manager D(curly hair, wearing a red T-shirt) prepares to mediate, focusing on a whiteboard. Bright overhead lighting highlights their expressions, with a blurred whiteboard and glass wall in the background.",
|
||||
# "A low-angle close-up shot captures the figure from below, framing him from the chest up. His face appears resolute and commanding, his eyes piercing as he speaks passionately. Flecks of saliva are visible, emphasizing his intensity. The overcast sky breaks with occasional light, casting him as a heroic, almost monumental figure against the gloom.",
|
||||
# "An extremely close-up of an old, motionless pocket watch. Soft light highlights scratches on its brass case and the enamel dial with Roman numerals. The second hand remains fixed at 'VIII', casting a sharp shadow. A wrinkled finger gently touches the glass surface, evoking a tangible sense of stillness and time.",
|
||||
# "An over-the-shoulder shot at eye level, positioned behind Character A(red hair, wearing a white T-shirt). The foreground, including A's shoulder and head, is softly blurred, directing focus onto Character B(with a beard, wearing a white T-shirt)'s face. B's subtle reactions—shifting from surprise to confusion, then to a glimmer of understanding—are clearly visible. The café background is gently blurred with warm lighting.",
|
||||
# ]
|
||||
)
|
||||
ff_vis_char_idxs: List[int] = Field(
|
||||
description="A list of indices of characters that are visible in the first frame of the shot, corresponding to the character list provided in the input.",
|
||||
examples=[[0], [1], [0, 1], []]
|
||||
)
|
||||
lf_desc: str = Field(
|
||||
description="A detailed description of the last frame of the shot, capturing the concluding visual elements and composition.",
|
||||
)
|
||||
lf_vis_char_idxs: List[int] = Field(
|
||||
description="A list of indices of characters that are visible in the last frame of the shot, corresponding to the character list provided in the input.",
|
||||
examples=[[0], [1], [0, 1], []]
|
||||
)
|
||||
motion_desc: str = Field(
|
||||
description="The motion description of the shot. Describe the dynamic visual changes within the shot (camera movement and the movement of elements within the frame)",
|
||||
examples=[
|
||||
"Static camera. Alice (short hair, wearing a green dress) is walking towards the camera.",
|
||||
"Dolly in from meidum shot to close-up. Bob (with a beard, wearing a white T-shirt) smiles to the camera.",
|
||||
]
|
||||
)
|
||||
variation_type: Literal["large", "medium", "small"] = Field(
|
||||
description="Indicates the degree of change between the first frame and the last frame.",
|
||||
)
|
||||
variation_reason: str = Field(
|
||||
description="The reason for the variation type of the shot.",
|
||||
examples=[
|
||||
"This is a smooth transition shot from the sky to the ground. The content of the shot has changed significantly, so the variation type is large.",
|
||||
"Compared to the first frame, a new character appears in the last frame, and there are no significant changes in the composition. So the variation type is medium.",
|
||||
"Compared to the first frame, there are only minor changes in the composition. So the variation type is small.",
|
||||
"This shot only shows Alice speaking and the changes in her facial expressions, thus the variation type is small.",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
|
||||
class StoryboardArtist:
|
||||
def __init__(
|
||||
self,
|
||||
chat_model: BaseChatModel,
|
||||
):
|
||||
self.chat_model = chat_model
|
||||
|
||||
|
||||
@retry(stop=stop_after_attempt(3), after=after_func)
|
||||
async def design_storyboard(
|
||||
self,
|
||||
script: str,
|
||||
characters: List[CharacterInScene],
|
||||
user_requirement: Optional[str] = None,
|
||||
retry_timeout: int = 150,
|
||||
) -> List[ShotBriefDescription]:
|
||||
|
||||
class StoryboardResponse(BaseModel):
|
||||
storyboard: List[ShotBriefDescription] = Field(
|
||||
description="A complete storyboard of the scene, including the visual and audio description of each shot.",
|
||||
)
|
||||
|
||||
script_str = script.strip()
|
||||
characters_str = "\n".join([f"Character {index}: {char}" for index, char in enumerate(characters)])
|
||||
user_requirement_str = user_requirement.strip() if user_requirement else ""
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=StoryboardResponse)
|
||||
messages = [
|
||||
('system', system_prompt_template_design_storyboard.format(format_instructions=parser.get_format_instructions())),
|
||||
('human', human_prompt_template_design_storyboard.format(script_str=script_str, characters_str=characters_str, user_requirement_str=user_requirement_str)),
|
||||
]
|
||||
chain = self.chat_model | parser
|
||||
response: StoryboardResponse = await asyncio.wait_for(
|
||||
chain.ainvoke(messages),
|
||||
timeout=retry_timeout,
|
||||
)
|
||||
storyboard = response.storyboard
|
||||
|
||||
return storyboard
|
||||
|
||||
|
||||
|
||||
|
||||
@retry(stop=stop_after_attempt(3), after=after_func)
|
||||
async def decompose_visual_description(
|
||||
self,
|
||||
shot_brief_desc: ShotBriefDescription,
|
||||
characters: List[CharacterInScene],
|
||||
retry_timeout: int = 150,
|
||||
) -> ShotDescription:
|
||||
parser = PydanticOutputParser(pydantic_object=VisDescDecompositionResponse)
|
||||
prompt_template = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
('system', system_prompt_template_decompose_visual_description),
|
||||
('human', human_prompt_template_decompose_visual_description),
|
||||
]
|
||||
)
|
||||
chain = prompt_template | self.chat_model | parser
|
||||
|
||||
visual_desc = shot_brief_desc.visual_desc.strip()
|
||||
|
||||
characters_str = "\n".join([f"{char.identifier_in_scene}: (static) {char.static_features}; (dynamic) {char.dynamic_features}" for char in characters])
|
||||
|
||||
decomposition: VisDescDecompositionResponse = await asyncio.wait_for(
|
||||
chain.ainvoke(
|
||||
input={
|
||||
"format_instructions": parser.get_format_instructions(),
|
||||
"visual_desc": visual_desc,
|
||||
"characters_str": characters_str,
|
||||
},
|
||||
),
|
||||
timeout=retry_timeout,
|
||||
)
|
||||
|
||||
validate_char_idxs(decomposition.ff_vis_char_idxs, len(characters), "ff_vis_char_idxs")
|
||||
validate_char_idxs(decomposition.lf_vis_char_idxs, len(characters), "lf_vis_char_idxs")
|
||||
|
||||
return ShotDescription(
|
||||
idx=shot_brief_desc.idx,
|
||||
is_last=shot_brief_desc.is_last,
|
||||
cam_idx=shot_brief_desc.cam_idx,
|
||||
visual_desc=shot_brief_desc.visual_desc,
|
||||
variation_type=decomposition.variation_type,
|
||||
variation_reason=decomposition.variation_reason,
|
||||
ff_desc=decomposition.ff_desc,
|
||||
ff_vis_char_idxs=decomposition.ff_vis_char_idxs,
|
||||
lf_desc=decomposition.lf_desc,
|
||||
lf_vis_char_idxs=decomposition.lf_vis_char_idxs,
|
||||
motion_desc=decomposition.motion_desc,
|
||||
audio_desc=shot_brief_desc.audio_desc,
|
||||
)
|
||||
|
||||
|
||||
def validate_char_idxs(idxs, num_characters, field_name):
|
||||
"""Reject LLM-emitted character indices outside [0, num_characters).
|
||||
|
||||
Negative values would silently select the wrong character via Python
|
||||
indexing; out-of-range values would crash deep inside the render gather.
|
||||
Raising here lets the @retry on decompose_visual_description re-ask.
|
||||
"""
|
||||
invalid = [idx for idx in idxs if idx < 0 or idx >= num_characters]
|
||||
if invalid:
|
||||
raise ValueError(
|
||||
f"{field_name} contains invalid character indices {invalid}; "
|
||||
f"valid range is 0..{num_characters - 1}"
|
||||
)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.7 MiB |
@@ -0,0 +1,31 @@
|
||||
# ViMax agent runtime local configuration template.
|
||||
# Keep real keys out of commits. Prefer environment variables for shared or CI usage.
|
||||
|
||||
llm:
|
||||
model_provider: openai
|
||||
model: <YOUR_LLM_MODEL>
|
||||
base_url: <YOUR_LLM_BASE_URL>
|
||||
api_key: ''
|
||||
|
||||
image:
|
||||
model: <YOUR_IMAGE_MODEL>
|
||||
base_url: <YOUR_IMAGE_BASE_URL>
|
||||
api_key: ''
|
||||
|
||||
video:
|
||||
model: <YOUR_VIDEO_MODEL>
|
||||
base_url: <YOUR_VIDEO_BASE_URL>
|
||||
api_key: ''
|
||||
|
||||
# Optional. Fill these only when using novel2video planning.
|
||||
embedding:
|
||||
model_provider: openai
|
||||
model: <YOUR_EMBEDDING_MODEL>
|
||||
base_url: <YOUR_EMBEDDING_BASE_URL>
|
||||
api_key: ''
|
||||
|
||||
# Optional. Fill these only when using novel2video planning.
|
||||
reranker:
|
||||
model: <YOUR_RERANKER_MODEL>
|
||||
base_url: <YOUR_RERANKER_BASE_URL>
|
||||
api_key: ''
|
||||
@@ -0,0 +1,31 @@
|
||||
# ViMax agent runtime local configuration template.
|
||||
# Keep real keys out of commits. Prefer environment variables for shared or CI usage.
|
||||
|
||||
llm:
|
||||
model_provider: openai
|
||||
model: <YOUR_LLM_MODEL>
|
||||
base_url: <YOUR_LLM_BASE_URL>
|
||||
api_key: ''
|
||||
|
||||
image:
|
||||
model: <YOUR_IMAGE_MODEL>
|
||||
base_url: <YOUR_IMAGE_BASE_URL>
|
||||
api_key: ''
|
||||
|
||||
video:
|
||||
model: <YOUR_VIDEO_MODEL>
|
||||
base_url: <YOUR_VIDEO_BASE_URL>
|
||||
api_key: ''
|
||||
|
||||
# Optional. Fill these only when using novel2video planning.
|
||||
embedding:
|
||||
model_provider: openai
|
||||
model: <YOUR_EMBEDDING_MODEL>
|
||||
base_url: <YOUR_EMBEDDING_BASE_URL>
|
||||
api_key: ''
|
||||
|
||||
# Optional. Fill these only when using novel2video planning.
|
||||
reranker:
|
||||
model: <YOUR_RERANKER_MODEL>
|
||||
base_url: <YOUR_RERANKER_BASE_URL>
|
||||
api_key: ''
|
||||
@@ -0,0 +1,32 @@
|
||||
chat_model:
|
||||
init_args:
|
||||
model: google/gemini-2.5-flash-lite-preview-09-2025
|
||||
model_provider: openai
|
||||
api_key:
|
||||
base_url: https://openrouter.ai/api/v1
|
||||
# Rate limits for chat model API calls
|
||||
# Set to null to disable rate limiting for this service
|
||||
max_requests_per_minute: 500
|
||||
max_requests_per_day: 2000
|
||||
|
||||
image_generator:
|
||||
class_path: tools.ImageGeneratorNanobananaGoogleAPI
|
||||
init_args:
|
||||
api_key:
|
||||
# Rate limits for image generation API calls
|
||||
# Set to null to disable rate limiting for this service
|
||||
max_requests_per_minute: 10
|
||||
max_requests_per_day: 500
|
||||
|
||||
|
||||
video_generator:
|
||||
class_path: tools.VideoGeneratorVeoGoogleAPI
|
||||
init_args:
|
||||
api_key:
|
||||
# Rate limits for video generation API calls
|
||||
# Set to null to disable rate limiting for this service
|
||||
max_requests_per_minute: 2
|
||||
max_requests_per_day: 10
|
||||
|
||||
|
||||
working_dir: .working_dir/idea2video
|
||||
@@ -0,0 +1,33 @@
|
||||
# Example configuration using MiniMax as the chat model provider.
|
||||
# MiniMax M3 is used via its OpenAI-compatible API.
|
||||
#
|
||||
# Set your API key below or export MINIMAX_API_KEY in the environment.
|
||||
# Available models:
|
||||
# - MiniMax-M3 (latest, recommended)
|
||||
# - MiniMax-M2.7 (previous generation)
|
||||
# - MiniMax-M2.7-highspeed (fast variant)
|
||||
|
||||
chat_model:
|
||||
init_args:
|
||||
model: MiniMax-M3
|
||||
model_provider: minimax
|
||||
api_key: # leave empty to use the MINIMAX_API_KEY environment variable
|
||||
# base_url is auto-resolved to https://api.minimax.io/v1
|
||||
max_requests_per_minute: 500
|
||||
max_requests_per_day: 2000
|
||||
|
||||
image_generator:
|
||||
class_path: tools.ImageGeneratorNanobananaGoogleAPI
|
||||
init_args:
|
||||
api_key:
|
||||
max_requests_per_minute: 10
|
||||
max_requests_per_day: 500
|
||||
|
||||
video_generator:
|
||||
class_path: tools.VideoGeneratorVeoGoogleAPI
|
||||
init_args:
|
||||
api_key:
|
||||
max_requests_per_minute: 2
|
||||
max_requests_per_day: 10
|
||||
|
||||
working_dir: .working_dir/idea2video
|
||||
@@ -0,0 +1,33 @@
|
||||
chat_model:
|
||||
init_args:
|
||||
model: google/gemini-2.5-flash-lite-preview-09-2025
|
||||
model_provider: openai
|
||||
api_key:
|
||||
base_url: https://openrouter.ai/api/v1
|
||||
# Rate limits for chat model API calls
|
||||
# Set to null to disable rate limiting for this service
|
||||
max_requests_per_minute: null
|
||||
max_requests_per_day: null
|
||||
|
||||
|
||||
image_generator:
|
||||
class_path: tools.ImageGeneratorNanobananaGoogleAPI
|
||||
init_args:
|
||||
api_key:
|
||||
# Rate limits for image generation API calls
|
||||
# Set to null to disable rate limiting for this service
|
||||
max_requests_per_minute: 2
|
||||
max_requests_per_day: 50
|
||||
|
||||
|
||||
video_generator:
|
||||
class_path: tools.VideoGeneratorVeoGoogleAPI
|
||||
init_args:
|
||||
api_key:
|
||||
# Rate limits for video generation API calls
|
||||
# Set to null to disable rate limiting for this service
|
||||
max_requests_per_minute: 2
|
||||
max_requests_per_day: 50
|
||||
|
||||
|
||||
working_dir: .working_dir/script2video
|
||||
@@ -0,0 +1,33 @@
|
||||
# Example configuration using MiniMax as the chat model provider.
|
||||
# MiniMax M3 is used via its OpenAI-compatible API.
|
||||
#
|
||||
# Set your API key below or export MINIMAX_API_KEY in the environment.
|
||||
# Available models:
|
||||
# - MiniMax-M3 (latest, recommended)
|
||||
# - MiniMax-M2.7 (previous generation)
|
||||
# - MiniMax-M2.7-highspeed (fast variant)
|
||||
|
||||
chat_model:
|
||||
init_args:
|
||||
model: MiniMax-M3
|
||||
model_provider: minimax
|
||||
api_key: # leave empty to use the MINIMAX_API_KEY environment variable
|
||||
# base_url is auto-resolved to https://api.minimax.io/v1
|
||||
max_requests_per_minute: null
|
||||
max_requests_per_day: null
|
||||
|
||||
image_generator:
|
||||
class_path: tools.ImageGeneratorNanobananaGoogleAPI
|
||||
init_args:
|
||||
api_key:
|
||||
max_requests_per_minute: 2
|
||||
max_requests_per_day: 50
|
||||
|
||||
video_generator:
|
||||
class_path: tools.VideoGeneratorVeoGoogleAPI
|
||||
init_args:
|
||||
api_key:
|
||||
max_requests_per_minute: 2
|
||||
max_requests_per_day: 50
|
||||
|
||||
working_dir: .working_dir/script2video
|
||||
@@ -0,0 +1,22 @@
|
||||
from .camera import Camera
|
||||
from .character import CharacterInScene, CharacterInEvent, CharacterInNovel
|
||||
from .event import Event
|
||||
from .frame import Frame
|
||||
from .image_output import ImageOutput
|
||||
from .scene import Scene
|
||||
from .shot_description import ShotDescription, ShotBriefDescription
|
||||
from .video_output import VideoOutput
|
||||
|
||||
__all__ = [
|
||||
"Camera",
|
||||
"CharacterInScene",
|
||||
"CharacterInEvent",
|
||||
"CharacterInNovel",
|
||||
"Event",
|
||||
"Frame",
|
||||
"ImageOutput",
|
||||
"Scene",
|
||||
"ShotBriefDescription",
|
||||
"ShotDescription",
|
||||
"VideoOutput",
|
||||
]
|
||||
@@ -0,0 +1,43 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Union, Dict, Tuple
|
||||
|
||||
|
||||
|
||||
class Camera(BaseModel):
|
||||
idx: int = Field(
|
||||
description="The index of the camera in the scene, starting from 0.",
|
||||
)
|
||||
|
||||
active_shot_idxs: List[int] = Field(
|
||||
description="The indices of the shots that the camera can film.",
|
||||
)
|
||||
|
||||
parent_cam_idx: Optional[int] = Field(
|
||||
default=None,
|
||||
description="The index of the parent camera. If the camera has no parent, set this to None.",
|
||||
)
|
||||
|
||||
parent_shot_idx: Optional[int] = Field(
|
||||
default=None,
|
||||
description="The index of the dependent shot. If the camera has no parent, set this to None.",
|
||||
)
|
||||
|
||||
reason: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The reason for the selection of the parent camera. If the camera has no parent, set this to None.",
|
||||
)
|
||||
|
||||
parent_shot_idx: Optional[int] = Field(
|
||||
default=None,
|
||||
description="The index of the dependent shot. If the camera has no parent, set this to None.",
|
||||
)
|
||||
|
||||
is_parent_fully_covers_child: Optional[bool] = Field(
|
||||
default=None,
|
||||
description="Whether the parent camera fully covers the child camera's content. If the camera has no parent, set this to None.",
|
||||
)
|
||||
|
||||
missing_info: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The missing information in the child shot that is not covered by the parent shot. If the parent shot fully covers the child shot, set this to None.",
|
||||
)
|
||||
@@ -0,0 +1,100 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Union, Dict
|
||||
from PIL import Image
|
||||
|
||||
|
||||
|
||||
|
||||
class CharacterInScene(BaseModel):
|
||||
idx: int = Field(
|
||||
description="The index of the character in the scene, starting from 0",
|
||||
)
|
||||
identifier_in_scene: str = Field(
|
||||
description="The identifier for the character in this specific scene, which may differ from the base identifier",
|
||||
examples=["Alice", "Bob the Builder"],
|
||||
)
|
||||
is_visible: bool = Field(
|
||||
description="Indicates whether the character is visible in this scene",
|
||||
examples=[True, False],
|
||||
)
|
||||
static_features: str = Field(
|
||||
description="The static features of the character in this specific scene, such as facial features and body shape that remain constant or are rarely changed. If the character is not visible, this field can be left empty.",
|
||||
examples=[
|
||||
"Alice has long blonde hair and blue eyes, and is of slender build.",
|
||||
"Bob the Builder is a middle-aged man with a sturdy build.",
|
||||
]
|
||||
)
|
||||
dynamic_features: str = Field(
|
||||
description="The dynamic features of the character in this specific scene, such as clothing and accessories that may change from scene to scene. If not mentioned, this field can be left empty. If the character is not visible, this field should be None.",
|
||||
examples=[
|
||||
"Wearing a red scarf and a black leather jacket",
|
||||
]
|
||||
)
|
||||
|
||||
def __str__(self):
|
||||
# Alice[visible]
|
||||
# static features: Alice has long blonde hair and blue eyes, and is of slender build.
|
||||
# dynamic features: Wearing a red scarf and a black leather jacket
|
||||
|
||||
s = f"{self.identifier_in_scene}"
|
||||
s += "[visible]" if self.is_visible else "[not visible]"
|
||||
s += "\n"
|
||||
s += f"static features: {self.static_features}\n"
|
||||
s += f"dynamic features: {self.dynamic_features}\n"
|
||||
|
||||
return s
|
||||
|
||||
|
||||
|
||||
class CharacterInEvent(BaseModel):
|
||||
index: int = Field(
|
||||
description="The index of the character in the event, starting from 0",
|
||||
)
|
||||
identifier_in_event: str = Field(
|
||||
description="The unique identifier for the character in the event",
|
||||
examples=["Alice", "Bob the Builder"],
|
||||
)
|
||||
|
||||
active_scenes: Dict[int, str] = Field(
|
||||
description="A dictionary mapping scene indices to their identifiers in specific scenes.",
|
||||
examples=[
|
||||
{0: "Alice", 2: "Alice in Wonderland", 5: "Alice"},
|
||||
{1: "Bob the Builder", 3: "Bob", 4: "Bob"},
|
||||
]
|
||||
)
|
||||
|
||||
static_features: str = Field(
|
||||
description="The static features of the character in the event, such as facial features and body shape that remain constant or are rarely changed.",
|
||||
examples=[
|
||||
"Alice has long blonde hair and blue eyes, and is of slender build. She often wears casual, comfortable clothing.",
|
||||
"Bob the Builder is a middle-aged man with a sturdy build. He typically wears a hard hat and work overalls.",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
|
||||
class CharacterInNovel(BaseModel):
|
||||
index: int = Field(
|
||||
description="The index of the character in the novel, starting from 0",
|
||||
)
|
||||
identifier_in_novel: str = Field(
|
||||
description="The unique identifier for the character in the novel",
|
||||
examples=["Alice", "Bob the Builder"],
|
||||
)
|
||||
|
||||
active_events: Dict[int, str] = Field(
|
||||
description="A dictionary mapping event indices to their identifiers in specific events.",
|
||||
examples=[
|
||||
{0: "Alice", 2: "Alice in Wonderland", 5: "Alice"},
|
||||
{1: "Bob the Builder", 3: "Bob", 4: "Bob"},
|
||||
]
|
||||
)
|
||||
|
||||
static_features: str = Field(
|
||||
description="The static features of the character in the novel, such as facial features and body shape that remain constant or are rarely changed.",
|
||||
examples=[
|
||||
"Alice has long blonde hair and blue eyes, and is of slender build. She often wears casual, comfortable clothing.",
|
||||
"Bob the Builder is a middle-aged man with a sturdy build. He typically wears a hard hat and work overalls.",
|
||||
]
|
||||
)
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Union, Dict
|
||||
from PIL import Image
|
||||
|
||||
|
||||
|
||||
class EnvironmentInScene(BaseModel):
|
||||
slugline: str = Field(
|
||||
description="The slugline of the scene, indicating the location and time of day",
|
||||
examples=[
|
||||
"INT. COFFEE SHOP - NIGHT",
|
||||
"EXT. PARK - DAY",
|
||||
]
|
||||
)
|
||||
description: str = Field(
|
||||
description="A detailed description of the environment in the specific scene. Don't describe any characters or actions here, just the setting.",
|
||||
examples=[
|
||||
"The warm yellow light glowed against the mottled brick wall, while raindrops streaked the glass window with blurred neon reflections. Among the empty booths sat a lone half-finished iced latte—its foam collapsed, a faint lipstick mark on the rim. beads of condensation gleamed on the stainless steel espresso machine, and the record player's turntable rotated slowly in the shadows. A patch of wet floor shimmered with hazy reflected light.",
|
||||
]
|
||||
)
|
||||
|
||||
def __str__(self):
|
||||
s = f"{self.slugline} -- {self.description}"
|
||||
return s
|
||||
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Union, Dict
|
||||
|
||||
|
||||
|
||||
class Event(BaseModel):
|
||||
index: int = Field(
|
||||
description="The index of the event, starting from 0",
|
||||
)
|
||||
|
||||
is_last: bool = Field(
|
||||
description="Indicates if this is the last event in the sequence"
|
||||
)
|
||||
|
||||
description: str = Field(
|
||||
description="A concise description of the event, capturing its essence in one sentence",
|
||||
examples=[
|
||||
"A thief who stole a gem from a museum was caught after a rooftop chase with guards, and the gem was recovered.",
|
||||
]
|
||||
)
|
||||
|
||||
process_chain: List[str] = Field(
|
||||
description="A list of steps or actions that make up the event's process chain, which constitutes a complete causal chain.",
|
||||
examples=[
|
||||
[
|
||||
"A thief steals a gem from a museum, triggering the alarm. Guards notice and begin the chase.",
|
||||
"The thief rushes out the museum's back door and dashes through narrow alleys, with guards closely pursuing and calling for backup.",
|
||||
"The thief climbs a fire escape to the rooftops; the guards follow using low platforms on adjacent buildings.",
|
||||
"The thief leaps across a 1.5-meter gap between two buildings. The guards hesitate but take the risky jump, nearly losing their footing.",
|
||||
"The thief knocks over stacked wooden planks to create an obstacle. The guards dodge but lose speed.",
|
||||
"The thief attempts to slide down a rope to the opposite rooftop, but a guard lunges and grabs their ankle. Both tumble and grapple.",
|
||||
"Backup arrives, subduing the thief and recovering the gem.",
|
||||
],
|
||||
]
|
||||
)
|
||||
|
||||
def __str__(self):
|
||||
s = f"<Event {self.index}>"
|
||||
s += f"\nDescription: {self.description}"
|
||||
s += f"\nProcess Chain:"
|
||||
for process in self.process_chain:
|
||||
s += f"\n- {process}"
|
||||
return s
|
||||
@@ -0,0 +1,20 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Union, Dict, Tuple, Literal
|
||||
|
||||
|
||||
class Frame(BaseModel):
|
||||
shot_idx: int = Field(
|
||||
description="The index of the shot in the sequence, starting from 0."
|
||||
)
|
||||
|
||||
frame_type: Literal["first", "last"] = Field(
|
||||
description="The type of the frame, 'first' for the first frame of the shot, 'last' for the last frame of the shot."
|
||||
)
|
||||
|
||||
cam_idx: int = Field(
|
||||
description="The index of the camera used for this frame, starting from 0."
|
||||
)
|
||||
|
||||
vis_char_idxs: List[int] = Field(
|
||||
description="A list of indices of characters that are visible in this frame, corresponding to the character list provided in the input."
|
||||
)
|
||||
@@ -0,0 +1,61 @@
|
||||
import base64
|
||||
import cv2
|
||||
from typing import List, Literal, Optional, Union
|
||||
from PIL import Image
|
||||
|
||||
from utils.image import download_image
|
||||
|
||||
|
||||
|
||||
class ImageOutput:
|
||||
fmt: Literal["b64", "url", "pil", "np"]
|
||||
ext: str = "png"
|
||||
data: Union[str, Image.Image]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fmt: Literal["b64", "url", "pil", "np"],
|
||||
ext: str,
|
||||
data: Union[str, Image.Image],
|
||||
):
|
||||
self.fmt = fmt
|
||||
self.ext = ext
|
||||
self.data = data
|
||||
|
||||
|
||||
def save_b64(self, path: str) -> None:
|
||||
"""Save a base64 encoded image to the specified path.
|
||||
|
||||
Args:
|
||||
path (str): Path where the image will be saved.
|
||||
"""
|
||||
with open(path, 'wb') as f:
|
||||
f.write(base64.b64decode(self.data))
|
||||
|
||||
def save_url(self, path: str) -> None:
|
||||
"""Download and save an image from a URL to the specified path.
|
||||
|
||||
Args:
|
||||
path (str): Path where the image will be saved.
|
||||
"""
|
||||
download_image(self.data, path)
|
||||
|
||||
def save_pil(self, path: str) -> None:
|
||||
"""Save a PIL Image to the specified path.
|
||||
|
||||
Args:
|
||||
path (str): Path where the image will be saved.
|
||||
"""
|
||||
self.data.save(path)
|
||||
|
||||
def save_np(self, path: str) -> None:
|
||||
"""Save a numpy array to the specified path.
|
||||
|
||||
Args:
|
||||
path (str): Path where the image will be saved.
|
||||
"""
|
||||
cv2.imencode('.png', self.data)[1].tofile(path)
|
||||
|
||||
def save(self, path: str) -> None:
|
||||
save_func = getattr(self, f"save_{self.fmt}")
|
||||
save_func(path)
|
||||
@@ -0,0 +1,57 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Literal, Tuple
|
||||
from interfaces.environment import EnvironmentInScene
|
||||
from interfaces.character import CharacterInScene
|
||||
|
||||
|
||||
class Scene(BaseModel):
|
||||
idx: int = Field(
|
||||
description="The scene index, starting from 0",
|
||||
examples=[0, 1, 2],
|
||||
)
|
||||
is_last: bool = Field(
|
||||
description="Indicates if this is the last scene",
|
||||
examples=[False, True],
|
||||
)
|
||||
environment: EnvironmentInScene = Field(
|
||||
description="The detailed scene setting, including location and time",
|
||||
)
|
||||
characters: List[CharacterInScene] = Field(
|
||||
description="A list of characters appearing in the scene, along with their dynamic features like clothing and accessories",
|
||||
)
|
||||
script: str = Field(
|
||||
description="The screenplay script for the scene, including character actions and dialogues. Character names in the script should be enclosed in <>, except for character names within dialogues.",
|
||||
examples=[
|
||||
"<Jane> paces nervously, clutching a letter. She turns to <John>.\n<Jane>: John, we need to leave tonight.\n<John> shakes his head, stepping toward the window.\n<John>: It's too dangerous.",
|
||||
"<Alice> sits quietly, observing the chaos around her. She whispers to <Bob>.\n<Alice>: Bob, do you think they'll find us here?\n<Bob> nods slowly, his expression grim."
|
||||
],
|
||||
)
|
||||
|
||||
def __str__(self):
|
||||
s = f"Scene {self.idx}:"
|
||||
s += f"\nEnvironment: {str(self.environment)}"
|
||||
s += f"\nCharacters: {', '.join([str(c) for c in self.characters])}"
|
||||
s += f"\nScript: \n{self.script}"
|
||||
return s
|
||||
|
||||
|
||||
|
||||
# class Scene(BaseModel):
|
||||
# index: int = Field(
|
||||
# description="The index of the scene within the event, starting from 0"
|
||||
# )
|
||||
# character_indices: List[int] = Field(
|
||||
# description="List of indices of characters appearing in this scene, including main characters, supporting characters, and extras.",
|
||||
# )
|
||||
# environment_index: int = Field(
|
||||
# description="The index of the environment where the scene takes place."
|
||||
# )
|
||||
# key_items_indices: List[int] = Field(
|
||||
# default=[],
|
||||
# description="List of indices of key items involved in this scene, if any.",
|
||||
# )
|
||||
# script: str = Field(
|
||||
# description="The script of the scene, including actions and dialogues"
|
||||
# )
|
||||
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Literal, Tuple
|
||||
|
||||
|
||||
class ShotBriefDescription(BaseModel):
|
||||
idx: int = Field(
|
||||
description="The index of the shot in the sequence, starting from 0.",
|
||||
examples=[0, 1, 2],
|
||||
)
|
||||
is_last: bool = Field(
|
||||
description="Whether this is the last shot. If True, the story of the script has ended and no more shots will be planned after this one.",
|
||||
examples=[False, True],
|
||||
)
|
||||
|
||||
# visual
|
||||
cam_idx: int = Field(
|
||||
description="The index of the camera in the scene.",
|
||||
examples=[0, 1, 2],
|
||||
)
|
||||
visual_desc: str = Field(
|
||||
description='''A vivid and detailed visual description of the shot that convey rich visual information through text. The character identifiers in the description must match those in the character list and be enclosed in angle brackets (e.g., <Alice>, <Bob>). All visible characters should be described.
|
||||
If there is a conversation, please write down the content of the conversation), when you meet some dialogue, you should write into the visual content description with :" " symbols and the character's features (eg. <SLING> (male, late 20s, Texan accent softened by military precision, confident and energetic.) says: "Gear retracted. Flaps transitioning. Flight path stable. You are clear to climb.").
|
||||
''',
|
||||
examples=[
|
||||
"An over-the-shoulder shot at eye level, positioned behind <Alice>. The foreground, including <Alice>'s shoulder and head, is softly blurred, directing focus onto <Bob>'s face. <Bob>'s subtle reactions—shifting from surprise to delight—are clearly visible. The supermarket background is gently blurred with cool fluorescent lighting.",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
# audio
|
||||
audio_desc: str = Field(
|
||||
description="A detailed description of the audio in the shot.",
|
||||
examples=[
|
||||
"[Sound Effect] Ambient sound (supermarket background noise, shopping cart wheels rolling)",
|
||||
"[Speaker] Alice (Happy): Hello, how are you?",
|
||||
None,
|
||||
],
|
||||
)
|
||||
|
||||
# sound_effect: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The sound effects used in the shot.",
|
||||
# examples=[
|
||||
# "Ambient sound (supermarket background noise, shopping cart wheels rolling)",
|
||||
# None,
|
||||
# ],
|
||||
# )
|
||||
# speaker: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The speaker in the shot, if applicable. If there is no speaker, this field should be set to None.",
|
||||
# examples=[
|
||||
# "Alice",
|
||||
# None,
|
||||
# ],
|
||||
# )
|
||||
# is_speaker_lip_visible: Optional[bool] = Field(
|
||||
# default=None,
|
||||
# description="Indicates whether the speaker's lips are visible in the shot. If there is no speaker, this field should be set to None.",
|
||||
# examples=[
|
||||
# True,
|
||||
# False,
|
||||
# None,
|
||||
# ],
|
||||
# )
|
||||
# line: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The dialogue or monologue in the shot, if applicable. If there is a speaker, there must be a line. If there is no speaker, this field should be set to None.",
|
||||
# examples=[
|
||||
# "Hello, how are you?",
|
||||
# None,
|
||||
# ],
|
||||
# )
|
||||
# emotion: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The emotion of the speaker when delivering the line, if applicable. If there is a speaker, there must be an emotion. If there is no speaker, this field should be set to None.",
|
||||
# examples=[
|
||||
# "Happy",
|
||||
# None,
|
||||
# ],
|
||||
# )
|
||||
|
||||
def __str__(self):
|
||||
s = f"Shot {self.idx}:\n"
|
||||
s += f"Camera Index: {self.cam_idx}\n"
|
||||
s += f"Visual: {self.visual_desc}\n"
|
||||
if self.audio_desc:
|
||||
s += f"Audio: {self.audio_desc}"
|
||||
return s
|
||||
|
||||
|
||||
class ShotDescription(BaseModel):
|
||||
idx: int = Field(
|
||||
description="The index of the shot in the sequence, starting from 0."
|
||||
)
|
||||
is_last: bool = Field(
|
||||
description="Whether this is the last shot in the sequence. If True, no more shots will be planned after this one."
|
||||
)
|
||||
|
||||
# visual
|
||||
cam_idx: int = Field(
|
||||
description="The index of the camera in the scene.",
|
||||
examples=[0, 1, 2],
|
||||
)
|
||||
visual_desc: str = Field(
|
||||
description='''A vivid and detailed visual description of the shot that convey rich visual information through text. The character identifiers in the description must match those in the character list and be enclosed in angle brackets (e.g., <Alice>, <Bob>).
|
||||
If there is a conversation, please write down the content of the conversation), when you meet some dialogue, you should write into the visual content description with :" " symbols and the character's features (eg. <SLING> (male, late 20s, Texan accent softened by military precision, confident and energetic.) says: "Gear retracted. Flaps transitioning. Flight path stable. You are clear to climb."). ''',
|
||||
examples=[
|
||||
"An over-the-shoulder shot at eye level, positioned behind <Alice>. The foreground, including <Alice>'s shoulder and head, is softly blurred, directing focus onto <Bob>'s face. <Bob>'s subtle reactions—shifting from surprise to delight—are clearly visible. The supermarket background is gently blurred with cool fluorescent lighting.",
|
||||
]
|
||||
)
|
||||
variation_type: Literal["large", "medium", "small"] = Field(
|
||||
description="Indicates the degree of change in the shot's content.",
|
||||
examples=["large", "medium", "small"],
|
||||
)
|
||||
variation_reason: str = Field(
|
||||
description="The reason for the variation type of the shot.",
|
||||
examples=[
|
||||
"This is a transition shot where the content of the first frame and the last frame differs dramatically. So the variation type is large.",
|
||||
"Compared to the first frame, a new character appears in the last frame, and there are no significant changes in the composition. So the variation type is medium.",
|
||||
"Compared to the first frame, there are only minor changes in the composition. So the variation type is small.",
|
||||
"This shot only shows Alice speaking and the changes in her facial expressions, thus the variation type is small.",
|
||||
],
|
||||
)
|
||||
|
||||
ff_desc: str = Field(
|
||||
description="The first frame of the shot.",
|
||||
examples=[
|
||||
"Medium shot of a supermarket aisle at eye level. Bob(a tall man wearing a blue shirt and jeans) is positioned on the right side of the frame, captured in profile and facing right, while Alice(a young woman with short hair, wearing a green dress) is on the left, shown pushing a shopping cart with her gaze lowered toward the ground. They are arranged in a front-to-back spatial relationship. Shelves line both sides of the frame, and cool-toned fluorescent lighting from above washes over the scene. The vibrant colors of product packaging contrast with the metallic gray of the shopping cart, all contained within a stable, horizontally balanced composition.",
|
||||
"Extreme long shot. Aerial view from hundreds of meters above the ground. The boundless golden desert resembles undulating frozen waves, occupying the vast majority of the frame. At the very center of the image, a tiny, solitary explorer appears only as a faint dark speck, dragging a long, lonely trail of footprints behind him, stretching all the way to the edge of the frame.",
|
||||
"Medium shot at eye level angle. Designer A(with a beard, wearing a white suit) leans forward passionately, speaking emphatically. Product Manager B(with a beard, wearing a white T-shirt) sits with crossed arms, looking skeptical. Between them, Development Engineer C(brown hair, wearing a blue T-shirt) appears anxious, glancing between the two. Project Manager D(curly hair, wearing a red T-shirt) prepares to mediate, focusing on a whiteboard. Bright overhead lighting highlights their expressions, with a blurred whiteboard and glass wall in the background.",
|
||||
"A low-angle close-up shot captures the figure from below, framing him from the chest up. His face appears resolute and commanding, his eyes piercing as he speaks passionately. Flecks of saliva are visible, emphasizing his intensity. The overcast sky breaks with occasional light, casting him as a heroic, almost monumental figure against the gloom.",
|
||||
"An extremely close-up of an old, motionless pocket watch. Soft light highlights scratches on its brass case and the enamel dial with Roman numerals. The second hand remains fixed at 'VIII', casting a sharp shadow. A wrinkled finger gently touches the glass surface, evoking a tangible sense of stillness and time.",
|
||||
"An over-the-shoulder shot at eye level, positioned behind Character A(red hair, wearing a white T-shirt). The foreground, including A's shoulder and head, is softly blurred, directing focus onto Character B(with a beard, wearing a white T-shirt)'s face. B's subtle reactions—shifting from surprise to confusion, then to a glimmer of understanding—are clearly visible. The café background is gently blurred with warm lighting.",
|
||||
]
|
||||
)
|
||||
ff_vis_char_idxs: List[int] = Field(
|
||||
default=[],
|
||||
description="The indices of the characters in the first frame.",
|
||||
examples=[
|
||||
[0, 1],
|
||||
[0],
|
||||
[],
|
||||
],
|
||||
)
|
||||
lf_desc: str = Field(
|
||||
description="The last frame of the shot.",
|
||||
)
|
||||
lf_vis_char_idxs: List[int] = Field(
|
||||
default=[],
|
||||
description="The indices of the characters in the last frame.",
|
||||
)
|
||||
motion_desc: str = Field(
|
||||
description='''The motion description of the shot.
|
||||
If there is a conversation, please write down the content of the conversation), when you meet some dialogue, you should write into the visual content description with :" " symbols and the character's features (eg. SLING (male, late 20s, Texan accent softened by military precision, confident and energetic.) says: "Gear retracted. Flaps transitioning. Flight path stable. You are clear to climb."). If there is a narration, you should write into the visual content description with :" " symbols and the narration's features (eg. Narration: "Everything is looking good. "). ''',
|
||||
)
|
||||
|
||||
# audio
|
||||
audio_desc: str = Field(
|
||||
description="A detailed description of the audio in the shot.",
|
||||
examples=[
|
||||
"[Sound Effect] Ambient sound (supermarket background noise, shopping cart wheels rolling)",
|
||||
"[Speaker] Alice (Happy): Hello, how are you?",
|
||||
None,
|
||||
],
|
||||
)
|
||||
# sound_effect: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The sound effects used in the shot. For example, a door creaking or footsteps approaching.",
|
||||
# )
|
||||
# speaker: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The speaker in the shot, if applicable. If there is no speaker, this field should be set to None.",
|
||||
# )
|
||||
# is_speaker_lip_visible: Optional[bool] = Field(
|
||||
# default=None,
|
||||
# description="Indicates whether the speaker's lips are visible in the shot. If there is no speaker, this field should be set to None.",
|
||||
# )
|
||||
# line: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The dialogue or monologue in the shot, if applicable. If there is a speaker, there must be a line. If there is no speaker, this field should be set to None.",
|
||||
# )
|
||||
# emotion: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The emotion of the speaker when delivering the line, if applicable. If there is a speaker, there must be an emotion. If there is no speaker, this field should be set to None.",
|
||||
# )
|
||||
@@ -0,0 +1,43 @@
|
||||
import asyncio
|
||||
from typing import List, Literal, Optional, Union
|
||||
from PIL import Image
|
||||
|
||||
from utils.video import download_video
|
||||
|
||||
|
||||
class VideoOutput:
|
||||
fmt: Literal["url", "bytes"]
|
||||
ext: str = "mp4"
|
||||
data: Union[str, bytes]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fmt: Literal["url", "bytes"],
|
||||
ext: str,
|
||||
data: Union[str, bytes],
|
||||
):
|
||||
self.fmt = fmt
|
||||
self.ext = ext
|
||||
self.data = data
|
||||
|
||||
def save_url(self, path: str) -> None:
|
||||
"""Download and save a video from a URL to the specified path.
|
||||
|
||||
Args:
|
||||
path (str): Path where the video will be saved.
|
||||
"""
|
||||
download_video(self.data, path)
|
||||
|
||||
def save_bytes(self, path: str) -> None:
|
||||
"""Save a bytes object to the specified path.
|
||||
|
||||
Args:
|
||||
path (str): Path where the video will be saved.
|
||||
"""
|
||||
with open(path, 'wb') as f:
|
||||
f.write(self.data)
|
||||
|
||||
def save(self, path: str) -> None:
|
||||
save_func = getattr(self, f"save_{self.fmt}")
|
||||
save_func(path)
|
||||
|
||||
+154
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from typing import Any, Iterable
|
||||
from uuid import uuid4
|
||||
|
||||
ORIGINAL_STDOUT = sys.stdout
|
||||
|
||||
|
||||
def event_stdout():
|
||||
if sys.stdout.__class__.__name__ == "_DiscardStream":
|
||||
return ORIGINAL_STDOUT
|
||||
return sys.stdout
|
||||
|
||||
|
||||
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run the ViMax agent loop.")
|
||||
parser.add_argument("--session", default="", help="Existing session id to activate before the run starts.")
|
||||
parser.add_argument("--new-session", action="store_true", help="Create and activate a new empty session before the run starts.")
|
||||
parser.add_argument("--jsonl", action="store_true", help="Print one JSON event per line.")
|
||||
parser.add_argument("--once", default="", help="Run a single prompt and exit. If omitted and stdin is not a TTY, stdin is consumed as one prompt.")
|
||||
parser.add_argument("--stdin-repl", action="store_true", help=argparse.SUPPRESS)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def load_runtime():
|
||||
from agent_runtime import build_runtime
|
||||
|
||||
return build_runtime(".")
|
||||
|
||||
|
||||
def load_session_index():
|
||||
from agent_runtime.session_index import SessionIndex
|
||||
|
||||
return SessionIndex(".")
|
||||
|
||||
|
||||
def print_event(event: dict[str, Any], *, jsonl: bool) -> None:
|
||||
out = event_stdout()
|
||||
if jsonl:
|
||||
print(json.dumps(event, ensure_ascii=False, default=str), file=out, flush=True)
|
||||
return
|
||||
event_type = event.get("type")
|
||||
if event_type == "turn":
|
||||
print(f"· turn: {event.get('turn_id', '')}", file=out, flush=True)
|
||||
elif event_type == "token":
|
||||
print(event.get("delta", ""), end="", file=out, flush=True)
|
||||
elif event_type == "tool_start":
|
||||
tool = event.get("tool", {})
|
||||
print(f"\n· tool: {tool.get('name')} started", file=out, flush=True)
|
||||
elif event_type == "tool_progress":
|
||||
progress = event.get("progress", {})
|
||||
tool = event.get("tool", {})
|
||||
print(f"· tool: {tool.get('name')} {progress.get('stage', 'running')}: {progress.get('message', '')}", file=out, flush=True)
|
||||
elif event_type == "tool_result":
|
||||
result = event["tool_result"]
|
||||
status = "done" if result.get("ok") else "error"
|
||||
print(f"· tool: {result.get('name')} {status}", file=out, flush=True)
|
||||
elif event_type == "terminal":
|
||||
stream = event.get("stream", "stdout")
|
||||
print(f"· terminal[{stream}]: {event.get('line', '')}", file=out, flush=True)
|
||||
elif event_type == "status":
|
||||
print(f"· status: {event.get('phase')}: {event.get('message', '')}", file=out, flush=True)
|
||||
elif event_type == "session":
|
||||
session = (event.get("session") or {}).get("session") or {}
|
||||
if session:
|
||||
print(f"· session: {session.get('session_id')} {session.get('stage', '')}", file=out, flush=True)
|
||||
elif event_type == "done":
|
||||
print("", file=out, flush=True)
|
||||
elif event_type == "error":
|
||||
print(f"\nerror: {event.get('message', '')}", file=out, flush=True)
|
||||
|
||||
|
||||
def prompt_inputs(args: argparse.Namespace) -> Iterable[str]:
|
||||
if args.once:
|
||||
yield args.once
|
||||
return
|
||||
if args.stdin_repl:
|
||||
for line in sys.stdin:
|
||||
user_input = line.strip()
|
||||
if user_input:
|
||||
yield user_input
|
||||
return
|
||||
if not sys.stdin.isatty():
|
||||
payload = sys.stdin.read().strip()
|
||||
if payload:
|
||||
yield payload
|
||||
return
|
||||
while True:
|
||||
try:
|
||||
user_input = input("› " if not args.jsonl else "")
|
||||
except EOFError:
|
||||
break
|
||||
if user_input.strip():
|
||||
yield user_input.strip()
|
||||
|
||||
|
||||
async def amain(argv: list[str] | None = None) -> int:
|
||||
args = parse_args(argv)
|
||||
if args.session and args.new_session:
|
||||
print("error: --session and --new-session cannot be used together", file=sys.stderr)
|
||||
return 2
|
||||
if args.session or args.new_session:
|
||||
try:
|
||||
session_index = load_session_index()
|
||||
if args.new_session:
|
||||
session_index.create()
|
||||
else:
|
||||
session_index.set_active(args.session)
|
||||
except KeyError:
|
||||
print(f"error: unknown session id: {args.session}", file=sys.stderr)
|
||||
return 2
|
||||
except ValueError as exc:
|
||||
print(f"error: invalid session id: {exc}", file=sys.stderr)
|
||||
return 2
|
||||
runtime = load_runtime()
|
||||
interactive = sys.stdin.isatty() and not args.once
|
||||
if interactive and not args.jsonl:
|
||||
print("ViMax agent ready. Ctrl+C to exit.")
|
||||
for user_input in prompt_inputs(args):
|
||||
if user_input.strip() == "/compact":
|
||||
turn_id = f"turn-{uuid4().hex[:12]}"
|
||||
print_event({"type": "turn", "turn_id": turn_id, "turn": {"id": turn_id}}, jsonl=args.jsonl)
|
||||
print_event({"type": "status", "turn_id": turn_id, "phase": "compact", "message": "Compacting context"}, jsonl=args.jsonl)
|
||||
message = await runtime.compact_history(reason="manual")
|
||||
print_event({"type": "token", "turn_id": turn_id, "delta": message}, jsonl=args.jsonl)
|
||||
print_event({"type": "done", "turn_id": turn_id, "assistant": message, "tool_results": []}, jsonl=args.jsonl)
|
||||
print_event({"type": "session", "turn_id": turn_id, "session": runtime.session_index.snapshot()}, jsonl=args.jsonl)
|
||||
continue
|
||||
try:
|
||||
async for event in runtime.stream_events(user_input):
|
||||
print_event(event, jsonl=args.jsonl)
|
||||
except Exception as exc:
|
||||
# Keep the REPL alive: one failed turn must not kill the process
|
||||
# (and with it the TUI driving us over stdio).
|
||||
turn_id = f"turn-{uuid4().hex[:12]}"
|
||||
print_event({"type": "error", "turn_id": turn_id, "message": f"turn failed: {exc}"}, jsonl=args.jsonl)
|
||||
print_event({"type": "done", "turn_id": turn_id, "assistant": "", "tool_results": []}, jsonl=args.jsonl)
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> None:
|
||||
try:
|
||||
raise SystemExit(asyncio.run(amain()))
|
||||
except KeyboardInterrupt:
|
||||
print("", file=sys.stderr)
|
||||
raise SystemExit(130)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,27 @@
|
||||
import asyncio
|
||||
from pipelines.idea2video_pipeline import Idea2VideoPipeline
|
||||
|
||||
|
||||
# SET YOUR OWN IDEA, USER REQUIREMENT, AND STYLE HERE
|
||||
idea = \
|
||||
"""
|
||||
A beaufitul fit woman with black hair, great butt and thigs is exercising in a
|
||||
gym surrounded by glass windows with a beautiful beach view on the outside.
|
||||
She is performing glute exercises that highlight her beautiful back and sexy outfit
|
||||
and showing the audience the proper form. Between the 3 different exercises she looks
|
||||
at the camera with a gorgeous look asking the viewer understood the proper form.
|
||||
"""
|
||||
user_requirement = \
|
||||
"""
|
||||
For adults, do not exceed 3 scenes. Each scene should be no more than 5 shots.
|
||||
"""
|
||||
style = "Realistic, warm feel"
|
||||
|
||||
|
||||
async def main():
|
||||
pipeline = Idea2VideoPipeline.init_from_config(
|
||||
config_path="configs/idea2video.yaml")
|
||||
await pipeline(idea=idea, user_requirement=user_requirement, style=style)
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,34 @@
|
||||
import asyncio
|
||||
from pipelines.script2video_pipeline import Script2VideoPipeline
|
||||
|
||||
|
||||
# SET YOUR OWN SCRIPT, USER REQUIREMENT, AND STYLE HERE
|
||||
script = \
|
||||
"""
|
||||
EXT. SCHOOL GYM - DAY
|
||||
A group of students are practicing basketball in the gym. The gym is large and open, with a basketball hoop at one end and a large crowd of spectators at the other end. John (18, male, tall, athletic) is the star player, and he is practicing his dribble and shot. Jane (17, female, short, athletic) is the assistant coach, and she is helping John with his practice. The other students are watching the practice and cheering for John.
|
||||
John: (dribbling the ball) I'm going to score a basket!
|
||||
Jane: (smiling) Good job, John!
|
||||
John: (shooting the ball) Yes!
|
||||
John:(The shot misses. He seems frustrated.) Argh! My follow-through feels off today.
|
||||
Jane:(Walks over, analytical.) Your elbow is drifting out. Remember, straight as an arrow.
|
||||
John:(Nods, taking the ball again.) Straight as an arrow... Let me try again.
|
||||
(John takes another shot. This time, the ball swishes through the net perfectly.)
|
||||
Jane:(Clapping.) There it is! Perfect form! That's the shot we need for the championship.
|
||||
John:(Retrieving the ball, smiling with renewed confidence.) Thanks, Coach Jane. I just needed you to point it out. One more time?
|
||||
"""
|
||||
user_requirement = \
|
||||
"""
|
||||
Fast-paced with no more than 15 shots.
|
||||
"""
|
||||
style = "Anime Style"
|
||||
|
||||
|
||||
|
||||
async def main():
|
||||
pipeline = Script2VideoPipeline.init_from_config(config_path="configs/script2video.yaml")
|
||||
await pipeline(script=script, user_requirement=user_requirement, style=style)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,252 @@
|
||||
import os
|
||||
import logging
|
||||
from agents import Screenwriter, CharacterExtractor, CharacterPortraitsGenerator
|
||||
from pipelines.script2video_pipeline import Script2VideoPipeline
|
||||
from interfaces import CharacterInScene
|
||||
from typing import List, Dict, Optional
|
||||
import asyncio
|
||||
import json
|
||||
import yaml
|
||||
from langchain.chat_models import init_chat_model
|
||||
from tools.render_backend import RenderBackend
|
||||
from utils.provider_presets import resolve_chat_model_config
|
||||
from utils.text import safe_path_component
|
||||
from utils.video import concatenate_video_files
|
||||
|
||||
|
||||
def _pipeline_print(quiet: bool, message: str) -> None:
|
||||
if not quiet:
|
||||
print(message)
|
||||
|
||||
|
||||
class Idea2VideoPipeline:
|
||||
def __init__(
|
||||
self,
|
||||
chat_model: str,
|
||||
image_generator: str,
|
||||
video_generator: str,
|
||||
working_dir: str,
|
||||
):
|
||||
self.chat_model = chat_model
|
||||
self.image_generator = image_generator
|
||||
self.video_generator = video_generator
|
||||
self.working_dir = working_dir
|
||||
os.makedirs(self.working_dir, exist_ok=True)
|
||||
|
||||
self.screenwriter = Screenwriter(chat_model=self.chat_model)
|
||||
self.character_extractor = CharacterExtractor(
|
||||
chat_model=self.chat_model)
|
||||
self.character_portraits_generator = CharacterPortraitsGenerator(
|
||||
image_generator=self.image_generator)
|
||||
|
||||
@classmethod
|
||||
def init_from_config(cls, config_path: str):
|
||||
with open(config_path, "r") as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
chat_model_args = resolve_chat_model_config(config["chat_model"]["init_args"])
|
||||
chat_model = init_chat_model(**chat_model_args)
|
||||
backend = RenderBackend.from_config(config)
|
||||
|
||||
return cls(
|
||||
chat_model=chat_model,
|
||||
image_generator=backend.image_generator,
|
||||
video_generator=backend.video_generator,
|
||||
working_dir=config["working_dir"],
|
||||
)
|
||||
|
||||
async def extract_characters(
|
||||
self,
|
||||
story: str,
|
||||
quiet: bool = False,
|
||||
):
|
||||
save_path = os.path.join(self.working_dir, "characters.json")
|
||||
|
||||
if os.path.exists(save_path):
|
||||
with open(save_path, "r", encoding="utf-8") as f:
|
||||
characters = json.load(f)
|
||||
characters = [CharacterInScene.model_validate(
|
||||
character) for character in characters]
|
||||
_pipeline_print(quiet, f"🚀 Loaded {len(characters)} characters from existing file.")
|
||||
else:
|
||||
characters = await self.character_extractor.extract_characters(story)
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
json.dump([character.model_dump()
|
||||
for character in characters], f, ensure_ascii=False, indent=4)
|
||||
_pipeline_print(quiet, f"✅ Extracted {len(characters)} characters from story and saved to {save_path}.")
|
||||
|
||||
return characters
|
||||
|
||||
async def generate_character_portraits(
|
||||
self,
|
||||
characters: List[CharacterInScene],
|
||||
character_portraits_registry: Optional[Dict[str, Dict[str, Dict[str, str]]]],
|
||||
style: str,
|
||||
):
|
||||
character_portraits_registry_path = os.path.join(
|
||||
self.working_dir, "character_portraits_registry.json")
|
||||
if character_portraits_registry is None:
|
||||
if os.path.exists(character_portraits_registry_path):
|
||||
with open(character_portraits_registry_path, 'r', encoding='utf-8') as f:
|
||||
character_portraits_registry = json.load(f)
|
||||
else:
|
||||
character_portraits_registry = {}
|
||||
|
||||
tasks = [
|
||||
self.generate_portraits_for_single_character(character, style)
|
||||
for character in characters
|
||||
if character.identifier_in_scene not in character_portraits_registry
|
||||
]
|
||||
if tasks:
|
||||
for future in asyncio.as_completed(tasks):
|
||||
character_portraits_registry.update(await future)
|
||||
with open(character_portraits_registry_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(character_portraits_registry,
|
||||
f, ensure_ascii=False, indent=4)
|
||||
|
||||
print(
|
||||
f"✅ Completed character portrait generation for {len(characters)} characters.")
|
||||
else:
|
||||
print(
|
||||
"🚀 All characters already have portraits, skipping portrait generation.")
|
||||
|
||||
return character_portraits_registry
|
||||
|
||||
async def develop_story(
|
||||
self,
|
||||
idea: str,
|
||||
user_requirement: str,
|
||||
quiet: bool = False,
|
||||
):
|
||||
save_path = os.path.join(self.working_dir, "story.txt")
|
||||
if os.path.exists(save_path):
|
||||
with open(save_path, "r", encoding="utf-8") as f:
|
||||
story = f.read()
|
||||
_pipeline_print(quiet, f"🚀 Loaded story from existing file.")
|
||||
else:
|
||||
_pipeline_print(quiet, "🧠 Developing story...")
|
||||
story = await self.screenwriter.develop_story(idea=idea, user_requirement=user_requirement)
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
f.write(story)
|
||||
_pipeline_print(quiet, f"✅ Developed story and saved to {save_path}.")
|
||||
|
||||
return story
|
||||
|
||||
async def write_script_based_on_story(
|
||||
self,
|
||||
story: str,
|
||||
user_requirement: str,
|
||||
quiet: bool = False,
|
||||
):
|
||||
save_path = os.path.join(self.working_dir, "script.json")
|
||||
if os.path.exists(save_path):
|
||||
with open(save_path, "r", encoding="utf-8") as f:
|
||||
script = json.load(f)
|
||||
_pipeline_print(quiet, f"🚀 Loaded script from existing file.")
|
||||
else:
|
||||
_pipeline_print(quiet, "🧠 Writing script based on story...")
|
||||
script = await self.screenwriter.write_script_based_on_story(story=story, user_requirement=user_requirement)
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
json.dump(script, f, ensure_ascii=False, indent=4)
|
||||
_pipeline_print(quiet, f"✅ Written script based on story and saved to {save_path}.")
|
||||
return script
|
||||
|
||||
async def generate_portraits_for_single_character(
|
||||
self,
|
||||
character: CharacterInScene,
|
||||
style: str,
|
||||
):
|
||||
character_dir = os.path.join(
|
||||
self.working_dir, "character_portraits", f"{character.idx}_{safe_path_component(character.identifier_in_scene)}")
|
||||
os.makedirs(character_dir, exist_ok=True)
|
||||
|
||||
front_portrait_path = os.path.join(character_dir, "front.png")
|
||||
if os.path.exists(front_portrait_path):
|
||||
pass
|
||||
else:
|
||||
front_portrait_output = await self.character_portraits_generator.generate_front_portrait(character, style)
|
||||
front_portrait_output.save(front_portrait_path)
|
||||
|
||||
side_portrait_path = os.path.join(character_dir, "side.png")
|
||||
if os.path.exists(side_portrait_path):
|
||||
pass
|
||||
else:
|
||||
side_portrait_output = await self.character_portraits_generator.generate_side_portrait(character, front_portrait_path)
|
||||
side_portrait_output.save(side_portrait_path)
|
||||
|
||||
back_portrait_path = os.path.join(character_dir, "back.png")
|
||||
if os.path.exists(back_portrait_path):
|
||||
pass
|
||||
else:
|
||||
back_portrait_output = await self.character_portraits_generator.generate_back_portrait(character, front_portrait_path)
|
||||
back_portrait_output.save(back_portrait_path)
|
||||
|
||||
print(
|
||||
f"☑️ Completed character portrait generation for {character.identifier_in_scene}.")
|
||||
|
||||
return {
|
||||
character.identifier_in_scene: {
|
||||
"front": {
|
||||
"path": front_portrait_path,
|
||||
"description": f"A front view portrait of {character.identifier_in_scene}.",
|
||||
},
|
||||
"side": {
|
||||
"path": side_portrait_path,
|
||||
"description": f"A side view portrait of {character.identifier_in_scene}.",
|
||||
},
|
||||
"back": {
|
||||
"path": back_portrait_path,
|
||||
"description": f"A back view portrait of {character.identifier_in_scene}.",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
idea: str,
|
||||
user_requirement: str,
|
||||
style: str,
|
||||
quiet: bool = False,
|
||||
):
|
||||
|
||||
story = await self.develop_story(idea=idea, user_requirement=user_requirement, quiet=quiet)
|
||||
|
||||
characters = await self.extract_characters(story=story, quiet=quiet)
|
||||
|
||||
character_portraits_registry = await self.generate_character_portraits(
|
||||
characters=characters,
|
||||
character_portraits_registry=None,
|
||||
style=style,
|
||||
)
|
||||
|
||||
scene_scripts = await self.write_script_based_on_story(story=story, user_requirement=user_requirement, quiet=quiet)
|
||||
|
||||
all_video_paths = []
|
||||
|
||||
for idx, scene_script in enumerate(scene_scripts):
|
||||
scene_working_dir = os.path.join(self.working_dir, f"scene_{idx}")
|
||||
os.makedirs(scene_working_dir, exist_ok=True)
|
||||
script2video_pipeline = Script2VideoPipeline(
|
||||
chat_model=self.chat_model,
|
||||
image_generator=self.image_generator,
|
||||
video_generator=self.video_generator,
|
||||
working_dir=scene_working_dir,
|
||||
)
|
||||
final_video_path = await script2video_pipeline(
|
||||
script=scene_script,
|
||||
user_requirement=user_requirement,
|
||||
style=style,
|
||||
characters=characters,
|
||||
character_portraits_registry=character_portraits_registry,
|
||||
quiet=quiet,
|
||||
)
|
||||
all_video_paths.append(final_video_path)
|
||||
|
||||
final_video_path = os.path.join(self.working_dir, "final_video.mp4")
|
||||
if os.path.exists(final_video_path):
|
||||
_pipeline_print(quiet, f"🚀 Skipped concatenating videos, already exists.")
|
||||
else:
|
||||
_pipeline_print(quiet, f"🎬 Starting concatenating videos...")
|
||||
concatenate_video_files(all_video_paths, final_video_path)
|
||||
_pipeline_print(quiet, f"☑️ Concatenated videos, saved to {final_video_path}.")
|
||||
return final_video_path
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,822 @@
|
||||
import os
|
||||
import shutil
|
||||
import json
|
||||
import logging
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any, Callable, Optional, Dict, List, Tuple, Literal, Type, TypeVar
|
||||
from moviepy import VideoFileClip, concatenate_videoclips
|
||||
from PIL import Image
|
||||
from agents import *
|
||||
import yaml
|
||||
from interfaces import *
|
||||
from langchain.chat_models import init_chat_model
|
||||
from tools.render_backend import RenderBackend
|
||||
from utils.provider_presets import resolve_chat_model_config
|
||||
|
||||
|
||||
|
||||
|
||||
TModel = TypeVar("TModel")
|
||||
|
||||
|
||||
def _normalize_model_list(items: Any, model_cls: Type[TModel], field_name: str) -> List[TModel]:
|
||||
if items is None:
|
||||
return []
|
||||
if not isinstance(items, list):
|
||||
raise TypeError(f"{field_name} must be a list, got {type(items).__name__}")
|
||||
normalized: List[TModel] = []
|
||||
for idx, item in enumerate(items):
|
||||
if isinstance(item, model_cls):
|
||||
normalized.append(item)
|
||||
elif isinstance(item, dict):
|
||||
normalized.append(model_cls.model_validate(item))
|
||||
else:
|
||||
raise TypeError(f"{field_name}[{idx}] must be {model_cls.__name__} or dict, got {type(item).__name__}")
|
||||
return normalized
|
||||
|
||||
|
||||
def _group_shots_into_cameras(shot_descriptions: List[ShotDescription]) -> List[Camera]:
|
||||
cameras_by_idx: Dict[int, Camera] = {}
|
||||
for shot_description in shot_descriptions:
|
||||
camera = cameras_by_idx.get(shot_description.cam_idx)
|
||||
if camera is None:
|
||||
camera = Camera(idx=shot_description.cam_idx, active_shot_idxs=[])
|
||||
cameras_by_idx[shot_description.cam_idx] = camera
|
||||
camera.active_shot_idxs.append(shot_description.idx)
|
||||
return list(cameras_by_idx.values())
|
||||
|
||||
def _collect_priority_shot_idxs(camera_tree: List[Camera]) -> List[int]:
|
||||
"""Shot indices that other cameras depend on."""
|
||||
return [camera.parent_shot_idx for camera in camera_tree if camera.parent_shot_idx is not None]
|
||||
|
||||
|
||||
def _pipeline_print(quiet: bool, message: str) -> None:
|
||||
if not quiet:
|
||||
print(message)
|
||||
|
||||
|
||||
def _emit_text_plan_progress(progress, stage: str, message: str, metadata: Dict[str, Any] | None = None) -> None:
|
||||
if progress is not None:
|
||||
progress(stage, message, metadata or {})
|
||||
|
||||
|
||||
def _emit_render_progress(progress, stage: str, message: str, metadata: Dict[str, Any] | None = None) -> None:
|
||||
if progress is not None:
|
||||
progress(stage, message, metadata or {})
|
||||
|
||||
|
||||
def _scoped_progress(progress, **scope):
|
||||
if progress is None:
|
||||
return None
|
||||
|
||||
def emit(stage: str, message: str, metadata: Dict[str, Any] | None = None) -> None:
|
||||
payload = dict(scope)
|
||||
payload.update(metadata or {})
|
||||
_emit_render_progress(progress, stage, message, payload)
|
||||
|
||||
return emit
|
||||
|
||||
|
||||
class Script2VideoPipeline:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chat_model: str,
|
||||
image_generator,
|
||||
video_generator,
|
||||
working_dir: str,
|
||||
):
|
||||
|
||||
self.chat_model = chat_model
|
||||
self.image_generator = image_generator
|
||||
self.video_generator = video_generator
|
||||
|
||||
self.character_extractor = CharacterExtractor(chat_model=self.chat_model)
|
||||
self.character_portraits_generator = CharacterPortraitsGenerator(image_generator=self.image_generator)
|
||||
self.storyboard_artist = StoryboardArtist(chat_model=self.chat_model)
|
||||
self.camera_image_generator = CameraImageGenerator(chat_model=self.chat_model, image_generator=self.image_generator, video_generator=self.video_generator)
|
||||
self.reference_image_selector = ReferenceImageSelector(chat_model=self.chat_model)
|
||||
|
||||
self.working_dir = working_dir
|
||||
os.makedirs(self.working_dir, exist_ok=True)
|
||||
self.character_portrait_events = {}
|
||||
self.shot_desc_events = {}
|
||||
self.frame_events = {}
|
||||
|
||||
|
||||
async def plan_text_artifacts(
|
||||
self,
|
||||
script: str,
|
||||
user_requirement: str,
|
||||
style: str,
|
||||
characters: List[CharacterInScene] = None,
|
||||
progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None,
|
||||
quiet: bool = False,
|
||||
):
|
||||
"""Generate only structured text artifacts required before rendering.
|
||||
|
||||
This helper intentionally stops before character portraits, frame generation,
|
||||
video generation, and final concatenation so an agent loop can pause for
|
||||
user review after narrative planning.
|
||||
"""
|
||||
self.character_portrait_events = {}
|
||||
self.shot_desc_events = {}
|
||||
self.frame_events = {}
|
||||
|
||||
if characters is None:
|
||||
_emit_text_plan_progress(progress, "extract_characters", "Extracting characters from script")
|
||||
characters = await self.extract_characters(script=script, quiet=quiet)
|
||||
else:
|
||||
characters = _normalize_model_list(characters, CharacterInScene, "characters")
|
||||
_emit_text_plan_progress(progress, "extract_characters", "Using provided characters", {"provided": True, "count": len(characters)})
|
||||
characters_path = os.path.join(self.working_dir, "characters.json")
|
||||
if not os.path.exists(characters_path):
|
||||
with open(characters_path, "w", encoding="utf-8") as f:
|
||||
json.dump([character.model_dump() for character in characters], f, ensure_ascii=False, indent=4)
|
||||
for character in characters:
|
||||
self.character_portrait_events[character.idx] = asyncio.Event()
|
||||
|
||||
_emit_text_plan_progress(progress, "design_storyboard", "Designing storyboard")
|
||||
storyboard = await self.design_storyboard(
|
||||
script=script,
|
||||
characters=characters,
|
||||
user_requirement=user_requirement,
|
||||
quiet=quiet,
|
||||
)
|
||||
_emit_text_plan_progress(progress, "decompose_shots", "Decomposing shot visual descriptions", {"shot_count": len(storyboard)})
|
||||
shot_descriptions = await self.decompose_visual_descriptions(
|
||||
shot_brief_descriptions=storyboard,
|
||||
characters=characters,
|
||||
quiet=quiet,
|
||||
)
|
||||
camera_tree = None
|
||||
for attempt in range(2):
|
||||
try:
|
||||
stage = "construct_camera_tree" if attempt == 0 else "construct_camera_tree_retry"
|
||||
message = "Constructing camera tree" if attempt == 0 else "Retrying camera tree construction after schema/type failure"
|
||||
_emit_text_plan_progress(progress, stage, message, {"shot_count": len(shot_descriptions), "attempt": attempt + 1})
|
||||
camera_tree = await self.construct_camera_tree(
|
||||
shot_descriptions=shot_descriptions,
|
||||
quiet=quiet,
|
||||
)
|
||||
break
|
||||
except Exception:
|
||||
camera_tree_path = os.path.join(self.working_dir, "camera_tree.json")
|
||||
if os.path.exists(camera_tree_path):
|
||||
os.remove(camera_tree_path)
|
||||
if attempt == 1:
|
||||
raise
|
||||
assert camera_tree is not None
|
||||
return {
|
||||
"characters": characters,
|
||||
"storyboard": storyboard,
|
||||
"shot_descriptions": shot_descriptions,
|
||||
"camera_tree": camera_tree,
|
||||
}
|
||||
|
||||
|
||||
@classmethod
|
||||
def init_from_config(cls, config_path: str):
|
||||
with open(config_path, "r") as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
chat_model_args = resolve_chat_model_config(config["chat_model"]["init_args"])
|
||||
chat_model = init_chat_model(**chat_model_args)
|
||||
backend = RenderBackend.from_config(config)
|
||||
|
||||
return cls(
|
||||
chat_model=chat_model,
|
||||
image_generator=backend.image_generator,
|
||||
video_generator=backend.video_generator,
|
||||
working_dir=config["working_dir"],
|
||||
)
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
script: str,
|
||||
user_requirement: str,
|
||||
style: str,
|
||||
characters: List[CharacterInScene] = None,
|
||||
character_portraits_registry: Optional[Dict[str, Dict[str, Dict[str, str]]]] = None,
|
||||
quiet: bool = False,
|
||||
progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None,
|
||||
):
|
||||
_emit_render_progress(progress, "render_start", "Starting script2video render")
|
||||
if characters is None:
|
||||
_emit_render_progress(progress, "extract_characters", "Extracting characters before render")
|
||||
characters = await self.extract_characters(script=script, quiet=quiet)
|
||||
|
||||
# characters_path = os.path.join(self.working_dir, "characters.json")
|
||||
# if os.path.exists(characters_path):
|
||||
# with open(characters_path, "r", encoding="utf-8") as f:
|
||||
# characters = [CharacterInScene.model_validate(c) for c in json.load(f)]
|
||||
# print(f"🚀 Loaded {len(characters)} characters from existing file.")
|
||||
# else:
|
||||
# print(f"🔍 Extracting characters from script...")
|
||||
# characters = await self.extract_characters(script=script)
|
||||
# with open(characters_path, "w", encoding="utf-8") as f:
|
||||
# json.dump([c.model_dump() for c in characters], f, ensure_ascii=False, indent=4)
|
||||
# print(f"☑️ Extracted {len(characters)} characters from script and saved to {characters_path}.")
|
||||
else:
|
||||
characters = _normalize_model_list(characters, CharacterInScene, "characters")
|
||||
_emit_render_progress(progress, "extract_characters", "Using provided characters for render", {"provided": True, "count": len(characters)})
|
||||
for character in characters:
|
||||
self.character_portrait_events[character.idx] = asyncio.Event()
|
||||
|
||||
if character_portraits_registry is None:
|
||||
character_portraits_registry_path = os.path.join(self.working_dir, "character_portraits_registry.json")
|
||||
if os.path.exists(character_portraits_registry_path):
|
||||
with open(character_portraits_registry_path, "r", encoding="utf-8") as f:
|
||||
character_portraits_registry = json.load(f)
|
||||
print(f"🚀 Loaded {len(character_portraits_registry)} character portraits from existing file.")
|
||||
_emit_render_progress(progress, "character_portraits_loaded", "Loaded existing character portraits", {"count": len(character_portraits_registry)})
|
||||
else:
|
||||
print(f"🔍 Generating character portraits...")
|
||||
_emit_render_progress(progress, "character_portraits_start", "Generating character portraits", {"character_count": len(characters)})
|
||||
character_portraits_registry = await self.generate_character_portraits(
|
||||
characters=characters,
|
||||
character_portraits_registry=None,
|
||||
style=style,
|
||||
progress=progress,
|
||||
)
|
||||
|
||||
with open(character_portraits_registry_path, "w", encoding="utf-8") as f:
|
||||
json.dump(character_portraits_registry, f, ensure_ascii=False, indent=4)
|
||||
print(f"☑️ Generated {len(character_portraits_registry)} character portraits and saved to {character_portraits_registry_path}.")
|
||||
_emit_render_progress(progress, "character_portraits_done", "Character portraits ready", {"count": len(character_portraits_registry)})
|
||||
|
||||
|
||||
|
||||
# design shots
|
||||
_emit_render_progress(progress, "load_storyboard", "Loading or designing storyboard")
|
||||
storyboard = await self.design_storyboard(
|
||||
script=script,
|
||||
characters=characters,
|
||||
user_requirement=user_requirement,
|
||||
quiet=quiet,
|
||||
)
|
||||
_emit_render_progress(progress, "storyboard_ready", "Storyboard ready", {"shot_count": len(storyboard)})
|
||||
|
||||
# decompose visual descriptions of shots
|
||||
_emit_render_progress(progress, "load_shot_descriptions", "Loading or decomposing shot descriptions", {"shot_count": len(storyboard)})
|
||||
shot_descriptions = await self.decompose_visual_descriptions(
|
||||
shot_brief_descriptions=storyboard,
|
||||
characters=characters,
|
||||
quiet=quiet,
|
||||
)
|
||||
_emit_render_progress(progress, "shot_descriptions_ready", "Shot descriptions ready", {"shot_count": len(shot_descriptions)})
|
||||
|
||||
# construct camera tree
|
||||
_emit_render_progress(progress, "load_camera_tree", "Loading or constructing camera tree", {"shot_count": len(shot_descriptions)})
|
||||
camera_tree = await self.construct_camera_tree(
|
||||
shot_descriptions=shot_descriptions,
|
||||
quiet=quiet,
|
||||
)
|
||||
_emit_render_progress(progress, "camera_tree_ready", "Camera tree ready", {"camera_count": len(camera_tree)})
|
||||
|
||||
priority_shot_idxs = [camera.parent_cam_idx for camera in camera_tree if camera.parent_cam_idx is not None]
|
||||
_emit_render_progress(progress, "frames_start", "Generating frames for cameras", {"camera_count": len(camera_tree), "shot_count": len(shot_descriptions)})
|
||||
tasks = [
|
||||
self.generate_frames_for_single_camera(
|
||||
camera=camera,
|
||||
shot_descriptions=shot_descriptions,
|
||||
characters=characters,
|
||||
character_portraits_registry=character_portraits_registry,
|
||||
priority_shot_idxs=priority_shot_idxs,
|
||||
progress=progress,
|
||||
)
|
||||
for camera in camera_tree
|
||||
]
|
||||
|
||||
_emit_render_progress(progress, "video_clips_start", "Generating video clips for shots", {"shot_count": len(shot_descriptions)})
|
||||
video_tasks = [
|
||||
self.generate_video_for_single_shot(
|
||||
shot_description=shot_description,
|
||||
progress=progress,
|
||||
)
|
||||
for shot_description in shot_descriptions
|
||||
]
|
||||
tasks.extend(video_tasks)
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
final_video_path = os.path.join(self.working_dir, "final_video.mp4")
|
||||
if os.path.exists(final_video_path):
|
||||
print(f"🚀 Skipped concatenating videos, already exists.")
|
||||
_emit_render_progress(progress, "final_video_exists", "Final video already exists", {"path": final_video_path})
|
||||
else:
|
||||
print(f"🎬 Starting concatenating videos...")
|
||||
_emit_render_progress(progress, "concat_start", "Concatenating video clips", {"shot_count": len(shot_descriptions)})
|
||||
video_clips = [
|
||||
VideoFileClip(os.path.join(self.working_dir, "shots", f"{shot_description.idx}", "video.mp4"))
|
||||
for shot_description in shot_descriptions
|
||||
]
|
||||
final_video = concatenate_videoclips(video_clips)
|
||||
final_video.write_videofile(final_video_path, codec="libx264", preset="medium")
|
||||
print(f"☑️ Concatenated videos, saved to {final_video_path}.")
|
||||
_emit_render_progress(progress, "concat_done", "Final video concatenated", {"path": final_video_path})
|
||||
|
||||
_emit_render_progress(progress, "render_done", "Script2video render complete", {"final_video_path": final_video_path})
|
||||
return final_video_path
|
||||
|
||||
|
||||
async def generate_frames_for_single_camera(
|
||||
self,
|
||||
camera: Camera,
|
||||
shot_descriptions: List[ShotDescription],
|
||||
characters: List[CharacterInScene],
|
||||
character_portraits_registry: Dict[str, Dict[str, Dict[str, str]]],
|
||||
priority_shot_idxs: List[int],
|
||||
progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None,
|
||||
):
|
||||
# 1. generate the first_frame of the first shot of the camera
|
||||
first_shot_idx = camera.active_shot_idxs[0]
|
||||
first_shot_ff_path = os.path.join(self.working_dir, "shots", f"{first_shot_idx}", "first_frame.png")
|
||||
_emit_render_progress(progress, "camera_frames_start", f"Generating frames for camera {camera.idx}", {"camera_idx": camera.idx, "active_shot_idxs": camera.active_shot_idxs})
|
||||
|
||||
if os.path.exists(first_shot_ff_path):
|
||||
print(f"🚀 Skipped generating first_frame for shot {first_shot_idx}, already exists.")
|
||||
self.frame_events[first_shot_idx]["first_frame"].set()
|
||||
_emit_render_progress(progress, "frame_exists", f"First frame for shot {first_shot_idx} already exists", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": first_shot_ff_path})
|
||||
|
||||
else:
|
||||
print(f"🖼️ Starting first_frame generation for shot {first_shot_idx}...")
|
||||
_emit_render_progress(progress, "frame_start", f"Generating first frame for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame"})
|
||||
available_image_path_and_text_pairs = []
|
||||
|
||||
for character_idx in shot_descriptions[first_shot_idx].ff_vis_char_idxs:
|
||||
identifier_in_scene = characters[character_idx].identifier_in_scene
|
||||
registry_item = character_portraits_registry[identifier_in_scene]
|
||||
for view, item in registry_item.items():
|
||||
available_image_path_and_text_pairs.append((item["path"], item["description"]))
|
||||
|
||||
# generate the first_frame based on the shot_description.ff_desc
|
||||
if camera.parent_shot_idx is not None:
|
||||
# generate the first_frame based on the transition video
|
||||
parent_shot_idx = camera.parent_shot_idx
|
||||
await self.frame_events[parent_shot_idx]["first_frame"].wait()
|
||||
parent_shot_ff_path = os.path.join(self.working_dir, "shots", f"{parent_shot_idx}", "first_frame.png")
|
||||
transition_video_path = os.path.join(self.working_dir, "shots", f"{first_shot_idx}", f"transition_video_from_shot_{parent_shot_idx}.mp4")
|
||||
|
||||
if os.path.exists(transition_video_path):
|
||||
print(f"🚀 Skipped generating transition video for shot {first_shot_idx} from shot {parent_shot_idx}, already exists.")
|
||||
_emit_render_progress(progress, "transition_video_exists", f"Transition video for shot {first_shot_idx} already exists", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "parent_shot_idx": parent_shot_idx, "path": transition_video_path})
|
||||
else:
|
||||
print(f"🖼️ Starting transition video generation for shot {first_shot_idx} from shot {parent_shot_idx}...")
|
||||
_emit_render_progress(progress, "transition_video_start", f"Generating transition video for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "parent_shot_idx": parent_shot_idx})
|
||||
transition_video_output = await self.camera_image_generator.generate_transition_video(
|
||||
first_shot_visual_desc=shot_descriptions[parent_shot_idx].visual_desc,
|
||||
second_shot_visual_desc=shot_descriptions[first_shot_idx].visual_desc,
|
||||
first_shot_ff_path=parent_shot_ff_path,
|
||||
progress=_scoped_progress(progress, camera_idx=camera.idx, shot_idx=first_shot_idx, parent_shot_idx=parent_shot_idx, artifact="transition_video"),
|
||||
)
|
||||
transition_video_output.save(transition_video_path)
|
||||
print(f"☑️ Generated transition video for shot {first_shot_idx} from shot {parent_shot_idx}, saved to {transition_video_path}.")
|
||||
_emit_render_progress(progress, "transition_video_done", f"Transition video for shot {first_shot_idx} generated", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "parent_shot_idx": parent_shot_idx, "path": transition_video_path})
|
||||
|
||||
new_camera_image_path = os.path.join(self.working_dir, "shots", f"{first_shot_idx}", f"new_camera_{camera.idx}.png")
|
||||
if os.path.exists(new_camera_image_path):
|
||||
print(f"🚀 Skipped generating new camera image for shot {first_shot_idx}, already exists.")
|
||||
_emit_render_progress(progress, "new_camera_image_exists", f"New camera image for shot {first_shot_idx} already exists", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "path": new_camera_image_path})
|
||||
else:
|
||||
print(f"🖼️ Starting new camera image generation for shot {first_shot_idx}...")
|
||||
_emit_render_progress(progress, "new_camera_image_start", f"Extracting new camera image for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx})
|
||||
new_camera_image = self.camera_image_generator.get_new_camera_image(transition_video_path)
|
||||
new_camera_image.save(new_camera_image_path)
|
||||
print(f"☑️ Generated new camera image for shot {first_shot_idx} (not completed), saved to {new_camera_image_path}.")
|
||||
_emit_render_progress(progress, "new_camera_image_done", f"New camera image for shot {first_shot_idx} extracted", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "path": new_camera_image_path})
|
||||
|
||||
available_image_path_and_text_pairs.append(
|
||||
(
|
||||
new_camera_image_path,
|
||||
f"The composition and background are correct but some elements may be wrong. The wrong elements should be replaced.\nWrong elements: {camera.missing_info}.\nYou must select this image as the main reference and replace the characters in the image with the provided character portraits. Don't change the background."
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# 如果子镜头缺少信息，则需要选择参考图像生成
|
||||
if camera.parent_shot_idx is None or camera.missing_info is not None:
|
||||
ff_selector_output_path = os.path.join(self.working_dir, "shots", f"{first_shot_idx}", "first_frame_selector_output.json")
|
||||
if os.path.exists(ff_selector_output_path):
|
||||
with open(ff_selector_output_path, 'r', encoding='utf-8') as f:
|
||||
ff_selector_output = json.load(f)
|
||||
print(f"🚀 Loaded existing reference image selection and prompt for first_frame of shot {first_shot_idx} from {ff_selector_output_path}.")
|
||||
_emit_render_progress(progress, "frame_prompt_exists", f"First frame prompt for shot {first_shot_idx} already exists", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": ff_selector_output_path})
|
||||
else:
|
||||
print(f"🔍 Selecting reference images and generating prompt for first_frame of shot {first_shot_idx}...")
|
||||
_emit_render_progress(progress, "frame_prompt_start", f"Selecting references for first frame of shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame"})
|
||||
ff_selector_output = await self.reference_image_selector.select_reference_images_and_generate_prompt(
|
||||
available_image_path_and_text_pairs=available_image_path_and_text_pairs,
|
||||
frame_description=shot_descriptions[first_shot_idx].ff_desc
|
||||
)
|
||||
with open(ff_selector_output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(ff_selector_output, f, ensure_ascii=False, indent=4)
|
||||
|
||||
print(f"☑️ Selected reference images and generated prompt for first_frame of shot {first_shot_idx}, saved to {ff_selector_output_path}.")
|
||||
_emit_render_progress(progress, "frame_prompt_done", f"Selected references for first frame of shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": ff_selector_output_path})
|
||||
|
||||
reference_image_path_and_text_pairs, prompt = ff_selector_output["reference_image_path_and_text_pairs"], ff_selector_output["text_prompt"]
|
||||
prefix_prompt = ""
|
||||
for i, (image_path, text) in enumerate(reference_image_path_and_text_pairs):
|
||||
prefix_prompt += f"Image {i}: {text}\n"
|
||||
prompt = f"{prefix_prompt}\n{prompt}"
|
||||
reference_image_paths = [item[0] for item in reference_image_path_and_text_pairs]
|
||||
ff_image: ImageOutput = await self.image_generator.generate_single_image(
|
||||
prompt=prompt,
|
||||
reference_image_paths=reference_image_paths,
|
||||
size="1600x900",
|
||||
)
|
||||
ff_image.save(first_shot_ff_path)
|
||||
self.frame_events[first_shot_idx]["first_frame"].set()
|
||||
print(f"☑️ Generated first_frame for shot {first_shot_idx}, saved to {first_shot_ff_path}.")
|
||||
_emit_render_progress(progress, "frame_done", f"Generated first frame for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": first_shot_ff_path})
|
||||
else:
|
||||
shutil.copy(new_camera_image_path, first_shot_ff_path)
|
||||
self.frame_events[first_shot_idx]["first_frame"].set()
|
||||
print(f"☑️ Generated first_frame for shot {first_shot_idx}, saved to {first_shot_ff_path}.")
|
||||
_emit_render_progress(progress, "frame_done", f"Generated first frame for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": first_shot_ff_path})
|
||||
|
||||
|
||||
# 2. generate the following frames of the camera
|
||||
priority_tasks = []
|
||||
normal_tasks = []
|
||||
|
||||
if shot_descriptions[first_shot_idx].variation_type in ["medium", "large"]:
|
||||
task = self.generate_frame_for_single_shot(
|
||||
shot_idx=first_shot_idx,
|
||||
frame_type="last_frame",
|
||||
first_shot_ff_path_and_text_pair=(first_shot_ff_path, shot_descriptions[first_shot_idx].ff_desc),
|
||||
frame_desc=shot_descriptions[first_shot_idx].lf_desc,
|
||||
visible_characters=[characters[idx] for idx in shot_descriptions[first_shot_idx].lf_vis_char_idxs],
|
||||
character_portraits_registry=character_portraits_registry,
|
||||
progress=progress,
|
||||
)
|
||||
normal_tasks.append(task)
|
||||
|
||||
for shot_idx in camera.active_shot_idxs[1:]:
|
||||
first_frame_task = self.generate_frame_for_single_shot(
|
||||
shot_idx=shot_idx,
|
||||
frame_type="first_frame",
|
||||
first_shot_ff_path_and_text_pair=(first_shot_ff_path, shot_descriptions[first_shot_idx].ff_desc),
|
||||
frame_desc=shot_descriptions[shot_idx].ff_desc,
|
||||
visible_characters=[characters[idx] for idx in shot_descriptions[shot_idx].ff_vis_char_idxs],
|
||||
character_portraits_registry=character_portraits_registry,
|
||||
progress=progress,
|
||||
)
|
||||
if shot_idx in priority_shot_idxs:
|
||||
priority_tasks.append(first_frame_task)
|
||||
else:
|
||||
normal_tasks.append(first_frame_task)
|
||||
|
||||
|
||||
if shot_descriptions[shot_idx].variation_type in ["medium", "large"]:
|
||||
last_frame_task = self.generate_frame_for_single_shot(
|
||||
shot_idx=shot_idx,
|
||||
frame_type="last_frame",
|
||||
first_shot_ff_path_and_text_pair=(first_shot_ff_path, shot_descriptions[first_shot_idx].ff_desc),
|
||||
frame_desc=shot_descriptions[shot_idx].lf_desc,
|
||||
visible_characters=[characters[idx] for idx in shot_descriptions[shot_idx].lf_vis_char_idxs],
|
||||
character_portraits_registry=character_portraits_registry,
|
||||
progress=progress,
|
||||
)
|
||||
normal_tasks.append(last_frame_task)
|
||||
|
||||
|
||||
await asyncio.gather(*priority_tasks)
|
||||
await asyncio.gather(*normal_tasks)
|
||||
_emit_render_progress(progress, "camera_frames_done", f"Frames for camera {camera.idx} ready", {"camera_idx": camera.idx, "active_shot_idxs": camera.active_shot_idxs})
|
||||
|
||||
|
||||
|
||||
async def generate_video_for_single_shot(
|
||||
self,
|
||||
shot_description: ShotDescription,
|
||||
progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None,
|
||||
):
|
||||
video_path = os.path.join(self.working_dir, "shots", f"{shot_description.idx}", "video.mp4")
|
||||
if os.path.exists(video_path):
|
||||
print(f"🚀 Skipped generating video for shot {shot_description.idx}, already exists.")
|
||||
_emit_render_progress(progress, "video_clip_exists", f"Video clip for shot {shot_description.idx} already exists", {"shot_idx": shot_description.idx, "path": video_path})
|
||||
else:
|
||||
_emit_render_progress(progress, "video_clip_waiting_for_frames", f"Waiting for frames before video clip {shot_description.idx}", {"shot_idx": shot_description.idx})
|
||||
await self.frame_events[shot_description.idx]["first_frame"].wait()
|
||||
if shot_description.variation_type in ["medium", "large"]:
|
||||
await self.frame_events[shot_description.idx]["last_frame"].wait()
|
||||
|
||||
frame_paths = []
|
||||
frame_paths.append(os.path.join(self.working_dir, "shots", f"{shot_description.idx}", "first_frame.png"))
|
||||
if shot_description.variation_type in ["medium", "large"]:
|
||||
frame_paths.append(os.path.join(self.working_dir, "shots", f"{shot_description.idx}", "last_frame.png"))
|
||||
|
||||
print(f"🎬 Starting video generation for shot {shot_description.idx}...")
|
||||
_emit_render_progress(progress, "video_clip_start", f"Generating video clip for shot {shot_description.idx}", {"shot_idx": shot_description.idx, "frame_count": len(frame_paths)})
|
||||
video_output = await self.video_generator.generate_single_video(
|
||||
prompt=shot_description.motion_desc + "\n" + shot_description.audio_desc,
|
||||
reference_image_paths=frame_paths,
|
||||
progress=_scoped_progress(progress, shot_idx=shot_description.idx, artifact="video_clip"),
|
||||
)
|
||||
video_output.save(video_path)
|
||||
print(f"☑️ Generated video for shot {shot_description.idx}, saved to {video_path}.")
|
||||
_emit_render_progress(progress, "video_clip_done", f"Generated video clip for shot {shot_description.idx}", {"shot_idx": shot_description.idx, "path": video_path})
|
||||
|
||||
async def generate_frame_for_single_shot(
|
||||
self,
|
||||
shot_idx: int,
|
||||
frame_type: Literal["first_frame", "last_frame"],
|
||||
first_shot_ff_path_and_text_pair: Tuple[str, str],
|
||||
frame_desc: str,
|
||||
visible_characters: List[CharacterInScene],
|
||||
character_portraits_registry: Dict[str, Dict[str, Dict[str, str]]],
|
||||
progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None,
|
||||
) -> ImageOutput:
|
||||
|
||||
frame_image_path = os.path.join(self.working_dir, "shots", f"{shot_idx}", f"{frame_type}.png")
|
||||
|
||||
if os.path.exists(frame_image_path):
|
||||
print(f"🚀 Skipped generating {frame_type} for shot {shot_idx}, already exists.")
|
||||
_emit_render_progress(progress, "frame_exists", f"{frame_type} for shot {shot_idx} already exists", {"shot_idx": shot_idx, "frame_type": frame_type, "path": frame_image_path})
|
||||
|
||||
else:
|
||||
print(f"🖼️ Starting {frame_type} generation for shot {shot_idx}...")
|
||||
_emit_render_progress(progress, "frame_start", f"Generating {frame_type} for shot {shot_idx}", {"shot_idx": shot_idx, "frame_type": frame_type})
|
||||
available_image_path_and_text_pairs = []
|
||||
for visible_character in visible_characters:
|
||||
identifier_in_scene = visible_character.identifier_in_scene
|
||||
registry_item = character_portraits_registry[identifier_in_scene]
|
||||
for view, item in registry_item.items():
|
||||
available_image_path_and_text_pairs.append((item["path"], item["description"]))
|
||||
|
||||
available_image_path_and_text_pairs.append(first_shot_ff_path_and_text_pair)
|
||||
|
||||
selector_output_path = os.path.join(self.working_dir, "shots", f"{shot_idx}", f"{frame_type}_selector_output.json")
|
||||
if os.path.exists(selector_output_path):
|
||||
with open(selector_output_path, 'r', encoding='utf-8') as f:
|
||||
selector_output = json.load(f)
|
||||
print(f"🚀 Loaded existing reference image selection and prompt for {frame_type} frame of shot {shot_idx} from {selector_output_path}.")
|
||||
_emit_render_progress(progress, "frame_prompt_exists", f"Prompt for {frame_type} of shot {shot_idx} already exists", {"shot_idx": shot_idx, "frame_type": frame_type, "path": selector_output_path})
|
||||
else:
|
||||
print(f"🔍 Selecting reference images and generating prompt for {frame_type} frame of shot {shot_idx}...")
|
||||
_emit_render_progress(progress, "frame_prompt_start", f"Selecting references for {frame_type} of shot {shot_idx}", {"shot_idx": shot_idx, "frame_type": frame_type})
|
||||
selector_output = await self.reference_image_selector.select_reference_images_and_generate_prompt(
|
||||
available_image_path_and_text_pairs=available_image_path_and_text_pairs,
|
||||
frame_description=frame_desc
|
||||
)
|
||||
with open(selector_output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(selector_output, f, ensure_ascii=False, indent=4)
|
||||
print(f"☑️ Selected reference images and generated prompt for {frame_type} frame of shot {shot_idx}, saved to {selector_output_path}.")
|
||||
_emit_render_progress(progress, "frame_prompt_done", f"Selected references for {frame_type} of shot {shot_idx}", {"shot_idx": shot_idx, "frame_type": frame_type, "path": selector_output_path})
|
||||
|
||||
reference_image_path_and_text_pairs, prompt = selector_output["reference_image_path_and_text_pairs"], selector_output["text_prompt"]
|
||||
prefix_prompt = ""
|
||||
for i, (image_path, text) in enumerate(reference_image_path_and_text_pairs):
|
||||
prefix_prompt += f"Image {i}: {text}\n"
|
||||
prompt = f"{prefix_prompt}\n{prompt}"
|
||||
reference_image_paths = [item[0] for item in reference_image_path_and_text_pairs]
|
||||
|
||||
frame_image: ImageOutput = await self.image_generator.generate_single_image(
|
||||
prompt=prompt,
|
||||
reference_image_paths=reference_image_paths,
|
||||
size="1600x900",
|
||||
)
|
||||
frame_image.save(frame_image_path)
|
||||
print(f"☑️ Generated {frame_type} frame for shot {shot_idx}, saved to {frame_image_path}.")
|
||||
_emit_render_progress(progress, "frame_done", f"Generated {frame_type} for shot {shot_idx}", {"shot_idx": shot_idx, "frame_type": frame_type, "path": frame_image_path})
|
||||
|
||||
|
||||
self.frame_events[shot_idx][frame_type].set()
|
||||
return frame_image_path
|
||||
|
||||
|
||||
async def construct_camera_tree(
|
||||
self,
|
||||
shot_descriptions: List[ShotDescription],
|
||||
quiet: bool = False,
|
||||
):
|
||||
camera_tree_path = os.path.join(self.working_dir, "camera_tree.json")
|
||||
|
||||
if os.path.exists(camera_tree_path):
|
||||
with open(camera_tree_path, "r", encoding="utf-8") as f:
|
||||
camera_tree = json.load(f)
|
||||
camera_tree = [Camera.model_validate(camera) for camera in camera_tree]
|
||||
_pipeline_print(quiet, f"🚀 Loaded {len(camera_tree)} cameras from existing file.")
|
||||
return camera_tree
|
||||
|
||||
shot_descriptions = _normalize_model_list(shot_descriptions, ShotDescription, "shot_descriptions")
|
||||
cameras = _group_shots_into_cameras(shot_descriptions)
|
||||
|
||||
camera_tree = await self.camera_image_generator.construct_camera_tree(cameras=cameras, shot_descs=shot_descriptions)
|
||||
camera_tree = _normalize_model_list(camera_tree, Camera, "camera_tree")
|
||||
with open(camera_tree_path, "w", encoding="utf-8") as f:
|
||||
json.dump([camera.model_dump() for camera in camera_tree], f, ensure_ascii=False, indent=4)
|
||||
_pipeline_print(quiet, f"✅ Constructed camera tree and saved to {camera_tree_path}.")
|
||||
return camera_tree
|
||||
|
||||
|
||||
|
||||
|
||||
async def extract_characters(
|
||||
self,
|
||||
script: str,
|
||||
quiet: bool = False,
|
||||
):
|
||||
save_path = os.path.join(self.working_dir, "characters.json")
|
||||
|
||||
if os.path.exists(save_path):
|
||||
with open(save_path, "r", encoding="utf-8") as f:
|
||||
characters = json.load(f)
|
||||
characters = [CharacterInScene.model_validate(character) for character in characters]
|
||||
_pipeline_print(quiet, f"🚀 Loaded {len(characters)} characters from existing file.")
|
||||
else:
|
||||
characters = await self.character_extractor.extract_characters(script)
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
json.dump([character.model_dump() for character in characters], f, ensure_ascii=False, indent=4)
|
||||
_pipeline_print(quiet, f"✅ Extracted {len(characters)} characters from script and saved to {save_path}.")
|
||||
|
||||
for character in characters:
|
||||
self.character_portrait_events[character.idx] = asyncio.Event()
|
||||
|
||||
return characters
|
||||
|
||||
|
||||
async def generate_character_portraits(
|
||||
self,
|
||||
characters: List[CharacterInScene],
|
||||
character_portraits_registry: Optional[Dict[str, Dict[str, Dict[str, str]]]],
|
||||
style: str,
|
||||
progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None,
|
||||
):
|
||||
character_portraits_registry_path = os.path.join(self.working_dir, "character_portraits_registry.json")
|
||||
if character_portraits_registry is None:
|
||||
if os.path.exists(character_portraits_registry_path):
|
||||
with open(character_portraits_registry_path, 'r', encoding='utf-8') as f:
|
||||
character_portraits_registry = json.load(f)
|
||||
else:
|
||||
character_portraits_registry = {}
|
||||
|
||||
|
||||
tasks = [
|
||||
self.generate_portraits_for_single_character(character, style, progress=progress)
|
||||
for character in characters
|
||||
if character.identifier_in_scene not in character_portraits_registry
|
||||
]
|
||||
if tasks:
|
||||
for future in asyncio.as_completed(tasks):
|
||||
character_portraits_registry.update(await future)
|
||||
with open(character_portraits_registry_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(character_portraits_registry, f, ensure_ascii=False, indent=4)
|
||||
|
||||
print(f"✅ Completed character portrait generation for {len(characters)} characters.")
|
||||
_emit_render_progress(progress, "character_portraits_done", "Completed character portrait generation", {"character_count": len(characters)})
|
||||
else:
|
||||
print("🚀 All characters already have portraits, skipping portrait generation.")
|
||||
_emit_render_progress(progress, "character_portraits_exist", "All character portraits already exist", {"character_count": len(characters)})
|
||||
return character_portraits_registry
|
||||
|
||||
|
||||
async def generate_portraits_for_single_character(
|
||||
self,
|
||||
character: CharacterInScene,
|
||||
style: str,
|
||||
progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None,
|
||||
):
|
||||
character_dir = os.path.join(self.working_dir, "character_portraits", f"{character.idx}_{character.identifier_in_scene}")
|
||||
os.makedirs(character_dir, exist_ok=True)
|
||||
_emit_render_progress(progress, "character_portrait_start", f"Generating portraits for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene})
|
||||
|
||||
front_portrait_path = os.path.join(character_dir, "front.png")
|
||||
if os.path.exists(front_portrait_path):
|
||||
pass
|
||||
else:
|
||||
_emit_render_progress(progress, "character_portrait_front_start", f"Generating front portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene})
|
||||
front_portrait_output = await self.character_portraits_generator.generate_front_portrait(character, style)
|
||||
front_portrait_output.save(front_portrait_path)
|
||||
_emit_render_progress(progress, "character_portrait_front_done", f"Generated front portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene, "path": front_portrait_path})
|
||||
|
||||
|
||||
side_portrait_path = os.path.join(character_dir, "side.png")
|
||||
if os.path.exists(side_portrait_path):
|
||||
pass
|
||||
else:
|
||||
_emit_render_progress(progress, "character_portrait_side_start", f"Generating side portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene})
|
||||
side_portrait_output = await self.character_portraits_generator.generate_side_portrait(character, front_portrait_path)
|
||||
side_portrait_output.save(side_portrait_path)
|
||||
_emit_render_progress(progress, "character_portrait_side_done", f"Generated side portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene, "path": side_portrait_path})
|
||||
|
||||
back_portrait_path = os.path.join(character_dir, "back.png")
|
||||
if os.path.exists(back_portrait_path):
|
||||
pass
|
||||
else:
|
||||
_emit_render_progress(progress, "character_portrait_back_start", f"Generating back portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene})
|
||||
back_portrait_output = await self.character_portraits_generator.generate_back_portrait(character, front_portrait_path)
|
||||
back_portrait_output.save(back_portrait_path)
|
||||
_emit_render_progress(progress, "character_portrait_back_done", f"Generated back portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene, "path": back_portrait_path})
|
||||
|
||||
self.character_portrait_events[character.idx].set()
|
||||
|
||||
print(f"☑️ Completed character portrait generation for {character.identifier_in_scene}.")
|
||||
_emit_render_progress(progress, "character_portrait_done", f"Portraits for {character.identifier_in_scene} ready", {"character_idx": character.idx, "identifier": character.identifier_in_scene})
|
||||
|
||||
return {
|
||||
character.identifier_in_scene: {
|
||||
"front": {
|
||||
"path": front_portrait_path,
|
||||
"description": f"A front view portrait of {character.identifier_in_scene}.",
|
||||
},
|
||||
"side": {
|
||||
"path": side_portrait_path,
|
||||
"description": f"A side view portrait of {character.identifier_in_scene}.",
|
||||
},
|
||||
"back": {
|
||||
"path": back_portrait_path,
|
||||
"description": f"A back view portrait of {character.identifier_in_scene}.",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
async def design_storyboard(
|
||||
self,
|
||||
script: str,
|
||||
characters: List[CharacterInScene],
|
||||
user_requirement: str,
|
||||
quiet: bool = False,
|
||||
):
|
||||
storyboard_path = os.path.join(self.working_dir, "storyboard.json")
|
||||
if os.path.exists(storyboard_path):
|
||||
with open(storyboard_path, 'r', encoding='utf-8') as f:
|
||||
storyboard = json.load(f)
|
||||
storyboard = [ShotBriefDescription.model_validate(shot) for shot in storyboard]
|
||||
_pipeline_print(quiet, f"🚀 Loaded {len(storyboard)} shot brief descriptions from existing file.")
|
||||
else:
|
||||
_pipeline_print(quiet, f"🔍 Designing storyboard...")
|
||||
storyboard = await self.storyboard_artist.design_storyboard(
|
||||
script=script,
|
||||
characters=characters,
|
||||
user_requirement=user_requirement,
|
||||
retry_timeout=150,
|
||||
)
|
||||
storyboard = _normalize_model_list(storyboard, ShotBriefDescription, "storyboard")
|
||||
with open(storyboard_path, 'w', encoding='utf-8') as f:
|
||||
json.dump([shot.model_dump() for shot in storyboard], f, ensure_ascii=False, indent=4)
|
||||
_pipeline_print(quiet, f"✅ Designed storyboard and saved to {storyboard_path}.")
|
||||
|
||||
for shot_brief_description in storyboard:
|
||||
self.shot_desc_events[shot_brief_description.idx] = asyncio.Event()
|
||||
|
||||
return storyboard
|
||||
|
||||
|
||||
|
||||
async def decompose_visual_descriptions(
|
||||
self,
|
||||
shot_brief_descriptions: List[ShotBriefDescription],
|
||||
characters: List[CharacterInScene],
|
||||
quiet: bool = False,
|
||||
):
|
||||
tasks = [
|
||||
self.decompose_visual_description_for_single_shot_brief_description(shot_brief_description, characters, quiet=quiet)
|
||||
for shot_brief_description in shot_brief_descriptions
|
||||
]
|
||||
|
||||
shot_descriptions = await asyncio.gather(*tasks)
|
||||
return shot_descriptions
|
||||
|
||||
|
||||
async def decompose_visual_description_for_single_shot_brief_description(
|
||||
self,
|
||||
shot_brief_description: ShotBriefDescription,
|
||||
characters: List[CharacterInScene],
|
||||
quiet: bool = False,
|
||||
):
|
||||
shot_description_path = os.path.join(self.working_dir, "shots", f"{shot_brief_description.idx}", "shot_description.json")
|
||||
os.makedirs(os.path.dirname(shot_description_path), exist_ok=True)
|
||||
|
||||
if os.path.exists(shot_description_path):
|
||||
with open(shot_description_path, 'r', encoding='utf-8') as f:
|
||||
shot_description = ShotDescription.model_validate(json.load(f))
|
||||
_pipeline_print(quiet, f"🚀 Loaded shot {shot_brief_description.idx} description from existing file.")
|
||||
else:
|
||||
shot_description = await self.storyboard_artist.decompose_visual_description(
|
||||
shot_brief_desc=shot_brief_description,
|
||||
characters=characters,
|
||||
retry_timeout=120,
|
||||
)
|
||||
shot_description = _normalize_model_list([shot_description], ShotDescription, "shot_description")[0]
|
||||
with open(shot_description_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(shot_description.model_dump(), f, ensure_ascii=False, indent=4)
|
||||
_pipeline_print(quiet, f"✅ Decomposed visual description for shot {shot_brief_description.idx} and saved to {shot_description_path}.")
|
||||
|
||||
self.shot_desc_events[shot_brief_description.idx].set()
|
||||
|
||||
if shot_description.variation_type in ["medium", "large"]:
|
||||
self.frame_events[shot_brief_description.idx] = {
|
||||
"first_frame": asyncio.Event(),
|
||||
"last_frame": asyncio.Event(),
|
||||
}
|
||||
else:
|
||||
self.frame_events[shot_brief_description.idx] = {
|
||||
"first_frame": asyncio.Event(),
|
||||
}
|
||||
|
||||
return shot_description
|
||||
@@ -0,0 +1,5 @@
|
||||
You are the ViMax Agent, a multimodal generation agent.
|
||||
|
||||
Core loop contract:
|
||||
- Do not claim that planning, rendering, or file edits happened unless a tool result or `.working_dir` state proves it.
|
||||
- Do not claim render has started unless `vimax_render_video` reports that it started or completed.
|
||||
@@ -0,0 +1,63 @@
|
||||
ViMax supports three separate workflows: `idea2video`, `script2video`, and `novel2video`.
|
||||
|
||||
Idea2Video workflow DAG:
|
||||
|
||||
```text
|
||||
input_idea
|
||||
-> project_brief
|
||||
-> characters
|
||||
-> script
|
||||
-> storyboard
|
||||
-> shot_decomposition
|
||||
-> camera_tree
|
||||
-> frame_prompts
|
||||
-> keyframes
|
||||
-> video_clips
|
||||
-> final_video
|
||||
```
|
||||
|
||||
`.working_dir/<session_id-or-run_id>/` is the artifact authority. `.vimax/sessions.json` is only a session index. `.vimax/memory.md` stores user preferences only.
|
||||
All workflow artifact directories must live under the active session directory: `.working_dir/<session_id-or-run_id>/idea2video/`, `.working_dir/<session_id-or-run_id>/script2video/`, and `.working_dir/<session_id-or-run_id>/novel2video/`. Never read from or write to `.working_dir/idea2video/`, `.working_dir/script2video/`, or `.working_dir/novel2video/` at the root level.
|
||||
|
||||
Workflow confirmation gate: before calling any planning tool, the user must explicitly confirm which workflow to run: `idea2video`, `script2video`, or `novel2video`. Do not treat a vague idea, a request to "make a short film", or a request to "plan a script" as workflow confirmation. If the current user request does not explicitly name the workflow, do not call a planning tool; ask a concise clarification question first, for example: "Which workflow do you prefer: `idea2video`, `script2video`, or `novel2video`?" Only proceed to a planning tool after the user explicitly chooses one workflow in the current session. Source requirements still apply: `script2video` needs explicit script text for `script2video/script.txt`; `novel2video` needs explicit novel prose for `novel2video/novel/novel.txt`; vague ideas belong to `idea2video` only after the user confirms `idea2video`.
|
||||
|
||||
You may help the user draft, rewrite, or discuss a script in normal assistant text before planning. Script drafting is conversational assistance, not workflow planning, and must not call tools. If you draft a script and the user wants to use it for `script2video`, ask the user to confirm that exact script before calling `vimax_narrative_planning` with the `script` argument.
|
||||
Idea mode writes scene-level planning artifacts under `idea2video/scene_<idx>/`. Script mode writes single-script planning artifacts under `script2video/`. Use `vimax_narrative_planning` to create or revise structured text artifacts. Use `vimax_render_video` only when narrative planning dependencies exist.
|
||||
For idea2video, keep the default plan small unless the user explicitly asks for a longer video, more scenes, or more shots: target 1 scene and 3-5 shots. Do not expand a vague idea into many scenes or many shots by default.
|
||||
|
||||
Script2Video workflow DAG:
|
||||
|
||||
```text
|
||||
input_script
|
||||
-> characters
|
||||
-> storyboard
|
||||
-> shot_decomposition
|
||||
-> camera_tree
|
||||
-> frame_prompts
|
||||
-> keyframes
|
||||
-> video_clips
|
||||
-> final_video
|
||||
```
|
||||
|
||||
Script2Video requires an explicit source script. Only use script mode when the user provides concrete script text, a screenplay, a shot list, or says to use "this script". In that case, call `vimax_narrative_planning` with the `script` argument, not `idea`. Script mode stores the exact source script at `script2video/script.txt` and writes planning artifacts under `script2video/`. Do not infer or fabricate `script2video/script.txt` from a vague idea; use idea2video for vague ideas. Do not expand a supplied script into an idea2video story first unless the user explicitly asks to rewrite or develop it as an idea.
|
||||
|
||||
When the user asks to continue an existing project or fill missing text planning nodes, call `vimax_narrative_planning` for the active session. You may omit `idea` and `script`; the tool will reuse the active session source and existing cached artifacts. Do not use fake `revision_target` values such as `missing_structured_text_artifacts`; revision targets must be real relative file paths.
|
||||
|
||||
After project_brief, characters, script, storyboard, shot_decomposition, and camera_tree exist, if the user did not ask for end-to-end generation or render, do not call another tool. Reply that text planning is complete and ask whether to revise or enter render.
|
||||
|
||||
If the user explicitly asks for end-to-end generation, continue from planning into render tools.
|
||||
|
||||
|
||||
Novel workflow DAG:
|
||||
|
||||
```text
|
||||
novel_text
|
||||
-> compressed_novel
|
||||
-> events
|
||||
-> relevant_chunks
|
||||
-> scenes
|
||||
-> global_characters
|
||||
-> scene_scripts
|
||||
```
|
||||
|
||||
Novel2Video requires explicit novel prose. Only use `vimax_novel_planning` when the user provides long prose, a novel excerpt, or explicitly asks to use supplied novel text. Novel planning stores the source at `novel2video/novel/novel.txt`, then produces `novel2video/novel/novel_compressed.txt` and downstream novel artifacts. Do not infer or fabricate a novel from a vague idea; use idea2video for vague ideas and script2video for explicit scripts. `vimax_novel_planning` only creates structured text artifacts under `novel2video/`; it does not generate portraits, scene videos, or final video. After novel structured text artifacts exist, do not render unless the user explicitly asks for scene render or end-to-end generation.
|
||||
@@ -0,0 +1,43 @@
|
||||
[project]
|
||||
name = "autolongvideogeneration"
|
||||
version = "1.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "readme.md"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"aiohttp>=3.12.14",
|
||||
"chardet>=5.2.0",
|
||||
"faiss-cpu>=1.12.0",
|
||||
"google-genai>=1.47.0",
|
||||
"langchain>=0.3.26",
|
||||
"langchain-community>=0.3.27",
|
||||
"langchain-openai>=0.3.27",
|
||||
"moviepy>=2.2.1",
|
||||
"openai>=1.95.0",
|
||||
"opencv-python",
|
||||
"pillow>=11.3.0",
|
||||
"pyyaml>=6.0.2",
|
||||
"requests>=2.32.4",
|
||||
"scenedetect[opencv]>=0.6.7.1",
|
||||
"tenacity>=9.1.2",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
torch = [
|
||||
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
torchaudio = [
|
||||
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
|
||||
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cu128"
|
||||
url = "https://download.pytorch.org/whl/cu128"
|
||||
explicit = true
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest>=8",
|
||||
]
|
||||
@@ -0,0 +1,517 @@
|
||||
<div align="center">
|
||||
<img src="./assets/vimax.png">
|
||||
<br>
|
||||
<br>
|
||||
<a href="https://trendshift.io/repositories/15299" target="_blank"><img src="https://trendshift.io/api/badge/repositories/15299" alt="HKUDS%2FViMax | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
|
||||
<h1 align="center">ViMax: Agentic Video Generation</h1>
|
||||
|
||||
<div align="center">
|
||||
</div>
|
||||
|
||||
<p align="center">
|
||||
<img src="https://img.shields.io/badge/🐍Python-3.12-00d9ff?style=for-the-badge&logo=python&logoColor=white&labelColor=1a1a2e">
|
||||
<a href="https://github.com/astral-sh/uv"><img src="https://img.shields.io/badge/⚡uv-Ready-ff6b6b?style=for-the-badge&logo=python&logoColor=white&labelColor=1a1a2e"></a>
|
||||
<img src="https://img.shields.io/badge/License-MIT-4ecdc4?style=for-the-badge&logo=opensourceinitiative&logoColor=white" alt="MIT License">
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="./Communication.md"><img src="https://img.shields.io/badge/💬Feishu-Group-07c160?style=for-the-badge&logoColor=white&labelColor=1a1a2e"></a>
|
||||
<a href="./Communication.md"><img src="https://img.shields.io/badge/WeChat-Group-07c160?style=for-the-badge&logo=wechat&logoColor=white&labelColor=1a1a2e"></a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href='https://www.youtube.com/@AI-Creator-is-here'><img src='https://img.shields.io/badge/YouTube-ff0000?style=for-the-badge&logo=youtube&logoColor=white&labelColor=1a1a2e' /></a>
|
||||
<a href='https://arxiv.org/abs/2606.07649'><img src='https://img.shields.io/badge/arXiv-2606.07649-b31b1b?style=for-the-badge&logo=arxiv&logoColor=white&labelColor=1a1a2e' /></a>
|
||||
</p>
|
||||
|
||||
</div>
|
||||
<div align="center">
|
||||
|
||||
<p>
|
||||
<a href="readme.md"><img src="https://img.shields.io/badge/English-1a1a2e?style=for-the-badge"></a>
|
||||
<a href="README_ZH.md"><img src="https://img.shields.io/badge/中文版-1a1a2e?style=for-the-badge"></a>
|
||||
</p>
|
||||
<a href="#quick-start" style="text-decoration: none;">
|
||||
<img src="https://img.shields.io/badge/Quick%20Start-Get%20Started%20Now-FFC107?style=for-the-badge&logo=rocket&logoColor=white&labelColor=1a1a2e">
|
||||
</a>
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
### 🚨 Current Video Generation Limitations:
|
||||
- ❌ **Limited to Short Clips** - Most AI tools generate only seconds of footage. <br>
|
||||
- ❌ **Consistency Chaos** - Characters and scenes change unpredictably across frames. <br>
|
||||
- ❌ **Visual-Only Focus** - Missing scripts, audio, narrative structure, and storytelling depth. <br>
|
||||
|
||||
### 💡 ViMax Solution:
|
||||
🎬 **Director**, **Screenwriter**, **Producer**, and **Video Generator** **All-in-One**! We're exploring a future where AI becomes a complete creative powerhouse. 💡 Simply input your concept. ViMax autonomously handles the rest. It orchestrates scriptwriting, storyboarding, character creation, and final video generation—all end-to-end. 🚀
|
||||
|
||||
https://github.com/user-attachments/assets/5bad46b2-8276-4e1d-9480-3522640744b2
|
||||
|
||||
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
|
||||
### 📰 **News**
|
||||
|
||||
- **2026-06-28** 🛠️ Agent Loop and TUI stability update: stronger LLM retries, persistent render status, landscape image guards, and Script2Video resume fixes.
|
||||
- **2026-06-09** 📄 Technical report released.
|
||||
- **2026-06-08** 🤖 Agents Loop + TUI workflow integrated for interactive planning, revision, rendering control, session reuse, and context compaction.
|
||||
- **2026-06-07** 📖 Novel2Video workflow released.
|
||||
- **2026-06-01** 🎬 Google Omni video generator support added.
|
||||
- **2026-03-23** ⚡ MiniMax chat model provider support added.
|
||||
|
||||
---
|
||||
|
||||
|
||||
|
||||
## 📑 Table of Contents
|
||||
|
||||
- [💡 Key Features](#key-features)
|
||||
- [🔮 Demos](#Video-Demos-Generated-from-Scratch)
|
||||
- [🏗️ Architecture](#️-architecture)
|
||||
- [🚀 Quick Start](#quick-start)
|
||||
|
||||
---
|
||||
## 💡 Key Features
|
||||
|
||||
<br/>
|
||||
|
||||
<table align="center" width="100%" style="border: none; table-layout: fixed;">
|
||||
<tr>
|
||||
<td width="25%" align="center" style="vertical-align: top; padding: 20px;">
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<h3 style="margin: 0; padding: 0;">🌟 <strong>Idea2Video</strong></h3>
|
||||
</div>
|
||||
|
||||
<div align="center" style="margin: 15px 0;">
|
||||
<img src="https://img.shields.io/badge/IDEA-GENERATION-ff6b6b?style=for-the-badge&logo=algorithm&logoColor=white" alt="Algorithm Badge" />
|
||||
</div>
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>From Spark to Screen</strong></p>
|
||||
</div>
|
||||
|
||||
<div style="height: 60px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center">Transform <strong> raw ideas </strong> into complete video stories through intelligent multi-agent workflows automating <strong> storytelling, character design, and production </strong>.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
</td>
|
||||
<td width="25%" align="center" style="vertical-align: top; padding: 20px;">
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<h3 style="margin: 0; padding: 0;">🎨 <strong>Novel2Video</strong></h3>
|
||||
</div>
|
||||
|
||||
<div align="center" style="margin: 15px 0;">
|
||||
<img src="https://img.shields.io/badge/NOVEL-ADAPTATION-4ecdc4?style=for-the-badge&logo=book&logoColor=white" alt="Frontend Badge" />
|
||||
</div>
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>Smart Literary Adaptation Engine</strong></p>
|
||||
</div>
|
||||
|
||||
<div style="height: 60px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center">Transform <strong>complete novels</strong> into <strong>episodic video content</strong> with intelligent narrative compression, character tracking, and scene-by-scene visual adaptation</p>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
</td>
|
||||
<td width="25%" align="center" style="vertical-align: top; padding: 20px;">
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<h3 style="margin: 0; padding: 0;">⚙️ <strong>Script2Video</strong></h3>
|
||||
</div>
|
||||
|
||||
<div align="center" style="margin: 15px 0;">
|
||||
<img src="https://img.shields.io/badge/SCRIPT-ADAPTATION-9b59b6?style=for-the-badge&logo=server&logoColor=white" alt="Backend Badge" />
|
||||
</div>
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>Unlimited Screenplay Video Creation</strong></p>
|
||||
</div>
|
||||
|
||||
<div style="height: 60px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center">Unleash your creativity by writing <strong>any screenplay</strong> from personal stories to epic adventures, giving you complete control over every aspect of your visual storytelling.</p>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
</td>
|
||||
<td width="25%" align="center" style="vertical-align: top; padding: 20px;">
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<h3 style="margin: 0; padding: 0;">🤳 <strong>AutoCameo</strong></h3>
|
||||
</div>
|
||||
|
||||
<div align="center" style="margin: 15px 0;">
|
||||
<img src="https://img.shields.io/badge/INTERACTIVE-GENERATION-FFC107?style=for-the-badge&logo=server&logoColor=white" alt="Backend Badge" />
|
||||
</div>
|
||||
|
||||
<div style="height: 80px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"><strong>Generate Video from Your Photo</strong></p>
|
||||
</div>
|
||||
|
||||
<div style="height: 60px; display: flex; align-items: center; justify-content: center;">
|
||||
<p align="center"> <strong>Create your own cameo</strong> video, transforming yourself/pet into a guest star who appears across limitless creative scripts, cinematic sequences, and interactive storylines.</p>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<br/>
|
||||
|
||||
---
|
||||
|
||||
## 🔮Video Demos Generated from Scratch
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/c2fb27b0-218c-4976-b3d6-2abf8ea06be7" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/bfa566a8-688d-4d53-a9e2-6cedeb4a399d" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/49f61134-4f78-4285-9a9e-bb5e3e0c4abf" controls width="100%"></video>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/a950f449-a15c-449b-a1b8-c393951aa9be" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/bb3ff0fd-9433-4806-886a-3f77b61d06ec" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/2624a3f0-9f66-4fa4-b527-45c0ea0353fc" controls width="100%"></video>
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/5dbb80f7-aff0-4211-940c-a898f91fb80c" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/cc0b0bcd-e7db-4839-950b-0b03949637bd" controls width="100%"></video>
|
||||
</td>
|
||||
<td align="center" width="33%">
|
||||
<video src="https://github.com/user-attachments/assets/85919b59-80f0-461a-af7e-a93d3fb412fc" controls width="100%"></video>
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
</table>
|
||||
|
||||
---
|
||||
|
||||
|
||||
|
||||
### 🎯 **End-to-End Video Creation Engine**
|
||||
|
||||
**The Challenges**:
|
||||
|
||||
- 🌅 **Reference Images**: Time-consuming acquisition, organization, and alignment of reference frames that accurately capture characters, objects, positions, and environments.
|
||||
|
||||
- 🫠 **Consistency Check**: Sometimes, the image generator may generate unusable images even if it is given the correct characters, position, environment reference image and prompts.
|
||||
|
||||
- 📄 **Scripts Generation**: Professional and high-quality videos need to have rich information density and structured design.
|
||||
|
||||
- 📝 **Storyboard Design**: Converting stories into visual narratives requires expertise in cinematography, scene composition, and visual storytelling that most creators lack.
|
||||
|
||||
- 🎬 **Shot Design**: Creating coherent camera sequences with proper angles, transitions, and pacing while maintaining narrative flow across complex scenes.
|
||||
|
||||
- 🎨 **Development Delays**: Ensuring character appearances, environments, and artistic style remain consistent across hundreds of shots in long-form content.
|
||||
|
||||
- ⏱️ **Production Efficiency**: Traditional video creation involves multiple specialists and lengthy workflows, creating barriers for independent creators and rapid prototyping.
|
||||
|
||||
- 🎥 **Scaling AI Generated Video**: AI-generated videos are usually only a few seconds long, high-quality long videos at the minute or even hour level require complex cross-scene continuity and multi-storyboards design and processing capabilities.
|
||||
|
||||
|
||||
**ViMAX**: eliminates these production bottlenecks by automating the entire video creation pipeline from narrative input to final video output.
|
||||
|
||||
---
|
||||
|
||||
|
||||
### 🔥 **Why ViMax?**
|
||||
|
||||
| 🧠 **Effortless Production** | 🚀 **Complete Creative Freedom** | 🔊 **Audio and Video Binding** | 🎨 **Professional Quality** | 🤩 **Interactive Video**
|
||||
|:---:|:---:|:---:|:---:|:---:|
|
||||
| One-Prompt to Finished Video | From Any Narrative to Reality | Synchronized Storytelling | Movie-Grade Output | Make Your Own Cameo Video
|
||||
| Skip the technical complexity—just describe your vision and let ViMax handle script generation, storyboarding, shot design, reference management, and consistency validation | No creative limits—whether it's a trailer, short story, novel chapter, or original concept, ViMax intelligently structures narratives and designs cinematography to bring any idea to life | Seamlessly integrate character voice, and sound effects with visual content to create immersive experiences where audio and video work in perfect harmony | Automated quality control ensures character consistency, proper scene composition, and professional visual standards across every frame of your video | Interact in your own short stories by uploading your photo—ViMax intelligently integrates you as a character with consistent appearance and natural interactions throughout the entire video
|
||||
|
||||
ViMax now also includes an **Agents Loop + TUI** workflow for interactive planning, revision, rendering control, session reuse, and context compaction while preserving the original direct pipeline entrypoints.
|
||||
|
||||
|
||||
---
|
||||
|
||||
### ☄️ **RoadMap**
|
||||
|
||||
- 🚧 🖥️ **Web frontend interface**
|
||||
- 🚧 🎬 **Seedance 2.0 video generator support**
|
||||
- 🚧 🖼️ **GPT-Image 2 image generator support**
|
||||
|
||||
---
|
||||
|
||||
|
||||
|
||||
## 🏗️ Architecture
|
||||
|
||||
### 📊 **System Overview**
|
||||
|
||||
**ViMax** is a multi-agent video framework that enables automated multi-shot video generation while ensuring character and scene consistency. Our system seamlessly translates your ideas into corresponding videos, allowing you to focus on storytelling rather than technical implementation.
|
||||
|
||||
🎯 **Technical Capabilities**:
|
||||
|
||||
🧬 **Intelligent Long Script Generation**
|
||||
|
||||
RAG-based long script design engine that intelligently analyzes lengthy, novel-like stories and automatically segments them into a multi-scene script format. The process meticulously ensures that all key plot developments and character dialogues are accurately retained within the new structure.
|
||||
|
||||
🪄 **Expressive Storyboard Design**
|
||||
|
||||
Shot-level storyboard design system that create expressive storyboards through cinematography language based on user requirements and target audiences, which establishs the narrative rhythm for subsequent video generation.
|
||||
|
||||
🔮 **Multi-camera Filming Simulation**
|
||||
|
||||
Simulates multi-camera filming to deliver an immersive viewing experience while maintaining consistent character positioning and backgrounds within the same scene.
|
||||
|
||||
🧸 **Intelligent Reference Images Selection**
|
||||
|
||||
Intelligently select the reference image required for the first frame of the current video, including the storyboards that occurred in the previous timeline, to ensure the accuracy of multiple characters and environmental elements as the video becomes longer.
|
||||
|
||||
⚙️ **Automated Images Generation**
|
||||
|
||||
Based on the selected reference image and the visual logical order on the previous timeline, the prompt of the image generator is automatically generated to reasonably arrange the spatial interaction position between the character and the environment.
|
||||
|
||||
✅ **Automated Image Generation Consistency Check**
|
||||
|
||||
Generate multiple images in parallel and select the best consistent image as the first frame through MLLM/VLM to imitate the workflow of human creators.
|
||||
|
||||
⚡ **High-efficiency Parallel Shot Generation**
|
||||
|
||||
Parallel processing for sequential shots captured from the same camera enables highly efficient video production.
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
### 🤖 <strong>Multi-Agent Video Generation Pipeline</strong>
|
||||
|
||||
<div align="center">
|
||||
<table align="center" width="100%" style="border: none; border-collapse: collapse;">
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px; color: white; font-weight: bold;">
|
||||
🧠 <strong>INPUT LAYER</strong><br/>
|
||||
📝 Idea & Scripts & Novels • 💭 Natural Language Prompts • 🖼️ Reference Images • 🎨 Style Directives • 🧩 Configs
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="20"></td></tr>
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 15px; background: linear-gradient(135deg, #ff6b6b 0%, #ee5a24 100%); border-radius: 12px; color: white; font-weight: bold;">
|
||||
🧭 <strong>CENTRAL ORCHESTRATION</strong><br/>
|
||||
Agent Scheduling • Stage Transitions • Resource Management • Retry/Fallback Logic
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<td align="center" style="padding: 12px; background: linear-gradient(135deg, #3742fa 0%, #2f3542 100%); border-radius: 10px; color: white; width: 50%;">
|
||||
🧾 <strong>SCRIPT UNDERSTANDING</strong><br/>
|
||||
<small>Character/Environment Extraction • Scene Boundaries • Style Intent</small>
|
||||
</td>
|
||||
<td width="10"></td>
|
||||
<td align="center" style="padding: 12px; background: linear-gradient(135deg, #8c7ae6 0%, #9c88ff 100%); border-radius: 10px; color: white; width: 50%;">
|
||||
🎥 <strong>SCENE & SHOT PLANNING</strong><br/>
|
||||
<small>Storyboard Steps • Shot List • Key Frames & Beats</small>
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 15px; background: linear-gradient(135deg, #00d2d3 0%, #54a0ff 100%); border-radius: 12px; color: white; font-weight: bold;">
|
||||
🧪 <strong>VISUAL ASSET PLANNING</strong><br/>
|
||||
Reference Image Selection • Look/Style Guidance • Prompt Conditioning
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<!-- Swapped: ASSET INDEXING is now on the left -->
|
||||
<td align="center" style="padding: 12px; background: linear-gradient(135deg, #e056fd 0%, #f368e0 100%); border-radius: 10px; color: white; width: 50%;">
|
||||
🗂️ <strong>ASSET INDEXING</strong><br/>
|
||||
<small>Frames/Refs Catalog • Embeddings • Retrieval for Reuse</small>
|
||||
</td>
|
||||
<td width="10"></td>
|
||||
<!-- Swapped: CONSISTENCY & CONTINUITY is now on the right -->
|
||||
<td align="center" style="padding: 12px; background: linear-gradient(135deg, #ffa726 0%, #ff7043 100%); border-radius: 10px; color: white; width: 50%;">
|
||||
♻️ <strong>CONSISTENCY & CONTINUITY</strong><br/>
|
||||
<small>Character/Environment Tracking • Ref Matching • Temporal Coherence</small>
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 15px; background: linear-gradient(135deg, #26de81 0%, #20bf6b 100%); border-radius: 12px; color: white; font-weight: bold;">
|
||||
✂️ <strong>VISUAL SYNTHESIS & ASSEMBLY</strong><br/>
|
||||
Image Generation • Best-Frame Selection • First/Last-Frame→Video • Cut & Timeline Assembly
|
||||
</td>
|
||||
</tr>
|
||||
<tr><td colspan="3" height="15"></td></tr>
|
||||
<tr>
|
||||
<td colspan="3" align="center" style="padding: 20px; background: linear-gradient(135deg, #045de9 0%, #09c6f9 100%); border-radius: 15px; color: white; font-weight: bold;">
|
||||
🚀 <strong>OUTPUT LAYER</strong><br/>
|
||||
🖼️ Frames • 🎞️ Clips & Final Videos • 📜 Logs • 📦 Working Directory Artifacts
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
## 🚀Quick Start
|
||||
|
||||
### 🖥️ **Environment**
|
||||
|
||||
```
|
||||
OS: Linux, Windows
|
||||
```
|
||||
|
||||
### 📥 **Clone and Install**
|
||||
We use uv to manage the environment. For uv installation, please refer to the https://docs.astral.sh/uv/getting-started/installation/.
|
||||
```bash
|
||||
git clone https://github.com/HKUDS/ViMax.git
|
||||
cd ViMax
|
||||
uv sync
|
||||
```
|
||||
|
||||
|
||||
### 🧠 **Agent TUI**
|
||||
ViMax also provides a minimal TUI for interactive agent-based video creation. Configure the model and API key information in `configs/agent.local.yaml`, including the LLM, image generator, and video generator, as shown below.
|
||||
```yaml
|
||||
llm:
|
||||
model_provider: openai
|
||||
model: <YOUR_LLM_MODEL>
|
||||
base_url: <YOUR_LLM_BASE_URL>
|
||||
api_key: <YOUR_API_KEY>
|
||||
|
||||
image:
|
||||
model: <YOUR_IMAGE_MODEL>
|
||||
base_url: <YOUR_IMAGE_BASE_URL>
|
||||
api_key: <YOUR_API_KEY>
|
||||
|
||||
video:
|
||||
model: <YOUR_VIDEO_MODEL>
|
||||
base_url: <YOUR_VIDEO_BASE_URL>
|
||||
api_key: <YOUR_API_KEY>
|
||||
```
|
||||
|
||||
Then, start the TUI from the ViMax root directory:
|
||||
```bash
|
||||
vimax tui
|
||||
```
|
||||
|
||||
Start a new session or resume an existing one:
|
||||
```bash
|
||||
vimax tui new
|
||||
vimax tui resume
|
||||
vimax tui resume <session_id>
|
||||
```
|
||||
|
||||
You can also keep `configs/agent.local.yaml` empty and provide the same values through environment variables, such as `VIMAX_LLM_API_KEY`, `VIMAX_IMAGE_API_KEY`, and `VIMAX_VIDEO_API_KEY`.
|
||||
|
||||
### 🎯 **Usage**
|
||||
main_idea2video.py is used to convert your ideas into videos.
|
||||
You need to configure the model and API key information in the configs/idea2video.yaml file, including three parts—the chat model, the image generator, and the video generator, as shown below
|
||||
```yaml
|
||||
chat_model:
|
||||
init_args:
|
||||
model: google/gemini-2.5-flash-lite-preview-09-2025
|
||||
model_provider: openai
|
||||
api_key: <YOUR_API_KEY>
|
||||
base_url: https://openrouter.ai/api/v1
|
||||
|
||||
image_generator:
|
||||
class_path: tools.ImageGeneratorNanobananaGoogleAPI
|
||||
init_args:
|
||||
api_key: <YOUR_API_KEY>
|
||||
|
||||
video_generator:
|
||||
class_path: tools.VideoGeneratorVeoGoogleAPI
|
||||
init_args:
|
||||
api_key: <YOUR_API_KEY>
|
||||
|
||||
working_dir: .working_dir/idea2video
|
||||
```
|
||||
|
||||
Then, provide a simple yet thoughtful idea and the corresponding creative requirements in main_idea2video.py.
|
||||
```bash
|
||||
idea = \
|
||||
"""
|
||||
If a cat and a dog are best friends, what would happen when they meet a new cat?
|
||||
"""
|
||||
user_requirement = \
|
||||
"""
|
||||
For children, do not exceed 3 scenes.
|
||||
"""
|
||||
style = "Cartoon"
|
||||
```
|
||||
|
||||
main_script2video.py generates a video based on a specific script.
|
||||
You similarly need to set up the API configuration in configs/script2video.yaml file. Then, provide a scene script and the corresponding creative requirements in main_script2video.py, as shown below.
|
||||
```python
|
||||
script = \
|
||||
"""
|
||||
EXT. SCHOOL GYM - DAY
|
||||
A group of students are practicing basketball in the gym. The gym is large and open, with a basketball hoop at one end and a large crowd of spectators at the other end. John (18, male, tall, athletic) is the star player, and he is practicing his dribble and shot. Jane (17, female, short, athletic) is the assistant coach, and she is helping John with his practice. The other students are watching the practice and cheering for John.
|
||||
John: (dribbling the ball) I'm going to score a basket!
|
||||
Jane: (smiling) Good job, John!
|
||||
John: (shooting the ball) Yes!
|
||||
...
|
||||
"""
|
||||
user_requirement = \
|
||||
"""
|
||||
Fast-paced with no more than 20 shots.
|
||||
"""
|
||||
style = "Animate Style"
|
||||
```
|
||||
|
||||
|
||||
---
|
||||
|
||||
**🌟 If this project helps you, please give us a Star!**
|
||||
|
||||
<p align="center">
|
||||
<em> ❤️ Thanks for visiting ✨ ViMax!</em><br><br>
|
||||
</p>
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import yaml
|
||||
|
||||
from agent_runtime.config import (
|
||||
api_provider_from_base_url,
|
||||
embedding_api_key,
|
||||
embedding_base_url,
|
||||
embedding_model,
|
||||
embedding_model_provider,
|
||||
image_api_key,
|
||||
image_base_url,
|
||||
image_model,
|
||||
llm_api_key,
|
||||
llm_base_url,
|
||||
llm_model,
|
||||
llm_model_provider,
|
||||
load_agent_config,
|
||||
reranker_api_key,
|
||||
reranker_base_url,
|
||||
reranker_model,
|
||||
video_api_key,
|
||||
video_base_url,
|
||||
video_model,
|
||||
video_provider,
|
||||
)
|
||||
|
||||
|
||||
class AgentConfigTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
load_agent_config.cache_clear()
|
||||
|
||||
def tearDown(self):
|
||||
load_agent_config.cache_clear()
|
||||
|
||||
def test_reads_agent_local_config(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
config_dir = Path(tmp) / "configs"
|
||||
config_dir.mkdir()
|
||||
(config_dir / "agent.local.yaml").write_text(yaml.safe_dump({
|
||||
"llm": {"model_provider": "openai", "model": "config-llm", "base_url": "https://config.test/v1", "api_key": "config-key"},
|
||||
"image": {"model": "config-image", "base_url": "https://image.test", "api_key": "image-key"},
|
||||
"video": {"model": "config-video", "base_url": "https://openrouter.ai/api/v1", "api_key": "video-key"},
|
||||
"embedding": {"model_provider": "openai", "model": "config-embedding", "base_url": "https://embedding.test/v1", "api_key": "embedding-key"},
|
||||
"reranker": {"model": "config-reranker", "base_url": "https://reranker.test", "api_key": "reranker-key"},
|
||||
}), encoding="utf-8")
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
self.assertEqual(llm_model(tmp), "config-llm")
|
||||
self.assertEqual(llm_model_provider(tmp), "openai")
|
||||
self.assertEqual(llm_base_url(tmp), "https://config.test/v1")
|
||||
self.assertEqual(llm_api_key(tmp), "config-key")
|
||||
self.assertEqual(image_model(tmp), "config-image")
|
||||
self.assertEqual(image_base_url(tmp), "https://image.test")
|
||||
self.assertEqual(image_api_key(tmp), "image-key")
|
||||
self.assertEqual(video_model(tmp), "config-video")
|
||||
self.assertEqual(video_provider(tmp), "openrouter")
|
||||
self.assertEqual(video_base_url(tmp), "https://openrouter.ai/api/v1")
|
||||
self.assertEqual(video_api_key(tmp), "video-key")
|
||||
self.assertEqual(embedding_model_provider(tmp), "openai")
|
||||
self.assertEqual(embedding_model(tmp), "config-embedding")
|
||||
self.assertEqual(embedding_base_url(tmp), "https://embedding.test/v1")
|
||||
self.assertEqual(embedding_api_key(tmp), "embedding-key")
|
||||
self.assertEqual(reranker_model(tmp), "config-reranker")
|
||||
self.assertEqual(reranker_base_url(tmp), "https://reranker.test")
|
||||
self.assertEqual(reranker_api_key(tmp), "reranker-key")
|
||||
|
||||
def test_environment_overrides_agent_local_config(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
config_dir = Path(tmp) / "configs"
|
||||
config_dir.mkdir()
|
||||
(config_dir / "agent.local.yaml").write_text(yaml.safe_dump({"llm": {"model": "config-llm", "api_key": "config-key"}}), encoding="utf-8")
|
||||
with patch.dict(os.environ, {"VIMAX_LLM_MODEL": "env-llm", "VIMAX_LLM_MODEL_PROVIDER": "openai", "VIMAX_LLM_API_KEY": "env-key", "VIMAX_VIDEO_BASE_URL": "https://openrouter.ai/api/v1", "VIMAX_EMBEDDING_MODEL": "env-embedding", "VIMAX_EMBEDDING_BASE_URL": "https://env-embedding.test/v1", "VIMAX_EMBEDDING_API_KEY": "env-embedding-key", "VIMAX_RERANKER_MODEL": "env-reranker", "VIMAX_RERANKER_BASE_URL": "https://env-reranker.test", "VIMAX_RERANKER_API_KEY": "env-reranker-key"}, clear=True):
|
||||
self.assertEqual(llm_model(tmp), "env-llm")
|
||||
self.assertEqual(llm_model_provider(tmp), "openai")
|
||||
self.assertEqual(llm_api_key(tmp), "env-key")
|
||||
self.assertEqual(video_provider(tmp), "openrouter")
|
||||
self.assertEqual(video_base_url(tmp), "https://openrouter.ai/api/v1")
|
||||
self.assertEqual(embedding_model(tmp), "env-embedding")
|
||||
self.assertEqual(embedding_base_url(tmp), "https://env-embedding.test/v1")
|
||||
self.assertEqual(embedding_api_key(tmp), "env-embedding-key")
|
||||
self.assertEqual(reranker_model(tmp), "env-reranker")
|
||||
self.assertEqual(reranker_base_url(tmp), "https://env-reranker.test")
|
||||
self.assertEqual(reranker_api_key(tmp), "env-reranker-key")
|
||||
|
||||
def test_image_and_video_keys_fall_back_to_llm_key(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
config_dir = Path(tmp) / "configs"
|
||||
config_dir.mkdir()
|
||||
(config_dir / "agent.local.yaml").write_text(yaml.safe_dump({"llm": {"api_key": "shared-key"}}), encoding="utf-8")
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
self.assertEqual(image_api_key(tmp), "shared-key")
|
||||
self.assertEqual(video_api_key(tmp), "shared-key")
|
||||
|
||||
def test_video_provider_is_inferred_from_base_url(self):
|
||||
self.assertEqual(api_provider_from_base_url("https://openrouter.ai/api/v1"), "openrouter")
|
||||
self.assertEqual(api_provider_from_base_url("https://yunwu.ai/v1"), "yunwu")
|
||||
self.assertEqual(api_provider_from_base_url("https://example.com/v1"), "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,48 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from agent_runtime.llm import OpenAICompatibleLLM
|
||||
|
||||
|
||||
class AgentLLMTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_string_response_retries_before_clear_error(self):
|
||||
llm = OpenAICompatibleLLM(model="test", base_url="https://example.invalid/v1", api_key="test-key")
|
||||
create = AsyncMock(side_effect=["data: [DONE]", "bad response"])
|
||||
llm.client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
|
||||
with self.assertRaisesRegex(RuntimeError, "returned a string"):
|
||||
await llm.complete([], [])
|
||||
self.assertEqual(create.await_count, 2)
|
||||
|
||||
async def test_string_response_retry_can_recover(self):
|
||||
llm = OpenAICompatibleLLM(model="test", base_url="https://example.invalid/v1", api_key="test-key")
|
||||
create = AsyncMock(side_effect=["data: [DONE]", {"choices": [{"message": {"content": "recovered", "tool_calls": []}}]}])
|
||||
llm.client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
|
||||
message = await llm.complete([], [])
|
||||
self.assertEqual(message.text, "recovered")
|
||||
self.assertEqual(create.await_count, 2)
|
||||
|
||||
async def test_tool_request_falls_back_to_plain_chat_after_bad_tool_responses(self):
|
||||
llm = OpenAICompatibleLLM(model="test", base_url="https://example.invalid/v1", api_key="test-key")
|
||||
create = AsyncMock(side_effect=[
|
||||
"data: [DONE]",
|
||||
"data: [DONE]",
|
||||
{"choices": [{"message": {"content": "plain fallback", "tool_calls": []}}]},
|
||||
])
|
||||
llm.client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
|
||||
message = await llm.complete([], [{"type": "function", "function": {"name": "x", "parameters": {}}}])
|
||||
self.assertEqual(message.text, "plain fallback")
|
||||
self.assertEqual(create.await_count, 3)
|
||||
self.assertIsNone(create.await_args_list[-1].kwargs.get("tools"))
|
||||
|
||||
async def test_dict_response_is_accepted(self):
|
||||
llm = OpenAICompatibleLLM(model="test", base_url="https://example.invalid/v1", api_key="test-key")
|
||||
llm.client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=AsyncMock(return_value={
|
||||
"choices": [{"message": {"content": "hello", "tool_calls": []}}]
|
||||
}))))
|
||||
message = await llm.complete([], [])
|
||||
self.assertEqual(message.text, "hello")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,135 @@
|
||||
import asyncio
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from agent_runtime.context_compactor import ContextCompactor
|
||||
from agent_runtime.llm import AssistantMessage
|
||||
from agent_runtime.loop import AgentLoop
|
||||
from agent_runtime.models import ToolCall, ToolResult
|
||||
from agent_runtime.prompts import PromptBuilder
|
||||
from agent_runtime.session_index import SessionIndex
|
||||
from agent_runtime.tool_executor import ToolExecutor
|
||||
from agent_runtime.tools import ToolArgumentSchema, ToolRegistry, ToolSpec
|
||||
|
||||
|
||||
class FakeLLM:
|
||||
def __init__(self, replies):
|
||||
self.replies = list(replies)
|
||||
|
||||
async def complete(self, messages, tools):
|
||||
return self.replies.pop(0)
|
||||
|
||||
|
||||
class FailingLLM:
|
||||
async def complete(self, messages, tools):
|
||||
raise RuntimeError("provider returned invalid response shape")
|
||||
|
||||
|
||||
class AgentLoopTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_no_tool_call_finishes(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
registry = ToolRegistry([])
|
||||
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), FakeLLM([AssistantMessage(text="done")]))
|
||||
events = [event async for event in loop.stream_events("hi")]
|
||||
self.assertEqual(events[-2]["type"], "done")
|
||||
turn_id = events[0]["turn_id"]
|
||||
self.assertTrue(all(event.get("turn_id") == turn_id for event in events))
|
||||
log_text = (index.logs_dir / "loop_history.jsonl").read_text(encoding="utf-8")
|
||||
self.assertIn("assistant_finished_without_tools", log_text)
|
||||
|
||||
|
||||
async def test_turn_record_follows_session_created_by_tool(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
old = index.create(idea="old")
|
||||
|
||||
def create_actual(args):
|
||||
record = index.create(idea="actual")
|
||||
return ToolResult("create_actual", True, record["session_id"])
|
||||
|
||||
registry = ToolRegistry([ToolSpec("create_actual", "Create actual session", create_actual, schema={})])
|
||||
llm = FakeLLM([AssistantMessage(tool_calls=[ToolCall(name="create_actual", arguments={})]), AssistantMessage(text="finished")])
|
||||
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), llm)
|
||||
events = [event async for event in loop.stream_events("start new project")]
|
||||
active = index.active()
|
||||
self.assertNotEqual(active["session_id"], old["session_id"])
|
||||
self.assertEqual(len(index.get(active["session_id"])["recent_turn_records"]), 1)
|
||||
self.assertEqual(index.get(old["session_id"])["recent_turn_records"], [])
|
||||
self.assertEqual(events[-1]["session"]["active_session_id"], active["session_id"])
|
||||
|
||||
|
||||
async def test_tool_progress_streams_before_tool_result(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
release = asyncio.Event()
|
||||
|
||||
async def slow_tool(args, runtime):
|
||||
runtime.emit_progress("started", stage="running")
|
||||
await release.wait()
|
||||
return ToolResult("slow_tool", True, "done")
|
||||
|
||||
registry = ToolRegistry([ToolSpec("slow_tool", "Slow tool", slow_tool, schema={})])
|
||||
llm = FakeLLM([AssistantMessage(tool_calls=[ToolCall(name="slow_tool", arguments={})]), AssistantMessage(text="finished")])
|
||||
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), llm)
|
||||
agen = loop.stream_events("start")
|
||||
seen = []
|
||||
while True:
|
||||
event = await asyncio.wait_for(anext(agen), timeout=1)
|
||||
seen.append(event["type"])
|
||||
if event["type"] == "tool_progress":
|
||||
self.assertFalse(release.is_set())
|
||||
break
|
||||
release.set()
|
||||
async for event in agen:
|
||||
seen.append(event["type"])
|
||||
self.assertLess(seen.index("tool_progress"), seen.index("tool_result"))
|
||||
|
||||
|
||||
async def test_preflight_compact_summarizes_old_history(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
registry = ToolRegistry([])
|
||||
compactor = ContextCompactor(None, token_threshold=200, buffer_tokens=0, preserve_last_n=2, summary_max_chars=2000)
|
||||
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), FakeLLM([AssistantMessage(text="after compact")]), compactor)
|
||||
loop.history = [
|
||||
{"role": "user", "content": "old request " + "x" * 1200},
|
||||
{"role": "assistant", "content": "old answer " + "y" * 1200},
|
||||
{"role": "user", "content": "recent request"},
|
||||
{"role": "assistant", "content": "recent answer"},
|
||||
]
|
||||
events = [event async for event in loop.stream_events("continue")]
|
||||
self.assertIn("compact", [event.get("phase") for event in events if event["type"] == "status"])
|
||||
session = index.active()
|
||||
self.assertIn("Reference Context Only", session["compacted_summary"])
|
||||
self.assertGreaterEqual(session["compacted_turns"], 1)
|
||||
self.assertTrue(session["compaction_snapshots"])
|
||||
self.assertEqual(loop.history[0]["role"], "system")
|
||||
self.assertIn("after compact", loop.history[-1]["content"])
|
||||
self.assertNotIn("old request", index.memory_text())
|
||||
|
||||
|
||||
async def test_llm_sampling_error_yields_error_without_crashing_loop(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
registry = ToolRegistry([])
|
||||
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), FailingLLM())
|
||||
events = [event async for event in loop.stream_events("start")]
|
||||
self.assertTrue(any(event["type"] == "error" and event.get("metadata", {}).get("error_type") == "llm_sampling_failed" for event in events))
|
||||
self.assertEqual(events[-2]["type"], "done")
|
||||
self.assertEqual(events[-1]["type"], "session")
|
||||
self.assertEqual(index.active()["recent_turn_records"][-1]["status"], "failed")
|
||||
|
||||
async def test_tool_call_continues_then_finishes(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
|
||||
def hello(args):
|
||||
return ToolResult("hello", True, "hello result")
|
||||
|
||||
registry = ToolRegistry([ToolSpec("hello", "Say hello", hello, schema={"name": ToolArgumentSchema(str, False, "x")})])
|
||||
llm = FakeLLM([AssistantMessage(tool_calls=[ToolCall(name="hello", arguments={})]), AssistantMessage(text="finished")])
|
||||
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), llm)
|
||||
events = [event async for event in loop.stream_events("start")]
|
||||
self.assertTrue(any(event["type"] == "tool_result" for event in events))
|
||||
self.assertEqual(events[-2]["assistant"], "finished")
|
||||
@@ -0,0 +1,70 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from agent_runtime.prompts import PromptBuilder
|
||||
from agent_runtime.session_index import SessionIndex
|
||||
from agent_runtime.tools import build_builtin_registry
|
||||
|
||||
|
||||
class PromptBuilderTests(unittest.TestCase):
|
||||
def test_prompt_injects_context_and_tool_manifest(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp)
|
||||
(root / "prompts").mkdir()
|
||||
(root / "prompts" / "agent.md").write_text("agent rules", encoding="utf-8")
|
||||
(root / "prompts" / "workflow.md").write_text("workflow rules", encoding="utf-8")
|
||||
index = SessionIndex(root)
|
||||
index.create(idea="cat")
|
||||
registry = build_builtin_registry(root, index)
|
||||
builder = PromptBuilder(root / "prompts", index, registry)
|
||||
messages = builder.build_messages("start")
|
||||
self.assertIn("Available tools", messages[0]["content"])
|
||||
self.assertIn("当前 working_dir 尚未完成结构化文本文件", messages[0]["content"])
|
||||
self.assertIn("read_file", messages[0]["content"])
|
||||
trace = builder.trace(builder.build_parts("start"))
|
||||
self.assertGreater(trace["total_estimated_tokens"], 0)
|
||||
|
||||
|
||||
def test_prompt_injects_compacted_session_summary_as_reference(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp)
|
||||
(root / "prompts").mkdir()
|
||||
(root / "prompts" / "agent.md").write_text("agent rules", encoding="utf-8")
|
||||
(root / "prompts" / "workflow.md").write_text("workflow rules", encoding="utf-8")
|
||||
index = SessionIndex(root)
|
||||
record = index.create(idea="cat")
|
||||
index.update_compaction(record["session_id"], {"summary": "## Reference Context Only\n- user wants moon cat", "compacted_message_count": 4, "preserved_message_count": 2})
|
||||
registry = build_builtin_registry(root, index)
|
||||
builder = PromptBuilder(root / "prompts", index, registry)
|
||||
message = builder.build_messages("continue")[0]["content"]
|
||||
self.assertIn("Session context summary", message)
|
||||
self.assertIn("reference context only", message)
|
||||
self.assertIn("user wants moon cat", message)
|
||||
trace = builder.trace(builder.build_parts("continue"))
|
||||
self.assertGreater(trace["totals"]["dynamic_tokens"], 0)
|
||||
|
||||
def test_prompt_treats_novel_text_artifacts_as_text_stage_complete(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp)
|
||||
(root / "prompts").mkdir()
|
||||
(root / "prompts" / "agent.md").write_text("agent rules", encoding="utf-8")
|
||||
(root / "prompts" / "workflow.md").write_text("workflow rules", encoding="utf-8")
|
||||
index = SessionIndex(root)
|
||||
record = index.create(idea="novel")
|
||||
session_root = root / record["working_dir"] / "novel2video"
|
||||
(session_root / "novel").mkdir(parents=True)
|
||||
(session_root / "novel" / "novel_compressed.txt").write_text("compressed", encoding="utf-8")
|
||||
(session_root / "events").mkdir()
|
||||
(session_root / "events" / "event_0.json").write_text("{}", encoding="utf-8")
|
||||
(session_root / "relevant_chunks" / "event_0").mkdir(parents=True)
|
||||
(session_root / "relevant_chunks" / "event_0" / "chunk.txt").write_text("chunk", encoding="utf-8")
|
||||
(session_root / "scenes" / "event_0").mkdir(parents=True)
|
||||
(session_root / "scenes" / "event_0" / "scene_0.json").write_text("{}", encoding="utf-8")
|
||||
(session_root / "global_information" / "characters" / "novel_level").mkdir(parents=True)
|
||||
(session_root / "global_information" / "characters" / "novel_level" / "novel_characters_after_event_0.json").write_text("[]", encoding="utf-8")
|
||||
registry = build_builtin_registry(root, index)
|
||||
builder = PromptBuilder(root / "prompts", index, registry)
|
||||
messages = builder.build_messages("continue")
|
||||
self.assertIn("文本规划阶段已完成", messages[0]["content"])
|
||||
self.assertIn("novel2video/events/event_*.json: present", messages[0]["content"])
|
||||
@@ -0,0 +1,69 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from agent_runtime.session_index import SessionIndex
|
||||
|
||||
|
||||
class SessionIndexTests(unittest.TestCase):
|
||||
def test_generated_session_id_round_trips_after_slug_truncation(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="A red ball rolls across a white table.")
|
||||
self.assertIsNotNone(index.get(record["session_id"]))
|
||||
self.assertEqual(index.working_dir(record["session_id"]).name, record["session_id"])
|
||||
|
||||
def test_create_session_and_checklist(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="Moon cat", user_requirement="short", style="anime")
|
||||
self.assertEqual(index.active()["session_id"], record["session_id"])
|
||||
working_dir = Path(tmp) / record["working_dir"]
|
||||
self.assertTrue((working_dir / "idea2video").exists())
|
||||
self.assertTrue((working_dir / "script2video").exists())
|
||||
checklist = index.artifact_checklist(record["session_id"])
|
||||
self.assertFalse(checklist["script2video/storyboard.json"])
|
||||
self.assertFalse(checklist["idea2video/scene_*/storyboard.json"])
|
||||
self.assertEqual(record["compacted_summary"], "")
|
||||
self.assertEqual(record["compaction_snapshots"], [])
|
||||
|
||||
|
||||
def test_session_id_is_sanitized_and_stays_under_working_dir(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(session_id="../../escaped-review")
|
||||
self.assertEqual(record["session_id"], "escaped-review")
|
||||
working_dir = (Path(tmp) / record["working_dir"]).resolve()
|
||||
self.assertTrue(str(working_dir).startswith(str((Path(tmp) / ".working_dir").resolve())))
|
||||
self.assertFalse((Path(tmp).parent / "escaped-review").exists())
|
||||
|
||||
|
||||
def test_update_compaction_writes_session_state_not_memory(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="compact")
|
||||
index.update_compaction(record["session_id"], {
|
||||
"summary": "## Reference Context Only\n- old context",
|
||||
"compacted_message_count": 4,
|
||||
"preserved_message_count": 2,
|
||||
"estimated_tokens_before": 1000,
|
||||
"estimated_tokens_after": 300,
|
||||
"reason": "manual",
|
||||
"mode": "fallback-local",
|
||||
})
|
||||
session = index.get(record["session_id"])
|
||||
self.assertIn("old context", session["compacted_summary"])
|
||||
self.assertEqual(session["compacted_turns"], 2)
|
||||
self.assertEqual(session["last_compaction_reason"], "manual")
|
||||
self.assertTrue(session["compaction_snapshots"])
|
||||
self.assertNotIn("old context", index.memory_text())
|
||||
|
||||
def test_memory_and_turn_record_boundaries(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create()
|
||||
index.write_memory("# User Preferences\n- 16:9\n")
|
||||
self.assertIn("16:9", index.memory_text())
|
||||
index.append_turn_record(record["session_id"], {"turn_id": "t1", "status": "completed", "tool_rounds": [], "final_assistant_text": "done"})
|
||||
self.assertTrue((Path(tmp) / ".vimax" / "logs" / "loop_history.jsonl").exists())
|
||||
self.assertEqual(len(index.get(record["session_id"])["recent_turn_records"]), 1)
|
||||
@@ -0,0 +1,127 @@
|
||||
import asyncio
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from agent_runtime.models import ToolCall, TurnControl
|
||||
from agent_runtime.session_index import SessionIndex
|
||||
from agent_runtime.tool_executor import ToolExecutor
|
||||
from agent_runtime.tools import build_builtin_registry
|
||||
|
||||
|
||||
class ToolRegistryTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_validation_default_write_json_and_logging(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
index.create()
|
||||
registry = build_builtin_registry(tmp, index)
|
||||
executor = ToolExecutor(registry, index)
|
||||
record = await executor.execute(ToolCall(name="write_json", arguments={"path": "data/a.json", "data": {"x": 1}}), TurnControl())
|
||||
self.assertTrue(record.result.ok)
|
||||
self.assertEqual(json.loads((Path(tmp) / "data" / "a.json").read_text())["x"], 1)
|
||||
self.assertTrue((Path(tmp) / ".vimax" / "logs" / "tool_calls.jsonl").exists())
|
||||
|
||||
async def test_unknown_and_missing_argument_return_tool_errors(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
registry = build_builtin_registry(tmp, index)
|
||||
executor = ToolExecutor(registry, index)
|
||||
missing = await executor.execute(ToolCall(name="read_file", arguments={}), TurnControl())
|
||||
self.assertFalse(missing.result.ok)
|
||||
unknown = await executor.execute(ToolCall(name="does_not_exist", arguments={}), TurnControl())
|
||||
self.assertFalse(unknown.result.ok)
|
||||
|
||||
|
||||
async def test_run_shell_is_disabled_by_default(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
registry = build_builtin_registry(tmp, index)
|
||||
executor = ToolExecutor(registry, index)
|
||||
record = await executor.execute(ToolCall(name="run_shell", arguments={"command": "pwd"}), TurnControl())
|
||||
self.assertFalse(record.result.ok)
|
||||
self.assertEqual(record.result.metadata["error_type"], "disabled")
|
||||
|
||||
|
||||
async def test_todo_read_returns_empty_items_by_default(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
registry = build_builtin_registry(tmp, index)
|
||||
executor = ToolExecutor(registry, index)
|
||||
record = await executor.execute(ToolCall(name="todo_read", arguments={}), TurnControl())
|
||||
self.assertTrue(record.result.ok)
|
||||
self.assertEqual(json.loads(record.result.content)["items"], [])
|
||||
|
||||
async def test_todo_write_then_read_persists_items_and_logs(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
index.create()
|
||||
registry = build_builtin_registry(tmp, index)
|
||||
executor = ToolExecutor(registry, index)
|
||||
items = [{"content": "实现 TUI", "status": "in_progress"}, {"content": "补测试"}]
|
||||
written = await executor.execute(ToolCall(name="todo_write", arguments={"items": items}), TurnControl())
|
||||
self.assertTrue(written.result.ok)
|
||||
todo_path = Path(tmp) / ".vimax" / "todo.json"
|
||||
self.assertTrue(todo_path.exists())
|
||||
payload = json.loads(todo_path.read_text())
|
||||
self.assertEqual(payload["items"][0]["content"], "实现 TUI")
|
||||
self.assertEqual(payload["items"][1]["status"], "pending")
|
||||
|
||||
read = await executor.execute(ToolCall(name="todo_read", arguments={}), TurnControl())
|
||||
self.assertTrue(read.result.ok)
|
||||
self.assertEqual(json.loads(read.result.content)["items"], payload["items"])
|
||||
logs = (Path(tmp) / ".vimax" / "logs" / "tool_calls.jsonl").read_text()
|
||||
self.assertIn('"tool": "todo_write"', logs)
|
||||
|
||||
async def test_todo_write_rejects_invalid_payload(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
registry = build_builtin_registry(tmp, index)
|
||||
executor = ToolExecutor(registry, index)
|
||||
not_list = await executor.execute(ToolCall(name="todo_write", arguments={"items": {"content": "x"}}), TurnControl())
|
||||
self.assertFalse(not_list.result.ok)
|
||||
missing_content = await executor.execute(ToolCall(name="todo_write", arguments={"items": [{"status": "pending"}]}), TurnControl())
|
||||
self.assertFalse(missing_content.result.ok)
|
||||
bad_status = await executor.execute(ToolCall(name="todo_write", arguments={"items": [{"content": "x", "status": "blocked"}]}), TurnControl())
|
||||
self.assertFalse(bad_status.result.ok)
|
||||
|
||||
async def test_read_json_supports_virtual_session_json_path(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(session_id="20260630-125442-vimax", idea="surfing")
|
||||
registry = build_builtin_registry(tmp, index)
|
||||
executor = ToolExecutor(registry, index)
|
||||
result = await executor.execute(
|
||||
ToolCall(name="read_json", arguments={"path": f"{record['working_dir']}/session.json"}),
|
||||
TurnControl(),
|
||||
)
|
||||
self.assertTrue(result.result.ok)
|
||||
payload = json.loads(result.result.content)
|
||||
self.assertEqual(payload["session"]["session_id"], "20260630-125442-vimax")
|
||||
self.assertEqual(payload["source"], ".vimax/sessions.json")
|
||||
self.assertTrue(result.result.metadata["virtual_path"])
|
||||
|
||||
async def test_read_file_supports_virtual_session_log_path(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(session_id="20260630-125442-vimax", idea="surfing")
|
||||
index.append_turn_record(record["session_id"], {"turn_id": "turn-1", "status": "completed", "final_assistant_text": "done"})
|
||||
registry = build_builtin_registry(tmp, index)
|
||||
executor = ToolExecutor(registry, index)
|
||||
result = await executor.execute(
|
||||
ToolCall(name="read_file", arguments={"path": ".vimax/logs/20260630-125442-vimax.log"}),
|
||||
TurnControl(),
|
||||
)
|
||||
self.assertTrue(result.result.ok)
|
||||
payload = json.loads(result.result.content)
|
||||
self.assertEqual(payload["session_id"], "20260630-125442-vimax")
|
||||
self.assertEqual(payload["source"], ".vimax/logs/*.jsonl")
|
||||
self.assertEqual(payload["records"][0]["turn_id"], "turn-1")
|
||||
self.assertTrue(result.result.metadata["virtual_path"])
|
||||
|
||||
def test_concurrency_partition_groups_read_tools(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
registry = build_builtin_registry(tmp, SessionIndex(tmp))
|
||||
batches = registry.partition_calls([ToolCall("read_file", {"path": "x"}), ToolCall("glob_files", {"pattern": "*"}), ToolCall("write_json", {"path": "x", "data": {}})])
|
||||
self.assertEqual(len(batches), 2)
|
||||
self.assertEqual(len(batches[0]), 2)
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Regression tests for small crash bugs in helper stringification paths."""
|
||||
|
||||
import unittest
|
||||
|
||||
from agent_runtime.context_compactor import ContextCompactor
|
||||
from interfaces.shot_description import ShotBriefDescription
|
||||
|
||||
|
||||
class TestContextCompactorToolCallPreview(unittest.TestCase):
|
||||
def test_fallback_summary_handles_tool_call_messages(self):
|
||||
compactor = ContextCompactor(None, token_threshold=200, buffer_tokens=0, preserve_last_n=2, summary_max_chars=2000)
|
||||
messages = [
|
||||
{"role": "user", "content": "list the files"},
|
||||
{"role": "assistant", "content": "", "tool_calls": [{"id": "c1", "function": {"name": "list_files", "arguments": "{}"}}]},
|
||||
]
|
||||
summary = compactor._fallback_summary(messages, [], "", "test")
|
||||
self.assertIn("[tool calls]", summary)
|
||||
self.assertIn("list_files", summary)
|
||||
|
||||
|
||||
class TestShotBriefDescriptionStr(unittest.TestCase):
|
||||
def test_str_uses_existing_fields(self):
|
||||
shot = ShotBriefDescription(
|
||||
idx=0,
|
||||
is_last=False,
|
||||
cam_idx=1,
|
||||
visual_desc="<Alice> waves at the camera.",
|
||||
audio_desc="[Speaker] Alice (Happy): Hello!",
|
||||
)
|
||||
text = str(shot)
|
||||
self.assertIn("Shot 0", text)
|
||||
self.assertIn("Camera Index: 1", text)
|
||||
self.assertIn("<Alice> waves at the camera.", text)
|
||||
self.assertIn("[Speaker] Alice (Happy): Hello!", text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Every video generator must satisfy the VideoGenerator protocol, which
|
||||
declares **kwargs: pipelines pass progress= callbacks, and generators that
|
||||
reject unknown kwargs crash mid-render (TypeError) on the transition path."""
|
||||
|
||||
import inspect
|
||||
import unittest
|
||||
|
||||
from tools.video_generator_doubao_seedance_yunwu_api import VideoGeneratorDoubaoSeedanceYunwuAPI
|
||||
from tools.video_generator_omni_yunwu_api import VideoGeneratorOmniYunwuAPI
|
||||
from tools.video_generator_openrouter_api import VideoGeneratorOpenRouterAPI
|
||||
from tools.video_generator_veo_google_api import VideoGeneratorVeoGoogleAPI
|
||||
from tools.video_generator_veo_yunwu_api import VideoGeneratorVeoYunwuAPI
|
||||
|
||||
|
||||
class TestVideoGeneratorProtocol(unittest.TestCase):
|
||||
GENERATORS = [
|
||||
VideoGeneratorDoubaoSeedanceYunwuAPI,
|
||||
VideoGeneratorOmniYunwuAPI,
|
||||
VideoGeneratorOpenRouterAPI,
|
||||
VideoGeneratorVeoGoogleAPI,
|
||||
VideoGeneratorVeoYunwuAPI,
|
||||
]
|
||||
|
||||
def test_generate_single_video_accepts_arbitrary_kwargs(self):
|
||||
for cls in self.GENERATORS:
|
||||
with self.subTest(cls=cls.__name__):
|
||||
params = inspect.signature(cls.generate_single_video).parameters
|
||||
accepts_var_kwargs = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
|
||||
self.assertTrue(
|
||||
accepts_var_kwargs,
|
||||
f"{cls.__name__}.generate_single_video must accept **kwargs per tools.protocols.VideoGenerator "
|
||||
"(pipelines pass progress=...)",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,229 @@
|
||||
"""Regression tests for unbounded retry/polling loops and LLM-trusted graphs.
|
||||
|
||||
The bugs under test previously looped forever. Each test's fakes succeed after
|
||||
N calls, so the buggy code produces a fast assertion failure (extra calls or a
|
||||
missing exception) rather than hanging the suite; the fixed code must give up
|
||||
before the fake ever succeeds.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import requests
|
||||
|
||||
from agents.camera_image_generator import _validate_camera_tree
|
||||
from agents.event_extractor import EventExtractor
|
||||
from interfaces.camera import Camera
|
||||
from interfaces.event import Event
|
||||
from pipelines.novel2movie_pipeline import _ensure_extraction_cap
|
||||
from tools.video_generator_doubao_seedance_yunwu_api import VideoGeneratorDoubaoSeedanceYunwuAPI
|
||||
from tools.video_generator_omni_yunwu_api import VideoGeneratorOmniYunwuAPI
|
||||
from utils.image import download_image
|
||||
from utils.video import download_video
|
||||
|
||||
|
||||
def _no_sleep(fn):
|
||||
retrying = getattr(fn, "retry", None)
|
||||
if retrying is not None:
|
||||
retrying.sleep = lambda seconds: None
|
||||
|
||||
|
||||
class TestDownloadRetries(unittest.TestCase):
|
||||
def setUp(self):
|
||||
_no_sleep(download_image)
|
||||
_no_sleep(download_video)
|
||||
|
||||
def test_download_image_gives_up_on_persistent_network_error(self):
|
||||
calls = {"n": 0}
|
||||
|
||||
def flaky_get(url, **kwargs):
|
||||
calls["n"] += 1
|
||||
if calls["n"] < 10:
|
||||
raise requests.ConnectionError("connection refused")
|
||||
return MagicMock()
|
||||
|
||||
with patch("utils.image.requests.get", side_effect=flaky_get):
|
||||
with self.assertRaises(requests.ConnectionError):
|
||||
download_image("http://example.com/a.png", "/tmp/a.png")
|
||||
self.assertLessEqual(calls["n"], 5, "retry must be bounded, not retry-until-success")
|
||||
|
||||
def test_download_image_does_not_retry_client_errors(self):
|
||||
calls = {"n": 0}
|
||||
gone = MagicMock()
|
||||
gone.raise_for_status.side_effect = requests.HTTPError(
|
||||
"404", response=MagicMock(status_code=404)
|
||||
)
|
||||
|
||||
def expired_then_ok(url, **kwargs):
|
||||
calls["n"] += 1
|
||||
if calls["n"] < 3:
|
||||
return gone
|
||||
return MagicMock()
|
||||
|
||||
with patch("utils.image.requests.get", side_effect=expired_then_ok):
|
||||
with self.assertRaises(requests.HTTPError):
|
||||
download_image("http://example.com/expired.png", "/tmp/a.png")
|
||||
self.assertEqual(calls["n"], 1, "4xx responses must fail fast, not be retried")
|
||||
|
||||
def test_download_image_sets_a_timeout(self):
|
||||
captured = {}
|
||||
|
||||
def record_get(url, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return MagicMock()
|
||||
|
||||
with patch("utils.image.requests.get", side_effect=record_get):
|
||||
download_image("http://example.com/a.png", "/tmp/a.png")
|
||||
self.assertIsNotNone(captured.get("timeout"), "requests.get must not wait forever")
|
||||
|
||||
def test_download_video_gives_up_on_persistent_network_error(self):
|
||||
calls = {"n": 0}
|
||||
|
||||
def flaky_get(url, **kwargs):
|
||||
calls["n"] += 1
|
||||
if calls["n"] < 10:
|
||||
raise requests.ConnectionError("connection refused")
|
||||
return MagicMock()
|
||||
|
||||
with patch("utils.video.requests.get", side_effect=flaky_get):
|
||||
with self.assertRaises(requests.ConnectionError):
|
||||
download_video("http://example.com/a.mp4", "/tmp/a.mp4")
|
||||
self.assertLessEqual(calls["n"], 5, "retry must be bounded, not retry-until-success")
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def json(self):
|
||||
return self.payload
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
"""Returns each scripted (payload, status) in turn, repeating the last one."""
|
||||
|
||||
def __init__(self, scripted):
|
||||
self.scripted = list(scripted)
|
||||
self.calls = 0
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def _next(self):
|
||||
response = self.scripted[min(self.calls, len(self.scripted) - 1)]
|
||||
self.calls += 1
|
||||
return _FakeResponse(*response)
|
||||
|
||||
def post(self, url, **kwargs):
|
||||
return self._next()
|
||||
|
||||
def get(self, url, **kwargs):
|
||||
return self._next()
|
||||
|
||||
|
||||
class TestSeedanceClientBounds(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_create_task_fails_fast_on_auth_error(self):
|
||||
session = _FakeSession([
|
||||
({"error": "invalid api key"}, 401),
|
||||
({"error": "invalid api key"}, 401),
|
||||
({"id": "task-1"}, 200),
|
||||
])
|
||||
generator = VideoGeneratorDoubaoSeedanceYunwuAPI(api_key="bad-key")
|
||||
with patch("tools.video_generator_doubao_seedance_yunwu_api.aiohttp.ClientSession", return_value=session), \
|
||||
patch("tools.video_generator_doubao_seedance_yunwu_api.asyncio.sleep", new=AsyncMock()):
|
||||
with self.assertRaises(RuntimeError):
|
||||
await generator.create_video_generation_task("a prompt", [])
|
||||
self.assertEqual(session.calls, 1, "4xx must not be retried")
|
||||
|
||||
async def test_query_task_polling_is_bounded(self):
|
||||
session = _FakeSession([({"status": "queued"}, 200)])
|
||||
generator = VideoGeneratorDoubaoSeedanceYunwuAPI(api_key="key", max_poll_attempts=3)
|
||||
with patch("tools.video_generator_doubao_seedance_yunwu_api.aiohttp.ClientSession", return_value=session), \
|
||||
patch("tools.video_generator_doubao_seedance_yunwu_api.asyncio.sleep", new=AsyncMock()):
|
||||
with self.assertRaises(TimeoutError):
|
||||
await generator.query_video_generation_task("task-1")
|
||||
self.assertLessEqual(session.calls, 3)
|
||||
|
||||
|
||||
class TestOmniClientBounds(unittest.IsolatedAsyncioTestCase):
|
||||
def test_polling_is_bounded_by_default(self):
|
||||
generator = VideoGeneratorOmniYunwuAPI(api_key="key")
|
||||
self.assertIsNotNone(generator.max_poll_attempts, "default polling must have a deadline")
|
||||
|
||||
async def test_create_task_fails_fast_on_auth_error(self):
|
||||
session = _FakeSession([
|
||||
({"error": "invalid api key"}, 401),
|
||||
({"error": "invalid api key"}, 401),
|
||||
({"id": "task-1"}, 200),
|
||||
])
|
||||
generator = VideoGeneratorOmniYunwuAPI(api_key="bad-key")
|
||||
with patch("tools.video_generator_omni_yunwu_api.aiohttp.ClientSession", return_value=session), \
|
||||
patch("tools.video_generator_omni_yunwu_api.asyncio.sleep", new=AsyncMock()):
|
||||
with self.assertRaises(RuntimeError):
|
||||
await generator.create_video_generation_task("a prompt", [])
|
||||
self.assertEqual(session.calls, 1, "4xx must not be retried")
|
||||
|
||||
|
||||
class TestEventExtractionCap(unittest.TestCase):
|
||||
def test_extraction_aborts_when_model_never_emits_is_last(self):
|
||||
extractor = object.__new__(EventExtractor)
|
||||
calls = {"n": 0}
|
||||
|
||||
def never_last(novel_text, extracted_events):
|
||||
calls["n"] += 1
|
||||
if calls["n"] > 200:
|
||||
raise AssertionError("loop was not capped")
|
||||
return Event(
|
||||
index=len(extracted_events),
|
||||
is_last=False,
|
||||
description="an event",
|
||||
process_chain=["something happens"],
|
||||
)
|
||||
|
||||
extractor.extract_next_event = never_last
|
||||
with self.assertRaisesRegex(RuntimeError, "[Mm]ax|[Cc]ap|exceed"):
|
||||
extractor("some novel text")
|
||||
|
||||
def test_pipeline_extraction_cap_helper(self):
|
||||
_ensure_extraction_cap(0, 50, "events")
|
||||
_ensure_extraction_cap(49, 50, "events")
|
||||
with self.assertRaisesRegex(RuntimeError, "events"):
|
||||
_ensure_extraction_cap(50, 50, "events")
|
||||
|
||||
|
||||
class TestCameraTreeValidation(unittest.TestCase):
|
||||
def _camera(self, idx, parent=None):
|
||||
return Camera(idx=idx, active_shot_idxs=[idx], parent_cam_idx=parent, parent_shot_idx=idx if parent is not None else None)
|
||||
|
||||
def test_valid_tree_passes(self):
|
||||
cameras = [self._camera(0), self._camera(1, parent=0), self._camera(2, parent=1)]
|
||||
_validate_camera_tree(cameras)
|
||||
|
||||
def test_cycle_is_rejected(self):
|
||||
cameras = [self._camera(0, parent=1), self._camera(1, parent=0)]
|
||||
with self.assertRaisesRegex(ValueError, "[Cc]ycle"):
|
||||
_validate_camera_tree(cameras)
|
||||
|
||||
def test_self_parent_is_rejected(self):
|
||||
cameras = [self._camera(0, parent=0)]
|
||||
with self.assertRaises(ValueError):
|
||||
_validate_camera_tree(cameras)
|
||||
|
||||
def test_unknown_parent_index_is_rejected(self):
|
||||
cameras = [self._camera(0), self._camera(1, parent=7)]
|
||||
with self.assertRaises(ValueError):
|
||||
_validate_camera_tree(cameras)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Regression tests for rate-limiter lock behavior, media resource cleanup,
|
||||
packaging metadata, config templates, and test-suite isolation."""
|
||||
|
||||
import asyncio
|
||||
import subprocess
|
||||
import sys
|
||||
import tomllib
|
||||
import unittest
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import yaml
|
||||
|
||||
from utils.rate_limiter import RateLimiter
|
||||
from utils.video import concatenate_video_files
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
class TestRateLimiterLocking(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_waiting_acquirer_does_not_hold_the_lock(self):
|
||||
limiter = RateLimiter(max_requests_per_minute=1)
|
||||
await limiter.acquire() # consume the only slot in the window
|
||||
waiter = asyncio.create_task(limiter.acquire())
|
||||
await asyncio.sleep(0.05) # waiter is now waiting for the window to free
|
||||
try:
|
||||
try:
|
||||
await asyncio.wait_for(limiter.lock.acquire(), timeout=0.25)
|
||||
limiter.lock.release()
|
||||
except asyncio.TimeoutError:
|
||||
self.fail("rate limiter sleeps while holding its lock, blocking every other caller")
|
||||
finally:
|
||||
waiter.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await waiter
|
||||
|
||||
async def test_min_delay_smoothing_still_applies(self):
|
||||
limiter = RateLimiter(max_requests_per_minute=600) # min delay 0.1s
|
||||
loop = asyncio.get_running_loop()
|
||||
start = loop.time()
|
||||
await limiter.acquire()
|
||||
await limiter.acquire()
|
||||
self.assertGreaterEqual(loop.time() - start, 0.08)
|
||||
|
||||
|
||||
class TestVideoConcatenationCleanup(unittest.TestCase):
|
||||
def test_concatenate_closes_all_clips_even_on_failure(self):
|
||||
clips = [MagicMock(), MagicMock()]
|
||||
final = MagicMock()
|
||||
with patch("utils.video.VideoFileClip", side_effect=clips), \
|
||||
patch("utils.video.concatenate_videoclips", return_value=final):
|
||||
concatenate_video_files(["a.mp4", "b.mp4"], "out.mp4")
|
||||
final.write_videofile.assert_called_once()
|
||||
final.close.assert_called_once()
|
||||
for clip in clips:
|
||||
clip.close.assert_called_once()
|
||||
|
||||
# And when writing fails, the ffmpeg readers must still be released.
|
||||
clips = [MagicMock(), MagicMock()]
|
||||
final = MagicMock()
|
||||
final.write_videofile.side_effect = OSError("disk full")
|
||||
with patch("utils.video.VideoFileClip", side_effect=clips), \
|
||||
patch("utils.video.concatenate_videoclips", return_value=final):
|
||||
with self.assertRaises(OSError):
|
||||
concatenate_video_files(["a.mp4", "b.mp4"], "out.mp4")
|
||||
final.close.assert_called_once()
|
||||
for clip in clips:
|
||||
clip.close.assert_called_once()
|
||||
|
||||
|
||||
class TestPackagingMetadata(unittest.TestCase):
|
||||
def test_pyproject_is_consistent(self):
|
||||
with open(REPO_ROOT / "pyproject.toml", "rb") as f:
|
||||
data = tomllib.load(f)
|
||||
readme = data["project"]["readme"]
|
||||
self.assertTrue((REPO_ROOT / readme).exists(), f"readme points at missing file: {readme}")
|
||||
self.assertNotIn("index", data, "top-level [[index]] is not valid pyproject TOML and is silently ignored")
|
||||
dev_group = data.get("dependency-groups", {}).get("dev", [])
|
||||
self.assertTrue(any("pytest" in str(item) for item in dev_group), "pytest must be a declared dev dependency so the suite runs from the venv")
|
||||
|
||||
|
||||
class TestProviderConfigTemplates(unittest.TestCase):
|
||||
def test_minimax_templates_do_not_ship_truthy_placeholders(self):
|
||||
for name in ("idea2video_minimax.yaml", "script2video_minimax.yaml"):
|
||||
with open(REPO_ROOT / "configs" / name, encoding="utf-8") as f:
|
||||
config = yaml.safe_load(f)
|
||||
chat_key = config["chat_model"]["init_args"].get("api_key")
|
||||
self.assertFalse(chat_key, f"{name}: a truthy api_key placeholder defeats the MINIMAX_API_KEY env fallback")
|
||||
for section in ("image_generator", "video_generator"):
|
||||
key = config[section]["init_args"].get("api_key")
|
||||
self.assertNotIn("<", str(key or ""), f"{name}: {section} ships an angle-bracket placeholder that would be sent as a bearer token")
|
||||
|
||||
|
||||
class TestSuiteIsolation(unittest.TestCase):
|
||||
def test_importing_minimax_tests_does_not_stub_global_modules(self):
|
||||
code = (
|
||||
"import sys; import tests.test_minimax_integration; "
|
||||
"mod = sys.modules.get('cv2'); "
|
||||
"from unittest.mock import MagicMock; "
|
||||
"exit(1 if isinstance(mod, MagicMock) else 0)"
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
cwd=REPO_ROOT,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
self.assertEqual(
|
||||
result.returncode, 0,
|
||||
"importing tests.test_minimax_integration replaces sys.modules entries at import time, "
|
||||
"making every later-collected test module see MagicMock stubs instead of real libraries",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,29 @@
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from tools.image_orientation import ensure_not_portrait, landscape_guard_requested
|
||||
|
||||
|
||||
class ImageOrientationTests(unittest.TestCase):
|
||||
def test_landscape_guard_requested_defaults_to_all_images(self):
|
||||
self.assertTrue(landscape_guard_requested(size="1600x900"))
|
||||
self.assertTrue(landscape_guard_requested(size="512x512"))
|
||||
self.assertTrue(landscape_guard_requested(size=None))
|
||||
self.assertTrue(landscape_guard_requested(aspect_ratio="16:9", enforce_landscape=False))
|
||||
self.assertFalse(landscape_guard_requested(allow_portrait=True))
|
||||
|
||||
def test_portrait_detection_allows_slightly_tall_images(self):
|
||||
ensure_not_portrait(Image.new("RGB", (1000, 1040)))
|
||||
with self.assertRaises(ValueError):
|
||||
ensure_not_portrait(Image.new("RGB", (1000, 1100)))
|
||||
|
||||
def test_portrait_tolerance_env_override(self):
|
||||
with patch.dict(os.environ, {"VIMAX_IMAGE_PORTRAIT_RETRY_TOLERANCE": "1.20"}):
|
||||
ensure_not_portrait(Image.new("RGB", (1000, 1100)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,29 @@
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from tools.image_response import image_from_response_part
|
||||
|
||||
|
||||
class ImageResponseTests(unittest.TestCase):
|
||||
def test_extracts_image_when_part_has_no_as_image_method(self):
|
||||
expected = Image.new("RGB", (1, 1), (255, 0, 0))
|
||||
with patch("tools.image_response.Image.open", return_value=expected):
|
||||
part = SimpleNamespace(inline_data=SimpleNamespace(data=b"fake-png"))
|
||||
image = image_from_response_part(part)
|
||||
self.assertEqual(image.size, (1, 1))
|
||||
|
||||
def test_extracts_base64_data_url(self):
|
||||
expected = Image.new("RGB", (1, 1), (255, 0, 0))
|
||||
payload = "data:image/png;base64," + base64.b64encode(b"fake-png").decode("ascii")
|
||||
with patch("tools.image_response.Image.open", return_value=expected):
|
||||
part = {"inline_data": {"data": payload}}
|
||||
image = image_from_response_part(part)
|
||||
self.assertEqual(image.size, (1, 1))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,162 @@
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import main_agent
|
||||
|
||||
|
||||
class FakeSessionIndex:
|
||||
def __init__(self, fail_session=False):
|
||||
self.fail_session = fail_session
|
||||
self.activated = ""
|
||||
self.created = 0
|
||||
|
||||
def set_active(self, session_id):
|
||||
if self.fail_session:
|
||||
raise KeyError(session_id)
|
||||
self.activated = session_id
|
||||
|
||||
def create(self):
|
||||
self.created += 1
|
||||
self.activated = f"new-{self.created}"
|
||||
return {"session_id": self.activated}
|
||||
|
||||
def snapshot(self):
|
||||
return {"active_session_id": self.activated, "session": {"session_id": self.activated, "stage": "created"}}
|
||||
|
||||
|
||||
class FakeRuntime:
|
||||
def __init__(self, fail_session=False):
|
||||
self.session_index = FakeSessionIndex(fail_session=fail_session)
|
||||
self.inputs = []
|
||||
|
||||
async def compact_history(self, reason="manual"):
|
||||
return "Compacted context 100 -> 50 (fallback-local)."
|
||||
|
||||
async def stream_events(self, user_input):
|
||||
self.inputs.append(user_input)
|
||||
turn_id = "turn-test"
|
||||
yield {"type": "turn", "turn_id": turn_id, "turn": {"id": turn_id}}
|
||||
yield {"type": "status", "turn_id": turn_id, "phase": "sampling_assistant", "message": "Sampling assistant"}
|
||||
yield {"type": "tool_progress", "turn_id": turn_id, "tool": {"name": "fake_tool"}, "progress": {"stage": "running", "message": "working"}}
|
||||
yield {"type": "terminal", "turn_id": turn_id, "stream": "stdout", "line": "pipeline output"}
|
||||
yield {"type": "token", "turn_id": turn_id, "delta": "done"}
|
||||
yield {"type": "done", "turn_id": turn_id, "assistant": "done", "tool_results": []}
|
||||
yield {"type": "session", "turn_id": turn_id, "session": {"session": {"session_id": "s1", "stage": "narrative_planned"}}}
|
||||
|
||||
|
||||
class MainAgentCliTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def run_cli(self, argv, runtime=None, stdin_text="", session_index=None, load_runtime_side_effect=None):
|
||||
runtime = runtime or FakeRuntime()
|
||||
session_index = session_index or runtime.session_index
|
||||
stdout = io.StringIO()
|
||||
stderr = io.StringIO()
|
||||
stdin = io.StringIO(stdin_text)
|
||||
load_runtime_patch = patch.object(main_agent, "load_runtime", return_value=runtime)
|
||||
if load_runtime_side_effect is not None:
|
||||
load_runtime_patch = patch.object(main_agent, "load_runtime", side_effect=load_runtime_side_effect)
|
||||
with load_runtime_patch, \
|
||||
patch.object(main_agent, "load_session_index", return_value=session_index), \
|
||||
patch.object(sys, "stdin", stdin), \
|
||||
contextlib.redirect_stdout(stdout), \
|
||||
contextlib.redirect_stderr(stderr):
|
||||
code = await main_agent.amain(argv)
|
||||
return code, stdout.getvalue(), stderr.getvalue(), runtime
|
||||
|
||||
def test_help_parser_contains_once(self):
|
||||
stdout = io.StringIO()
|
||||
with contextlib.redirect_stdout(stdout), self.assertRaises(SystemExit) as ctx:
|
||||
main_agent.parse_args(["--help"])
|
||||
self.assertEqual(ctx.exception.code, 0)
|
||||
self.assertIn("--once", stdout.getvalue())
|
||||
|
||||
async def test_jsonl_once_outputs_valid_events_with_turn_id(self):
|
||||
code, stdout, stderr, runtime = await self.run_cli(["--jsonl", "--once", "hello"])
|
||||
self.assertEqual(code, 0)
|
||||
self.assertEqual(stderr, "")
|
||||
self.assertEqual(runtime.inputs, ["hello"])
|
||||
lines = [json.loads(line) for line in stdout.splitlines()]
|
||||
self.assertTrue(lines)
|
||||
self.assertEqual({event["turn_id"] for event in lines}, {"turn-test"})
|
||||
self.assertEqual(lines[0]["type"], "turn")
|
||||
self.assertIn("terminal", [event["type"] for event in lines])
|
||||
self.assertNotIn("›", stdout)
|
||||
|
||||
async def test_stdin_non_tty_is_single_prompt(self):
|
||||
code, stdout, stderr, runtime = await self.run_cli(["--jsonl"], stdin_text="from stdin\n")
|
||||
self.assertEqual(code, 0)
|
||||
self.assertEqual(runtime.inputs, ["from stdin"])
|
||||
self.assertTrue(stdout.strip())
|
||||
|
||||
|
||||
async def test_stdin_repl_reads_each_line_as_a_turn(self):
|
||||
code, stdout, stderr, runtime = await self.run_cli(["--jsonl", "--stdin-repl"], stdin_text="first\nsecond\n")
|
||||
self.assertEqual(code, 0)
|
||||
self.assertEqual(stderr, "")
|
||||
self.assertEqual(runtime.inputs, ["first", "second"])
|
||||
events = [json.loads(line) for line in stdout.splitlines()]
|
||||
self.assertEqual([event["type"] for event in events if event["type"] == "done"], ["done", "done"])
|
||||
|
||||
async def test_session_error_is_clear_before_runtime_load(self):
|
||||
runtime = FakeRuntime()
|
||||
failing_index = FakeSessionIndex(fail_session=True)
|
||||
code, stdout, stderr, runtime = await self.run_cli(
|
||||
["--session", "missing", "--once", "hello"],
|
||||
runtime=runtime,
|
||||
session_index=failing_index,
|
||||
load_runtime_side_effect=AssertionError("runtime should not load"),
|
||||
)
|
||||
self.assertEqual(code, 2)
|
||||
self.assertEqual(stdout, "")
|
||||
self.assertIn("unknown session id", stderr)
|
||||
self.assertEqual(runtime.inputs, [])
|
||||
|
||||
|
||||
async def test_new_session_is_created_before_runtime_load(self):
|
||||
runtime = FakeRuntime()
|
||||
session_index = FakeSessionIndex()
|
||||
code, stdout, stderr, runtime = await self.run_cli(
|
||||
["--new-session", "--jsonl", "--once", "hello"],
|
||||
runtime=runtime,
|
||||
session_index=session_index,
|
||||
)
|
||||
self.assertEqual(code, 0)
|
||||
self.assertEqual(stderr, "")
|
||||
self.assertEqual(session_index.created, 1)
|
||||
self.assertEqual(session_index.activated, "new-1")
|
||||
self.assertEqual(runtime.inputs, ["hello"])
|
||||
|
||||
async def test_new_session_and_session_are_mutually_exclusive(self):
|
||||
runtime = FakeRuntime()
|
||||
code, stdout, stderr, runtime = await self.run_cli(
|
||||
["--new-session", "--session", "s1", "--once", "hello"],
|
||||
runtime=runtime,
|
||||
load_runtime_side_effect=AssertionError("runtime should not load"),
|
||||
)
|
||||
self.assertEqual(code, 2)
|
||||
self.assertEqual(stdout, "")
|
||||
self.assertIn("cannot be used together", stderr)
|
||||
self.assertEqual(runtime.inputs, [])
|
||||
|
||||
|
||||
async def test_compact_command_outputs_jsonl_without_llm_turn(self):
|
||||
code, stdout, stderr, runtime = await self.run_cli(["--jsonl", "--once", "/compact"])
|
||||
self.assertEqual(code, 0)
|
||||
self.assertEqual(stderr, "")
|
||||
self.assertEqual(runtime.inputs, [])
|
||||
events = [json.loads(line) for line in stdout.splitlines()]
|
||||
self.assertEqual([event["type"] for event in events], ["turn", "status", "token", "done", "session"])
|
||||
turn_ids = {event["turn_id"] for event in events}
|
||||
self.assertEqual(len(turn_ids), 1)
|
||||
self.assertTrue(next(iter(turn_ids)).startswith("turn-"))
|
||||
self.assertIn("Compacted context", events[2]["delta"])
|
||||
|
||||
async def test_plain_mode_prints_progress_terminal_and_session(self):
|
||||
code, stdout, stderr, _ = await self.run_cli(["--once", "hello"])
|
||||
self.assertEqual(code, 0)
|
||||
self.assertIn("tool: fake_tool running: working", stdout)
|
||||
self.assertIn("terminal[stdout]: pipeline output", stdout)
|
||||
self.assertIn("session: s1 narrative_planned", stdout)
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Integration tests for MiniMax provider support.
|
||||
|
||||
These tests verify provider preset resolution and default pipeline config
|
||||
loading. They mock the LangChain factory so no real API calls are made.
|
||||
|
||||
Heavy multimedia dependencies (moviepy, scenedetect, cv2, google-genai,
|
||||
etc.) are stubbed in setUpModule and restored in tearDownModule. Stubbing
|
||||
at import time leaked the MagicMocks into sys.modules for every test module
|
||||
collected after this one, making suite results import-order dependent.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
_STUB_MODULES = [
|
||||
"moviepy", "cv2", "scenedetect", "scenedetect.detectors",
|
||||
"PIL", "PIL.Image",
|
||||
"faiss",
|
||||
"google", "google.genai", "google.genai.types", "google.genai.errors",
|
||||
"langchain_community", "langchain_community.vectorstores",
|
||||
"langchain_community.vectorstores.FAISS",
|
||||
]
|
||||
_saved = {}
|
||||
_modules_before_stubs = set()
|
||||
|
||||
|
||||
def setUpModule():
|
||||
_modules_before_stubs.update(sys.modules)
|
||||
for _mod in _STUB_MODULES:
|
||||
_saved[_mod] = sys.modules.get(_mod)
|
||||
mock = MagicMock()
|
||||
# Give stub a __spec__ so importlib.util.find_spec() works
|
||||
mock.__spec__ = importlib.machinery.ModuleSpec(_mod, None)
|
||||
mock.__path__ = []
|
||||
sys.modules[_mod] = mock
|
||||
|
||||
|
||||
def tearDownModule():
|
||||
# Drop project modules that were first imported while the stubs were
|
||||
# active, so later test modules import them fresh against real libraries.
|
||||
for name in list(sys.modules):
|
||||
if name in _modules_before_stubs:
|
||||
continue
|
||||
if name.split(".")[0] in {"pipelines", "agents", "tools", "interfaces"}:
|
||||
del sys.modules[name]
|
||||
for _mod, original in _saved.items():
|
||||
if original is None:
|
||||
sys.modules.pop(_mod, None)
|
||||
else:
|
||||
sys.modules[_mod] = original
|
||||
_saved.clear()
|
||||
_modules_before_stubs.clear()
|
||||
|
||||
|
||||
from utils.provider_presets import resolve_chat_model_config
|
||||
|
||||
|
||||
class TestPipelineConfigResolution(unittest.TestCase):
|
||||
"""Integration: config dict -> resolve -> init_chat_model kwargs."""
|
||||
|
||||
def _make_minimax_config(self, **overrides):
|
||||
base = {
|
||||
"model": "MiniMax-M3",
|
||||
"model_provider": "minimax",
|
||||
"api_key": "test-key",
|
||||
}
|
||||
base.update(overrides)
|
||||
return base
|
||||
|
||||
def test_full_minimax_config_resolution(self):
|
||||
config = self._make_minimax_config()
|
||||
resolved = resolve_chat_model_config(config)
|
||||
self.assertEqual(resolved["model_provider"], "openai")
|
||||
self.assertEqual(resolved["base_url"], "https://api.minimax.io/v1")
|
||||
self.assertEqual(resolved["model"], "MiniMax-M3")
|
||||
self.assertEqual(resolved["api_key"], "test-key")
|
||||
|
||||
def test_minimax_highspeed_model(self):
|
||||
config = self._make_minimax_config(model="MiniMax-M2.7-highspeed")
|
||||
resolved = resolve_chat_model_config(config)
|
||||
self.assertEqual(resolved["model"], "MiniMax-M2.7-highspeed")
|
||||
self.assertEqual(resolved["model_provider"], "openai")
|
||||
|
||||
def test_minimax_m27_model(self):
|
||||
config = self._make_minimax_config(model="MiniMax-M2.7")
|
||||
resolved = resolve_chat_model_config(config)
|
||||
self.assertEqual(resolved["model"], "MiniMax-M2.7")
|
||||
|
||||
@patch.dict(os.environ, {"MINIMAX_API_KEY": "env-api-key"})
|
||||
def test_env_key_fallback_in_config(self):
|
||||
config = {
|
||||
"model": "MiniMax-M3",
|
||||
"model_provider": "minimax",
|
||||
}
|
||||
resolved = resolve_chat_model_config(config)
|
||||
self.assertEqual(resolved["api_key"], "env-api-key")
|
||||
|
||||
def test_openrouter_config_unchanged(self):
|
||||
"""Existing OpenRouter configs must not be affected."""
|
||||
config = {
|
||||
"model": "google/gemini-2.5-flash-lite-preview-09-2025",
|
||||
"model_provider": "openai",
|
||||
"api_key": "or-key",
|
||||
"base_url": "https://openrouter.ai/api/v1",
|
||||
}
|
||||
resolved = resolve_chat_model_config(config)
|
||||
self.assertEqual(resolved["model_provider"], "openai")
|
||||
self.assertEqual(resolved["base_url"], "https://openrouter.ai/api/v1")
|
||||
self.assertEqual(resolved["model"], "google/gemini-2.5-flash-lite-preview-09-2025")
|
||||
|
||||
def test_init_chat_model_receives_openai_provider(self):
|
||||
"""Verify that resolved kwargs have model_provider='openai'."""
|
||||
config = self._make_minimax_config()
|
||||
resolved = resolve_chat_model_config(config)
|
||||
self.assertEqual(resolved["model_provider"], "openai")
|
||||
self.assertEqual(resolved["base_url"], "https://api.minimax.io/v1")
|
||||
self.assertEqual(resolved["model"], "MiniMax-M3")
|
||||
|
||||
def test_temperature_clamping_in_pipeline_flow(self):
|
||||
config = self._make_minimax_config(temperature=2.0)
|
||||
resolved = resolve_chat_model_config(config)
|
||||
self.assertEqual(resolved["temperature"], 1.0)
|
||||
|
||||
def test_extra_kwargs_preserved(self):
|
||||
config = self._make_minimax_config(max_tokens=4096, top_p=0.9)
|
||||
resolved = resolve_chat_model_config(config)
|
||||
self.assertEqual(resolved["max_tokens"], 4096)
|
||||
self.assertEqual(resolved["top_p"], 0.9)
|
||||
|
||||
|
||||
class TestPipelineInitFromConfig(unittest.TestCase):
|
||||
"""Integration: full pipeline init_from_config with provider configs."""
|
||||
|
||||
@patch("pipelines.idea2video_pipeline.init_chat_model")
|
||||
@patch("pipelines.idea2video_pipeline.RenderBackend.from_config")
|
||||
def test_idea2video_pipeline_minimax_config(self, mock_backend, mock_init):
|
||||
mock_model = MagicMock()
|
||||
mock_init.return_value = mock_model
|
||||
mock_backend.return_value = MagicMock(image_generator=MagicMock(), video_generator=MagicMock())
|
||||
|
||||
from pipelines.idea2video_pipeline import Idea2VideoPipeline
|
||||
pipeline = Idea2VideoPipeline.init_from_config("configs/idea2video_minimax.yaml")
|
||||
|
||||
mock_init.assert_called_once()
|
||||
call_kwargs = mock_init.call_args[1]
|
||||
self.assertEqual(call_kwargs["model_provider"], "openai")
|
||||
self.assertEqual(call_kwargs["base_url"], "https://api.minimax.io/v1")
|
||||
self.assertEqual(call_kwargs["model"], "MiniMax-M3")
|
||||
|
||||
@patch("pipelines.script2video_pipeline.init_chat_model")
|
||||
@patch("pipelines.script2video_pipeline.RenderBackend.from_config")
|
||||
def test_script2video_pipeline_minimax_config(self, mock_backend, mock_init):
|
||||
mock_model = MagicMock()
|
||||
mock_init.return_value = mock_model
|
||||
mock_backend.return_value = MagicMock(image_generator=MagicMock(), video_generator=MagicMock())
|
||||
|
||||
from pipelines.script2video_pipeline import Script2VideoPipeline
|
||||
pipeline = Script2VideoPipeline.init_from_config("configs/script2video_minimax.yaml")
|
||||
|
||||
mock_init.assert_called_once()
|
||||
call_kwargs = mock_init.call_args[1]
|
||||
self.assertEqual(call_kwargs["model_provider"], "openai")
|
||||
self.assertEqual(call_kwargs["base_url"], "https://api.minimax.io/v1")
|
||||
self.assertEqual(call_kwargs["model"], "MiniMax-M3")
|
||||
"""Integration: full pipeline init_from_config with default configs."""
|
||||
|
||||
@patch("pipelines.idea2video_pipeline.init_chat_model")
|
||||
@patch("pipelines.idea2video_pipeline.RenderBackend.from_config")
|
||||
def test_existing_openrouter_config_still_works(self, mock_backend, mock_init):
|
||||
mock_model = MagicMock()
|
||||
mock_init.return_value = mock_model
|
||||
mock_backend.return_value = MagicMock(image_generator=MagicMock(), video_generator=MagicMock())
|
||||
|
||||
from pipelines.idea2video_pipeline import Idea2VideoPipeline
|
||||
pipeline = Idea2VideoPipeline.init_from_config("configs/idea2video.yaml")
|
||||
|
||||
mock_init.assert_called_once()
|
||||
call_kwargs = mock_init.call_args[1]
|
||||
self.assertEqual(call_kwargs["model_provider"], "openai")
|
||||
self.assertEqual(call_kwargs["base_url"], "https://openrouter.ai/api/v1")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Tests for Novel2MoviePipeline initialization."""
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
import unittest
|
||||
|
||||
|
||||
class TestNovel2MoviePipelineInit(unittest.TestCase):
|
||||
def _class_node(self):
|
||||
source = Path("pipelines/novel2movie_pipeline.py").read_text(encoding="utf-8")
|
||||
tree = ast.parse(source)
|
||||
return next(
|
||||
node
|
||||
for node in tree.body
|
||||
if isinstance(node, ast.ClassDef) and node.name == "Novel2MoviePipeline"
|
||||
)
|
||||
|
||||
def test_initializes_runtime_dependencies(self):
|
||||
class_node = self._class_node()
|
||||
init_node = next(
|
||||
node
|
||||
for node in class_node.body
|
||||
if isinstance(node, ast.FunctionDef) and node.name == "__init__"
|
||||
)
|
||||
assigned = {
|
||||
target.attr
|
||||
for node in ast.walk(init_node)
|
||||
if isinstance(node, ast.Assign)
|
||||
for target in node.targets
|
||||
if isinstance(target, ast.Attribute)
|
||||
and isinstance(target.value, ast.Name)
|
||||
and target.value.id == "self"
|
||||
}
|
||||
|
||||
self.assertTrue(
|
||||
{
|
||||
"working_dir",
|
||||
"novel_compressor",
|
||||
"event_extractor",
|
||||
"embeddings",
|
||||
"rerank_model",
|
||||
"scene_extractor",
|
||||
"global_information_planner",
|
||||
"image_generator",
|
||||
"rewriter",
|
||||
"script2video_pipeline",
|
||||
}.issubset(assigned)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,277 @@
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from interfaces import CharacterInEvent, CharacterInNovel, CharacterInScene, Event, Scene
|
||||
from interfaces.environment import EnvironmentInScene
|
||||
from agent_runtime.session_index import SessionIndex
|
||||
from agent_runtime.tools import ToolRuntimeContext
|
||||
from agent_runtime.vimax_adapters import ViMaxAdapters, _run_planning_step
|
||||
from agents.global_information_planner import GlobalInformationPlanner, MergeCharactersAcrossScenesInEventResponse
|
||||
from pipelines.novel2movie_pipeline import Novel2MoviePipeline
|
||||
|
||||
|
||||
class FakeCompressor:
|
||||
def split(self, novel_text):
|
||||
return [novel_text]
|
||||
|
||||
async def compress_single_novel_chunk(self, semaphore, index, novel_chunk):
|
||||
return index, f"compressed {novel_chunk}"
|
||||
|
||||
def aggregate(self, chunks):
|
||||
return "\n".join(chunks)
|
||||
|
||||
|
||||
class FakeEventExtractor:
|
||||
def extract_next_event(self, novel_text, extracted_events):
|
||||
return Event(index=len(extracted_events), is_last=True, description="Hero leaves home", process_chain=["Hero opens the door"])
|
||||
|
||||
|
||||
class FakeKnowledgeBase:
|
||||
def similarity_search(self, process, k=10):
|
||||
return [SimpleNamespace(page_content="Hero opens the old wooden door.")]
|
||||
|
||||
|
||||
class FakeReranker:
|
||||
async def __call__(self, documents, query, top_n):
|
||||
return [(documents[0], 0.95)]
|
||||
|
||||
|
||||
class FakeSceneExtractor:
|
||||
async def get_next_scene(self, relevant_chunks, event, previous_scenes):
|
||||
return Scene(
|
||||
idx=len(previous_scenes),
|
||||
is_last=True,
|
||||
environment=EnvironmentInScene(slugline="INT. HOUSE - DAY", description="A quiet room."),
|
||||
characters=[CharacterInScene(idx=0, identifier_in_scene="Hero", is_visible=True, static_features="adult", dynamic_features="coat")],
|
||||
script="<Hero> opens the door.",
|
||||
)
|
||||
|
||||
|
||||
class FakeGlobalPlanner:
|
||||
async def merge_characters_across_scenes_in_event(self, event_idx, scenes):
|
||||
return [CharacterInEvent(index=0, identifier_in_event="Hero", active_scenes={0: "Hero"}, static_features="adult")]
|
||||
|
||||
def merge_characters_to_existing_characters_in_novel(self, event_idx, existing_characters_in_novel, characters_in_event):
|
||||
return [CharacterInNovel(index=0, identifier_in_novel="Hero", active_events={event_idx: "Hero"}, static_features="adult")]
|
||||
|
||||
|
||||
class FakeNovelPipeline:
|
||||
def __init__(self, working_dir: Path):
|
||||
self.working_dir = working_dir
|
||||
|
||||
async def plan_text_artifacts(self, novel_text, user_requirement="", style="", progress=None, quiet=False):
|
||||
if progress:
|
||||
for stage in ["save_novel", "compress_novel", "extract_events", "retrieve_chunks", "extract_scenes", "merge_characters", "completed"]:
|
||||
progress(stage, stage, {})
|
||||
novel = self.working_dir / "novel"
|
||||
novel.mkdir(parents=True, exist_ok=True)
|
||||
(novel / "novel.txt").write_text(novel_text, encoding="utf-8")
|
||||
(novel / "novel_compressed.txt").write_text("compressed", encoding="utf-8")
|
||||
events = self.working_dir / "events"
|
||||
events.mkdir(parents=True, exist_ok=True)
|
||||
(events / "event_0.json").write_text(json.dumps(Event(index=0, is_last=True, description="d", process_chain=["p"]).model_dump()), encoding="utf-8")
|
||||
chunks = self.working_dir / "relevant_chunks" / "event_0"
|
||||
chunks.mkdir(parents=True, exist_ok=True)
|
||||
(chunks / "chunk_0-score_0.95.txt").write_text("chunk", encoding="utf-8")
|
||||
scenes = self.working_dir / "scenes" / "event_0"
|
||||
scenes.mkdir(parents=True, exist_ok=True)
|
||||
scene = Scene(idx=0, is_last=True, environment=EnvironmentInScene(slugline="INT. ROOM - DAY", description="room"), characters=[CharacterInScene(idx=0, identifier_in_scene="Hero", is_visible=True, static_features="adult", dynamic_features="coat")], script="<Hero> walks.")
|
||||
(scenes / "scene_0.json").write_text(json.dumps(scene.model_dump()), encoding="utf-8")
|
||||
event_level = self.working_dir / "global_information" / "characters" / "event_level"
|
||||
event_level.mkdir(parents=True, exist_ok=True)
|
||||
event_char = CharacterInEvent(index=0, identifier_in_event="Hero", active_scenes={0: "Hero"}, static_features="adult")
|
||||
(event_level / "event_0_characters.json").write_text(json.dumps([event_char.model_dump()]), encoding="utf-8")
|
||||
novel_level = self.working_dir / "global_information" / "characters" / "novel_level"
|
||||
novel_level.mkdir(parents=True, exist_ok=True)
|
||||
novel_char = CharacterInNovel(index=0, identifier_in_novel="Hero", active_events={0: "Hero"}, static_features="adult")
|
||||
(novel_level / "novel_characters_after_event_0.json").write_text(json.dumps([novel_char.model_dump()]), encoding="utf-8")
|
||||
return {}
|
||||
|
||||
|
||||
class FakeMergeChain:
|
||||
async def ainvoke(self, messages):
|
||||
return MergeCharactersAcrossScenesInEventResponse(
|
||||
characters=[CharacterInEvent(index=0, identifier_in_event="Hero", active_scenes={0: "Hero"}, static_features="adult")]
|
||||
)
|
||||
|
||||
|
||||
class FakeMergeChatModel:
|
||||
def __or__(self, parser):
|
||||
return FakeMergeChain()
|
||||
|
||||
|
||||
class GlobalInformationPlannerCompatibilityTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_merge_event_characters_uses_scene_character_idx(self):
|
||||
planner = GlobalInformationPlanner.__new__(GlobalInformationPlanner)
|
||||
planner.chat_model = FakeMergeChatModel()
|
||||
scene = Scene(
|
||||
idx=0,
|
||||
is_last=True,
|
||||
environment=EnvironmentInScene(slugline="INT. ROOM - DAY", description="room"),
|
||||
characters=[CharacterInScene(idx=0, identifier_in_scene="Hero", is_visible=True, static_features="adult", dynamic_features="coat")],
|
||||
script="<Hero> walks.",
|
||||
)
|
||||
characters = await planner.merge_characters_across_scenes_in_event(event_idx=0, scenes=[scene])
|
||||
self.assertEqual(characters[0].identifier_in_event, "Hero")
|
||||
|
||||
|
||||
class PlanningStepOutputSuppressionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_run_planning_step_suppresses_stdout_stderr_and_warnings(self):
|
||||
async def noisy_step():
|
||||
print("NOISE_STDOUT")
|
||||
logging.warning("NOISE_WARNING")
|
||||
return "ok"
|
||||
|
||||
stdout = io.StringIO()
|
||||
stderr = io.StringIO()
|
||||
with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
|
||||
result = await _run_planning_step("message", "stage", noisy_step(), runtime=None)
|
||||
self.assertEqual(result, "ok")
|
||||
self.assertNotIn("NOISE_STDOUT", stdout.getvalue())
|
||||
self.assertNotIn("NOISE_WARNING", stderr.getvalue())
|
||||
|
||||
|
||||
class Novel2MoviePlanningTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_plan_text_artifacts_writes_structured_text_and_progress(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
pipeline = Novel2MoviePipeline(
|
||||
novel_compressor=FakeCompressor(),
|
||||
event_extractor=FakeEventExtractor(),
|
||||
embeddings=SimpleNamespace(model="fake-embedding"),
|
||||
rerank_model=FakeReranker(),
|
||||
scene_extractor=FakeSceneExtractor(),
|
||||
global_information_planner=FakeGlobalPlanner(),
|
||||
image_generator=object(),
|
||||
rewriter=object(),
|
||||
script2video_pipeline=object(),
|
||||
working_dir=tmp,
|
||||
)
|
||||
events = []
|
||||
with patch("pipelines.novel2movie_pipeline.CacheBackedEmbeddings.from_bytes_store", return_value=object()), \
|
||||
patch("pipelines.novel2movie_pipeline.FAISS.from_texts", return_value=FakeKnowledgeBase()):
|
||||
result = await pipeline.plan_text_artifacts("Hero opens a door.", progress=lambda stage, message, metadata=None: events.append(stage), quiet=True)
|
||||
self.assertEqual(events, ["save_novel", "compress_novel", "extract_events", "retrieve_chunks", "extract_scenes", "merge_characters", "completed"])
|
||||
root = Path(tmp)
|
||||
self.assertTrue((root / "novel" / "novel_compressed.txt").exists())
|
||||
self.assertTrue((root / "events" / "event_0.json").exists())
|
||||
self.assertTrue((root / "relevant_chunks" / "event_0" / "chunk_0-score_0.95.txt").exists())
|
||||
self.assertTrue((root / "scenes" / "event_0" / "scene_0.json").exists())
|
||||
self.assertTrue((root / "global_information" / "characters" / "novel_level" / "novel_characters_after_event_0.json").exists())
|
||||
self.assertFalse((root / "character_portraits").exists())
|
||||
self.assertFalse((root / "videos").exists())
|
||||
self.assertEqual(len(result["events"]), 1)
|
||||
|
||||
|
||||
class FakeNovelRenderPipeline:
|
||||
def __init__(self, working_dir: Path):
|
||||
self.working_dir = Path(working_dir)
|
||||
|
||||
async def render_video_artifacts(self, style, user_requirement="", progress=None, quiet=False):
|
||||
if progress:
|
||||
progress("novel_portraits_start", "portraits", {})
|
||||
progress("novel_scene_render_start", "scene", {"event_idx": 0, "scene_idx": 0})
|
||||
progress("novel_render_completed", "done", {"scene_count": 1})
|
||||
scene_dir = self.working_dir / "videos" / "event_0" / "scene_0"
|
||||
scene_dir.mkdir(parents=True, exist_ok=True)
|
||||
(scene_dir / "final_video.mp4").write_text("video", encoding="utf-8")
|
||||
return {
|
||||
"character_portraits_dir": str(self.working_dir / "character_portraits"),
|
||||
"scene_videos_dir": str(self.working_dir / "videos"),
|
||||
"scene_video_dirs": [str(scene_dir)],
|
||||
"scene_count": 1,
|
||||
}
|
||||
|
||||
|
||||
def write_minimal_novel_artifacts(root: Path):
|
||||
novel = root / "novel2video"
|
||||
(novel / "novel").mkdir(parents=True, exist_ok=True)
|
||||
(novel / "novel" / "novel.txt").write_text("novel", encoding="utf-8")
|
||||
(novel / "novel" / "novel_compressed.txt").write_text("compressed", encoding="utf-8")
|
||||
events = novel / "events"
|
||||
events.mkdir(parents=True, exist_ok=True)
|
||||
(events / "event_0.json").write_text(json.dumps(Event(index=0, is_last=True, description="d", process_chain=["p"]).model_dump()), encoding="utf-8")
|
||||
chunks = novel / "relevant_chunks" / "event_0"
|
||||
chunks.mkdir(parents=True, exist_ok=True)
|
||||
(chunks / "chunk_0-score_0.95.txt").write_text("chunk", encoding="utf-8")
|
||||
scenes = novel / "scenes" / "event_0"
|
||||
scenes.mkdir(parents=True, exist_ok=True)
|
||||
scene = Scene(idx=0, is_last=True, environment=EnvironmentInScene(slugline="INT. ROOM - DAY", description="room"), characters=[CharacterInScene(idx=0, identifier_in_scene="Hero", is_visible=True, static_features="adult", dynamic_features="coat")], script="<Hero> walks.")
|
||||
(scenes / "scene_0.json").write_text(json.dumps(scene.model_dump()), encoding="utf-8")
|
||||
event_level = novel / "global_information" / "characters" / "event_level"
|
||||
event_level.mkdir(parents=True, exist_ok=True)
|
||||
event_char = CharacterInEvent(index=0, identifier_in_event="Hero", active_scenes={0: "Hero"}, static_features="adult")
|
||||
(event_level / "event_0_characters.json").write_text(json.dumps([event_char.model_dump()]), encoding="utf-8")
|
||||
novel_level = novel / "global_information" / "characters" / "novel_level"
|
||||
novel_level.mkdir(parents=True, exist_ok=True)
|
||||
novel_char = CharacterInNovel(index=0, identifier_in_novel="Hero", active_events={0: "Hero"}, static_features="adult")
|
||||
(novel_level / "novel_characters_after_event_0.json").write_text(json.dumps([novel_char.model_dump()]), encoding="utf-8")
|
||||
|
||||
|
||||
class NovelAdapterTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_missing_rag_config_returns_tool_error_and_marks_session_error(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch.dict("os.environ", {"VIMAX_LLM_API_KEY": "llm-key", "VIMAX_LLM_BASE_URL": "https://llm.test/v1"}, clear=True), \
|
||||
patch("agent_runtime.vimax_adapters._build_embedding_model", side_effect=RuntimeError("embedding config missing")):
|
||||
result = await adapter.vimax_novel_planning({"novel_text": "Hero opens a door."})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertIn("embedding", result.content.lower())
|
||||
self.assertEqual(index.active()["stage"], "error")
|
||||
|
||||
async def test_success_writes_novel2video_artifacts_and_marks_novel_planned(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
progress_events = []
|
||||
runtime = ToolRuntimeContext("vimax_novel_planning", "vimax_novel_planning", turn_id="turn-test", progress_callback=progress_events.append)
|
||||
with patch("agent_runtime.vimax_adapters._build_novel_pipeline", side_effect=lambda working_dir: FakeNovelPipeline(Path(working_dir))):
|
||||
result = await adapter.vimax_novel_planning({"novel_text": "Hero opens a door.", "style": "noir"}, runtime)
|
||||
self.assertTrue(result.ok)
|
||||
payload = json.loads(result.content)
|
||||
root = Path(tmp) / payload["working_dir"]
|
||||
self.assertTrue((root / "novel2video" / "novel" / "novel_compressed.txt").exists())
|
||||
self.assertTrue((root / "novel2video" / "events" / "event_0.json").exists())
|
||||
self.assertIn("novel2video/scenes/event_*/scene_*.json", payload["generated"])
|
||||
self.assertFalse(payload["ready_for_scene_render"])
|
||||
self.assertEqual(index.active()["stage"], "novel_planned")
|
||||
stages = [event["progress"]["stage"] for event in progress_events if event.get("type") == "tool_progress"]
|
||||
self.assertIn("novel_plan_text_artifacts", stages)
|
||||
self.assertIn("merge_characters", stages)
|
||||
|
||||
async def test_render_video_routes_novel2video_with_mock_pipeline(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="novel", user_requirement="scene render", style="noir")
|
||||
root = Path(tmp) / record["working_dir"]
|
||||
write_minimal_novel_artifacts(root)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
progress_events = []
|
||||
runtime = ToolRuntimeContext("vimax_render_video", "vimax_render_video", turn_id="turn-test", progress_callback=progress_events.append)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters._build_image_generator", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters._build_video_generator", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters._build_novel_render_pipeline", side_effect=lambda working_dir, chat_model, image_generator, video_generator: FakeNovelRenderPipeline(Path(working_dir))):
|
||||
result = await adapter.vimax_render_video({}, runtime)
|
||||
self.assertTrue(result.ok)
|
||||
payload = json.loads(result.content)
|
||||
self.assertEqual(payload["render_mode"], "novel2video")
|
||||
self.assertTrue(payload["scene_render_completed"])
|
||||
self.assertIsNone(payload["final_video_path"])
|
||||
self.assertEqual(payload["scene_count"], 1)
|
||||
self.assertEqual(index.get(record["session_id"])["stage"], "novel_scene_rendered")
|
||||
self.assertTrue((root / "novel2video" / "videos" / "event_0" / "scene_0" / "final_video.mp4").exists())
|
||||
stages = [event["progress"]["stage"] for event in progress_events if event.get("type") == "tool_progress"]
|
||||
self.assertIn("novel_scene_render_start", stages)
|
||||
self.assertIn("novel_render_completed", stages)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,136 @@
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from tools.video_generator_omni_yunwu_api import (
|
||||
VideoGeneratorOmniYunwuAPI,
|
||||
VideoGeneratorOminiYunwuAPI,
|
||||
)
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload):
|
||||
self.payload = payload
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def json(self):
|
||||
return self.payload
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, payload):
|
||||
self.payload = payload
|
||||
self.calls = []
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def get(self, url, **kwargs):
|
||||
self.calls.append(("get", url, kwargs))
|
||||
return _FakeResponse(self.payload)
|
||||
|
||||
|
||||
class TestVideoGeneratorOmniYunwuAPI(unittest.IsolatedAsyncioTestCase):
|
||||
def test_text_to_video_payload(self):
|
||||
generator = VideoGeneratorOmniYunwuAPI(api_key="test-key")
|
||||
|
||||
payload = generator._build_payload(
|
||||
prompt="hello world",
|
||||
reference_image_paths=[],
|
||||
aspect_ratio="16:9",
|
||||
seconds=8,
|
||||
size=None,
|
||||
enable_upsample=False,
|
||||
enable_sample=None,
|
||||
)
|
||||
|
||||
self.assertEqual(payload["model"], "omni-flash")
|
||||
self.assertEqual(payload["type"], 1)
|
||||
self.assertEqual(payload["seconds"], "8")
|
||||
self.assertEqual(payload["aspect_ratio"], "16:9")
|
||||
self.assertFalse(payload["enable_upsample"])
|
||||
self.assertNotIn("images", payload)
|
||||
|
||||
def test_first_last_frame_payload(self):
|
||||
generator = VideoGeneratorOmniYunwuAPI(api_key="test-key")
|
||||
|
||||
payload = generator._build_payload(
|
||||
prompt="transition",
|
||||
reference_image_paths=["https://example.com/first.png", "https://example.com/last.png"],
|
||||
aspect_ratio="9:16",
|
||||
seconds=6,
|
||||
size=None,
|
||||
enable_upsample=None,
|
||||
enable_sample=True,
|
||||
)
|
||||
|
||||
self.assertEqual(payload["type"], 2)
|
||||
self.assertEqual(payload["images"], ["https://example.com/first.png", "https://example.com/last.png"])
|
||||
self.assertEqual(payload["seconds"], "6")
|
||||
self.assertTrue(payload["enable_sample"])
|
||||
|
||||
def test_three_reference_images_use_reference_mode(self):
|
||||
generator = VideoGeneratorOmniYunwuAPI(api_key="test-key")
|
||||
|
||||
payload = generator._build_payload(
|
||||
prompt="blend references",
|
||||
reference_image_paths=[
|
||||
"https://example.com/1.png",
|
||||
"https://example.com/2.png",
|
||||
"https://example.com/3.png",
|
||||
],
|
||||
aspect_ratio="16:9",
|
||||
seconds=None,
|
||||
size=None,
|
||||
enable_upsample=None,
|
||||
enable_sample=None,
|
||||
)
|
||||
|
||||
self.assertEqual(payload["type"], 3)
|
||||
self.assertEqual(len(payload["images"]), 3)
|
||||
|
||||
def test_too_many_reference_images_raise(self):
|
||||
generator = VideoGeneratorOmniYunwuAPI(api_key="test-key")
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
generator._build_payload(
|
||||
prompt="too many",
|
||||
reference_image_paths=["1", "2", "3", "4"],
|
||||
aspect_ratio="16:9",
|
||||
seconds=None,
|
||||
size=None,
|
||||
enable_upsample=None,
|
||||
enable_sample=None,
|
||||
)
|
||||
|
||||
async def test_query_completed_uses_top_level_video_url(self):
|
||||
generator = VideoGeneratorOmniYunwuAPI(api_key="test-key", poll_interval=0, max_poll_attempts=1)
|
||||
session = _FakeSession({"status": "completed", "video_url": "https://example.com/out.mp4"})
|
||||
|
||||
with patch("tools.video_generator_omni_yunwu_api.aiohttp.ClientSession", return_value=session):
|
||||
video_url = await generator.query_video_generation_task("task-1", "omni-flash")
|
||||
|
||||
self.assertEqual(video_url, "https://example.com/out.mp4")
|
||||
self.assertEqual(session.calls[0][2]["params"], {"id": "task-1", "model": "omni-flash"})
|
||||
|
||||
async def test_query_failed_raises(self):
|
||||
generator = VideoGeneratorOmniYunwuAPI(api_key="test-key", poll_interval=0, max_poll_attempts=1)
|
||||
session = _FakeSession({"status": "failed", "error": "视频生成失败"})
|
||||
|
||||
with patch("tools.video_generator_omni_yunwu_api.aiohttp.ClientSession", return_value=session):
|
||||
with self.assertRaises(RuntimeError):
|
||||
await generator.query_video_generation_task("task-1", "omni-flash")
|
||||
|
||||
def test_omini_alias(self):
|
||||
self.assertTrue(issubclass(VideoGeneratorOminiYunwuAPI, VideoGeneratorOmniYunwuAPI))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,40 @@
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from tools.video_generator_openrouter_api import VideoGeneratorOpenRouterAPI
|
||||
from interfaces.video_output import VideoOutput
|
||||
|
||||
|
||||
class OpenRouterVideoGeneratorTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_default_duration_is_eight_seconds(self):
|
||||
captured = {}
|
||||
|
||||
async def fake_post_json(url, *, headers, payload, timeout, hard_timeout_seconds):
|
||||
captured["payload"] = payload
|
||||
return 200, {"id": "job-1", "polling_url": "/videos/job-1", "status": "queued"}
|
||||
|
||||
async def fake_get_json(url, *, headers, timeout, hard_timeout_seconds):
|
||||
return 200, {"status": "completed", "unsigned_urls": ["https://cdn.example/out.mp4"]}
|
||||
|
||||
async def fake_get_bytes(url, *, headers, timeout, hard_timeout_seconds):
|
||||
return 200, b"video"
|
||||
|
||||
async def fake_sleep(seconds):
|
||||
return None
|
||||
|
||||
generator = VideoGeneratorOpenRouterAPI(api_key="test-key", model="google/veo-3.1-lite")
|
||||
with patch.dict(os.environ, {}, clear=True), \
|
||||
patch("tools.video_generator_openrouter_api._post_json", fake_post_json), \
|
||||
patch("tools.video_generator_openrouter_api._get_json", fake_get_json), \
|
||||
patch("tools.video_generator_openrouter_api._get_bytes", fake_get_bytes), \
|
||||
patch("tools.video_generator_openrouter_api.asyncio.sleep", fake_sleep):
|
||||
output = await generator.generate_single_video(prompt="hello")
|
||||
|
||||
self.assertIsInstance(output, VideoOutput)
|
||||
self.assertEqual(captured["payload"]["duration"], 8)
|
||||
self.assertEqual(captured["payload"]["model"], "google/veo-3.1-lite")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,175 @@
|
||||
"""Unit tests for utils.provider_presets."""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from utils.provider_presets import (
|
||||
PROVIDER_PRESETS,
|
||||
resolve_chat_model_config,
|
||||
detect_provider_from_env,
|
||||
)
|
||||
|
||||
|
||||
class TestProviderPresets(unittest.TestCase):
|
||||
"""Tests for the PROVIDER_PRESETS registry."""
|
||||
|
||||
def test_minimax_preset_exists(self):
|
||||
self.assertIn("minimax", PROVIDER_PRESETS)
|
||||
|
||||
def test_minimax_preset_base_url(self):
|
||||
self.assertEqual(
|
||||
PROVIDER_PRESETS["minimax"]["base_url"],
|
||||
"https://api.minimax.io/v1",
|
||||
)
|
||||
|
||||
def test_minimax_preset_env_key(self):
|
||||
self.assertEqual(PROVIDER_PRESETS["minimax"]["env_key"], "MINIMAX_API_KEY")
|
||||
|
||||
def test_minimax_preset_default_model(self):
|
||||
self.assertEqual(PROVIDER_PRESETS["minimax"]["default_model"], "MiniMax-M3")
|
||||
|
||||
def test_minimax_preset_has_models_list(self):
|
||||
models = PROVIDER_PRESETS["minimax"]["models"]
|
||||
self.assertIn("MiniMax-M3", models)
|
||||
self.assertIn("MiniMax-M2.7", models)
|
||||
self.assertIn("MiniMax-M2.7-highspeed", models)
|
||||
|
||||
def test_minimax_preset_temperature_range(self):
|
||||
lo, hi = PROVIDER_PRESETS["minimax"]["temperature_range"]
|
||||
self.assertEqual(lo, 0.0)
|
||||
self.assertEqual(hi, 1.0)
|
||||
|
||||
|
||||
class TestResolveChatModelConfig(unittest.TestCase):
|
||||
"""Tests for resolve_chat_model_config()."""
|
||||
|
||||
def test_unknown_provider_passes_through(self):
|
||||
args = {"model_provider": "openai", "model": "gpt-4", "base_url": "https://example.com"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["model_provider"], "openai")
|
||||
self.assertEqual(result["model"], "gpt-4")
|
||||
self.assertEqual(result["base_url"], "https://example.com")
|
||||
|
||||
def test_no_model_provider_passes_through(self):
|
||||
args = {"model": "gpt-4"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["model"], "gpt-4")
|
||||
|
||||
def test_minimax_rewrites_provider_to_openai(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk-test"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["model_provider"], "openai")
|
||||
|
||||
def test_minimax_sets_base_url(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk-test"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["base_url"], "https://api.minimax.io/v1")
|
||||
|
||||
def test_minimax_preserves_custom_base_url(self):
|
||||
args = {
|
||||
"model_provider": "minimax",
|
||||
"model": "MiniMax-M3",
|
||||
"api_key": "sk-test",
|
||||
"base_url": "https://custom-proxy.example.com/v1",
|
||||
}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["base_url"], "https://custom-proxy.example.com/v1")
|
||||
|
||||
def test_minimax_defaults_model(self):
|
||||
args = {"model_provider": "minimax", "api_key": "sk-test"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["model"], "MiniMax-M3")
|
||||
|
||||
def test_minimax_preserves_explicit_model(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M2.7-highspeed", "api_key": "sk-test"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["model"], "MiniMax-M2.7-highspeed")
|
||||
|
||||
@patch.dict(os.environ, {"MINIMAX_API_KEY": "env-key-123"})
|
||||
def test_minimax_reads_api_key_from_env(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["api_key"], "env-key-123")
|
||||
|
||||
def test_minimax_prefers_explicit_api_key_over_env(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "explicit-key"}
|
||||
with patch.dict(os.environ, {"MINIMAX_API_KEY": "env-key"}):
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["api_key"], "explicit-key")
|
||||
|
||||
def test_minimax_clamps_temperature_above_max(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk", "temperature": 1.5}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["temperature"], 1.0)
|
||||
|
||||
def test_minimax_clamps_temperature_below_min(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk", "temperature": -0.5}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["temperature"], 0.0)
|
||||
|
||||
def test_minimax_passes_valid_temperature(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk", "temperature": 0.7}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["temperature"], 0.7)
|
||||
|
||||
def test_minimax_temperature_zero_allowed(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk", "temperature": 0.0}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["temperature"], 0.0)
|
||||
|
||||
def test_minimax_no_temperature_key(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertNotIn("temperature", result)
|
||||
|
||||
def test_minimax_temperature_none_ignored(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk", "temperature": None}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertIsNone(result["temperature"])
|
||||
|
||||
def test_original_dict_not_mutated(self):
|
||||
args = {"model_provider": "minimax", "model": "MiniMax-M3", "api_key": "sk"}
|
||||
resolve_chat_model_config(args)
|
||||
self.assertEqual(args["model_provider"], "minimax")
|
||||
|
||||
def test_empty_model_string_gets_default(self):
|
||||
args = {"model_provider": "minimax", "model": "", "api_key": "sk"}
|
||||
result = resolve_chat_model_config(args)
|
||||
self.assertEqual(result["model"], "MiniMax-M3")
|
||||
|
||||
|
||||
class TestDetectProviderFromEnv(unittest.TestCase):
|
||||
"""Tests for detect_provider_from_env()."""
|
||||
|
||||
@patch.dict(os.environ, {"MINIMAX_API_KEY": "test-key"}, clear=False)
|
||||
def test_detects_minimax(self):
|
||||
self.assertEqual(detect_provider_from_env(), "minimax")
|
||||
|
||||
@patch.dict(os.environ, {}, clear=True)
|
||||
def test_returns_none_when_no_keys(self):
|
||||
self.assertIsNone(detect_provider_from_env())
|
||||
|
||||
|
||||
class TestConfigYAMLLoading(unittest.TestCase):
|
||||
"""Test that MiniMax example config files are valid YAML."""
|
||||
|
||||
def test_idea2video_minimax_yaml(self):
|
||||
import yaml
|
||||
path = os.path.join(os.path.dirname(__file__), "..", "configs", "idea2video_minimax.yaml")
|
||||
with open(path) as f:
|
||||
config = yaml.safe_load(f)
|
||||
self.assertEqual(config["chat_model"]["init_args"]["model_provider"], "minimax")
|
||||
self.assertEqual(config["chat_model"]["init_args"]["model"], "MiniMax-M3")
|
||||
|
||||
def test_script2video_minimax_yaml(self):
|
||||
import yaml
|
||||
path = os.path.join(os.path.dirname(__file__), "..", "configs", "script2video_minimax.yaml")
|
||||
with open(path) as f:
|
||||
config = yaml.safe_load(f)
|
||||
self.assertEqual(config["chat_model"]["init_args"]["model_provider"], "minimax")
|
||||
self.assertEqual(config["chat_model"]["init_args"]["model"], "MiniMax-M3")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Regression tests for error-boundary and durability fixes.
|
||||
|
||||
Covers: LLM client retry/empty-choices handling, agent-loop turn error
|
||||
boundary, session index corruption/atomicity/concurrency, and bounded
|
||||
retry policies with backoff across agents and API clients.
|
||||
"""
|
||||
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from tenacity.stop import stop_never
|
||||
from tenacity.wait import wait_none
|
||||
|
||||
from agent_runtime.llm import OpenAICompatibleLLM
|
||||
from agent_runtime.loop import AgentLoop
|
||||
from agent_runtime.prompts import PromptBuilder
|
||||
from agent_runtime.session_index import SessionIndex
|
||||
from agent_runtime.tool_executor import ToolExecutor
|
||||
from agent_runtime.tools import ToolRegistry
|
||||
from agents.screenwriter import Screenwriter
|
||||
from agents.script_planner import ScriptPlanner
|
||||
from tools.image_generator_doubao_seedream_yunwu_api import ImageGeneratorDoubaoSeedreamYunwuAPI
|
||||
from tools.image_generator_nanobanana_google_api import ImageGeneratorNanobananaGoogleAPI
|
||||
from tools.image_generator_nanobanana_yunwu_api import ImageGeneratorNanobananaYunwuAPI
|
||||
from tools.reranker_bge_silicon_api import RerankerBgeSiliconapi
|
||||
|
||||
|
||||
class FakeStatusError(Exception):
|
||||
def __init__(self, status_code):
|
||||
self.status_code = status_code
|
||||
super().__init__(f"http status {status_code}")
|
||||
|
||||
|
||||
def _fake_completion(text="ok"):
|
||||
message = MagicMock()
|
||||
message.content = text
|
||||
message.tool_calls = None
|
||||
message.model_dump.return_value = {}
|
||||
return MagicMock(choices=[MagicMock(message=message)])
|
||||
|
||||
|
||||
class TestLLMClient(unittest.IsolatedAsyncioTestCase):
|
||||
def _llm(self, create):
|
||||
llm = OpenAICompatibleLLM(model="m", base_url="http://localhost:1", api_key="k")
|
||||
llm.client = MagicMock(chat=MagicMock(completions=MagicMock(create=create)))
|
||||
return llm
|
||||
|
||||
async def test_retries_rate_limit_then_succeeds(self):
|
||||
create = AsyncMock(side_effect=[FakeStatusError(429), _fake_completion("recovered")])
|
||||
llm = self._llm(create)
|
||||
result = await llm.complete([{"role": "user", "content": "x"}], tools=[])
|
||||
self.assertEqual(result.text, "recovered")
|
||||
self.assertEqual(create.await_count, 2)
|
||||
|
||||
async def test_does_not_retry_auth_errors(self):
|
||||
create = AsyncMock(side_effect=FakeStatusError(401))
|
||||
llm = self._llm(create)
|
||||
with self.assertRaises(FakeStatusError):
|
||||
await llm.complete([{"role": "user", "content": "x"}], tools=[])
|
||||
self.assertEqual(create.await_count, 1)
|
||||
|
||||
async def test_gives_up_after_bounded_attempts(self):
|
||||
create = AsyncMock(side_effect=FakeStatusError(500))
|
||||
llm = self._llm(create)
|
||||
with self.assertRaises(FakeStatusError):
|
||||
await llm.complete([{"role": "user", "content": "x"}], tools=[])
|
||||
self.assertLessEqual(create.await_count, 4)
|
||||
self.assertGreater(create.await_count, 1)
|
||||
|
||||
async def test_empty_choices_raises_clear_error(self):
|
||||
create = AsyncMock(return_value=MagicMock(choices=[]))
|
||||
llm = self._llm(create)
|
||||
with self.assertRaisesRegex(RuntimeError, "choice"):
|
||||
await llm.complete([{"role": "user", "content": "x"}], tools=[])
|
||||
|
||||
|
||||
class BoomLLM:
|
||||
async def complete(self, messages, tools):
|
||||
raise RuntimeError("boom-llm")
|
||||
|
||||
|
||||
class TestLoopErrorBoundary(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_llm_failure_emits_error_and_persists_failed_turn(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
registry = ToolRegistry([])
|
||||
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), BoomLLM())
|
||||
events = [event async for event in loop.stream_events("hi")]
|
||||
kinds = [event["type"] for event in events]
|
||||
self.assertIn("error", kinds)
|
||||
error_event = next(event for event in events if event["type"] == "error")
|
||||
self.assertIn("boom-llm", error_event["message"])
|
||||
self.assertEqual(events[-2]["type"], "done")
|
||||
self.assertEqual(events[-1]["type"], "session")
|
||||
active = index.active()
|
||||
records = index.get(active["session_id"])["recent_turn_records"]
|
||||
self.assertEqual(records[-1]["status"], "failed")
|
||||
|
||||
|
||||
class TestSessionIndexDurability(unittest.TestCase):
|
||||
def test_corrupt_sessions_file_is_backed_up_not_silently_replaced(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
index.create(idea="precious work", session_id="keep-me")
|
||||
index.sessions_path.write_text("{ definitely not json", encoding="utf-8")
|
||||
data = index.load()
|
||||
self.assertEqual(data["sessions"], {})
|
||||
backups = list(index.vimax_dir.glob("sessions.json.corrupt-*"))
|
||||
self.assertEqual(len(backups), 1, "corrupt state must be preserved for recovery, not discarded")
|
||||
self.assertIn("definitely not json", backups[0].read_text(encoding="utf-8"))
|
||||
|
||||
def test_save_is_atomic_and_leaves_no_temp_files(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
index.create(session_id="roundtrip")
|
||||
self.assertEqual(list(index.vimax_dir.glob("*.tmp")), [])
|
||||
self.assertIn("roundtrip", index.load()["sessions"])
|
||||
|
||||
def test_concurrent_creates_do_not_lose_sessions(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index_a = SessionIndex(tmp)
|
||||
index_b = SessionIndex(tmp)
|
||||
|
||||
def worker(index, tag):
|
||||
for i in range(40):
|
||||
index.create(session_id=f"s-{tag}-{i}")
|
||||
|
||||
threads = [
|
||||
threading.Thread(target=worker, args=(index_a, "a")),
|
||||
threading.Thread(target=worker, args=(index_b, "b")),
|
||||
]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
sessions = index_a.load()["sessions"]
|
||||
self.assertEqual(len(sessions), 80, "concurrent read-modify-write must not lose sessions")
|
||||
|
||||
|
||||
class TestBoundedRetryPolicies(unittest.TestCase):
|
||||
CASES = [
|
||||
("Screenwriter.write_script_based_on_story", Screenwriter.write_script_based_on_story),
|
||||
("ScriptPlanner.plan_script", ScriptPlanner.plan_script),
|
||||
("RerankerBgeSiliconapi.__call__", RerankerBgeSiliconapi.__call__),
|
||||
("ImageGeneratorDoubaoSeedreamYunwuAPI.generate_single_image", ImageGeneratorDoubaoSeedreamYunwuAPI.generate_single_image),
|
||||
("ImageGeneratorNanobananaGoogleAPI.generate_single_image", ImageGeneratorNanobananaGoogleAPI.generate_single_image),
|
||||
("ImageGeneratorNanobananaYunwuAPI.generate_single_image", ImageGeneratorNanobananaYunwuAPI.generate_single_image),
|
||||
]
|
||||
|
||||
def test_every_retry_is_bounded_with_backoff(self):
|
||||
for name, fn in self.CASES:
|
||||
with self.subTest(name=name):
|
||||
retrying = getattr(fn, "retry", None)
|
||||
self.assertIsNotNone(retrying, f"{name} must have a retry policy")
|
||||
self.assertIsNot(retrying.stop, stop_never, f"{name} must not retry forever")
|
||||
self.assertNotIsInstance(retrying.wait, wait_none, f"{name} must back off between attempts")
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def json(self):
|
||||
return self.payload
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, scripted):
|
||||
self.scripted = list(scripted)
|
||||
self.calls = 0
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def _next(self):
|
||||
response = self.scripted[min(self.calls, len(self.scripted) - 1)]
|
||||
self.calls += 1
|
||||
return _FakeResponse(*response)
|
||||
|
||||
def post(self, url, **kwargs):
|
||||
return self._next()
|
||||
|
||||
def get(self, url, **kwargs):
|
||||
return self._next()
|
||||
|
||||
|
||||
class TestClientHttpErrors(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_reranker_surfaces_http_error_without_retry(self):
|
||||
session = _FakeSession([
|
||||
({"message": "invalid api key"}, 401),
|
||||
({"results": []}, 200),
|
||||
])
|
||||
reranker = RerankerBgeSiliconapi(api_key="bad", base_url="http://x")
|
||||
with patch("tools.reranker_bge_silicon_api.aiohttp.ClientSession", return_value=session):
|
||||
with self.assertRaisesRegex(RuntimeError, "401"):
|
||||
await reranker(documents=["doc"], query="q", top_n=1)
|
||||
self.assertEqual(session.calls, 1, "4xx must fail fast with the real error, not retry into KeyError")
|
||||
|
||||
async def test_seedream_surfaces_http_error_without_retry(self):
|
||||
session = _FakeSession([
|
||||
({"error": {"message": "invalid api key"}}, 401),
|
||||
({"data": [{"url": "http://img"}]}, 200),
|
||||
])
|
||||
generator = ImageGeneratorDoubaoSeedreamYunwuAPI(api_key="bad")
|
||||
with patch("tools.image_generator_doubao_seedream_yunwu_api.aiohttp.ClientSession", return_value=session):
|
||||
with self.assertRaisesRegex(RuntimeError, "401"):
|
||||
await generator.generate_single_image(prompt="p")
|
||||
self.assertEqual(session.calls, 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,62 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from interfaces import Camera, ShotBriefDescription, ShotDescription
|
||||
from pipelines.script2video_pipeline import Script2VideoPipeline, _group_shots_into_cameras
|
||||
|
||||
|
||||
class FlakyCameraImageGenerator:
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
|
||||
async def construct_camera_tree(self, cameras, shot_descs):
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
return ["not-a-camera"]
|
||||
return cameras
|
||||
|
||||
|
||||
class Script2VideoPipelineGuardTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_group_shots_into_cameras_does_not_use_camera_idx_as_list_index(self):
|
||||
shots = [
|
||||
ShotDescription(idx=0, is_last=False, cam_idx=2, visual_desc="a", variation_type="small", variation_reason="same", ff_desc="a", ff_vis_char_idxs=[], lf_desc="a", lf_vis_char_idxs=[], motion_desc="a", audio_desc="none"),
|
||||
ShotDescription(idx=1, is_last=True, cam_idx=5, visual_desc="b", variation_type="small", variation_reason="same", ff_desc="b", ff_vis_char_idxs=[], lf_desc="b", lf_vis_char_idxs=[], motion_desc="b", audio_desc="none"),
|
||||
ShotDescription(idx=2, is_last=True, cam_idx=2, visual_desc="c", variation_type="small", variation_reason="same", ff_desc="c", ff_vis_char_idxs=[], lf_desc="c", lf_vis_char_idxs=[], motion_desc="c", audio_desc="none"),
|
||||
]
|
||||
cameras = _group_shots_into_cameras(shots)
|
||||
self.assertEqual([camera.idx for camera in cameras], [2, 5])
|
||||
self.assertEqual(cameras[0].active_shot_idxs, [0, 2])
|
||||
self.assertEqual(cameras[1].active_shot_idxs, [1])
|
||||
|
||||
async def test_plan_text_artifacts_retries_bad_camera_tree_schema(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
pipeline = Script2VideoPipeline(chat_model=object(), image_generator=object(), video_generator=object(), working_dir=tmp)
|
||||
pipeline.camera_image_generator = FlakyCameraImageGenerator()
|
||||
|
||||
async def design_storyboard(script, characters, user_requirement, quiet=False):
|
||||
return [{"idx": 0, "is_last": True, "cam_idx": 3, "visual_desc": "wide shot", "audio_desc": "waves"}]
|
||||
|
||||
async def decompose_visual_descriptions(shot_brief_descriptions, characters, quiet=False):
|
||||
return [{"idx": 0, "is_last": True, "cam_idx": 3, "visual_desc": "wide shot", "variation_type": "small", "variation_reason": "simple", "ff_desc": "start", "ff_vis_char_idxs": [], "lf_desc": "end", "lf_vis_char_idxs": [], "motion_desc": "walk", "audio_desc": "waves"}]
|
||||
|
||||
pipeline.design_storyboard = design_storyboard
|
||||
pipeline.decompose_visual_descriptions = decompose_visual_descriptions
|
||||
events = []
|
||||
result = await pipeline.plan_text_artifacts(
|
||||
"script",
|
||||
"req",
|
||||
"style",
|
||||
characters=[{"idx": 0, "identifier_in_scene": "Man", "is_visible": True, "static_features": "adult", "dynamic_features": "coat"}],
|
||||
progress=lambda stage, message, metadata=None: events.append(stage),
|
||||
quiet=True,
|
||||
)
|
||||
|
||||
self.assertEqual(pipeline.camera_image_generator.calls, 2)
|
||||
self.assertIn("construct_camera_tree_retry", events)
|
||||
self.assertEqual(result["camera_tree"][0].idx, 3)
|
||||
self.assertTrue((Path(tmp) / "camera_tree.json").exists())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,452 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from interfaces import Camera, CharacterInScene, ShotBriefDescription, ShotDescription
|
||||
from agent_runtime.session_index import SessionIndex
|
||||
from agent_runtime.vimax_adapters import ViMaxAdapters
|
||||
from agent_runtime.tools import ToolRuntimeContext
|
||||
from pipelines.idea2video_pipeline import Idea2VideoPipeline
|
||||
from pipelines.script2video_pipeline import Script2VideoPipeline
|
||||
|
||||
|
||||
class FakeIdeaPipeline:
|
||||
def __init__(self, chat_model, image_generator, video_generator, working_dir):
|
||||
self.working_dir = Path(working_dir)
|
||||
self.working_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
async def develop_story(self, idea, user_requirement, quiet=False):
|
||||
path = self.working_dir / "story.txt"
|
||||
path.write_text("story", encoding="utf-8")
|
||||
return "story"
|
||||
|
||||
async def extract_characters(self, story, quiet=False):
|
||||
chars = [CharacterInScene(idx=0, identifier_in_scene="Cat", is_visible=True, static_features="black cat", dynamic_features="helmet")]
|
||||
(self.working_dir / "characters.json").write_text(json.dumps([c.model_dump() for c in chars]), encoding="utf-8")
|
||||
return chars
|
||||
|
||||
async def write_script_based_on_story(self, story, user_requirement, quiet=False):
|
||||
script = [{"scene": "cat jumps"}]
|
||||
(self.working_dir / "script.json").write_text(json.dumps(script), encoding="utf-8")
|
||||
return script
|
||||
|
||||
|
||||
|
||||
|
||||
class HangingIdeaPipeline(FakeIdeaPipeline):
|
||||
async def develop_story(self, idea, user_requirement, quiet=False):
|
||||
await asyncio.sleep(10)
|
||||
return "story"
|
||||
|
||||
|
||||
|
||||
class FakeRevisionModel:
|
||||
async def ainvoke(self, prompt):
|
||||
return SimpleNamespace(content='[{"idx": 0, "description": "more oppressive"}]')
|
||||
|
||||
|
||||
class FailRenderIdeaPipeline(FakeIdeaPipeline):
|
||||
async def __call__(self, idea, user_requirement, style, quiet=False):
|
||||
raise RuntimeError("render failed")
|
||||
|
||||
|
||||
class FailRender403IdeaPipeline(FakeIdeaPipeline):
|
||||
async def __call__(self, idea, user_requirement, style, quiet=False):
|
||||
raise RuntimeError("OpenRouter video create failed with HTTP 403: {'error': {'message': 'Key limit exceeded (total limit). Manage it using token sk-short', 'code': 403}}")
|
||||
|
||||
|
||||
class NoisyRenderIdeaPipeline(FakeIdeaPipeline):
|
||||
async def __call__(self, idea, user_requirement, style, quiet=False):
|
||||
print("NOISE_FROM_RENDER_PIPELINE")
|
||||
final = self.working_dir / "final_video.mp4"
|
||||
final.write_text("video", encoding="utf-8")
|
||||
return str(final)
|
||||
|
||||
|
||||
class FakeScriptPipeline:
|
||||
def __init__(self, chat_model, image_generator, video_generator, working_dir):
|
||||
self.working_dir = Path(working_dir)
|
||||
self.working_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
async def plan_text_artifacts(self, script, user_requirement, style, characters=None, progress=None, quiet=False):
|
||||
if progress:
|
||||
progress("design_storyboard", "Designing storyboard", {})
|
||||
progress("decompose_shots", "Decomposing shot visual descriptions", {"shot_count": 1})
|
||||
progress("construct_camera_tree", "Constructing camera tree", {"shot_count": 1})
|
||||
(self.working_dir / "storyboard.json").write_text("[]", encoding="utf-8")
|
||||
(self.working_dir / "camera_tree.json").write_text("[]", encoding="utf-8")
|
||||
shot_dir = self.working_dir / "shots" / "0"
|
||||
shot_dir.mkdir(parents=True, exist_ok=True)
|
||||
(shot_dir / "shot_description.json").write_text("{}", encoding="utf-8")
|
||||
if characters:
|
||||
(self.working_dir / "characters.json").write_text(json.dumps([c.model_dump() for c in characters]), encoding="utf-8")
|
||||
return {}
|
||||
|
||||
|
||||
|
||||
|
||||
class FailingScriptPipeline(FakeScriptPipeline):
|
||||
async def plan_text_artifacts(self, script, user_requirement, style, characters=None, progress=None, quiet=False):
|
||||
if progress:
|
||||
progress("design_storyboard", "Designing storyboard", {})
|
||||
raise RuntimeError("storyboard failed")
|
||||
|
||||
|
||||
class FakeInitChatModel:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
self.calls.append(kwargs)
|
||||
return object()
|
||||
|
||||
|
||||
class Script2VideoPlanningProgressTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_plan_text_artifacts_emits_progress_in_order(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
pipeline = Script2VideoPipeline(chat_model=object(), image_generator=object(), video_generator=object(), working_dir=tmp)
|
||||
chars = [CharacterInScene(idx=0, identifier_in_scene="Cat", is_visible=True, static_features="black cat", dynamic_features="helmet")]
|
||||
storyboard = [ShotBriefDescription(idx=0, is_last=True, cam_idx=0, visual_desc="cat jumps", audio_desc="wind")]
|
||||
shot = ShotDescription(idx=0, is_last=True, cam_idx=0, visual_desc="cat jumps", variation_type="small", variation_reason="simple motion", ff_desc="cat starts", ff_vis_char_idxs=[0], lf_desc="cat lands", lf_vis_char_idxs=[0], motion_desc="cat jumps", audio_desc="wind")
|
||||
camera = [Camera(idx=0, active_shot_idxs=[0])]
|
||||
|
||||
async def design_storyboard(script, characters, user_requirement, quiet=False):
|
||||
return storyboard
|
||||
|
||||
async def decompose_visual_descriptions(shot_brief_descriptions, characters, quiet=False):
|
||||
return [shot]
|
||||
|
||||
async def construct_camera_tree(shot_descriptions, quiet=False):
|
||||
return camera
|
||||
|
||||
pipeline.design_storyboard = design_storyboard
|
||||
pipeline.decompose_visual_descriptions = decompose_visual_descriptions
|
||||
pipeline.construct_camera_tree = construct_camera_tree
|
||||
events = []
|
||||
await pipeline.plan_text_artifacts("script", "req", "style", characters=chars, progress=lambda stage, message, metadata=None: events.append(stage))
|
||||
self.assertEqual(events, ["extract_characters", "design_storyboard", "decompose_shots", "construct_camera_tree"])
|
||||
|
||||
|
||||
async def test_idea_pipeline_quiet_suppresses_text_planning_prints(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
pipeline = Idea2VideoPipeline(chat_model=object(), image_generator=object(), video_generator=object(), working_dir=tmp)
|
||||
|
||||
async def develop_story(idea, user_requirement):
|
||||
return "story"
|
||||
|
||||
pipeline.screenwriter = SimpleNamespace(develop_story=develop_story)
|
||||
stdout = io.StringIO()
|
||||
with contextlib.redirect_stdout(stdout):
|
||||
result = await pipeline.develop_story("idea", "req", quiet=True)
|
||||
self.assertEqual(result, "story")
|
||||
self.assertEqual(stdout.getvalue(), "")
|
||||
|
||||
|
||||
class ViMaxAdapterTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_build_chat_model_uses_bounded_init_chat_model_kwargs(self):
|
||||
fake = FakeInitChatModel()
|
||||
with patch.dict("os.environ", {
|
||||
"VIMAX_LLM_API_KEY": "test-key",
|
||||
"VIMAX_LLM_MODEL": "test-model",
|
||||
"VIMAX_LLM_BASE_URL": "https://example.invalid/v1",
|
||||
"VIMAX_LLM_REQUEST_TIMEOUT_SECONDS": "12",
|
||||
"VIMAX_NARRATIVE_MAX_TOKENS": "1234",
|
||||
}), patch("agent_runtime.vimax_adapters.init_chat_model", fake):
|
||||
from agent_runtime.vimax_adapters import _build_chat_model
|
||||
|
||||
_build_chat_model()
|
||||
|
||||
self.assertEqual(fake.calls[0]["model"], "test-model")
|
||||
self.assertEqual(fake.calls[0]["base_url"], "https://example.invalid/v1")
|
||||
self.assertEqual(fake.calls[0]["timeout"], 12.0)
|
||||
self.assertEqual(fake.calls[0]["max_retries"], 0)
|
||||
self.assertEqual(fake.calls[0]["max_completion_tokens"], 1234)
|
||||
|
||||
|
||||
async def test_narrative_planning_uses_text_only_pipeline(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FakeIdeaPipeline), \
|
||||
patch("agent_runtime.vimax_adapters.Script2VideoPipeline", FakeScriptPipeline):
|
||||
result = await adapter.vimax_narrative_planning({"idea": "moon cat", "user_requirement": "short", "style": "anime"})
|
||||
self.assertTrue(result.ok)
|
||||
payload = json.loads(result.content)
|
||||
self.assertTrue(payload["ready_for_render"])
|
||||
root = Path(tmp) / payload["working_dir"]
|
||||
self.assertTrue((root / "idea2video" / "scene_0" / "storyboard.json").exists())
|
||||
self.assertTrue((root / "idea2video" / "scene_0" / "camera_tree.json").exists())
|
||||
self.assertTrue((root / "idea2video" / "scene_0" / "shots" / "0" / "shot_description.json").exists())
|
||||
self.assertFalse((root / "script2video" / "storyboard.json").exists())
|
||||
self.assertFalse((root / "script2video" / "final_video.mp4").exists())
|
||||
|
||||
|
||||
async def test_script_mode_persists_source_script_for_render(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
script = "A red ball rolls across a white table."
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Script2VideoPipeline", FakeScriptPipeline):
|
||||
result = await adapter.vimax_narrative_planning({"script": script, "user_requirement": "one shot"})
|
||||
self.assertTrue(result.ok)
|
||||
payload = json.loads(result.content)
|
||||
root = Path(tmp) / payload["working_dir"]
|
||||
self.assertEqual((root / "script2video" / "script.txt").read_text(encoding="utf-8"), script)
|
||||
self.assertEqual(index.artifact_checklist(payload["session_id"])["script2video/script.txt"], True)
|
||||
from agent_runtime.vimax_adapters import _load_script_text
|
||||
self.assertEqual(_load_script_text(root), script)
|
||||
|
||||
|
||||
async def test_narrative_planning_forwards_pipeline_progress(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
events = []
|
||||
runtime = ToolRuntimeContext("vimax_narrative_planning", "vimax_narrative_planning", turn_id="turn-test", progress_callback=events.append)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FakeIdeaPipeline), \
|
||||
patch("agent_runtime.vimax_adapters.Script2VideoPipeline", FakeScriptPipeline):
|
||||
result = await adapter.vimax_narrative_planning({"idea": "moon cat"}, runtime)
|
||||
self.assertTrue(result.ok)
|
||||
stages = [event["progress"]["stage"] for event in events if event.get("type") == "tool_progress"]
|
||||
self.assertIn("initializing_llm", stages)
|
||||
self.assertIn("develop_story", stages)
|
||||
self.assertIn("design_storyboard", stages)
|
||||
self.assertIn("decompose_shots", stages)
|
||||
self.assertIn("construct_camera_tree", stages)
|
||||
|
||||
|
||||
async def test_plan_scene_failure_marks_session_error(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FakeIdeaPipeline), \
|
||||
patch("agent_runtime.vimax_adapters.Script2VideoPipeline", FailingScriptPipeline):
|
||||
result = await adapter.vimax_narrative_planning({"idea": "moon cat"})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.metadata["error_type"], "recoverable_planning_step_failed")
|
||||
self.assertTrue(result.metadata["retryable"])
|
||||
session = index.active()
|
||||
self.assertEqual(session["stage"], "error")
|
||||
self.assertIn("storyboard failed", session["summary"])
|
||||
|
||||
|
||||
async def test_narrative_planning_timeout_marks_session_error(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch.dict("os.environ", {"VIMAX_NARRATIVE_STEP_TIMEOUT_SECONDS": "0.01"}), \
|
||||
patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", HangingIdeaPipeline):
|
||||
result = await adapter.vimax_narrative_planning({"idea": "moon cat"})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.metadata["error_type"], "recoverable_planning_step_failed")
|
||||
session = index.active()
|
||||
self.assertIsNotNone(session)
|
||||
self.assertEqual(session["stage"], "error")
|
||||
self.assertIn("timed out", session["summary"])
|
||||
|
||||
|
||||
|
||||
async def test_active_session_without_new_input_continues_existing_idea(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="moon cat", user_requirement="short", style="anime")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FakeIdeaPipeline), patch("agent_runtime.vimax_adapters.Script2VideoPipeline", FakeScriptPipeline):
|
||||
result = await adapter.vimax_narrative_planning({})
|
||||
self.assertTrue(result.ok)
|
||||
payload = json.loads(result.content)
|
||||
self.assertEqual(payload["session_id"], record["session_id"])
|
||||
self.assertEqual(index.active()["session_id"], record["session_id"])
|
||||
|
||||
|
||||
async def test_active_session_continuation_preserves_existing_style(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="moon cat", user_requirement="short", style="anime")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FakeIdeaPipeline), patch("agent_runtime.vimax_adapters.Script2VideoPipeline", FakeScriptPipeline):
|
||||
result = await adapter.vimax_narrative_planning({"session_id": record["session_id"]})
|
||||
self.assertTrue(result.ok)
|
||||
self.assertEqual(index.get(record["session_id"])["style"], "anime")
|
||||
|
||||
async def test_new_idea_creates_new_session_instead_of_reusing_active(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FakeIdeaPipeline), \
|
||||
patch("agent_runtime.vimax_adapters.Script2VideoPipeline", FakeScriptPipeline):
|
||||
first = await adapter.vimax_narrative_planning({"idea": "moon cat"})
|
||||
second = await adapter.vimax_narrative_planning({"idea": "ocean robot"})
|
||||
self.assertNotEqual(json.loads(first.content)["session_id"], json.loads(second.content)["session_id"])
|
||||
|
||||
|
||||
async def test_explicit_session_with_different_idea_creates_new_session(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
old = index.create(idea="old cat")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FakeIdeaPipeline), \
|
||||
patch("agent_runtime.vimax_adapters.Script2VideoPipeline", FakeScriptPipeline):
|
||||
result = await adapter.vimax_narrative_planning({"session_id": old["session_id"], "idea": "new robot"})
|
||||
self.assertTrue(result.ok)
|
||||
payload = json.loads(result.content)
|
||||
self.assertNotEqual(payload["session_id"], old["session_id"])
|
||||
self.assertEqual(index.get(payload["session_id"])["idea"], "new robot")
|
||||
|
||||
async def test_revision_mode_rewrites_existing_artifact_and_logs(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="x")
|
||||
target = Path(tmp) / record["working_dir"] / "idea2video" / "scene_0" / "storyboard.json"
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text('[{"idx": 0, "description": "calm"}]', encoding="utf-8")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=FakeRevisionModel()):
|
||||
result = await adapter.vimax_narrative_planning({"revision_target": "idea2video/scene_0/storyboard.json", "revision_instruction": "make it oppressive"})
|
||||
self.assertTrue(result.ok)
|
||||
self.assertIn("more oppressive", target.read_text(encoding="utf-8"))
|
||||
self.assertTrue((Path(tmp) / ".vimax" / "logs" / "revisions.jsonl").exists())
|
||||
self.assertTrue(index.get(record["session_id"])["stale"]["final_video"])
|
||||
|
||||
|
||||
async def test_revision_missing_instruction_marks_error(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="x")
|
||||
target = Path(tmp) / record["working_dir"] / "idea2video" / "scene_0" / "storyboard.json"
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text('[]', encoding="utf-8")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
result = await adapter.vimax_narrative_planning({"revision_target": "idea2video/scene_0/storyboard.json"})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.metadata["error_type"], "missing_revision_instruction")
|
||||
self.assertEqual(index.get(record["session_id"])["stage"], "error")
|
||||
|
||||
|
||||
async def test_revision_missing_target_marks_error(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="x")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
result = await adapter.vimax_narrative_planning({"revision_target": "idea2video/scene_0/missing.json", "revision_instruction": "change it"})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.metadata["error_type"], "dependency_missing")
|
||||
self.assertEqual(index.get(record["session_id"])["stage"], "error")
|
||||
|
||||
async def test_render_setup_failure_marks_session_error(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="x")
|
||||
root = Path(tmp) / record["working_dir"] / "idea2video"
|
||||
(root / "scene_0" / "shots" / "0").mkdir(parents=True, exist_ok=True)
|
||||
(root / "story.txt").write_text("story", encoding="utf-8")
|
||||
(root / "characters.json").write_text("[]", encoding="utf-8")
|
||||
(root / "script.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "storyboard.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "camera_tree.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "shots" / "0" / "shot_description.json").write_text("{}", encoding="utf-8")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", side_effect=RuntimeError("missing key")):
|
||||
result = await adapter.vimax_render_video({})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.metadata["error_type"], "render_failed")
|
||||
self.assertIn("missing key", result.content)
|
||||
self.assertEqual(index.get(record["session_id"])["stage"], "error")
|
||||
|
||||
async def test_render_failure_marks_session_error(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="x")
|
||||
root = Path(tmp) / record["working_dir"] / "idea2video"
|
||||
(root / "scene_0" / "shots" / "0").mkdir(parents=True, exist_ok=True)
|
||||
(root / "story.txt").write_text("story", encoding="utf-8")
|
||||
(root / "characters.json").write_text("[]", encoding="utf-8")
|
||||
(root / "script.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "storyboard.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "camera_tree.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "shots" / "0" / "shot_description.json").write_text("{}", encoding="utf-8")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters._build_image_generator", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters._build_video_generator", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FailRenderIdeaPipeline):
|
||||
result = await adapter.vimax_render_video({})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.metadata["error_type"], "render_failed")
|
||||
self.assertIn("render failed", result.content)
|
||||
self.assertEqual(index.get(record["session_id"])["stage"], "error")
|
||||
status_path = Path(tmp) / record["working_dir"] / "render_status.json"
|
||||
events_path = Path(tmp) / record["working_dir"] / "render_events.jsonl"
|
||||
self.assertTrue(status_path.exists())
|
||||
self.assertTrue(events_path.exists())
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(status["status"], "error")
|
||||
self.assertEqual(status["error_type"], "render_failed")
|
||||
|
||||
async def test_render_403_key_limit_is_non_retryable_and_sanitized(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="x")
|
||||
root = Path(tmp) / record["working_dir"] / "idea2video"
|
||||
(root / "scene_0" / "shots" / "0").mkdir(parents=True, exist_ok=True)
|
||||
(root / "story.txt").write_text("story", encoding="utf-8")
|
||||
(root / "characters.json").write_text("[]", encoding="utf-8")
|
||||
(root / "script.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "storyboard.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "camera_tree.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "shots" / "0" / "shot_description.json").write_text("{}", encoding="utf-8")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters._build_image_generator", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters._build_video_generator", return_value=object()), \
|
||||
patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", FailRender403IdeaPipeline):
|
||||
result = await adapter.vimax_render_video({})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertFalse(result.metadata["retryable"])
|
||||
self.assertIn("<redacted>", result.metadata["error"])
|
||||
self.assertNotIn("sk-short", result.metadata["error"])
|
||||
status = json.loads((Path(tmp) / record["working_dir"] / "render_status.json").read_text(encoding="utf-8"))
|
||||
self.assertFalse(status["retryable"])
|
||||
self.assertNotIn("sk-short", status["error"])
|
||||
|
||||
|
||||
async def test_render_pipeline_stdout_is_suppressed(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
record = index.create(idea="x", style="anime")
|
||||
root = Path(tmp) / record["working_dir"] / "idea2video"
|
||||
(root / "scene_0" / "shots" / "0").mkdir(parents=True, exist_ok=True)
|
||||
(root / "story.txt").write_text("story", encoding="utf-8")
|
||||
(root / "characters.json").write_text("[]", encoding="utf-8")
|
||||
(root / "script.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "storyboard.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "camera_tree.json").write_text("[]", encoding="utf-8")
|
||||
(root / "scene_0" / "shots" / "0" / "shot_description.json").write_text("{}", encoding="utf-8")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
stdout = io.StringIO()
|
||||
with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), patch("agent_runtime.vimax_adapters._build_image_generator", return_value=object()), patch("agent_runtime.vimax_adapters._build_video_generator", return_value=object()), patch("agent_runtime.vimax_adapters.Idea2VideoPipeline", NoisyRenderIdeaPipeline), contextlib.redirect_stdout(stdout):
|
||||
result = await adapter.vimax_render_video({})
|
||||
self.assertTrue(result.ok)
|
||||
self.assertNotIn("NOISE_FROM_RENDER_PIPELINE", stdout.getvalue())
|
||||
|
||||
async def test_render_dependency_missing(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
index = SessionIndex(tmp)
|
||||
index.create(idea="x")
|
||||
adapter = ViMaxAdapters(Path(tmp), index)
|
||||
result = await adapter.vimax_render_video({})
|
||||
self.assertFalse(result.ok)
|
||||
self.assertEqual(result.metadata["error_type"], "dependency_missing")
|
||||
@@ -0,0 +1,184 @@
|
||||
"""Regression tests for silent wrong-output bugs in the script2video render path."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from agents.reference_image_selector import select_pairs_by_indices
|
||||
from agents.storyboard_artist import validate_char_idxs
|
||||
from interfaces.camera import Camera
|
||||
from interfaces.shot_description import ShotDescription
|
||||
from pipelines.script2video_pipeline import (
|
||||
Script2VideoPipeline,
|
||||
_collect_priority_shot_idxs,
|
||||
_group_shots_into_cameras,
|
||||
)
|
||||
from utils.text import safe_path_component
|
||||
|
||||
|
||||
def _shot(idx, cam_idx, variation_type="small", ff_chars=None, lf_chars=None):
|
||||
return ShotDescription(
|
||||
idx=idx,
|
||||
is_last=False,
|
||||
cam_idx=cam_idx,
|
||||
visual_desc=f"shot {idx}",
|
||||
variation_type=variation_type,
|
||||
variation_reason="r",
|
||||
ff_desc=f"first frame {idx}",
|
||||
ff_vis_char_idxs=ff_chars or [],
|
||||
lf_desc=f"last frame {idx}",
|
||||
lf_vis_char_idxs=lf_chars or [],
|
||||
motion_desc="m",
|
||||
audio_desc="a",
|
||||
)
|
||||
|
||||
|
||||
class TestCameraGrouping(unittest.TestCase):
|
||||
def test_out_of_order_camera_indices_group_correctly(self):
|
||||
# Shot 0 uses camera 1, shot 1 uses camera 0, shot 2 uses camera 1 again.
|
||||
shots = [_shot(0, cam_idx=1), _shot(1, cam_idx=0), _shot(2, cam_idx=1)]
|
||||
cameras = _group_shots_into_cameras(shots)
|
||||
by_idx = {camera.idx: camera for camera in cameras}
|
||||
self.assertEqual(by_idx[1].active_shot_idxs, [0, 2])
|
||||
self.assertEqual(by_idx[0].active_shot_idxs, [1])
|
||||
|
||||
|
||||
class TestPriorityShotIdxs(unittest.TestCase):
|
||||
def test_priorities_are_shot_indices_not_camera_indices(self):
|
||||
# Camera 2 depends on shot 7 of camera 0: shot 7 must be prioritized.
|
||||
camera_tree = [
|
||||
Camera(idx=0, active_shot_idxs=[7, 8]),
|
||||
Camera(idx=2, active_shot_idxs=[9], parent_cam_idx=0, parent_shot_idx=7),
|
||||
]
|
||||
self.assertEqual(_collect_priority_shot_idxs(camera_tree), [7])
|
||||
|
||||
|
||||
class TestEventDictsAreInstanceState(unittest.TestCase):
|
||||
def _pipeline(self, working_dir):
|
||||
return Script2VideoPipeline(
|
||||
chat_model=MagicMock(),
|
||||
image_generator=MagicMock(),
|
||||
video_generator=MagicMock(),
|
||||
working_dir=working_dir,
|
||||
)
|
||||
|
||||
def test_two_pipelines_do_not_share_events(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
p1 = self._pipeline(os.path.join(tmp, "a"))
|
||||
p2 = self._pipeline(os.path.join(tmp, "b"))
|
||||
p1.frame_events[0] = {"first_frame": asyncio.Event()}
|
||||
p1.shot_desc_events[0] = asyncio.Event()
|
||||
p1.character_portrait_events[0] = asyncio.Event()
|
||||
self.assertEqual(p2.frame_events, {})
|
||||
self.assertEqual(p2.shot_desc_events, {})
|
||||
self.assertEqual(p2.character_portrait_events, {})
|
||||
|
||||
def test_no_class_level_mutable_event_dicts(self):
|
||||
for name in ("frame_events", "shot_desc_events", "character_portrait_events"):
|
||||
self.assertNotIsInstance(
|
||||
Script2VideoPipeline.__dict__.get(name), dict,
|
||||
f"{name} must not be shared class state",
|
||||
)
|
||||
|
||||
|
||||
class TestResumeIncludesNewCameraReference(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_existing_new_camera_image_is_still_offered_to_selector(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
pipeline = Script2VideoPipeline(
|
||||
chat_model=MagicMock(),
|
||||
image_generator=MagicMock(),
|
||||
video_generator=MagicMock(),
|
||||
working_dir=tmp,
|
||||
)
|
||||
shots = [_shot(0, cam_idx=0), _shot(1, cam_idx=1)]
|
||||
camera = Camera(
|
||||
idx=1, active_shot_idxs=[1],
|
||||
parent_cam_idx=0, parent_shot_idx=0,
|
||||
missing_info="wrong background",
|
||||
)
|
||||
parent_done = asyncio.Event()
|
||||
parent_done.set()
|
||||
pipeline.frame_events = {
|
||||
0: {"first_frame": parent_done},
|
||||
1: {"first_frame": asyncio.Event()},
|
||||
}
|
||||
|
||||
# Resume state: transition video and new-camera image already on disk.
|
||||
shot_dir = os.path.join(tmp, "shots", "1")
|
||||
os.makedirs(shot_dir, exist_ok=True)
|
||||
new_camera_path = os.path.join(shot_dir, "new_camera_1.png")
|
||||
open(os.path.join(shot_dir, "transition_video_from_shot_0.mp4"), "wb").close()
|
||||
open(new_camera_path, "wb").close()
|
||||
|
||||
selector = AsyncMock(return_value={"reference_image_path_and_text_pairs": [], "text_prompt": "p"})
|
||||
pipeline.reference_image_selector = MagicMock(select_reference_images_and_generate_prompt=selector)
|
||||
fake_image = MagicMock()
|
||||
pipeline.image_generator.generate_single_image = AsyncMock(return_value=fake_image)
|
||||
|
||||
await pipeline.generate_frames_for_single_camera(
|
||||
camera=camera,
|
||||
shot_descriptions=shots,
|
||||
characters=[],
|
||||
character_portraits_registry={},
|
||||
priority_shot_idxs=[],
|
||||
)
|
||||
|
||||
selector.assert_awaited_once()
|
||||
offered = selector.await_args.kwargs["available_image_path_and_text_pairs"]
|
||||
offered_paths = [pair[0] for pair in offered]
|
||||
self.assertIn(new_camera_path, offered_paths,
|
||||
"resumed runs must offer the new-camera reference image to the selector")
|
||||
|
||||
|
||||
class TestCharIdxValidation(unittest.TestCase):
|
||||
def test_valid_indices_pass(self):
|
||||
validate_char_idxs([0, 1], 2, "ff_vis_char_idxs")
|
||||
|
||||
def test_out_of_range_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
validate_char_idxs([0, 2], 2, "ff_vis_char_idxs")
|
||||
|
||||
def test_negative_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
validate_char_idxs([-1], 2, "lf_vis_char_idxs")
|
||||
|
||||
|
||||
class TestReferenceSelectorIndices(unittest.TestCase):
|
||||
def test_valid_selection(self):
|
||||
pairs = [("a.png", "a"), ("b.png", "b")]
|
||||
self.assertEqual(select_pairs_by_indices(pairs, [1]), [("b.png", "b")])
|
||||
|
||||
def test_negative_index_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
select_pairs_by_indices([("a.png", "a")], [-1])
|
||||
|
||||
def test_out_of_range_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
select_pairs_by_indices([("a.png", "a")], [3])
|
||||
|
||||
|
||||
class TestSafePathComponent(unittest.TestCase):
|
||||
def test_clean_names_unchanged(self):
|
||||
self.assertEqual(safe_path_component("Alice"), "Alice")
|
||||
self.assertEqual(safe_path_component("Bob_2"), "Bob_2")
|
||||
|
||||
def test_unicode_names_preserved(self):
|
||||
self.assertEqual(safe_path_component("李雷"), "李雷")
|
||||
|
||||
def test_path_separators_removed(self):
|
||||
self.assertNotIn("/", safe_path_component("a/b"))
|
||||
self.assertNotIn("\\", safe_path_component("a\\b"))
|
||||
|
||||
def test_traversal_neutralized(self):
|
||||
cleaned = safe_path_component("../../etc/passwd")
|
||||
self.assertNotIn("/", cleaned)
|
||||
self.assertFalse(cleaned.startswith("."))
|
||||
|
||||
def test_empty_becomes_placeholder(self):
|
||||
self.assertEqual(safe_path_component(""), "unnamed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,35 @@
|
||||
# rendering abstraction
|
||||
from .protocols import ImageGenerator, VideoGenerator
|
||||
from .render_backend import RenderBackend
|
||||
|
||||
# image generators
|
||||
from .image_generator_doubao_seedream_yunwu_api import ImageGeneratorDoubaoSeedreamYunwuAPI
|
||||
from .image_generator_nanobanana_google_api import ImageGeneratorNanobananaGoogleAPI
|
||||
from .image_generator_nanobanana_yunwu_api import ImageGeneratorNanobananaYunwuAPI
|
||||
|
||||
# reranker for rag
|
||||
from .reranker_bge_silicon_api import RerankerBgeSiliconapi
|
||||
|
||||
# video generators
|
||||
from .video_generator_doubao_seedance_yunwu_api import VideoGeneratorDoubaoSeedanceYunwuAPI
|
||||
from .video_generator_omni_yunwu_api import VideoGeneratorOmniYunwuAPI, VideoGeneratorOminiYunwuAPI
|
||||
from .video_generator_openrouter_api import VideoGeneratorOpenRouterAPI
|
||||
from .video_generator_veo_google_api import VideoGeneratorVeoGoogleAPI
|
||||
from .video_generator_veo_yunwu_api import VideoGeneratorVeoYunwuAPI
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ImageGenerator",
|
||||
"VideoGenerator",
|
||||
"RenderBackend",
|
||||
"ImageGeneratorDoubaoSeedreamYunwuAPI",
|
||||
"ImageGeneratorNanobananaGoogleAPI",
|
||||
"ImageGeneratorNanobananaYunwuAPI",
|
||||
"RerankerBgeSiliconapi",
|
||||
"VideoGeneratorDoubaoSeedanceYunwuAPI",
|
||||
"VideoGeneratorOmniYunwuAPI",
|
||||
"VideoGeneratorOminiYunwuAPI",
|
||||
"VideoGeneratorOpenRouterAPI",
|
||||
"VideoGeneratorVeoGoogleAPI",
|
||||
"VideoGeneratorVeoYunwuAPI",
|
||||
]
|
||||
@@ -0,0 +1,74 @@
|
||||
# https://yunwu.apifox.cn/api-347960869
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import aiohttp
|
||||
from typing import List, Optional
|
||||
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
|
||||
from utils.retry import after_func
|
||||
from utils.image import image_path_to_b64
|
||||
from interfaces.image_output import ImageOutput
|
||||
|
||||
|
||||
class ImageGeneratorDoubaoSeedreamYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = "doubao-seedream-4-0-250828",
|
||||
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.base_url = "https://yunwu.ai/v1/images/generations"
|
||||
self.model = model
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, max=30),
|
||||
retry=retry_if_exception_type((aiohttp.ClientError, asyncio.TimeoutError)),
|
||||
reraise=True,
|
||||
after=after_func,
|
||||
)
|
||||
async def generate_single_image(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str] = [],
|
||||
size: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> ImageOutput:
|
||||
"""
|
||||
size: [1024x1024, 4096x4096]
|
||||
"""
|
||||
|
||||
logging.info(f"Calling {self.model} to generate image...")
|
||||
|
||||
image = [
|
||||
image_path_to_b64(path, mime=True) for path in reference_image_paths
|
||||
]
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"prompt": prompt,
|
||||
"sequential_image_generation": "disabled", # "auto" or "disabled"
|
||||
# "sequential_image_generation_options": {
|
||||
# "max_images": 1
|
||||
# },
|
||||
"response_format": "url",
|
||||
"size": size if size is not None else "1024x1024",
|
||||
}
|
||||
if len(image) > 0:
|
||||
payload["image"] = image
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(self.base_url, json=payload, headers=headers) as response:
|
||||
response_json = await response.json()
|
||||
if response.status >= 400:
|
||||
raise RuntimeError(f"Image generation failed with HTTP {response.status}: {response_json}")
|
||||
|
||||
data = response_json['data'][0]['url']
|
||||
return ImageOutput(fmt="url", ext="png", data=data)
|
||||
@@ -0,0 +1,97 @@
|
||||
# https://ai.google.dev/gemini-api/docs/image-generation
|
||||
|
||||
import logging
|
||||
import asyncio
|
||||
from PIL import Image
|
||||
from typing import List, Optional
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from google.genai.errors import ClientError
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
from interfaces.image_output import ImageOutput
|
||||
from tools.image_orientation import ensure_not_portrait, landscape_guard_requested
|
||||
from tools.image_response import image_from_response_part
|
||||
from utils.retry import after_func
|
||||
from utils.rate_limiter import RateLimiter
|
||||
|
||||
|
||||
class ImageGeneratorNanobananaGoogleAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
rate_limiter: Optional[RateLimiter] = None,
|
||||
):
|
||||
self.model = "gemini-2.5-flash-image"
|
||||
self.rate_limiter = rate_limiter
|
||||
self.client = genai.Client(
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10), after=after_func, reraise=True)
|
||||
async def generate_single_image(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str] = [],
|
||||
aspect_ratio: Optional[str] = "16:9",
|
||||
**kwargs,
|
||||
) -> ImageOutput:
|
||||
|
||||
"""
|
||||
aspect_ratio: The aspect ratio of the image.
|
||||
"""
|
||||
|
||||
logging.info(f"Calling {self.model} to generate image...")
|
||||
|
||||
# Apply rate limiting if configured
|
||||
if self.rate_limiter:
|
||||
await self.rate_limiter.acquire()
|
||||
|
||||
reference_images = [Image.open(path) for path in reference_image_paths]
|
||||
|
||||
# Retry logic for rate limit errors
|
||||
max_retries = 3
|
||||
retry_delay = 5
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = await self.client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=reference_images + [prompt],
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=["IMAGE"],
|
||||
image_config=types.ImageConfig(
|
||||
aspect_ratio=aspect_ratio,
|
||||
),
|
||||
),
|
||||
)
|
||||
break
|
||||
except ClientError as e:
|
||||
if e.status_code == 429 and attempt < max_retries - 1:
|
||||
wait_time = retry_delay * (2 ** attempt)
|
||||
logging.warning(f"Rate limit hit (429), retrying in {wait_time}s... (attempt {attempt + 1}/{max_retries})")
|
||||
await asyncio.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
|
||||
image = None
|
||||
text = ""
|
||||
for part in response.candidates[0].content.parts:
|
||||
if part.text is not None:
|
||||
text += part.text
|
||||
elif part.inline_data is not None:
|
||||
image = image_from_response_part(part)
|
||||
|
||||
if image is None:
|
||||
logging.error(f"No image generated. The response text is: {text}")
|
||||
raise ValueError("No image generated")
|
||||
|
||||
if landscape_guard_requested(
|
||||
size=kwargs.get("size"),
|
||||
aspect_ratio=aspect_ratio,
|
||||
enforce_landscape=kwargs.get("enforce_landscape", True),
|
||||
allow_portrait=kwargs.get("allow_portrait", False),
|
||||
):
|
||||
ensure_not_portrait(image)
|
||||
|
||||
return ImageOutput(fmt="pil", ext="png", data=image)
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
# https://ai.google.dev/gemini-api/docs/image-generation?hl=zh-cn
|
||||
|
||||
import logging
|
||||
from PIL import Image
|
||||
from typing import List, Optional
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
from interfaces.image_output import ImageOutput
|
||||
from tools.image_orientation import ensure_not_portrait, landscape_guard_requested
|
||||
from tools.image_response import image_from_response_part
|
||||
from utils.retry import after_func
|
||||
|
||||
|
||||
class ImageGeneratorNanobananaYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = "gemini-2.5-flash-image-preview",
|
||||
base_url: str = "https://yunwu.ai",
|
||||
):
|
||||
self.client = genai.Client(
|
||||
api_key=api_key,
|
||||
http_options=types.HttpOptions(
|
||||
base_url=base_url.rstrip("/"),
|
||||
api_version="v1beta",
|
||||
),
|
||||
)
|
||||
self.model = model
|
||||
|
||||
|
||||
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10), after=after_func, reraise=True)
|
||||
async def generate_single_image(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str] = [],
|
||||
aspect_ratio: Optional[str] = "16:9",
|
||||
**kwargs,
|
||||
) -> ImageOutput:
|
||||
"""
|
||||
aspect_ratio: The aspect ratio of the image.
|
||||
"""
|
||||
|
||||
logging.info(f"Calling {self.model} to generate image...")
|
||||
|
||||
reference_images = [Image.open(path) for path in reference_image_paths]
|
||||
|
||||
response = await self.client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=reference_images + [prompt],
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=["TEXT", "IMAGE"],
|
||||
image_config=types.ImageConfig(
|
||||
aspect_ratio=aspect_ratio,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
image = None
|
||||
text = ""
|
||||
for part in response.candidates[0].content.parts:
|
||||
if part.text is not None:
|
||||
text += part.text
|
||||
elif part.inline_data is not None:
|
||||
image = image_from_response_part(part)
|
||||
|
||||
if image is None:
|
||||
logging.error(f"No image generated. The response text is: {text}")
|
||||
raise ValueError(f"Error occurred while generating image.")
|
||||
|
||||
if landscape_guard_requested(
|
||||
size=kwargs.get("size"),
|
||||
aspect_ratio=aspect_ratio,
|
||||
enforce_landscape=kwargs.get("enforce_landscape", True),
|
||||
allow_portrait=kwargs.get("allow_portrait", False),
|
||||
):
|
||||
ensure_not_portrait(image)
|
||||
|
||||
return ImageOutput(fmt="pil", ext="png", data=image)
|
||||
@@ -0,0 +1,53 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def landscape_guard_requested(*, size: Any = None, aspect_ratio: Any = None, enforce_landscape: Any = True, allow_portrait: Any = False) -> bool:
|
||||
if bool(allow_portrait):
|
||||
return False
|
||||
if bool(enforce_landscape):
|
||||
return True
|
||||
parsed = _parse_size(size)
|
||||
if parsed and parsed[0] > parsed[1]:
|
||||
return True
|
||||
parsed_ratio = _parse_size(aspect_ratio)
|
||||
return bool(parsed_ratio and parsed_ratio[0] > parsed_ratio[1])
|
||||
|
||||
|
||||
def ensure_not_portrait(image: Image.Image, *, tolerance: float | None = None) -> None:
|
||||
width, height = image.size
|
||||
if width <= 0 or height <= 0:
|
||||
return
|
||||
threshold = tolerance if tolerance is not None else _portrait_tolerance()
|
||||
if height > width * threshold:
|
||||
raise ValueError(f"Generated image is portrait-oriented ({width}x{height}); retrying for a landscape frame")
|
||||
|
||||
|
||||
def _portrait_tolerance() -> float:
|
||||
raw = os.environ.get("VIMAX_IMAGE_PORTRAIT_RETRY_TOLERANCE", "1.05")
|
||||
try:
|
||||
return max(1.0, float(raw))
|
||||
except ValueError:
|
||||
return 1.05
|
||||
|
||||
|
||||
def _parse_size(size: Any) -> tuple[int, int] | None:
|
||||
if not isinstance(size, str):
|
||||
return None
|
||||
normalized = size.lower()
|
||||
separator = "x" if "x" in normalized else ":" if ":" in normalized else ""
|
||||
if not separator:
|
||||
return None
|
||||
left, right = normalized.split(separator, 1)
|
||||
try:
|
||||
width = int(left.strip())
|
||||
height = int(right.strip())
|
||||
except ValueError:
|
||||
return None
|
||||
if width <= 0 or height <= 0:
|
||||
return None
|
||||
return width, height
|
||||
@@ -0,0 +1,40 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def image_from_response_part(part: Any) -> Image.Image | None:
|
||||
inline_data = getattr(part, "inline_data", None)
|
||||
if inline_data is None and isinstance(part, dict):
|
||||
inline_data = part.get("inline_data")
|
||||
if inline_data is None:
|
||||
return None
|
||||
|
||||
as_image = getattr(part, "as_image", None)
|
||||
if callable(as_image):
|
||||
image = as_image()
|
||||
if isinstance(image, Image.Image):
|
||||
return image
|
||||
|
||||
data = _value(inline_data, "data")
|
||||
if data is None:
|
||||
return None
|
||||
if isinstance(data, str):
|
||||
if data.startswith("data:") and "," in data:
|
||||
data = data.split(",", 1)[1]
|
||||
data = base64.b64decode(data)
|
||||
if isinstance(data, bytearray):
|
||||
data = bytes(data)
|
||||
if not isinstance(data, bytes):
|
||||
return None
|
||||
return Image.open(BytesIO(data)).convert("RGB")
|
||||
|
||||
|
||||
def _value(obj: Any, key: str) -> Any:
|
||||
if isinstance(obj, dict):
|
||||
return obj.get(key)
|
||||
return getattr(obj, key, None)
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Structural typing contracts for rendering backends.
|
||||
|
||||
Any class that exposes the right method signatures satisfies these
|
||||
protocols -- no inheritance required. Existing generators (Google,
|
||||
Yunwu/Doubao, Yunwu/Veo) are already compliant by duck typing.
|
||||
"""
|
||||
|
||||
from typing import List, Protocol, runtime_checkable
|
||||
|
||||
from interfaces.image_output import ImageOutput
|
||||
from interfaces.video_output import VideoOutput
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ImageGenerator(Protocol):
|
||||
"""Generates a single image from a text prompt and optional reference images."""
|
||||
|
||||
async def generate_single_image(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
**kwargs,
|
||||
) -> ImageOutput: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VideoGenerator(Protocol):
|
||||
"""Generates a single video from a text prompt and optional reference images."""
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
**kwargs,
|
||||
) -> VideoOutput: ...
|
||||
@@ -0,0 +1,62 @@
|
||||
"""RenderBackend: config-driven factory for image and video generators.
|
||||
|
||||
Reads the ``image_generator`` and ``video_generator`` sections from a
|
||||
ViMax YAML config, instantiates the concrete classes via *class_path*,
|
||||
and wires up rate limiters.
|
||||
|
||||
Usage::
|
||||
|
||||
backend = RenderBackend.from_config(config)
|
||||
image = await backend.image_generator.generate_single_image(...)
|
||||
video = await backend.video_generator.generate_single_video(...)
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict
|
||||
|
||||
from utils.rate_limiter import RateLimiter
|
||||
|
||||
|
||||
@dataclass
|
||||
class RenderBackend:
|
||||
"""Bundles an image generator and a video generator."""
|
||||
|
||||
image_generator: Any
|
||||
video_generator: Any
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> "RenderBackend":
|
||||
"""Build a RenderBackend from a parsed YAML config dict.
|
||||
|
||||
Rate limiters are created from ``max_requests_per_minute`` /
|
||||
``max_requests_per_day`` if present in each generator section.
|
||||
"""
|
||||
img_cfg = config["image_generator"]
|
||||
vid_cfg = config["video_generator"]
|
||||
|
||||
image_gen = _instantiate(img_cfg, _build_rate_limiter(img_cfg))
|
||||
video_gen = _instantiate(vid_cfg, _build_rate_limiter(vid_cfg))
|
||||
|
||||
logging.info("RenderBackend: image=%s, video=%s",
|
||||
img_cfg["class_path"], vid_cfg["class_path"])
|
||||
|
||||
return cls(image_generator=image_gen, video_generator=video_gen)
|
||||
|
||||
|
||||
def _build_rate_limiter(section: Dict[str, Any]) -> RateLimiter | None:
|
||||
rpm = section.get("max_requests_per_minute")
|
||||
rpd = section.get("max_requests_per_day")
|
||||
if rpm or rpd:
|
||||
return RateLimiter(max_requests_per_minute=rpm, max_requests_per_day=rpd)
|
||||
return None
|
||||
|
||||
|
||||
def _instantiate(section: Dict[str, Any], rate_limiter: RateLimiter | None) -> Any:
|
||||
module_path, cls_name = section["class_path"].rsplit(".", 1)
|
||||
cls = getattr(importlib.import_module(module_path), cls_name)
|
||||
init_args = dict(section.get("init_args", {}))
|
||||
if rate_limiter is not None:
|
||||
init_args["rate_limiter"] = rate_limiter
|
||||
return cls(**init_args)
|
||||
@@ -0,0 +1,83 @@
|
||||
from typing import List
|
||||
import aiohttp
|
||||
import asyncio
|
||||
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
|
||||
import logging
|
||||
|
||||
|
||||
class RerankerBgeSiliconapi:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str = "BAAI/bge-reranker-v2-m3",
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
# return_documents: bool = True,
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, max=30),
|
||||
retry=retry_if_exception_type((aiohttp.ClientError, asyncio.TimeoutError)),
|
||||
reraise=True,
|
||||
after=lambda retry_state: logging.warning(f"Retrying SiliconReranker due to error: {retry_state.outcome.exception()}"),
|
||||
)
|
||||
async def __call__(
|
||||
self,
|
||||
documents: List[str],
|
||||
query: str,
|
||||
top_n: int,
|
||||
) -> List[str]:
|
||||
|
||||
url = f"{self.base_url}/rerank"
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
"top_n": top_n,
|
||||
"return_documents": True,
|
||||
}
|
||||
|
||||
|
||||
headers = {
|
||||
'Accept': 'application/json',
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, json=payload, headers=headers) as resp:
|
||||
response = await resp.json()
|
||||
if resp.status >= 400:
|
||||
raise RuntimeError(f"Rerank request failed with HTTP {resp.status}: {response}")
|
||||
|
||||
|
||||
"""
|
||||
{
|
||||
"id": "<string>",
|
||||
"results": [
|
||||
{
|
||||
"document": {
|
||||
"text": "<string>"
|
||||
},
|
||||
"index": 123,
|
||||
"relevance_score": 123
|
||||
}
|
||||
],
|
||||
"tokens": {
|
||||
"input_tokens": 123,
|
||||
"output_tokens": 123
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
results = []
|
||||
|
||||
for result in response["results"]:
|
||||
results.append((result["document"]["text"], result["relevance_score"]))
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,212 @@
|
||||
import logging
|
||||
from typing import List, Literal
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.image import image_path_to_b64
|
||||
|
||||
|
||||
class VideoGeneratorDoubaoSeedanceYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
t2v_model: str = "doubao-seedance-1-0-lite-t2v-250428",
|
||||
ff2v_model: str = "doubao-seedance-1-0-lite-i2v-250428",
|
||||
flf2v_model: str = "doubao-seedance-1-0-lite-i2v-250428",
|
||||
max_create_attempts: int = 3,
|
||||
poll_interval: int = 2,
|
||||
max_poll_attempts: int = 300,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.t2v_model = t2v_model
|
||||
self.ff2v_model = ff2v_model
|
||||
self.flf2v_model = flf2v_model
|
||||
self.max_create_attempts = max_create_attempts
|
||||
self.poll_interval = poll_interval
|
||||
self.max_poll_attempts = max_poll_attempts
|
||||
|
||||
|
||||
async def create_video_generation_task(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
resolution: Literal["480p", "720p", "1080p"] = "720p",
|
||||
aspect_ratio: str = "16:9",
|
||||
fps: Literal[16, 24] = 16,
|
||||
duration: Literal[5, 10] = 5,
|
||||
) -> str:
|
||||
"""
|
||||
Create a video generation task and return the task ID.
|
||||
|
||||
Args:
|
||||
prompt: Text prompt for video generation
|
||||
reference_image_paths: List of 1 or 2 reference images
|
||||
|
||||
Returns:
|
||||
Task ID string
|
||||
"""
|
||||
if len(reference_image_paths) == 0:
|
||||
model = self.t2v_model
|
||||
elif len(reference_image_paths) == 1:
|
||||
model = self.ff2v_model
|
||||
elif len(reference_image_paths) == 2:
|
||||
model = self.flf2v_model
|
||||
else:
|
||||
raise ValueError("reference_image_paths must contain 1 or 2 images.")
|
||||
|
||||
logging.info(f"Calling {model} to generate video...")
|
||||
|
||||
url = "https://yunwu.ai/volc/v1/contents/generations/tasks"
|
||||
|
||||
|
||||
content = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": prompt + f" --rs {resolution} --rt {aspect_ratio} --dur {duration} --fps {fps} --wm false --seed -1 --cf false"
|
||||
}
|
||||
]
|
||||
if len(reference_image_paths) >= 1:
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": image_path_to_b64(reference_image_paths[0])
|
||||
},
|
||||
"role": "first_frame",
|
||||
}
|
||||
)
|
||||
if len(reference_image_paths) >= 2:
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": image_path_to_b64(reference_image_paths[1])
|
||||
},
|
||||
"role": "last_frame",
|
||||
}
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"content": content
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
last_error = None
|
||||
for attempt in range(1, self.max_create_attempts + 1):
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, headers=headers, json=payload) as response:
|
||||
response_json = await response.json()
|
||||
http_status = response.status
|
||||
logging.debug(f"Response: {response_json}")
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
logging.error(f"Error occurred while creating video generation task (attempt {attempt}/{self.max_create_attempts}): {e}")
|
||||
if attempt < self.max_create_attempts:
|
||||
await asyncio.sleep(attempt)
|
||||
continue
|
||||
|
||||
if http_status >= 400:
|
||||
message = f"Video generation task creation failed with HTTP {http_status}: {response_json}"
|
||||
if http_status < 500:
|
||||
raise RuntimeError(message)
|
||||
last_error = RuntimeError(message)
|
||||
logging.error(f"{message} (attempt {attempt}/{self.max_create_attempts})")
|
||||
if attempt < self.max_create_attempts:
|
||||
await asyncio.sleep(attempt)
|
||||
continue
|
||||
|
||||
task_id = response_json.get("id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Video generation task creation returned no task id: {response_json}")
|
||||
logging.info(f"Video generation task created successfully. Task ID: {task_id}")
|
||||
return task_id
|
||||
|
||||
raise RuntimeError(f"Failed to create video generation task after {self.max_create_attempts} attempts.") from last_error
|
||||
|
||||
async def query_video_generation_task(
|
||||
self,
|
||||
task_id: str,
|
||||
) -> str:
|
||||
"""
|
||||
Query the video generation task until completion and return the video URL.
|
||||
|
||||
Args:
|
||||
task_id: Task ID to query
|
||||
|
||||
Returns:
|
||||
Video URL string
|
||||
"""
|
||||
url = f"https://yunwu.ai/volc/v1/contents/generations/tasks/{task_id}"
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
}
|
||||
|
||||
attempts = 0
|
||||
consecutive_errors = 0
|
||||
while True:
|
||||
if attempts >= self.max_poll_attempts:
|
||||
raise TimeoutError(f"Video generation did not complete after {attempts} polls.")
|
||||
attempts += 1
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url, headers=headers) as response:
|
||||
response_json = await response.json()
|
||||
http_status = response.status
|
||||
except Exception as e:
|
||||
consecutive_errors += 1
|
||||
if consecutive_errors >= 5:
|
||||
raise RuntimeError(f"Querying video generation task failed {consecutive_errors} times in a row.") from e
|
||||
logging.error(f"Error occurred while querying video generation task: {e}. Retrying in {self.poll_interval} seconds...")
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
continue
|
||||
consecutive_errors = 0
|
||||
|
||||
if http_status >= 400:
|
||||
raise RuntimeError(f"Querying video generation task failed with HTTP {http_status}: {response_json}")
|
||||
|
||||
status = response_json.get("status")
|
||||
if status == "succeeded":
|
||||
video_url = response_json["content"]["video_url"]
|
||||
logging.info(f"Video generation completed successfully. Video URL: {video_url}")
|
||||
return video_url
|
||||
elif status == "failed":
|
||||
logging.error(f"Video generation failed. Response: {response_json}")
|
||||
raise ValueError("Video generation failed.")
|
||||
else:
|
||||
logging.info(f"Video generation is still in progress. Checking again in {self.poll_interval} seconds...")
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
resolution: Literal["480p", "720p", "1080p"] = "720p",
|
||||
aspect_ratio: str = "16:9",
|
||||
fps: Literal[16, 24] = 16,
|
||||
duration: Literal[5, 10] = 5,
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
"""
|
||||
Generate a single video by creating a task and waiting for completion.
|
||||
|
||||
Args:
|
||||
prompt: Text prompt for video generation
|
||||
reference_image_paths: List of 1 or 2 reference images
|
||||
resolution: Resolution of the video
|
||||
aspect_ratio: Aspect ratio of the video
|
||||
fps: Frames per second of the video
|
||||
duration: Duration of the video
|
||||
Returns:
|
||||
VideoOutput containing the video URL
|
||||
"""
|
||||
task_id = await self.create_video_generation_task(prompt, reference_image_paths, resolution, aspect_ratio, fps, duration)
|
||||
video_url = await self.query_video_generation_task(task_id)
|
||||
return VideoOutput(fmt="url", ext="mp4", data=video_url)
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.image import image_path_to_b64
|
||||
from utils.rate_limiter import RateLimiter
|
||||
|
||||
|
||||
class VideoGeneratorOmniYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
t2v_model: str = "omni-flash",
|
||||
i2v_model: str = "omni-flash",
|
||||
base_url: str = "https://yunwu.ai",
|
||||
seconds: int = 8,
|
||||
enable_upsample: bool = False,
|
||||
enable_sample: Optional[bool] = None,
|
||||
poll_interval: int = 2,
|
||||
max_poll_attempts: Optional[int] = 300,
|
||||
max_create_attempts: int = 3,
|
||||
rate_limiter: Optional[RateLimiter] = None,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.t2v_model = t2v_model
|
||||
self.i2v_model = i2v_model
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.seconds = seconds
|
||||
self.enable_upsample = enable_upsample
|
||||
self.enable_sample = enable_sample
|
||||
self.poll_interval = poll_interval
|
||||
self.max_poll_attempts = max_poll_attempts
|
||||
self.max_create_attempts = max_create_attempts
|
||||
self.rate_limiter = rate_limiter
|
||||
|
||||
def _headers(self) -> dict:
|
||||
return {
|
||||
"Accept": "application/json",
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def _image_uri(self, image_path: str) -> str:
|
||||
if image_path.startswith(("http://", "https://", "data:")):
|
||||
return image_path
|
||||
return image_path_to_b64(image_path, mime=True)
|
||||
|
||||
def _build_payload(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
aspect_ratio: str,
|
||||
seconds: Optional[int],
|
||||
size: Optional[str],
|
||||
enable_upsample: Optional[bool],
|
||||
enable_sample: Optional[bool],
|
||||
) -> dict:
|
||||
if len(reference_image_paths) > 3:
|
||||
raise ValueError("The number of reference images must be no more than 3")
|
||||
|
||||
payload = {
|
||||
"model": self.t2v_model if len(reference_image_paths) == 0 else self.i2v_model,
|
||||
"prompt": prompt,
|
||||
"seconds": str(seconds or self.seconds),
|
||||
}
|
||||
|
||||
if len(reference_image_paths) == 0:
|
||||
payload["type"] = 1
|
||||
elif len(reference_image_paths) <= 2:
|
||||
payload["type"] = 2
|
||||
payload["images"] = [self._image_uri(path) for path in reference_image_paths]
|
||||
else:
|
||||
payload["type"] = 3
|
||||
payload["images"] = [self._image_uri(path) for path in reference_image_paths]
|
||||
|
||||
if aspect_ratio:
|
||||
payload["aspect_ratio"] = aspect_ratio
|
||||
if size:
|
||||
payload["size"] = size
|
||||
if enable_upsample is not None:
|
||||
payload["enable_upsample"] = enable_upsample
|
||||
if enable_sample is not None:
|
||||
payload["enable_sample"] = enable_sample
|
||||
|
||||
return payload
|
||||
|
||||
async def create_video_generation_task(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
aspect_ratio: str = "16:9",
|
||||
seconds: Optional[int] = None,
|
||||
size: Optional[str] = None,
|
||||
enable_upsample: Optional[bool] = None,
|
||||
enable_sample: Optional[bool] = None,
|
||||
) -> tuple[str, str]:
|
||||
payload = self._build_payload(
|
||||
prompt=prompt,
|
||||
reference_image_paths=reference_image_paths,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
enable_upsample=self.enable_upsample if enable_upsample is None else enable_upsample,
|
||||
enable_sample=self.enable_sample if enable_sample is None else enable_sample,
|
||||
)
|
||||
|
||||
logging.info("Calling %s to generate video...", payload["model"])
|
||||
|
||||
if self.rate_limiter:
|
||||
await self.rate_limiter.acquire()
|
||||
|
||||
url = f"{self.base_url}/v1/video/create"
|
||||
last_error = None
|
||||
for attempt in range(1, self.max_create_attempts + 1):
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, headers=self._headers(), json=payload) as response:
|
||||
response_json = await response.json()
|
||||
http_status = response.status
|
||||
logging.debug("Response: %s", response_json)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
logging.error(
|
||||
"Error occurred while creating video generation task (attempt %s/%s): %s",
|
||||
attempt,
|
||||
self.max_create_attempts,
|
||||
e,
|
||||
)
|
||||
if attempt < self.max_create_attempts:
|
||||
await asyncio.sleep(attempt)
|
||||
continue
|
||||
|
||||
if http_status >= 400:
|
||||
message = f"Video generation task creation failed with HTTP {http_status}: {response_json}"
|
||||
if http_status < 500:
|
||||
raise RuntimeError(message)
|
||||
last_error = RuntimeError(message)
|
||||
logging.error("%s (attempt %s/%s)", message, attempt, self.max_create_attempts)
|
||||
if attempt < self.max_create_attempts:
|
||||
await asyncio.sleep(attempt)
|
||||
continue
|
||||
|
||||
task_id = response_json.get("id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Video generation task creation returned no task id: {response_json}")
|
||||
logging.info("Video generation task created successfully. Task ID: %s", task_id)
|
||||
return task_id, payload["model"]
|
||||
|
||||
raise RuntimeError(
|
||||
f"Failed to create video generation task after {self.max_create_attempts} attempts."
|
||||
) from last_error
|
||||
|
||||
async def query_video_generation_task(self, task_id: str, model: str) -> str:
|
||||
url = f"{self.base_url}/v1/video/query"
|
||||
params = {"id": task_id, "model": model}
|
||||
|
||||
attempts = 0
|
||||
while True:
|
||||
if self.max_poll_attempts is not None and attempts >= self.max_poll_attempts:
|
||||
raise TimeoutError(f"Video generation did not complete after {attempts} polls.")
|
||||
attempts += 1
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url, headers=self._headers(), params=params) as response:
|
||||
response_json = await response.json()
|
||||
logging.debug("Response: %s", response_json)
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
"Error occurred while querying video generation task: %s. Retrying in %s seconds...",
|
||||
e,
|
||||
self.poll_interval,
|
||||
)
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
continue
|
||||
|
||||
status = response_json.get("status")
|
||||
if status == "completed":
|
||||
detail = response_json.get("detail") or {}
|
||||
video_url = (
|
||||
response_json.get("video_url")
|
||||
or detail.get("upsample_video_url")
|
||||
or detail.get("video_url")
|
||||
)
|
||||
if not video_url:
|
||||
raise RuntimeError(f"Video generation completed without a video URL: {response_json}")
|
||||
logging.info("Video generation completed successfully. Video URL: %s", video_url)
|
||||
return video_url
|
||||
|
||||
if status in {"failed", "error"}:
|
||||
raise RuntimeError(f"Video generation failed: {response_json}")
|
||||
|
||||
logging.info("Video generation status: %s, waiting %s seconds...", status, self.poll_interval)
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
aspect_ratio: str = "16:9",
|
||||
seconds: Optional[int] = None,
|
||||
size: Optional[str] = None,
|
||||
enable_upsample: Optional[bool] = None,
|
||||
enable_sample: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
task_id, model = await self.create_video_generation_task(
|
||||
prompt=prompt,
|
||||
reference_image_paths=reference_image_paths,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
enable_upsample=enable_upsample,
|
||||
enable_sample=enable_sample,
|
||||
)
|
||||
video_url = await self.query_video_generation_task(task_id, model)
|
||||
return VideoOutput(fmt="url", ext="mp4", data=video_url)
|
||||
|
||||
|
||||
class VideoGeneratorOminiYunwuAPI(VideoGeneratorOmniYunwuAPI):
|
||||
"""Backward-compatible alias for the common "omini" spelling."""
|
||||
@@ -0,0 +1,201 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from typing import List
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import aiohttp
|
||||
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.image import image_path_to_b64
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
try:
|
||||
return max(0, int(os.environ.get(name, str(default))))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _env_float(name: str, default: float) -> float:
|
||||
try:
|
||||
return max(0.0, float(os.environ.get(name, str(default))))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _env_bool(name: str, default: bool) -> bool:
|
||||
raw = os.environ.get(name)
|
||||
if raw is None:
|
||||
return default
|
||||
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _emit_progress(progress, stage: str, message: str, metadata: dict | None = None) -> None:
|
||||
if progress is not None:
|
||||
progress(stage, message, metadata or {})
|
||||
|
||||
|
||||
class VideoGeneratorOpenRouterAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = "google/veo-3.1-lite",
|
||||
base_url: str = "https://openrouter.ai/api/v1",
|
||||
http_referer: str = "",
|
||||
app_title: str = "ViMax",
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.http_referer = http_referer
|
||||
self.app_title = app_title
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str = "",
|
||||
reference_image_paths: List[str] = [],
|
||||
aspect_ratio: str = "16:9",
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
progress = kwargs.get("progress")
|
||||
request_timeout_seconds = _env_float("VIMAX_VIDEO_REQUEST_TIMEOUT_SECONDS", 60.0)
|
||||
query_timeout_seconds = _env_float("VIMAX_VIDEO_QUERY_TIMEOUT_SECONDS", 600.0)
|
||||
poll_interval_seconds = _env_float("VIMAX_VIDEO_POLL_INTERVAL_SECONDS", 10.0)
|
||||
duration = _env_int("VIMAX_OPENROUTER_VIDEO_DURATION", 8)
|
||||
resolution = os.environ.get("VIMAX_OPENROUTER_VIDEO_RESOLUTION", "720p")
|
||||
generate_audio = _env_bool("VIMAX_OPENROUTER_GENERATE_AUDIO", True)
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"duration": duration,
|
||||
"resolution": resolution,
|
||||
"generate_audio": generate_audio,
|
||||
}
|
||||
frame_images = _frame_images(reference_image_paths)
|
||||
if frame_images:
|
||||
payload["frame_images"] = frame_images
|
||||
|
||||
headers = self._headers()
|
||||
timeout = aiohttp.ClientTimeout(total=request_timeout_seconds)
|
||||
_emit_progress(progress, "video_create", f"Creating OpenRouter video generation task with {self.model}", {"model": self.model, "duration": duration, "resolution": resolution, "frame_count": len(frame_images)})
|
||||
|
||||
create_status, create_payload = await _post_json(
|
||||
f"{self.base_url}/videos",
|
||||
headers=headers,
|
||||
payload=payload,
|
||||
timeout=timeout,
|
||||
hard_timeout_seconds=request_timeout_seconds,
|
||||
)
|
||||
if create_status >= 400:
|
||||
raise RuntimeError(f"OpenRouter video create failed with HTTP {create_status}: {create_payload}")
|
||||
job_id = create_payload.get("id")
|
||||
polling_url = create_payload.get("polling_url")
|
||||
if not job_id or not polling_url:
|
||||
raise RuntimeError(f"OpenRouter video create response missing id or polling_url: {create_payload}")
|
||||
_emit_progress(progress, "video_task_created", "OpenRouter video generation task created", {"model": self.model, "job_id": job_id, "status": create_payload.get("status")})
|
||||
|
||||
poll_url = _absolute_url(self.base_url, polling_url)
|
||||
deadline = asyncio.get_running_loop().time() + query_timeout_seconds if query_timeout_seconds > 0 else None
|
||||
last_status = create_payload.get("status")
|
||||
last_payload = create_payload
|
||||
while deadline is None or asyncio.get_running_loop().time() < deadline:
|
||||
await asyncio.sleep(poll_interval_seconds)
|
||||
poll_status, poll_payload = await _get_json(
|
||||
poll_url,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
hard_timeout_seconds=request_timeout_seconds,
|
||||
)
|
||||
if poll_status >= 400:
|
||||
raise RuntimeError(f"OpenRouter video poll failed with HTTP {poll_status}: {poll_payload}")
|
||||
last_payload = poll_payload
|
||||
status = poll_payload.get("status")
|
||||
last_status = status
|
||||
_emit_progress(progress, "video_status", f"OpenRouter video generation status: {status}", {"model": self.model, "job_id": job_id, "status": status})
|
||||
|
||||
if status == "completed":
|
||||
urls = poll_payload.get("unsigned_urls") or []
|
||||
if urls:
|
||||
content_url = urls[0]
|
||||
else:
|
||||
content_url = f"{self.base_url}/videos/{job_id}/content?index=0"
|
||||
_emit_progress(progress, "video_download_start", "Downloading OpenRouter video output", {"model": self.model, "job_id": job_id})
|
||||
download_status, data = await _get_bytes(
|
||||
content_url,
|
||||
headers=headers if _needs_authorization(content_url) else {},
|
||||
timeout=timeout,
|
||||
hard_timeout_seconds=request_timeout_seconds,
|
||||
)
|
||||
if download_status >= 400:
|
||||
raise RuntimeError(f"OpenRouter video content download failed with HTTP {download_status}: {data[:500]!r}")
|
||||
_emit_progress(progress, "video_completed", "OpenRouter video generation completed and downloaded", {"model": self.model, "job_id": job_id})
|
||||
return VideoOutput(fmt="bytes", ext="mp4", data=data)
|
||||
if status in {"failed", "cancelled", "expired"}:
|
||||
raise RuntimeError(f"OpenRouter video generation {status} for job {job_id}: {poll_payload.get('error') or poll_payload}")
|
||||
|
||||
raise RuntimeError(f"OpenRouter video generation timed out after {query_timeout_seconds:g}s for job {job_id}; last_status={last_status}; last_payload={last_payload}")
|
||||
|
||||
def _headers(self) -> dict[str, str]:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if self.http_referer:
|
||||
headers["HTTP-Referer"] = self.http_referer
|
||||
if self.app_title:
|
||||
headers["X-OpenRouter-Title"] = self.app_title
|
||||
return headers
|
||||
|
||||
|
||||
def _frame_images(reference_image_paths: List[str]) -> list[dict]:
|
||||
if len(reference_image_paths) > 2:
|
||||
raise ValueError("OpenRouter video generation supports at most first and last frame images")
|
||||
frame_types = ["first_frame", "last_frame"]
|
||||
return [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_path_to_b64(path, mime=True)},
|
||||
"frame_type": frame_types[index],
|
||||
}
|
||||
for index, path in enumerate(reference_image_paths)
|
||||
]
|
||||
|
||||
|
||||
def _absolute_url(base_url: str, url: str) -> str:
|
||||
if url.startswith("http://") or url.startswith("https://"):
|
||||
return url
|
||||
return urljoin(f"{base_url.rstrip('/')}/", url.lstrip("/"))
|
||||
|
||||
|
||||
def _needs_authorization(url: str) -> bool:
|
||||
return url.startswith("https://openrouter.ai/api/")
|
||||
|
||||
|
||||
async def _post_json(url: str, *, headers: dict[str, str], payload: dict, timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, dict]:
|
||||
async def request() -> tuple[int, dict]:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.post(url, headers=headers, json=payload) as response:
|
||||
return response.status, await response.json(content_type=None)
|
||||
|
||||
return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5)
|
||||
|
||||
|
||||
async def _get_json(url: str, *, headers: dict[str, str], timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, dict]:
|
||||
async def request() -> tuple[int, dict]:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(url, headers=headers) as response:
|
||||
return response.status, await response.json(content_type=None)
|
||||
|
||||
return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5)
|
||||
|
||||
|
||||
async def _get_bytes(url: str, *, headers: dict[str, str], timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, bytes]:
|
||||
async def request() -> tuple[int, bytes]:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(url, headers=headers) as response:
|
||||
return response.status, await response.read()
|
||||
|
||||
return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5)
|
||||
@@ -0,0 +1,116 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
import asyncio
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from google.genai.errors import ClientError
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.rate_limiter import RateLimiter
|
||||
|
||||
# https://ai.google.dev/gemini-api/docs/video-generation?hl=zh-cn
|
||||
|
||||
|
||||
class VideoGeneratorVeoGoogleAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
t2v_model: str = "veo-3.1-generate-preview",
|
||||
ff2v_model: str = "veo-3.1-generate-preview",
|
||||
flf2v_model: str = "veo-3.1-generate-preview",
|
||||
rate_limiter: Optional[RateLimiter] = None,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.t2v_model = t2v_model
|
||||
self.ff2v_model = ff2v_model
|
||||
self.flf2v_model = flf2v_model
|
||||
self.rate_limiter = rate_limiter
|
||||
|
||||
self.client = genai.Client(
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
resolution: str = "1080p",
|
||||
aspect_ratio: str = "16:9",
|
||||
duration: int = 8,
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
|
||||
params = {
|
||||
"prompt": prompt,
|
||||
}
|
||||
config_params = {
|
||||
"resolution": resolution,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"duration_seconds": duration,
|
||||
}
|
||||
if len(reference_image_paths) == 0:
|
||||
params["model"] = self.t2v_model
|
||||
elif len(reference_image_paths) == 1:
|
||||
params["model"] = self.ff2v_model
|
||||
params["image"] = types.Image.from_file(location=reference_image_paths[0])
|
||||
elif len(reference_image_paths) == 2:
|
||||
params["model"] = self.flf2v_model
|
||||
params["image"] = types.Image.from_file(location=reference_image_paths[0])
|
||||
config_params["last_frame"] = types.Image.from_file(location=reference_image_paths[1])
|
||||
else:
|
||||
raise ValueError("The number of reference images must be no more than 2")
|
||||
|
||||
logging.info(f"Calling {params['model']} to generate video...")
|
||||
|
||||
# Apply rate limiting if configured
|
||||
if self.rate_limiter:
|
||||
await self.rate_limiter.acquire()
|
||||
|
||||
# Retry logic for rate limit errors
|
||||
max_retries = 3
|
||||
retry_delay = 5
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
operation = self.client.models.generate_videos(
|
||||
**params,
|
||||
config=types.GenerateVideosConfig(**config_params),
|
||||
)
|
||||
break
|
||||
except ClientError as e:
|
||||
if e.status_code == 429 and attempt < max_retries - 1:
|
||||
wait_time = retry_delay * (2 ** attempt)
|
||||
logging.warning(f"Rate limit hit (429), retrying in {wait_time}s... (attempt {attempt + 1}/{max_retries})")
|
||||
await asyncio.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
|
||||
while not operation.done:
|
||||
await asyncio.sleep(2)
|
||||
operation = self.client.operations.get(operation)
|
||||
logging.info(f"Video generation not completed, waiting 2 seconds...")
|
||||
|
||||
# Check if operation completed successfully
|
||||
if operation.error:
|
||||
error_msg = f"Video generation failed: {operation.error}"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
if not operation.response:
|
||||
error_msg = "Video generation completed but no response received"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
if not hasattr(operation.response, 'generated_videos') or not operation.response.generated_videos:
|
||||
error_msg = "Video generation completed but no videos were generated"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
generated_video = operation.response.generated_videos[0]
|
||||
self.client.files.download(file=generated_video.video)
|
||||
|
||||
video_output = VideoOutput(
|
||||
fmt="bytes",
|
||||
ext="mp4",
|
||||
data=generated_video.video.video_bytes,
|
||||
)
|
||||
return video_output
|
||||
@@ -0,0 +1,177 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from PIL import Image
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import os
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.image import image_path_to_b64
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
try:
|
||||
return max(0, int(os.environ.get(name, str(default))))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _env_float(name: str, default: float) -> float:
|
||||
try:
|
||||
return max(0.0, float(os.environ.get(name, str(default))))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _emit_progress(progress, stage: str, message: str, metadata: dict | None = None) -> None:
|
||||
if progress is not None:
|
||||
progress(stage, message, metadata or {})
|
||||
|
||||
|
||||
class VideoGeneratorVeoYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
t2v_model: str = "veo3.1-fast", # text to video
|
||||
ff2v_model: str = "veo3.1-fast", # first frame to video
|
||||
flf2v_model: str = "veo2-fast-frames", # first and last frame to video
|
||||
base_url: str = "https://yunwu.ai",
|
||||
):
|
||||
"""
|
||||
all models:
|
||||
veo2
|
||||
veo2-fast
|
||||
veo2-fast-frames
|
||||
veo2-fast-components
|
||||
veo2-pro
|
||||
veo3
|
||||
veo3-fast
|
||||
veo3-pro
|
||||
veo3-pro-frames
|
||||
veo3-fast-frames
|
||||
veo3-frames
|
||||
|
||||
NOTE: veo3 does not support first and last frame to video generation.
|
||||
"""
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.t2v_model = t2v_model
|
||||
self.ff2v_model = ff2v_model
|
||||
self.flf2v_model = flf2v_model
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str = "",
|
||||
reference_image_paths: List[Image.Image] = [],
|
||||
aspect_ratio: str = "16:9",
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
progress = kwargs.get("progress")
|
||||
create_retries = _env_int("VIMAX_VIDEO_CREATE_RETRIES", 3)
|
||||
query_timeout_seconds = _env_float("VIMAX_VIDEO_QUERY_TIMEOUT_SECONDS", 600.0)
|
||||
request_timeout_seconds = _env_float("VIMAX_VIDEO_REQUEST_TIMEOUT_SECONDS", 60.0)
|
||||
poll_interval_seconds = _env_float("VIMAX_VIDEO_POLL_INTERVAL_SECONDS", 5.0)
|
||||
max_query_errors = _env_int("VIMAX_VIDEO_MAX_QUERY_ERRORS", 5)
|
||||
if len(reference_image_paths) == 0:
|
||||
model = self.t2v_model
|
||||
elif len(reference_image_paths) == 1:
|
||||
model = self.ff2v_model
|
||||
elif len(reference_image_paths) == 2:
|
||||
model = self.flf2v_model
|
||||
else:
|
||||
raise ValueError("The number of reference images must be no more than 2")
|
||||
|
||||
logging.info(f"Calling {model} to generate video...")
|
||||
|
||||
# 1. Create video generation task
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"images": [image_path_to_b64(image_path, mime=True) for image_path in reference_image_paths],
|
||||
"enhance_prompt": True,
|
||||
}
|
||||
# only veo3 supports aspect ratio setting
|
||||
if model.startswith("veo3"):
|
||||
payload["aspect_ratio"] = aspect_ratio
|
||||
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
url = f"{self.base_url}/v1/video/create"
|
||||
task_id = None
|
||||
last_create_error = None
|
||||
timeout = aiohttp.ClientTimeout(total=request_timeout_seconds)
|
||||
for attempt in range(1, create_retries + 1):
|
||||
try:
|
||||
_emit_progress(progress, "video_create", f"Creating video generation task with {model}", {"model": model, "attempt": attempt, "max_attempts": create_retries})
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.post(url, headers=headers, json=payload) as response:
|
||||
response_payload = await response.json(content_type=None)
|
||||
logging.debug(f"Response: {response_payload}")
|
||||
if response.status >= 400:
|
||||
raise RuntimeError(f"Video create failed with HTTP {response.status}: {response_payload}")
|
||||
task_id = response_payload.get("id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Video create response missing id: {response_payload}")
|
||||
logging.info(f"Video generation task created successfully. Task ID: {task_id}")
|
||||
_emit_progress(progress, "video_task_created", "Video generation task created", {"model": model, "task_id": task_id})
|
||||
break
|
||||
except Exception as e:
|
||||
last_create_error = e
|
||||
logging.error(f"Error occurred while creating video generation task: {e}.")
|
||||
_emit_progress(progress, "video_create_error", f"Video create attempt {attempt} failed", {"model": model, "attempt": attempt, "error": str(e)})
|
||||
if attempt < create_retries:
|
||||
await asyncio.sleep(1)
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Video create failed after {create_retries} attempts: {last_create_error}")
|
||||
|
||||
|
||||
# 2. Query the video generation task until the video generation is completed
|
||||
headers = {
|
||||
'Accept': 'application/json',
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
}
|
||||
|
||||
deadline = asyncio.get_running_loop().time() + query_timeout_seconds if query_timeout_seconds > 0 else None
|
||||
query_errors = 0
|
||||
last_status = None
|
||||
while deadline is None or asyncio.get_running_loop().time() < deadline:
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(f"{self.base_url}/v1/video/query?id={task_id}", headers=headers) as response:
|
||||
payload = await response.json(content_type=None)
|
||||
logging.debug(f"Response: {payload}")
|
||||
if response.status >= 400:
|
||||
raise RuntimeError(f"Video query failed with HTTP {response.status}: {payload}")
|
||||
status = payload.get("status")
|
||||
if not status:
|
||||
raise RuntimeError(f"Video query response missing status: {payload}")
|
||||
query_errors = 0
|
||||
except Exception as e:
|
||||
query_errors += 1
|
||||
logging.error(f"Error occurred while querying video generation task: {e}.")
|
||||
_emit_progress(progress, "video_query_error", "Video query failed", {"model": model, "task_id": task_id, "error": str(e), "query_errors": query_errors, "max_query_errors": max_query_errors})
|
||||
if query_errors >= max_query_errors:
|
||||
raise RuntimeError(f"Video query failed {query_errors} times for task {task_id}: {e}")
|
||||
await asyncio.sleep(poll_interval_seconds)
|
||||
continue
|
||||
|
||||
if status == "completed":
|
||||
logging.info(f"Video generation completed successfully")
|
||||
video_url = payload.get("video_url")
|
||||
if not video_url:
|
||||
raise RuntimeError(f"Video task completed without video_url: {payload}")
|
||||
_emit_progress(progress, "video_completed", "Video generation completed", {"model": model, "task_id": task_id})
|
||||
return VideoOutput(fmt="url", ext="mp4", data=video_url)
|
||||
elif status == "failed":
|
||||
logging.error(f"Video generation failed: \n{payload}")
|
||||
raise RuntimeError(f"Video generation failed for task {task_id}: {payload}")
|
||||
else:
|
||||
logging.info(f"Video generation status: {status}, waiting 1 second...")
|
||||
last_status = status
|
||||
_emit_progress(progress, "video_status", f"Video generation status: {status}", {"model": model, "task_id": task_id, "status": status})
|
||||
await asyncio.sleep(poll_interval_seconds)
|
||||
continue
|
||||
raise RuntimeError(f"Video generation timed out after {query_timeout_seconds:g}s for task {task_id}; last_status={last_status}")
|
||||
Generated
+1196
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"name": "vimax-ui",
|
||||
"private": true,
|
||||
"version": "0.1.0",
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"tui": "tsx src/cli.tsx",
|
||||
"test": "tsx src/lineMapping.test.ts && tsx src/slashCommands.test.ts && tsx src/workspaceMeta.test.ts"
|
||||
},
|
||||
"dependencies": {
|
||||
"ink": "^4.4.1",
|
||||
"ink-text-input": "^5.0.1",
|
||||
"react": "^18.3.1",
|
||||
"react-dom": "^18.3.1"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^24.5.2",
|
||||
"@types/react": "^18.3.12",
|
||||
"@types/react-dom": "^18.3.1",
|
||||
"tsx": "^4.20.5",
|
||||
"typescript": "^5.6.3"
|
||||
}
|
||||
}
|
||||
+601
@@ -0,0 +1,601 @@
|
||||
import React, {useEffect, useMemo, useRef, useState} from 'react';
|
||||
import {render, Box, Text, useApp, useInput, useStdout} from 'ink';
|
||||
import stringWidth from 'string-width';
|
||||
import {spawn, type ChildProcessWithoutNullStreams} from 'node:child_process';
|
||||
import {existsSync} from 'node:fs';
|
||||
import path from 'node:path';
|
||||
import process from 'node:process';
|
||||
import {fileURLToPath} from 'node:url';
|
||||
import {applyStreamEvent, createMappingState} from './lineMapping.js';
|
||||
import {matchingSlashCommands, shouldShowSlashCommands} from './slashCommands.js';
|
||||
import {compactionLabel, compactTargetFromEnv, resolveWorkspacePath, type WorkspaceMeta} from './workspaceMeta.js';
|
||||
import type {MappingState, StreamEvent, WorkspaceLine} from './types.js';
|
||||
|
||||
const __dirname = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(__dirname, '..', '..');
|
||||
|
||||
const THINKING_FRAMES = ['', '.', '..', '...'];
|
||||
|
||||
const WORKSPACE_BORDER_COLORS = ['blue', 'blueBright', 'cyan', 'blueBright', 'blue'];
|
||||
|
||||
type CliOptions = {
|
||||
agentArgs: string[];
|
||||
};
|
||||
|
||||
const cliOptions = parseCliArgs(process.argv.slice(2));
|
||||
|
||||
function parseCliArgs(argv: string[]): CliOptions {
|
||||
const agentArgs: string[] = [];
|
||||
for (let index = 0; index < argv.length; index += 1) {
|
||||
const arg = argv[index];
|
||||
if (arg === '--new-session') {
|
||||
agentArgs.push('--new-session');
|
||||
continue;
|
||||
}
|
||||
if (arg === '--session') {
|
||||
const sessionId = argv[index + 1];
|
||||
if (!sessionId) throw new Error('--session requires a session id');
|
||||
agentArgs.push('--session', sessionId);
|
||||
index += 1;
|
||||
continue;
|
||||
}
|
||||
if (arg === '--help' || arg === '-h') {
|
||||
printHelpAndExit();
|
||||
}
|
||||
throw new Error(`Unknown TUI argument: ${arg}`);
|
||||
}
|
||||
return {agentArgs};
|
||||
}
|
||||
|
||||
function printHelpAndExit(): never {
|
||||
console.log(`Usage:
|
||||
./vimax tui
|
||||
./vimax tui new
|
||||
./vimax tui resume [session_id]
|
||||
|
||||
Direct TUI args:
|
||||
--new-session create and activate a new empty session
|
||||
--session <id> activate an existing session`);
|
||||
process.exit(0);
|
||||
}
|
||||
|
||||
function gradientColor(index: number, total: number): string {
|
||||
if (total <= 1) return WORKSPACE_BORDER_COLORS[0] ?? 'blue';
|
||||
const scaled = (index / (total - 1)) * (WORKSPACE_BORDER_COLORS.length - 1);
|
||||
return WORKSPACE_BORDER_COLORS[Math.min(WORKSPACE_BORDER_COLORS.length - 1, Math.max(0, Math.round(scaled)))] ?? 'blue';
|
||||
}
|
||||
|
||||
function useThinkingFrame(active: boolean): string {
|
||||
const [frame, setFrame] = useState(0);
|
||||
useEffect(() => {
|
||||
if (!active) {
|
||||
setFrame(0);
|
||||
return;
|
||||
}
|
||||
const timer = setInterval(() => setFrame((value) => (value + 1) % THINKING_FRAMES.length), 220);
|
||||
return () => clearInterval(timer);
|
||||
}, [active]);
|
||||
return THINKING_FRAMES[frame];
|
||||
}
|
||||
|
||||
function useTerminalWidth(stdout: NodeJS.WriteStream): number {
|
||||
const [terminal, setTerminal] = useState({width: Math.max(20, stdout.columns || 100), revision: 0});
|
||||
useEffect(() => {
|
||||
let resizeTimer: NodeJS.Timeout | null = null;
|
||||
const redraw = (clear: boolean) => {
|
||||
if (clear) {
|
||||
// Ink does not always erase cells from the previous frame when the
|
||||
// terminal is resized quickly. Clear only after resize, then force a
|
||||
// render even when the new width equals the previous width.
|
||||
stdout.write('\u001b[2J\u001b[3J\u001b[H');
|
||||
}
|
||||
setTerminal((current) => ({width: Math.max(20, stdout.columns || 100), revision: current.revision + 1}));
|
||||
};
|
||||
const update = () => {
|
||||
if (resizeTimer) clearTimeout(resizeTimer);
|
||||
resizeTimer = setTimeout(() => redraw(true), 60);
|
||||
};
|
||||
redraw(false);
|
||||
stdout.on('resize', update);
|
||||
return () => {
|
||||
if (resizeTimer) clearTimeout(resizeTimer);
|
||||
stdout.off('resize', update);
|
||||
};
|
||||
}, [stdout]);
|
||||
return terminal.width;
|
||||
}
|
||||
|
||||
|
||||
function baseAgentArgs(): string[] {
|
||||
return ['main_agent.py', '--jsonl', '--stdin-repl', ...cliOptions.agentArgs];
|
||||
}
|
||||
|
||||
function agentCommand(): {command: string; args: string[]} {
|
||||
if (process.env.VIMAX_AGENT_COMMAND) {
|
||||
return {command: process.env.VIMAX_AGENT_COMMAND, args: splitArgs(process.env.VIMAX_AGENT_ARGS ?? '')};
|
||||
}
|
||||
if (process.env.VIMAX_PYTHON_CMD) {
|
||||
return {command: process.env.VIMAX_PYTHON_CMD, args: baseAgentArgs()};
|
||||
}
|
||||
const bundledUv = process.env.VIMAX_UV_CMD ?? '/home/xavierhuang/.local/bin/uv';
|
||||
if (existsSync(bundledUv)) {
|
||||
return {command: bundledUv, args: ['run', 'python', ...baseAgentArgs()]};
|
||||
}
|
||||
const venvPython = path.join(repoRoot, '.venv', 'bin', 'python3');
|
||||
if (existsSync(venvPython)) {
|
||||
return {command: venvPython, args: baseAgentArgs()};
|
||||
}
|
||||
return {command: 'uv', args: ['run', 'python', ...baseAgentArgs()]};
|
||||
}
|
||||
|
||||
|
||||
function splitArgs(value: string): string[] {
|
||||
return value.split(/\s+/).map((part) => part.trim()).filter(Boolean);
|
||||
}
|
||||
|
||||
function App() {
|
||||
const {exit} = useApp();
|
||||
const {stdout} = useStdout();
|
||||
const terminalWidth = useTerminalWidth(stdout);
|
||||
const [lines, setLines] = useState<WorkspaceLine[]>([]);
|
||||
const [input, setInput] = useState('');
|
||||
const [cursor, setCursor] = useState(0);
|
||||
const inputRef = useRef('');
|
||||
const cursorRef = useRef(0);
|
||||
const [busy, setBusy] = useState(false);
|
||||
const [activityText, setActivityText] = useState('ViMax thinking');
|
||||
const [workspaceMeta, setWorkspaceMeta] = useState<WorkspaceMeta>({
|
||||
workspacePath: '.working_dir',
|
||||
sessionId: '',
|
||||
stage: '',
|
||||
compactionUsed: 0,
|
||||
compactionTarget: compactTargetFromEnv(process.env),
|
||||
});
|
||||
const stateRef = useRef<MappingState>(createMappingState());
|
||||
const childRef = useRef<ChildProcessWithoutNullStreams | null>(null);
|
||||
const bufferRef = useRef('');
|
||||
const responseIdleTimerRef = useRef<NodeJS.Timeout | null>(null);
|
||||
|
||||
const width = useMemo(() => Math.max(20, terminalWidth - 6), [terminalWidth]);
|
||||
const thinkingFrame = useThinkingFrame(busy);
|
||||
const slashMatches = useMemo(() => matchingSlashCommands(input), [input]);
|
||||
const showSlashPopup = shouldShowSlashCommands(input, busy);
|
||||
|
||||
useEffect(() => {
|
||||
inputRef.current = input;
|
||||
const length = Array.from(input).length;
|
||||
if (cursorRef.current > length) {
|
||||
cursorRef.current = length;
|
||||
setCursor(length);
|
||||
}
|
||||
}, [input]);
|
||||
|
||||
useEffect(() => {
|
||||
cursorRef.current = cursor;
|
||||
}, [cursor]);
|
||||
|
||||
function updateInput(next: string, nextCursor: number) {
|
||||
const length = Array.from(next).length;
|
||||
const boundedCursor = Math.max(0, Math.min(nextCursor, length));
|
||||
inputRef.current = next;
|
||||
cursorRef.current = boundedCursor;
|
||||
setInput(next);
|
||||
setCursor(boundedCursor);
|
||||
}
|
||||
|
||||
useInput((value, key) => {
|
||||
if (key.ctrl && value === 'c') {
|
||||
childRef.current?.kill();
|
||||
exit();
|
||||
return;
|
||||
}
|
||||
if (busy) return;
|
||||
const currentChars = Array.from(inputRef.current);
|
||||
const currentCursor = Math.max(0, Math.min(cursorRef.current, currentChars.length));
|
||||
if (key.leftArrow) {
|
||||
updateInput(inputRef.current, currentCursor - 1);
|
||||
return;
|
||||
}
|
||||
if (key.rightArrow) {
|
||||
updateInput(inputRef.current, currentCursor + 1);
|
||||
return;
|
||||
}
|
||||
if ((key as {home?: boolean}).home) {
|
||||
updateInput(inputRef.current, 0);
|
||||
return;
|
||||
}
|
||||
if ((key as {end?: boolean}).end) {
|
||||
updateInput(inputRef.current, currentChars.length);
|
||||
return;
|
||||
}
|
||||
if (value.includes('\r') || value.includes('\n')) {
|
||||
const [beforeBreak] = value.split(/[\r\n]/, 1);
|
||||
const pasted = Array.from(beforeBreak ?? '');
|
||||
const next = [...currentChars.slice(0, currentCursor), ...pasted, ...currentChars.slice(currentCursor)].join('');
|
||||
submit(next);
|
||||
return;
|
||||
}
|
||||
if (key.return) {
|
||||
submit(inputRef.current);
|
||||
return;
|
||||
}
|
||||
const isBackspace = key.backspace || value === '\u007f' || value === '\b' || value === '\u001b\u007f';
|
||||
const isDelete = key.delete || value === '\u001b[3~' || value === '\u001b[P';
|
||||
if (isBackspace) {
|
||||
if (currentCursor === 0) return;
|
||||
const next = [...currentChars.slice(0, currentCursor - 1), ...currentChars.slice(currentCursor)].join('');
|
||||
updateInput(next, currentCursor - 1);
|
||||
return;
|
||||
}
|
||||
if (isDelete) {
|
||||
if (currentCursor < currentChars.length) {
|
||||
const next = [...currentChars.slice(0, currentCursor), ...currentChars.slice(currentCursor + 1)].join('');
|
||||
updateInput(next, currentCursor);
|
||||
return;
|
||||
}
|
||||
if (currentCursor > 0) {
|
||||
const next = [...currentChars.slice(0, currentCursor - 1), ...currentChars.slice(currentCursor)].join('');
|
||||
updateInput(next, currentCursor - 1);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (!key.ctrl && !key.meta && value) {
|
||||
const inserted = Array.from(value);
|
||||
const next = [...currentChars.slice(0, currentCursor), ...inserted, ...currentChars.slice(currentCursor)].join('');
|
||||
updateInput(next, currentCursor + inserted.length);
|
||||
}
|
||||
});
|
||||
|
||||
useEffect(() => {
|
||||
const {command, args} = agentCommand();
|
||||
const child = spawn(command, args, {cwd: repoRoot, env: process.env});
|
||||
childRef.current = child;
|
||||
|
||||
child.stdout.setEncoding('utf8');
|
||||
child.stdout.on('data', (chunk: string) => {
|
||||
bufferRef.current += chunk;
|
||||
const parts = bufferRef.current.split('\n');
|
||||
bufferRef.current = parts.pop() ?? '';
|
||||
for (const part of parts) {
|
||||
consumeJsonLine(part);
|
||||
}
|
||||
});
|
||||
|
||||
child.stderr.setEncoding('utf8');
|
||||
child.stderr.on('data', (chunk: string) => {
|
||||
for (const line of chunk.split('\n')) {
|
||||
if (!line.trim()) continue;
|
||||
appendLine({kind: 'terminal', text: `[stderr]: ${line}`});
|
||||
}
|
||||
});
|
||||
|
||||
child.on('error', (error) => {
|
||||
appendLine({kind: 'error', text: `agent process error: ${error.message}`});
|
||||
});
|
||||
|
||||
child.on('exit', (code, signal) => {
|
||||
childRef.current = null;
|
||||
setBusy(false);
|
||||
if (code && code !== 0) {
|
||||
appendLine({kind: 'error', text: `agent process exited with code ${code}`});
|
||||
} else if (signal) {
|
||||
appendLine({kind: 'status', text: `agent process stopped by ${signal}`});
|
||||
}
|
||||
});
|
||||
|
||||
return () => {
|
||||
if (responseIdleTimerRef.current) clearTimeout(responseIdleTimerRef.current);
|
||||
child.kill();
|
||||
};
|
||||
}, []);
|
||||
|
||||
function appendLine(line: WorkspaceLine) {
|
||||
stateRef.current = createMappingState();
|
||||
setLines((current) => [...current, line]);
|
||||
}
|
||||
|
||||
function stripThinking(lines: WorkspaceLine[]): WorkspaceLine[] {
|
||||
return lines.filter((line) => line.kind !== 'thinking');
|
||||
}
|
||||
|
||||
function consumeJsonLine(line: string) {
|
||||
const trimmed = line.trim();
|
||||
if (!trimmed) return;
|
||||
let event: StreamEvent;
|
||||
try {
|
||||
event = JSON.parse(trimmed) as StreamEvent;
|
||||
} catch (error) {
|
||||
appendLine({kind: 'error', text: `invalid JSONL event: ${trimmed}`});
|
||||
return;
|
||||
}
|
||||
updateWorkspaceMeta(event);
|
||||
updateActivity(event);
|
||||
if (event.type === 'done' || event.type === 'error' || event.type === 'session') {
|
||||
clearResponseIdleTimer();
|
||||
setBusy(false);
|
||||
}
|
||||
setLines((current) => {
|
||||
const mapped = applyStreamEvent(stripThinking(current), stateRef.current, event);
|
||||
stateRef.current = mapped.state;
|
||||
return mapped.lines;
|
||||
});
|
||||
}
|
||||
|
||||
function clearResponseIdleTimer() {
|
||||
if (!responseIdleTimerRef.current) return;
|
||||
clearTimeout(responseIdleTimerRef.current);
|
||||
responseIdleTimerRef.current = null;
|
||||
}
|
||||
|
||||
function scheduleResponseIdleClear() {
|
||||
clearResponseIdleTimer();
|
||||
responseIdleTimerRef.current = setTimeout(() => {
|
||||
responseIdleTimerRef.current = null;
|
||||
setBusy(false);
|
||||
}, 1500);
|
||||
}
|
||||
|
||||
function updateActivity(event: StreamEvent) {
|
||||
if (event.type === 'tool_start') {
|
||||
clearResponseIdleTimer();
|
||||
setActivityText(`tool ${event.tool?.name ?? 'unknown'} running`);
|
||||
return;
|
||||
}
|
||||
if (event.type === 'tool_progress') {
|
||||
clearResponseIdleTimer();
|
||||
const stage = event.progress?.stage;
|
||||
setActivityText(stage ? `tool ${event.tool?.name ?? 'unknown'}: ${stage}` : `tool ${event.tool?.name ?? 'unknown'} running`);
|
||||
return;
|
||||
}
|
||||
if (event.type === 'tool_result') {
|
||||
clearResponseIdleTimer();
|
||||
setActivityText('ViMax thinking');
|
||||
return;
|
||||
}
|
||||
if (event.type === 'token') {
|
||||
setActivityText('ViMax responding');
|
||||
scheduleResponseIdleClear();
|
||||
return;
|
||||
}
|
||||
if (event.type === 'status') {
|
||||
clearResponseIdleTimer();
|
||||
setActivityText(statusActivityLabel(event.phase, event.message));
|
||||
return;
|
||||
}
|
||||
if (event.type === 'done' || event.type === 'error' || event.type === 'session') {
|
||||
clearResponseIdleTimer();
|
||||
setActivityText('ViMax thinking');
|
||||
return;
|
||||
}
|
||||
if (event.type === 'turn') {
|
||||
clearResponseIdleTimer();
|
||||
setActivityText('ViMax thinking');
|
||||
}
|
||||
}
|
||||
|
||||
function updateWorkspaceMeta(event: StreamEvent) {
|
||||
if (event.type === 'prompt_trace') {
|
||||
const used = event.prompt_trace?.totals?.total_tokens ?? event.prompt_trace?.totals?.total_estimated_tokens ?? event.prompt_trace?.total_estimated_tokens;
|
||||
if (typeof used === 'number' && Number.isFinite(used)) {
|
||||
setWorkspaceMeta((current) => {
|
||||
const nextUsed = Math.max(0, Math.round(used));
|
||||
const currentPercent = current.compactionTarget > 0 ? Math.round((current.compactionUsed / current.compactionTarget) * 100) : 0;
|
||||
const nextPercent = current.compactionTarget > 0 ? Math.round((nextUsed / current.compactionTarget) * 100) : 0;
|
||||
if (currentPercent === nextPercent && Math.abs(nextUsed - current.compactionUsed) < 100) return current;
|
||||
return {...current, compactionUsed: nextUsed};
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (event.type === 'session') {
|
||||
const session = event.session?.session;
|
||||
if (!session) return;
|
||||
setWorkspaceMeta((current) => ({
|
||||
...current,
|
||||
workspacePath: resolveWorkspacePath(repoRoot, session.working_dir),
|
||||
sessionId: session.session_id ?? current.sessionId,
|
||||
stage: session.stage ?? current.stage,
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
function submit(value: string) {
|
||||
const prompt = value.trim();
|
||||
if (!prompt || busy) return;
|
||||
const child = childRef.current;
|
||||
if (!child || child.killed || !child.stdin.writable) {
|
||||
appendLine({kind: 'error', text: 'agent process is not available'});
|
||||
return;
|
||||
}
|
||||
setLines((current) => [...stripThinking(current), {kind: 'user', text: prompt}]);
|
||||
clearResponseIdleTimer();
|
||||
setActivityText('ViMax thinking');
|
||||
stateRef.current = createMappingState();
|
||||
updateInput('', 0);
|
||||
setBusy(true);
|
||||
child.stdin.write(`${prompt}\n`);
|
||||
}
|
||||
|
||||
return (
|
||||
<Box flexDirection="column" paddingX={1} width={Math.max(20, terminalWidth - 4)}>
|
||||
<WorkspacePanel lines={lines} width={width} thinkingFrame={thinkingFrame} meta={workspaceMeta} busy={busy} activityText={activityText} />
|
||||
{showSlashPopup && <SlashCommandPopup matches={slashMatches} width={width} />}
|
||||
<Box borderStyle="round" borderColor="white" paddingX={1} marginTop={1} width={width}>
|
||||
<Text color={busy ? 'gray' : 'white'}>{busy ? '· ' : '› '}</Text>
|
||||
<InputText value={input} cursor={cursor} busy={busy} />
|
||||
</Box>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
function InputText({value, cursor, busy}: {value: string; cursor: number; busy: boolean}) {
|
||||
const chars = Array.from(value);
|
||||
const boundedCursor = Math.max(0, Math.min(cursor, chars.length));
|
||||
const before = chars.slice(0, boundedCursor).join('');
|
||||
const current = chars[boundedCursor] ?? ' ';
|
||||
const after = chars.slice(boundedCursor + 1).join('');
|
||||
if (busy) {
|
||||
return <Text color="gray">{value}</Text>;
|
||||
}
|
||||
return (
|
||||
<Text>
|
||||
<Text color="white">{before}</Text>
|
||||
<Text color="black" backgroundColor="white">{current}</Text>
|
||||
<Text color="white">{after}</Text>
|
||||
</Text>
|
||||
);
|
||||
}
|
||||
|
||||
function SlashCommandPopup({matches, width}: {matches: ReturnType<typeof matchingSlashCommands>; width: number}) {
|
||||
const panelWidth = Math.max(20, width);
|
||||
const visibleMatches = matches.slice(0, 6);
|
||||
return (
|
||||
<Box flexDirection="column" borderStyle="round" borderColor="blueBright" paddingX={1} marginTop={1} width={panelWidth}>
|
||||
{visibleMatches.length > 0 ? (
|
||||
visibleMatches.map((command) => (
|
||||
<Text key={command.name}>
|
||||
<Text color="cyanBright">{command.matchedPrefix}</Text>
|
||||
<Text color="blueBright">{command.unmatchedSuffix}</Text>
|
||||
<Text color="gray"> {command.description}</Text>
|
||||
</Text>
|
||||
))
|
||||
) : (
|
||||
<Text color="gray">No matching slash commands</Text>
|
||||
)}
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
function WorkspacePanel({lines, width, thinkingFrame, meta, busy, activityText}: {lines: WorkspaceLine[]; width: number; thinkingFrame: string; meta: WorkspaceMeta; busy: boolean; activityText: string}) {
|
||||
const panelWidth = Math.max(20, width);
|
||||
const contentWidth = Math.max(1, panelWidth - 4);
|
||||
return (
|
||||
<Box flexDirection="column" width={panelWidth}>
|
||||
<GradientBorderLine left="╭" fill="─" right="╮" width={panelWidth} />
|
||||
<WorkspaceContentLine text="ViMax Workspace" color="blueBright" width={panelWidth} />
|
||||
{workspaceHeaderLines(meta, contentWidth).map((line, index) => (
|
||||
<WorkspaceContentLine key={`header-${index}`} text={line.text} color={line.color} width={panelWidth} />
|
||||
))}
|
||||
{lines.flatMap((line, index) => {
|
||||
const rawText = `› ${line.text}`;
|
||||
return wrapText(rawText, contentWidth).map((part, partIndex) => (
|
||||
<WorkspaceContentLine key={`${line.kind}-${index}-${partIndex}`} text={part} color={lineColor(line)} width={panelWidth} />
|
||||
));
|
||||
})}
|
||||
{busy && wrapText(`› ${activityText}${thinkingFrame}`, contentWidth).map((part, index) => (
|
||||
<WorkspaceContentLine key={`activity-${index}`} text={part} color="cyanBright" width={panelWidth} />
|
||||
))}
|
||||
<GradientBorderLine left="╰" fill="─" right="╯" width={panelWidth} />
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
function workspaceHeaderLines(meta: WorkspaceMeta, width: number): Array<{text: string; color: string}> {
|
||||
const rows: Array<{text: string; color: string}> = [];
|
||||
for (const part of wrapText(`Path: ${meta.workspacePath}`, width)) {
|
||||
rows.push({text: part, color: 'gray'});
|
||||
}
|
||||
const session = [meta.sessionId, displayStage(meta.stage)].filter(Boolean).join(' · ');
|
||||
if (session) {
|
||||
for (const part of wrapText(`Session: ${session}`, width)) {
|
||||
rows.push({text: part, color: 'gray'});
|
||||
}
|
||||
}
|
||||
for (const part of wrapText(compactionLabel(meta.compactionUsed, meta.compactionTarget), width)) {
|
||||
rows.push({text: part, color: 'cyanBright'});
|
||||
}
|
||||
return rows;
|
||||
}
|
||||
|
||||
function statusActivityLabel(phase: string | undefined, message: string | undefined): string {
|
||||
if (phase === 'compact') return 'compacting context';
|
||||
if (phase === 'sampling_assistant') return 'ViMax thinking';
|
||||
if (phase === 'executing_tools') return 'running tools';
|
||||
const normalized = String(message ?? '').trim();
|
||||
return normalized || 'ViMax thinking';
|
||||
}
|
||||
|
||||
function displayStage(stage: string): string {
|
||||
const labels: Record<string, string> = {
|
||||
created: 'Created',
|
||||
narrative_planning: 'Planning text',
|
||||
narrative_planned: 'Text planned',
|
||||
novel_planning: 'Planning novel',
|
||||
novel_planned: 'Novel planned',
|
||||
rendering: 'Rendering',
|
||||
rendered: 'Rendered',
|
||||
error: 'Error',
|
||||
};
|
||||
return labels[stage] ?? stage.replace(/_/g, ' ');
|
||||
}
|
||||
|
||||
function GradientBorderLine({left, fill, right, width}: {left: string; fill: string; right: string; width: number}) {
|
||||
const fillWidth = Math.max(0, width - 2);
|
||||
return (
|
||||
<Text>
|
||||
<Text color={gradientColor(0, width)}>{left}</Text>
|
||||
{Array.from({length: fillWidth}, (_, index) => (
|
||||
<Text key={index} color={gradientColor(index + 1, width)}>{fill}</Text>
|
||||
))}
|
||||
<Text color={gradientColor(width - 1, width)}>{right}</Text>
|
||||
</Text>
|
||||
);
|
||||
}
|
||||
|
||||
function WorkspaceContentLine({text, color, width}: {text: string; color: string; width: number}) {
|
||||
const contentWidth = Math.max(1, width - 4);
|
||||
const padding = Math.max(0, contentWidth - stringWidth(text));
|
||||
return (
|
||||
<Text>
|
||||
<Text color={WORKSPACE_BORDER_COLORS[0]}>│</Text>
|
||||
<Text> </Text>
|
||||
<Text color={color}>{text}</Text>
|
||||
<Text>{' '.repeat(padding)}</Text>
|
||||
<Text> </Text>
|
||||
<Text color={WORKSPACE_BORDER_COLORS[WORKSPACE_BORDER_COLORS.length - 1]}>│</Text>
|
||||
</Text>
|
||||
);
|
||||
}
|
||||
|
||||
function wrapText(text: string, width: number): string[] {
|
||||
if (width <= 0) return [text];
|
||||
const rows: string[] = [];
|
||||
for (const segment of text.split(/\r?\n/)) {
|
||||
let current = '';
|
||||
let currentWidth = 0;
|
||||
for (const char of Array.from(segment)) {
|
||||
const charWidth = stringWidth(char);
|
||||
if (current && currentWidth + charWidth > width) {
|
||||
rows.push(current);
|
||||
current = char;
|
||||
currentWidth = charWidth;
|
||||
} else {
|
||||
current += char;
|
||||
currentWidth += charWidth;
|
||||
}
|
||||
}
|
||||
rows.push(current);
|
||||
}
|
||||
return rows;
|
||||
}
|
||||
|
||||
function lineColor(line: WorkspaceLine): string {
|
||||
if (line.kind === 'user') return 'yellow';
|
||||
if (line.kind === 'assistant') return 'white';
|
||||
if (line.kind === 'thinking') return 'cyanBright';
|
||||
if (line.kind === 'terminal') return 'cyan';
|
||||
if (line.kind === 'error') return 'red';
|
||||
if (line.kind === 'tool' && line.status === 'error') return 'red';
|
||||
if (line.kind === 'tool') return 'magenta';
|
||||
return 'gray';
|
||||
}
|
||||
|
||||
function clearTerminalForTuiStart() {
|
||||
if (process.env.VIMAX_TUI_NO_CLEAR === '1') return;
|
||||
if (!process.stdout.isTTY) return;
|
||||
process.stdout.write('\u001b[2J\u001b[3J\u001b[H');
|
||||
}
|
||||
|
||||
clearTerminalForTuiStart();
|
||||
render(<App />);
|
||||
@@ -0,0 +1,41 @@
|
||||
import assert from 'node:assert/strict';
|
||||
import {applyStreamEvent, createMappingState} from './lineMapping.js';
|
||||
import type {MappingState, WorkspaceLine} from './types.js';
|
||||
|
||||
let lines: WorkspaceLine[] = [];
|
||||
let state: MappingState = createMappingState();
|
||||
|
||||
lines = [...lines, {kind: 'user', text: 'start'}];
|
||||
assert.deepEqual(lines[0], {kind: 'user', text: 'start'});
|
||||
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'token', delta: 'hello'}));
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'token', delta: ' world'}));
|
||||
assert.equal(lines.length, 2);
|
||||
assert.deepEqual(lines[1], {kind: 'assistant', text: 'hello world'});
|
||||
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'tool_start', tool: {name: 'vimax_narrative_planning'}}));
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'tool_progress', tool: {name: 'vimax_narrative_planning'}, progress: {stage: 'running', message: 'planning'}}));
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'tool_result', tool_result: {name: 'vimax_narrative_planning', ok: true}}));
|
||||
assert.equal(lines.at(-3)?.kind, 'tool');
|
||||
assert.equal(lines.at(-2)?.kind, 'tool');
|
||||
assert.deepEqual(lines.at(-1), {kind: 'tool', status: 'done', text: 'tool vimax_narrative_planning done'});
|
||||
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'tool_result', tool_result: {name: 'vimax_narrative_planning', ok: false, content: 'Developing story failed: Request timed out.'}}));
|
||||
assert.deepEqual(lines.at(-1), {kind: 'tool', status: 'error', text: 'tool vimax_narrative_planning error: Developing story failed: Request timed out.'});
|
||||
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'terminal', stream: 'stdout', line: 'render log'}));
|
||||
assert.deepEqual(lines.at(-1), {kind: 'terminal', text: '[stdout]: render log'});
|
||||
|
||||
const lengthBeforeInternalEvents = lines.length;
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'turn', turn_id: 'turn-1'}));
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'status', phase: 'sampling_assistant', message: 'Sampling assistant'}));
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'session', session: {active_session_id: 's1', session: {session_id: 's1', stage: 'narrative_planned', working_dir: '.working_dir/s1'}}}));
|
||||
assert.equal(lines.length, lengthBeforeInternalEvents);
|
||||
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'error', message: 'bad'}));
|
||||
assert.deepEqual(lines.at(-1), {kind: 'error', text: 'bad'});
|
||||
|
||||
({lines, state} = applyStreamEvent(lines, state, {type: 'done', assistant: 'hello world'}));
|
||||
assert.equal(state.assistantStreaming, false);
|
||||
|
||||
console.log('lineMapping tests passed');
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user