vibe-aigc 0.6.3__tar.gz → 0.7.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {vibe_aigc-0.6.3/vibe_aigc.egg-info → vibe_aigc-0.7.0}/PKG-INFO +29 -1
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/README.md +28 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/pyproject.toml +1 -1
- vibe_aigc-0.7.0/tests/test_pipeline.py +421 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/__init__.py +46 -4
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/composer_general.py +408 -1
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/discovery.py +107 -2
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/knowledge.py +512 -24
- vibe_aigc-0.7.0/vibe_aigc/llm.py +447 -0
- vibe_aigc-0.7.0/vibe_aigc/models.py +148 -0
- vibe_aigc-0.7.0/vibe_aigc/pipeline.py +565 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/planner.py +145 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/tools.py +32 -0
- vibe_aigc-0.7.0/vibe_aigc/tools_audio.py +746 -0
- vibe_aigc-0.7.0/vibe_aigc/tools_comfyui.py +976 -0
- vibe_aigc-0.7.0/vibe_aigc/tools_utility.py +997 -0
- vibe_aigc-0.7.0/vibe_aigc/tools_video.py +799 -0
- vibe_aigc-0.7.0/vibe_aigc/tools_vision.py +1187 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/vibe_backend.py +11 -1
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/vlm_feedback.py +186 -7
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0/vibe_aigc.egg-info}/PKG-INFO +29 -1
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc.egg-info/SOURCES.txt +6 -0
- vibe_aigc-0.6.3/vibe_aigc/llm.py +0 -204
- vibe_aigc-0.6.3/vibe_aigc/models.py +0 -51
- vibe_aigc-0.6.3/vibe_aigc/tools_comfyui.py +0 -271
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/LICENSE +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/setup.cfg +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_adaptive_replanning.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_agents.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_assets.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_auto_checkpoint.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_automatic_checkpoints.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_checkpoint_serialization.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_error_handling.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_executor.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_feedback_system.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_integration.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_knowledge_base.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_metaplanner_resume.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_metaplanner_visualization.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_models.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_parallel_execution.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_planner.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_progress_callbacks.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_tools.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_visualization.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/tests/test_workflow_resume.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/agents.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/assets.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/audio.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/character.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/cli.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/comfyui.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/executor.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/fidelity.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/model_registry.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/mv_pipeline.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/persistence.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/tools_multimodal.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/video.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/visualization.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/workflow_backend.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/workflow_composer.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/workflow_executor.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/workflow_registry.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/workflow_strategies.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc/workflows.py +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc.egg-info/dependency_links.txt +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc.egg-info/entry_points.txt +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc.egg-info/requires.txt +0 -0
- {vibe_aigc-0.6.3 → vibe_aigc-0.7.0}/vibe_aigc.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: vibe-aigc
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.7.0
|
|
4
4
|
Summary: A New Paradigm for Content Generation via Agentic Orchestration
|
|
5
5
|
Author: Vibe AIGC Contributors
|
|
6
6
|
License-Expression: MIT
|
|
@@ -105,6 +105,34 @@ result = kb.query("Hitchcockian suspense")
|
|
|
105
105
|
- 📊 **Progress Tracking** — Real-time callbacks and visualization
|
|
106
106
|
- 🎨 **Workflow Visualization** — ASCII and Mermaid diagrams
|
|
107
107
|
|
|
108
|
+
## LLM Providers
|
|
109
|
+
|
|
110
|
+
Vibe AIGC supports multiple LLM backends with auto-detection:
|
|
111
|
+
|
|
112
|
+
| Provider | API Key | Model Example |
|
|
113
|
+
|----------|---------|---------------|
|
|
114
|
+
| **OpenAI** | `OPENAI_API_KEY` | `gpt-4` |
|
|
115
|
+
| **Anthropic** | `ANTHROPIC_API_KEY` | `claude-sonnet-4-20250514` |
|
|
116
|
+
| **Ollama** | None (local) | `qwen2.5:14b` |
|
|
117
|
+
|
|
118
|
+
```python
|
|
119
|
+
from vibe_aigc import MetaPlanner, LLMConfig
|
|
120
|
+
|
|
121
|
+
# Auto-detect (checks env vars, falls back to Ollama)
|
|
122
|
+
planner = MetaPlanner()
|
|
123
|
+
|
|
124
|
+
# Explicit Ollama (no API key needed!)
|
|
125
|
+
config = LLMConfig.for_ollama(
|
|
126
|
+
host="http://localhost:11434", # or your GPU server
|
|
127
|
+
model="qwen2.5:14b"
|
|
128
|
+
)
|
|
129
|
+
planner = MetaPlanner(llm_config=config)
|
|
130
|
+
|
|
131
|
+
# Explicit OpenAI
|
|
132
|
+
config = LLMConfig.for_openai(model="gpt-4o")
|
|
133
|
+
planner = MetaPlanner(llm_config=config)
|
|
134
|
+
```
|
|
135
|
+
|
|
108
136
|
## Installation
|
|
109
137
|
|
|
110
138
|
```bash
|
|
@@ -69,6 +69,34 @@ result = kb.query("Hitchcockian suspense")
|
|
|
69
69
|
- 📊 **Progress Tracking** — Real-time callbacks and visualization
|
|
70
70
|
- 🎨 **Workflow Visualization** — ASCII and Mermaid diagrams
|
|
71
71
|
|
|
72
|
+
## LLM Providers
|
|
73
|
+
|
|
74
|
+
Vibe AIGC supports multiple LLM backends with auto-detection:
|
|
75
|
+
|
|
76
|
+
| Provider | API Key | Model Example |
|
|
77
|
+
|----------|---------|---------------|
|
|
78
|
+
| **OpenAI** | `OPENAI_API_KEY` | `gpt-4` |
|
|
79
|
+
| **Anthropic** | `ANTHROPIC_API_KEY` | `claude-sonnet-4-20250514` |
|
|
80
|
+
| **Ollama** | None (local) | `qwen2.5:14b` |
|
|
81
|
+
|
|
82
|
+
```python
|
|
83
|
+
from vibe_aigc import MetaPlanner, LLMConfig
|
|
84
|
+
|
|
85
|
+
# Auto-detect (checks env vars, falls back to Ollama)
|
|
86
|
+
planner = MetaPlanner()
|
|
87
|
+
|
|
88
|
+
# Explicit Ollama (no API key needed!)
|
|
89
|
+
config = LLMConfig.for_ollama(
|
|
90
|
+
host="http://localhost:11434", # or your GPU server
|
|
91
|
+
model="qwen2.5:14b"
|
|
92
|
+
)
|
|
93
|
+
planner = MetaPlanner(llm_config=config)
|
|
94
|
+
|
|
95
|
+
# Explicit OpenAI
|
|
96
|
+
config = LLMConfig.for_openai(model="gpt-4o")
|
|
97
|
+
planner = MetaPlanner(llm_config=config)
|
|
98
|
+
```
|
|
99
|
+
|
|
72
100
|
## Installation
|
|
73
101
|
|
|
74
102
|
```bash
|
|
@@ -8,7 +8,7 @@ exclude = ["tests*", "docs*", "examples*", "landing*"]
|
|
|
8
8
|
|
|
9
9
|
[project]
|
|
10
10
|
name = "vibe-aigc"
|
|
11
|
-
version = "0.
|
|
11
|
+
version = "0.7.0"
|
|
12
12
|
description = "A New Paradigm for Content Generation via Agentic Orchestration"
|
|
13
13
|
authors = [{name = "Vibe AIGC Contributors"}]
|
|
14
14
|
license = "MIT"
|
|
@@ -0,0 +1,421 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Tests for the Pipeline workflow chaining system.
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
import pytest
|
|
7
|
+
from typing import Dict, Any, Optional
|
|
8
|
+
|
|
9
|
+
from vibe_aigc.pipeline import (
|
|
10
|
+
Pipeline,
|
|
11
|
+
PipelineStep,
|
|
12
|
+
PipelineResult,
|
|
13
|
+
PipelineStatus,
|
|
14
|
+
PipelineBuilder,
|
|
15
|
+
StepResult
|
|
16
|
+
)
|
|
17
|
+
from vibe_aigc.tools import BaseTool, ToolResult, ToolSpec, ToolCategory, ToolRegistry
|
|
18
|
+
from vibe_aigc.models import Vibe
|
|
19
|
+
from vibe_aigc.planner import MetaPlanner
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
# ==================== Mock Tools ====================
|
|
23
|
+
|
|
24
|
+
class MockImageTool(BaseTool):
|
|
25
|
+
"""Mock image generation tool for testing."""
|
|
26
|
+
|
|
27
|
+
@property
|
|
28
|
+
def spec(self) -> ToolSpec:
|
|
29
|
+
return ToolSpec(
|
|
30
|
+
name="image_generate",
|
|
31
|
+
description="Generate images (mock)",
|
|
32
|
+
category=ToolCategory.IMAGE,
|
|
33
|
+
input_schema={
|
|
34
|
+
"type": "object",
|
|
35
|
+
"required": ["prompt"],
|
|
36
|
+
"properties": {"prompt": {"type": "string"}, "size": {"type": "string"}}
|
|
37
|
+
},
|
|
38
|
+
output_schema={"type": "object"}
|
|
39
|
+
)
|
|
40
|
+
|
|
41
|
+
async def execute(self, inputs: Dict[str, Any], context: Optional[Dict[str, Any]] = None) -> ToolResult:
|
|
42
|
+
# Simulate image generation
|
|
43
|
+
await asyncio.sleep(0.01)
|
|
44
|
+
return ToolResult(
|
|
45
|
+
success=True,
|
|
46
|
+
output={
|
|
47
|
+
"url": f"https://example.com/image_{inputs.get('size', '1024x1024')}.png",
|
|
48
|
+
"prompt": inputs["prompt"],
|
|
49
|
+
"size": inputs.get("size", "1024x1024")
|
|
50
|
+
},
|
|
51
|
+
metadata={"model": "mock-dalle", "provider": "mock"}
|
|
52
|
+
)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class MockUpscaleTool(BaseTool):
|
|
56
|
+
"""Mock upscale tool for testing."""
|
|
57
|
+
|
|
58
|
+
@property
|
|
59
|
+
def spec(self) -> ToolSpec:
|
|
60
|
+
return ToolSpec(
|
|
61
|
+
name="image_upscale",
|
|
62
|
+
description="Upscale images (mock)",
|
|
63
|
+
category=ToolCategory.IMAGE,
|
|
64
|
+
input_schema={
|
|
65
|
+
"type": "object",
|
|
66
|
+
"required": ["url"],
|
|
67
|
+
"properties": {"url": {"type": "string"}, "scale": {"type": "integer"}}
|
|
68
|
+
},
|
|
69
|
+
output_schema={"type": "object"}
|
|
70
|
+
)
|
|
71
|
+
|
|
72
|
+
async def execute(self, inputs: Dict[str, Any], context: Optional[Dict[str, Any]] = None) -> ToolResult:
|
|
73
|
+
await asyncio.sleep(0.01)
|
|
74
|
+
scale = inputs.get("scale", 2)
|
|
75
|
+
original_url = inputs.get("url", "unknown")
|
|
76
|
+
return ToolResult(
|
|
77
|
+
success=True,
|
|
78
|
+
output={
|
|
79
|
+
"url": f"{original_url.replace('.png', '')}_upscaled_{scale}x.png",
|
|
80
|
+
"scale": scale,
|
|
81
|
+
"original_url": original_url
|
|
82
|
+
},
|
|
83
|
+
metadata={"upscaler": "mock-esrgan"}
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
class MockVideoTool(BaseTool):
|
|
88
|
+
"""Mock video generation tool for testing."""
|
|
89
|
+
|
|
90
|
+
@property
|
|
91
|
+
def spec(self) -> ToolSpec:
|
|
92
|
+
return ToolSpec(
|
|
93
|
+
name="video_generate",
|
|
94
|
+
description="Generate video (mock)",
|
|
95
|
+
category=ToolCategory.VIDEO,
|
|
96
|
+
input_schema={
|
|
97
|
+
"type": "object",
|
|
98
|
+
"required": ["prompt"],
|
|
99
|
+
"properties": {"prompt": {"type": "string"}, "frames": {"type": "integer"}}
|
|
100
|
+
},
|
|
101
|
+
output_schema={"type": "object"}
|
|
102
|
+
)
|
|
103
|
+
|
|
104
|
+
async def execute(self, inputs: Dict[str, Any], context: Optional[Dict[str, Any]] = None) -> ToolResult:
|
|
105
|
+
await asyncio.sleep(0.02)
|
|
106
|
+
return ToolResult(
|
|
107
|
+
success=True,
|
|
108
|
+
output={
|
|
109
|
+
"url": "https://example.com/video.mp4",
|
|
110
|
+
"frames": inputs.get("frames", 33),
|
|
111
|
+
"prompt": inputs["prompt"]
|
|
112
|
+
},
|
|
113
|
+
metadata={"model": "mock-ltx"}
|
|
114
|
+
)
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
class MockFailingTool(BaseTool):
|
|
118
|
+
"""Tool that fails for testing error handling."""
|
|
119
|
+
|
|
120
|
+
def __init__(self, fail_count: int = 1):
|
|
121
|
+
self.fail_count = fail_count
|
|
122
|
+
self.attempts = 0
|
|
123
|
+
|
|
124
|
+
@property
|
|
125
|
+
def spec(self) -> ToolSpec:
|
|
126
|
+
return ToolSpec(
|
|
127
|
+
name="failing_tool",
|
|
128
|
+
description="A tool that fails",
|
|
129
|
+
category=ToolCategory.UTILITY,
|
|
130
|
+
input_schema={"type": "object"},
|
|
131
|
+
output_schema={"type": "object"}
|
|
132
|
+
)
|
|
133
|
+
|
|
134
|
+
async def execute(self, inputs: Dict[str, Any], context: Optional[Dict[str, Any]] = None) -> ToolResult:
|
|
135
|
+
self.attempts += 1
|
|
136
|
+
if self.attempts <= self.fail_count:
|
|
137
|
+
return ToolResult(
|
|
138
|
+
success=False,
|
|
139
|
+
output=None,
|
|
140
|
+
error=f"Simulated failure (attempt {self.attempts})"
|
|
141
|
+
)
|
|
142
|
+
return ToolResult(
|
|
143
|
+
success=True,
|
|
144
|
+
output={"status": "recovered", "attempts": self.attempts}
|
|
145
|
+
)
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
# ==================== Fixtures ====================
|
|
149
|
+
|
|
150
|
+
@pytest.fixture
|
|
151
|
+
def mock_registry():
|
|
152
|
+
"""Create a registry with mock tools."""
|
|
153
|
+
registry = ToolRegistry()
|
|
154
|
+
registry.register(MockImageTool())
|
|
155
|
+
registry.register(MockUpscaleTool())
|
|
156
|
+
registry.register(MockVideoTool())
|
|
157
|
+
return registry
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
# ==================== Tests ====================
|
|
161
|
+
|
|
162
|
+
class TestPipelineBasic:
|
|
163
|
+
"""Basic pipeline functionality tests."""
|
|
164
|
+
|
|
165
|
+
@pytest.mark.asyncio
|
|
166
|
+
async def test_simple_pipeline(self, mock_registry):
|
|
167
|
+
"""Test a simple single-step pipeline."""
|
|
168
|
+
pipeline = Pipeline(
|
|
169
|
+
steps=[
|
|
170
|
+
PipelineStep(tool="image_generate", config={"size": "768x768"})
|
|
171
|
+
],
|
|
172
|
+
tool_registry=mock_registry
|
|
173
|
+
)
|
|
174
|
+
|
|
175
|
+
result = await pipeline.execute({"prompt": "a red apple"})
|
|
176
|
+
|
|
177
|
+
assert result.status == PipelineStatus.COMPLETED
|
|
178
|
+
assert result.success
|
|
179
|
+
assert len(result.step_results) == 1
|
|
180
|
+
assert "url" in result.final_output
|
|
181
|
+
assert "768x768" in result.final_output["url"]
|
|
182
|
+
|
|
183
|
+
@pytest.mark.asyncio
|
|
184
|
+
async def test_multi_step_pipeline(self, mock_registry):
|
|
185
|
+
"""Test chaining multiple steps."""
|
|
186
|
+
pipeline = Pipeline(
|
|
187
|
+
steps=[
|
|
188
|
+
PipelineStep(tool="image_generate", config={"size": "768x768"}, name="gen_image"),
|
|
189
|
+
PipelineStep(tool="image_upscale", config={"scale": 2}, name="upscale")
|
|
190
|
+
],
|
|
191
|
+
tool_registry=mock_registry
|
|
192
|
+
)
|
|
193
|
+
|
|
194
|
+
result = await pipeline.execute({"prompt": "mountain landscape"})
|
|
195
|
+
|
|
196
|
+
assert result.status == PipelineStatus.COMPLETED
|
|
197
|
+
assert len(result.step_results) == 2
|
|
198
|
+
# Verify chaining: upscale step got the URL from image step
|
|
199
|
+
assert "upscaled_2x" in result.final_output["url"]
|
|
200
|
+
assert "original_url" in result.final_output
|
|
201
|
+
|
|
202
|
+
@pytest.mark.asyncio
|
|
203
|
+
async def test_image_upscale_chain(self, mock_registry):
|
|
204
|
+
"""Test the specific image->upscale chain from requirements."""
|
|
205
|
+
pipeline = Pipeline([
|
|
206
|
+
PipelineStep(tool="image_generate", config={"size": "1024x1024"}),
|
|
207
|
+
PipelineStep(tool="image_upscale", config={"scale": 4})
|
|
208
|
+
], tool_registry=mock_registry)
|
|
209
|
+
|
|
210
|
+
result = await pipeline.execute({"prompt": "samurai in rain"})
|
|
211
|
+
|
|
212
|
+
assert result.success
|
|
213
|
+
assert result.step_results[0].status == PipelineStatus.COMPLETED
|
|
214
|
+
assert result.step_results[1].status == PipelineStatus.COMPLETED
|
|
215
|
+
# Verify the chain worked
|
|
216
|
+
final = result.final_output
|
|
217
|
+
assert "upscaled_4x" in final["url"]
|
|
218
|
+
print(f"Pipeline completed in {result.total_duration:.3f}s")
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
class TestPipelineErrorHandling:
|
|
222
|
+
"""Test error handling strategies."""
|
|
223
|
+
|
|
224
|
+
@pytest.mark.asyncio
|
|
225
|
+
async def test_fail_on_error(self, mock_registry):
|
|
226
|
+
"""Test default fail behavior."""
|
|
227
|
+
failing_tool = MockFailingTool(fail_count=10) # Always fails
|
|
228
|
+
mock_registry.register(failing_tool)
|
|
229
|
+
|
|
230
|
+
pipeline = Pipeline([
|
|
231
|
+
PipelineStep(tool="image_generate"),
|
|
232
|
+
PipelineStep(tool="failing_tool", on_error="fail")
|
|
233
|
+
], tool_registry=mock_registry)
|
|
234
|
+
|
|
235
|
+
result = await pipeline.execute({"prompt": "test"})
|
|
236
|
+
|
|
237
|
+
assert result.status == PipelineStatus.FAILED
|
|
238
|
+
assert not result.success
|
|
239
|
+
assert result.step_results[0].status == PipelineStatus.COMPLETED
|
|
240
|
+
assert result.step_results[1].status == PipelineStatus.FAILED
|
|
241
|
+
|
|
242
|
+
@pytest.mark.asyncio
|
|
243
|
+
async def test_skip_on_error(self, mock_registry):
|
|
244
|
+
"""Test skip behavior continues pipeline."""
|
|
245
|
+
failing_tool = MockFailingTool(fail_count=10)
|
|
246
|
+
mock_registry.register(failing_tool)
|
|
247
|
+
|
|
248
|
+
pipeline = Pipeline([
|
|
249
|
+
PipelineStep(tool="image_generate"),
|
|
250
|
+
PipelineStep(tool="failing_tool", on_error="skip"),
|
|
251
|
+
PipelineStep(tool="image_upscale", config={"scale": 2})
|
|
252
|
+
], tool_registry=mock_registry)
|
|
253
|
+
|
|
254
|
+
result = await pipeline.execute({"prompt": "test"})
|
|
255
|
+
|
|
256
|
+
assert result.status == PipelineStatus.COMPLETED
|
|
257
|
+
assert result.success
|
|
258
|
+
# All three steps ran (middle was skipped)
|
|
259
|
+
assert len(result.step_results) == 3
|
|
260
|
+
|
|
261
|
+
@pytest.mark.asyncio
|
|
262
|
+
async def test_retry_on_error(self, mock_registry):
|
|
263
|
+
"""Test retry behavior."""
|
|
264
|
+
failing_tool = MockFailingTool(fail_count=2) # Fails twice, then succeeds
|
|
265
|
+
mock_registry.register(failing_tool)
|
|
266
|
+
|
|
267
|
+
pipeline = Pipeline([
|
|
268
|
+
PipelineStep(tool="failing_tool", on_error="retry", max_retries=3)
|
|
269
|
+
], tool_registry=mock_registry)
|
|
270
|
+
|
|
271
|
+
result = await pipeline.execute({})
|
|
272
|
+
|
|
273
|
+
assert result.status == PipelineStatus.COMPLETED
|
|
274
|
+
assert result.step_results[0].retries == 2 # Took 2 retries to succeed
|
|
275
|
+
|
|
276
|
+
|
|
277
|
+
class TestPipelineBuilder:
|
|
278
|
+
"""Test fluent pipeline builder."""
|
|
279
|
+
|
|
280
|
+
@pytest.mark.asyncio
|
|
281
|
+
async def test_builder_pattern(self, mock_registry):
|
|
282
|
+
"""Test fluent builder creates working pipeline."""
|
|
283
|
+
pipeline = (PipelineBuilder("test_pipeline")
|
|
284
|
+
.add("image_generate", size="1024x1024")
|
|
285
|
+
.add("image_upscale", scale=2)
|
|
286
|
+
.add("video_generate", frames=33)
|
|
287
|
+
.build(mock_registry))
|
|
288
|
+
|
|
289
|
+
assert len(pipeline.steps) == 3
|
|
290
|
+
assert pipeline.name == "test_pipeline"
|
|
291
|
+
|
|
292
|
+
result = await pipeline.execute({"prompt": "epic scene"})
|
|
293
|
+
assert result.success
|
|
294
|
+
assert len(result.step_results) == 3
|
|
295
|
+
|
|
296
|
+
|
|
297
|
+
class TestMetaPlannerIntegration:
|
|
298
|
+
"""Test integration with MetaPlanner."""
|
|
299
|
+
|
|
300
|
+
@pytest.mark.asyncio
|
|
301
|
+
async def test_vibe_to_pipeline_image(self, mock_registry):
|
|
302
|
+
"""Test MetaPlanner converts image vibe to pipeline."""
|
|
303
|
+
planner = MetaPlanner(tool_registry=mock_registry)
|
|
304
|
+
|
|
305
|
+
vibe = Vibe(
|
|
306
|
+
description="A beautiful sunset over the ocean",
|
|
307
|
+
domain="art"
|
|
308
|
+
)
|
|
309
|
+
|
|
310
|
+
pipeline = planner.vibe_to_pipeline(vibe)
|
|
311
|
+
|
|
312
|
+
assert len(pipeline.steps) >= 1
|
|
313
|
+
assert pipeline.steps[0].name == "generate_image"
|
|
314
|
+
|
|
315
|
+
@pytest.mark.asyncio
|
|
316
|
+
async def test_vibe_to_pipeline_video(self, mock_registry):
|
|
317
|
+
"""Test MetaPlanner converts video vibe to pipeline."""
|
|
318
|
+
planner = MetaPlanner(tool_registry=mock_registry)
|
|
319
|
+
|
|
320
|
+
vibe = Vibe(
|
|
321
|
+
description="Animate a samurai drawing sword in the rain",
|
|
322
|
+
domain="video"
|
|
323
|
+
)
|
|
324
|
+
|
|
325
|
+
pipeline = planner.vibe_to_pipeline(vibe)
|
|
326
|
+
|
|
327
|
+
# Should have image generation + video generation
|
|
328
|
+
step_names = [s.name for s in pipeline.steps]
|
|
329
|
+
assert "generate_first_frame" in step_names or "generate_video" in step_names
|
|
330
|
+
|
|
331
|
+
@pytest.mark.asyncio
|
|
332
|
+
async def test_execute_pipeline_method(self, mock_registry):
|
|
333
|
+
"""Test MetaPlanner.execute_pipeline() works."""
|
|
334
|
+
planner = MetaPlanner(tool_registry=mock_registry)
|
|
335
|
+
|
|
336
|
+
vibe = Vibe(
|
|
337
|
+
description="Create an upscaled HD image of a dragon",
|
|
338
|
+
metadata={"upscale_factor": 2}
|
|
339
|
+
)
|
|
340
|
+
|
|
341
|
+
result = await planner.execute_pipeline(vibe)
|
|
342
|
+
|
|
343
|
+
assert result.success
|
|
344
|
+
assert len(result.step_results) >= 1
|
|
345
|
+
|
|
346
|
+
|
|
347
|
+
class TestPipelineResult:
|
|
348
|
+
"""Test PipelineResult functionality."""
|
|
349
|
+
|
|
350
|
+
@pytest.mark.asyncio
|
|
351
|
+
async def test_get_step_output_by_index(self, mock_registry):
|
|
352
|
+
"""Test retrieving step output by index."""
|
|
353
|
+
pipeline = Pipeline([
|
|
354
|
+
PipelineStep(tool="image_generate"),
|
|
355
|
+
PipelineStep(tool="image_upscale")
|
|
356
|
+
], tool_registry=mock_registry)
|
|
357
|
+
|
|
358
|
+
result = await pipeline.execute({"prompt": "test"})
|
|
359
|
+
|
|
360
|
+
step0_output = result.get_step_output(0)
|
|
361
|
+
step1_output = result.get_step_output(1)
|
|
362
|
+
|
|
363
|
+
assert step0_output is not None
|
|
364
|
+
assert step1_output is not None
|
|
365
|
+
assert "url" in step0_output
|
|
366
|
+
|
|
367
|
+
@pytest.mark.asyncio
|
|
368
|
+
async def test_get_step_output_by_name(self, mock_registry):
|
|
369
|
+
"""Test retrieving step output by name."""
|
|
370
|
+
pipeline = Pipeline([
|
|
371
|
+
PipelineStep(tool="image_generate", name="gen"),
|
|
372
|
+
PipelineStep(tool="image_upscale", name="upscale")
|
|
373
|
+
], tool_registry=mock_registry)
|
|
374
|
+
|
|
375
|
+
result = await pipeline.execute({"prompt": "test"})
|
|
376
|
+
|
|
377
|
+
gen_output = result.get_step_output("gen")
|
|
378
|
+
assert gen_output is not None
|
|
379
|
+
assert "url" in gen_output
|
|
380
|
+
|
|
381
|
+
@pytest.mark.asyncio
|
|
382
|
+
async def test_result_to_dict(self, mock_registry):
|
|
383
|
+
"""Test serialization to dict."""
|
|
384
|
+
pipeline = Pipeline([
|
|
385
|
+
PipelineStep(tool="image_generate")
|
|
386
|
+
], tool_registry=mock_registry)
|
|
387
|
+
|
|
388
|
+
result = await pipeline.execute({"prompt": "test"})
|
|
389
|
+
result_dict = result.to_dict()
|
|
390
|
+
|
|
391
|
+
assert "status" in result_dict
|
|
392
|
+
assert "step_results" in result_dict
|
|
393
|
+
assert "total_duration" in result_dict
|
|
394
|
+
|
|
395
|
+
|
|
396
|
+
# ==================== Main ====================
|
|
397
|
+
|
|
398
|
+
if __name__ == "__main__":
|
|
399
|
+
# Run a quick test
|
|
400
|
+
async def main():
|
|
401
|
+
registry = ToolRegistry()
|
|
402
|
+
registry.register(MockImageTool())
|
|
403
|
+
registry.register(MockUpscaleTool())
|
|
404
|
+
registry.register(MockVideoTool())
|
|
405
|
+
|
|
406
|
+
print("Testing image->upscale chain...")
|
|
407
|
+
pipeline = Pipeline([
|
|
408
|
+
PipelineStep(tool="image_generate", config={"size": "768x768"}),
|
|
409
|
+
PipelineStep(tool="image_upscale", config={"scale": 2})
|
|
410
|
+
], tool_registry=registry)
|
|
411
|
+
|
|
412
|
+
result = await pipeline.execute({"prompt": "samurai in rain"})
|
|
413
|
+
|
|
414
|
+
print(f"Status: {result.status.value}")
|
|
415
|
+
print(f"Duration: {result.total_duration:.3f}s")
|
|
416
|
+
print(f"Steps: {len(result.step_results)}")
|
|
417
|
+
for i, step in enumerate(result.step_results):
|
|
418
|
+
print(f" [{i}] {step.step_name}: {step.status.value} ({step.duration:.3f}s)")
|
|
419
|
+
print(f"Final output: {result.final_output}")
|
|
420
|
+
|
|
421
|
+
asyncio.run(main())
|
|
@@ -11,9 +11,9 @@ Architecture (Paper Section 5):
|
|
|
11
11
|
- AssetBank: Character and style consistency management
|
|
12
12
|
"""
|
|
13
13
|
|
|
14
|
-
from .models import Vibe, WorkflowPlan, WorkflowNode, WorkflowNodeType
|
|
14
|
+
from .models import Vibe, WorkflowPlan, WorkflowNode, WorkflowNodeType, GenerationRequest, CharacterProfile
|
|
15
15
|
from .planner import MetaPlanner
|
|
16
|
-
from .llm import LLMClient, LLMConfig
|
|
16
|
+
from .llm import LLMClient, LLMConfig, LLMProvider, list_ollama_models, check_ollama_available
|
|
17
17
|
from .executor import WorkflowExecutor, ExecutionStatus, ExecutionResult
|
|
18
18
|
|
|
19
19
|
# Paper Section 5.3: Domain-Specific Expert Knowledge Base
|
|
@@ -78,8 +78,15 @@ __version__ = "0.2.0"
|
|
|
78
78
|
__all__ = [
|
|
79
79
|
# Core models
|
|
80
80
|
"Vibe", "WorkflowPlan", "WorkflowNode", "WorkflowNodeType",
|
|
81
|
+
"GenerationRequest", "CharacterProfile",
|
|
82
|
+
# System Discovery
|
|
83
|
+
"SystemDiscovery", "SystemCapabilities", "Capability", "HardwareConstraints",
|
|
84
|
+
"AvailableNode", "AvailableModel", "discover_system",
|
|
85
|
+
# General Composer
|
|
86
|
+
"GeneralComposer", "NodeRequirement", "STANDARD_REQUIREMENTS", "create_composer",
|
|
81
87
|
# MetaPlanner (Section 5.2)
|
|
82
|
-
"MetaPlanner", "LLMClient", "LLMConfig",
|
|
88
|
+
"MetaPlanner", "LLMClient", "LLMConfig", "LLMProvider",
|
|
89
|
+
"list_ollama_models", "check_ollama_available",
|
|
83
90
|
# Executor
|
|
84
91
|
"WorkflowExecutor", "ExecutionStatus", "ExecutionResult",
|
|
85
92
|
# Knowledge Base (Section 5.3)
|
|
@@ -97,8 +104,31 @@ __all__ = [
|
|
|
97
104
|
"DesignerAgent", "ScreenwriterAgent", "ComposerAgent",
|
|
98
105
|
"create_default_agents",
|
|
99
106
|
# Asset Bank
|
|
100
|
-
"AssetBank", "Character", "StyleGuide", "Artifact", "create_asset_bank"
|
|
107
|
+
"AssetBank", "Character", "StyleGuide", "Artifact", "create_asset_bank",
|
|
108
|
+
# Pipeline chaining
|
|
109
|
+
"Pipeline", "PipelineStep", "PipelineResult", "PipelineStatus",
|
|
110
|
+
"PipelineBuilder", "StepResult",
|
|
111
|
+
"create_image_pipeline", "create_video_pipeline"
|
|
101
112
|
]
|
|
113
|
+
# System Discovery - Constraint-aware system discovery
|
|
114
|
+
from .discovery import (
|
|
115
|
+
SystemDiscovery,
|
|
116
|
+
SystemCapabilities,
|
|
117
|
+
Capability,
|
|
118
|
+
HardwareConstraints,
|
|
119
|
+
AvailableNode,
|
|
120
|
+
AvailableModel,
|
|
121
|
+
discover_system,
|
|
122
|
+
)
|
|
123
|
+
|
|
124
|
+
# General Workflow Composer - Builds workflows from discovered nodes
|
|
125
|
+
from .composer_general import (
|
|
126
|
+
GeneralComposer,
|
|
127
|
+
NodeRequirement,
|
|
128
|
+
STANDARD_REQUIREMENTS,
|
|
129
|
+
create_composer,
|
|
130
|
+
)
|
|
131
|
+
|
|
102
132
|
# Model Registry - Auto-detect available models
|
|
103
133
|
from .model_registry import ModelRegistry, ModelCapability, ModelFamily, ModelSpec
|
|
104
134
|
|
|
@@ -116,3 +146,15 @@ from .audio import MusicGenBackend, RiffusionBackend, ElevenLabsBackend, MusicGe
|
|
|
116
146
|
|
|
117
147
|
# MV Pipeline
|
|
118
148
|
from .mv_pipeline import MVPipeline, Shot, Storyboard, create_mv
|
|
149
|
+
|
|
150
|
+
# Pipeline chaining for workflow orchestration
|
|
151
|
+
from .pipeline import (
|
|
152
|
+
Pipeline,
|
|
153
|
+
PipelineStep,
|
|
154
|
+
PipelineResult,
|
|
155
|
+
PipelineStatus,
|
|
156
|
+
PipelineBuilder,
|
|
157
|
+
StepResult,
|
|
158
|
+
create_image_pipeline,
|
|
159
|
+
create_video_pipeline
|
|
160
|
+
)
|