返回 MoneyPrinterTurbo
twelvelabs.py
根目录 / app / services / twelvelabs.py
1 """
2 TwelveLabs (https://twelvelabs.io) integration — optional, opt-in helpers.
3
4 This module wraps two TwelveLabs models so MoneyPrinterTurbo can make better
5 use of the stock/B-roll footage it downloads:
6
7 * Marengo (multimodal embeddings, 512-dim) — used to *semantically reorder*
8 the LLM-generated search terms against the video subject, so that when the
9 timeline budget runs out the most on-topic footage is the footage that made
10 it in (instead of whatever the LLM happened to list first).
11
12 * Pegasus (video understanding) — used to QA / describe a generated clip from
13 a public URL, e.g. to sanity-check that a downloaded clip actually matches
14 the script before it ships.
15
16 The integration is fully opt-in and non-breaking:
17 * If `twelvelabs_api_keys` is not configured, every public function here is a
18 no-op that returns its input unchanged (or None), so default behavior is
19 identical to a build without TwelveLabs.
20 * The `twelvelabs` SDK is imported lazily, so the dependency is only required
21 when the feature is actually used.
22
23 Config (config.toml, [app] section):
24 twelvelabs_api_keys = ["tlk_xxx"] # required to enable
25 twelvelabs_rerank_terms = true # opt-in: reorder search terms by relevance
26 twelvelabs_marengo_model = "marengo3.0" # optional override
27 twelvelabs_pegasus_model = "pegasus1.5" # optional override
28
29 Configure a TwelveLabs API key from the TwelveLabs dashboard (https://twelvelabs.io) to enable this optional integration.
30 """
31
32 import math
33 from functools import lru_cache
34 from typing import List, Optional
35
36 from loguru import logger
37
38 from app.config import config
39 from app.services import material
40
41 DEFAULT_MARENGO_MODEL = "marengo3.0"
42 DEFAULT_PEGASUS_MODEL = "pegasus1.5"
43 # Pegasus requires max_tokens in [512, 98304]; 512 is plenty for a one-line QA.
44 _PEGASUS_MIN_MAX_TOKENS = 512
45
46
47 def is_enabled() -> bool:
48 """True only when at least one TwelveLabs API key is configured."""
49 keys = config.app.get("twelvelabs_api_keys")
50 return bool(keys)
51
52
53 def _client():
54 # Lazy import + rotated key reuse mirrors the other providers in
55 # material.py (get_api_key rotates across configured keys).
56 from twelvelabs import TwelveLabs
57
58 api_key = material.get_api_key("twelvelabs_api_keys")
59 return TwelveLabs(api_key=api_key)
60
61
62 def _cosine(a: List[float], b: List[float]) -> float:
63 dot = sum(x * y for x, y in zip(a, b))
64 na = math.sqrt(sum(x * x for x in a))
65 nb = math.sqrt(sum(x * x for x in b))
66 if na == 0 or nb == 0:
67 return 0.0
68 return dot / (na * nb)
69
70
71 def embed_text(text: str, model: Optional[str] = None) -> Optional[List[float]]:
72 """
73 Return a 512-dim Marengo text embedding, or None on failure / when disabled.
74
75 Cached so repeated terms across a session don't re-hit the API.
76 """
77 if not is_enabled() or not text or not text.strip():
78 return None
79 model = model or config.app.get("twelvelabs_marengo_model", DEFAULT_MARENGO_MODEL)
80 try:
81 # lru_cache only memoizes successful returns; a raised exception is not
82 # cached, so a transient API error never poisons the cache.
83 return _embed_text_cached(text.strip(), model)
84 except Exception as e: # noqa: BLE001 - never break the pipeline on TL errors
85 logger.warning(f"TwelveLabs embed_text failed, skipping rerank: {e}")
86 return None
87
88
89 @lru_cache(maxsize=512)
90 def _embed_text_cached(text: str, model: str) -> List[float]:
91 client = _client()
92 resp = client.embed.create(model_name=model, text=text)
93 # SDK aliases the raw JSON 'float' vector key to `float_`.
94 return list(resp.text_embedding.segments[0].float_)
95
96
97 def rerank_terms_by_subject(
98 video_subject: str,
99 search_terms: List[str],
100 model: Optional[str] = None,
101 ) -> List[str]:
102 """
103 Reorder `search_terms` so the terms most semantically relevant to
104 `video_subject` come first (Marengo cosine similarity).
105
106 Opt-in: only runs when TwelveLabs is enabled AND
107 `twelvelabs_rerank_terms` is truthy. Falls back to the original order on
108 any failure, so it can never make the pipeline worse.
109 """
110 if not is_enabled() or not config.app.get("twelvelabs_rerank_terms"):
111 return search_terms
112 if not video_subject or len(search_terms) < 2:
113 return search_terms
114
115 subject_vec = embed_text(video_subject, model)
116 if subject_vec is None:
117 return search_terms
118
119 scored = []
120 for term in search_terms:
121 vec = embed_text(term, model)
122 if vec is None:
123 # If any term can't be embedded, don't risk a partial reorder.
124 return search_terms
125 scored.append((term, _cosine(subject_vec, vec)))
126
127 ranked = [term for term, _ in sorted(scored, key=lambda x: x[1], reverse=True)]
128 logger.info(
129 f"TwelveLabs Marengo reranked {len(ranked)} search terms by relevance "
130 f"to subject '{video_subject}': {ranked}"
131 )
132 return ranked
133
134
135 def analyze_clip(
136 video_url: str,
137 prompt: str = "Describe what happens in this video in one sentence.",
138 model: Optional[str] = None,
139 max_tokens: int = _PEGASUS_MIN_MAX_TOKENS,
140 ) -> Optional[str]:
141 """
142 QA / describe a clip from a public URL with Pegasus, returning the model's
143 text answer (or None when disabled / on failure).
144
145 Notes (TwelveLabs API constraints):
146 * Pegasus needs a publicly reachable URL (or an uploaded asset), not a
147 bare local path; the analyzed window must be >= 4s.
148 * max_tokens must be >= 512 for this model.
149 """
150 if not is_enabled() or not video_url:
151 return None
152 model = model or config.app.get("twelvelabs_pegasus_model", DEFAULT_PEGASUS_MODEL)
153 try:
154 from twelvelabs.types import VideoContext_Url
155
156 client = _client()
157 resp = client.analyze(
158 model_name=model,
159 video=VideoContext_Url(url=video_url),
160 prompt=prompt,
161 max_tokens=max(max_tokens, _PEGASUS_MIN_MAX_TOKENS),
162 )
163 return resp.data
164 except Exception as e: # noqa: BLE001
165 logger.warning(f"TwelveLabs analyze_clip failed: {e}")
166 return None
167
167 lines PYTHON