chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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"))
|
||||
Reference in New Issue
Block a user