diff --git a/README.md b/README.md index 52f947c..5f18d70 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ The application templates provided in this repo scale up to **millions of pages | [`Adaptive RAG App`](templates/adaptive_rag/) | A RAG application using Adaptive RAG, a technique developed by Pathway to reduce token cost in RAG up to 4x while maintaining accuracy. | | [`Private RAG App with Mistral and Ollama`](templates/private_rag/) | A fully private (local) version of the `question_answering_rag` RAG pipeline using Pathway Live Data Framework, Mistral, and Ollama. | | [`Slides AI Search App`](templates/slides_ai_search/) | An indexing pipeline for retrieving slides. It performs multi-modal of PowerPoint and PDF and maintains live index of your slides."| +| [`Video RAG with TwelveLabs`](templates/video_rag_twelvelabs/) | A RAG pipeline over **video**. It uses [TwelveLabs](https://twelvelabs.io) Pegasus to turn videos into rich text descriptions and Marengo multimodal embeddings to index them, so you can ask questions about your videos on a live connected data source (files, Google Drive,...). | ## How do these AI Pipelines work? diff --git a/pyproject.toml b/pyproject.toml index 5e94869..3f62ae1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,7 @@ profile = "black" [tool.mypy] python_version = "3.10" -exclude = ["templates/adaptive_rag", "templates/private_rag", "templates/document_indexing", "templates/multimodal_rag", "templates/question_answering_rag"] +exclude = ["templates/adaptive_rag", "templates/private_rag", "templates/document_indexing", "templates/multimodal_rag", "templates/question_answering_rag", "templates/video_rag_twelvelabs"] ignore_missing_imports = true check_untyped_defs = true warn_redundant_casts = true diff --git a/templates/video_rag_twelvelabs/.env.example b/templates/video_rag_twelvelabs/.env.example new file mode 100644 index 0000000..f4af2bf --- /dev/null +++ b/templates/video_rag_twelvelabs/.env.example @@ -0,0 +1,6 @@ +# TwelveLabs API key, used by the Pegasus video parser and the Marengo embedder. +# Get a free key at https://twelvelabs.io +TWELVELABS_API_KEY= + +# OpenAI API key, used by the question-answering LLM. +OPENAI_API_KEY= diff --git a/templates/video_rag_twelvelabs/Dockerfile b/templates/video_rag_twelvelabs/Dockerfile new file mode 100644 index 0000000..957987d --- /dev/null +++ b/templates/video_rag_twelvelabs/Dockerfile @@ -0,0 +1,12 @@ +FROM pathwaycom/pathway:latest + +WORKDIR /app + +COPY requirements.txt . +RUN pip install -U --no-cache-dir -r requirements.txt + +COPY . . + +EXPOSE 8000 + +CMD ["python", "app.py"] diff --git a/templates/video_rag_twelvelabs/README.md b/templates/video_rag_twelvelabs/README.md new file mode 100644 index 0000000..ed49033 --- /dev/null +++ b/templates/video_rag_twelvelabs/README.md @@ -0,0 +1,122 @@ +

+ GCP Logo + Deploy with GCP | + AWS Logo + Deploy with AWS | + Azure Logo + Deploy with Azure | + Render Logo + Deploy with Render +

+ +# Video RAG with Pathway Live Data Framework and TwelveLabs + +## Overview + +This app template shows how to build a RAG application **over video** using the +Pathway Live Data Framework together with [TwelveLabs](https://twelvelabs.io). + +Videos are notoriously hard to put into a RAG pipeline: most stacks first +transcribe the audio and throw away everything that happens on screen. This +template indexes the *whole* video instead, using two TwelveLabs models: + +- **[Pegasus](https://docs.twelvelabs.io/docs/concepts/models/pegasus)** — a + video-understanding model that turns each video into a rich text description + (what happens, who and what appears, the setting, on-screen and spoken text, + the overall topic). Pathway indexes that text exactly like it would index a PDF. +- **[Marengo](https://docs.twelvelabs.io/docs/concepts/models/marengo)** — a + multimodal embedding model that produces 512-dimensional vectors in a shared + space for text, image, audio and video. It is used here as the retriever + embedder. + +Because the template uses the standard Pathway `DocumentStore` and +`BaseRAGQuestionAnswerer`, everything else works out of the box: live sync with +your data source, the in-memory vector index, caching, and the HTTP API. Drop a +new video into the connected folder and it becomes queryable automatically. + +## Architecture + +``` +video files ─▶ pw.io.fs.read (bytes) + ─▶ TwelveLabsVideoParser (Pegasus: video → text) + ─▶ TokenCountSplitter (chunking) + ─▶ DocumentStore + UsearchKnnFactory + with MarengoEmbedder (text → 512-dim vectors) + ─▶ BaseRAGQuestionAnswerer (OpenAI LLM over retrieved context) + ─▶ REST API on :8000 +``` + +The TwelveLabs components live in the local +[`pathway_twelvelabs`](pathway_twelvelabs/__init__.py) package: + +- `TwelveLabsVideoParser` — a Pathway parser (`pw.UDF`) that uploads the video + bytes as a TwelveLabs asset and asks Pegasus to describe it. +- `MarengoEmbedder` — a Pathway embedder (`BaseEmbedder`) that calls the Marengo + embedding endpoint. + +Both are wired in entirely through [`app.yaml`](app.yaml), so you can swap models, +prompts, the data source, or the LLM without touching any Python. + +## Customizing the pipeline + +- **Change what is extracted from the video.** Set the `prompt` field of + `$parser` in `app.yaml`, e.g. `"Describe this video, focusing on the products + that appear and any prices shown."` +- **Change the data source.** Replace the `!pw.io.fs.read` source with the Google + Drive, SharePoint, or S3 connector (a commented Google Drive example is + included in `app.yaml`). +- **Change the answering LLM.** Edit the `$llm` block — any + [Pathway LLM wrapper](https://pathway.com/developers/api-docs/pathway-xpacks-llm/llms) + works. + +## Running the app + +### Prerequisites + +- A TwelveLabs API key. Get a free one at [twelvelabs.io](https://twelvelabs.io) — + there is a generous free tier. +- An OpenAI API key for the question-answering LLM. + +Copy `.env.example` to `.env` and fill in your keys: + +```bash +cp .env.example .env +# edit .env and set TWELVELABS_API_KEY and OPENAI_API_KEY +``` + +Put one or more videos (e.g. `.mp4`, `.mov`) into the `data/` directory. + +### With Docker + +```bash +docker build -t video-rag-twelvelabs . +docker run -v $(pwd)/data:/app/data --env-file .env -p 8000:8000 video-rag-twelvelabs +``` + +### Locally + +```bash +pip install -r requirements.txt +python app.py +``` + +### Querying + +Once the server is up, ask questions about your videos: + +```bash +curl -X POST http://localhost:8000/v1/pw_ai_answer \ + -H "Content-Type: application/json" \ + -d '{"prompt": "What products are shown in the videos?"}' +``` + +## Tests + +A small test suite lives in [`test_twelvelabs.py`](test_twelvelabs.py). The +no-network tests run without any credentials; the live Marengo embedding test is +skipped unless `TWELVELABS_API_KEY` is set: + +```bash +pip install -r requirements.txt pytest +TWELVELABS_API_KEY=... pytest templates/video_rag_twelvelabs/test_twelvelabs.py +``` diff --git a/templates/video_rag_twelvelabs/app.py b/templates/video_rag_twelvelabs/app.py new file mode 100644 index 0000000..a5a8c20 --- /dev/null +++ b/templates/video_rag_twelvelabs/app.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python + +# Copyright © 2026 Pathway + +import logging +from pathlib import Path +from warnings import warn + +import pathway as pw + +# `pathway_twelvelabs` registers the TwelveLabsVideoParser and MarengoEmbedder +# classes so they can be referenced from `app.yaml` via the `!` YAML tags. +import pathway_twelvelabs # noqa: F401 +from dotenv import load_dotenv +from pathway.xpacks.llm.question_answering import BaseRAGQuestionAnswerer +from pathway.xpacks.llm.servers import QASummaryRestServer +from pydantic import BaseModel, ConfigDict, InstanceOf + +# To use advanced features with Pathway Live Data Framework Scale, get your free license key from +# https://pathway.com/features and paste it below. +# To use Pathway Live Data Framework Community, comment out the line below. +pw.set_license_key("demo-license-key-with-telemetry") + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(name)s %(levelname)s %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) + + +class App(BaseModel): + question_answerer: InstanceOf[BaseRAGQuestionAnswerer] + host: str = "0.0.0.0" + port: int = 8000 + + with_cache: bool | None = None # deprecated + persistence_backend: pw.persistence.Backend | None = None + persistence_mode: pw.PersistenceMode | None = pw.PersistenceMode.UDF_CACHING + terminate_on_error: bool = False + + def run(self) -> None: + server = QASummaryRestServer( # noqa: F841 + self.host, self.port, self.question_answerer + ) + + if self.persistence_mode is None: + if self.with_cache is True: + warn( + "`with_cache` is deprecated. Please use `persistence_mode` instead.", + DeprecationWarning, + ) + persistence_mode = pw.PersistenceMode.UDF_CACHING + else: + persistence_mode = None + else: + persistence_mode = self.persistence_mode + + if persistence_mode is not None: + if self.persistence_backend is None: + persistence_backend = pw.persistence.Backend.filesystem("./Cache") + else: + persistence_backend = self.persistence_backend + persistence_config = pw.persistence.Config( + persistence_backend, + persistence_mode=persistence_mode, + ) + else: + persistence_config = None + + pw.run( + persistence_config=persistence_config, + terminate_on_error=self.terminate_on_error, + monitoring_level=pw.MonitoringLevel.NONE, + ) + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + +if __name__ == "__main__": + base_dir = Path(__file__).resolve().parent + + load_dotenv(base_dir / ".env") + + with open(base_dir / "app.yaml") as f: + config = pw.load_yaml(f) + + app = App(**config) + app.run() diff --git a/templates/video_rag_twelvelabs/app.yaml b/templates/video_rag_twelvelabs/app.yaml new file mode 100644 index 0000000..f3c5663 --- /dev/null +++ b/templates/video_rag_twelvelabs/app.yaml @@ -0,0 +1,92 @@ +# This YAML configuration file sets up the TwelveLabs video RAG template. +# It indexes videos by turning them into text with the TwelveLabs Pegasus model, +# embedding that text with the TwelveLabs Marengo model, and serving a RAG +# question-answering endpoint on top. +# You can learn more about the YAML syntax here: +# https://pathway.com/developers/templates/configure-yaml + + +# $sources defines the data sources read and indexed in the RAG. +# Here we read video files from the local `data` directory. The connector emits +# raw bytes, which the TwelveLabs parser sends to Pegasus for analysis. +# You can learn more about configuring data sources here: +# https://pathway.com/developers/templates/yaml-examples/data-sources-examples +$sources: + - !pw.io.fs.read + path: data + format: binary + with_metadata: true + + # Uncomment to use the Google Drive connector to index videos from a Drive folder. + # - !pw.io.gdrive.read + # object_id: $DRIVE_ID + # service_user_credentials_file: gdrive_indexer.json + # file_name_pattern: + # - "*.mp4" + # - "*.mov" + # object_size_limit: null + # with_metadata: true + # refresh_interval: 30 + + +# The parser turns each video into text using the TwelveLabs Pegasus model. +# The TWELVELABS_API_KEY environment variable must be set (see .env.example). +# You can customize the prompt to extract exactly what your application needs. +$parser: !pathway_twelvelabs.TwelveLabsVideoParser + model: "pegasus1.5" + max_tokens: 2048 + cache_strategy: !pw.udfs.DefaultCache {} + # prompt: "Describe this video, focusing on the products that appear and any prices shown." + +# The embedder converts the parsed text into 512-dimensional multimodal +# embeddings using the TwelveLabs Marengo model. +$embedder: !pathway_twelvelabs.MarengoEmbedder + model: "marengo3.0" + cache_strategy: !pw.udfs.DefaultCache {} + retry_strategy: !pw.udfs.ExponentialBackoffRetryStrategy {} + +# Splits the parsed video descriptions into chunks for indexing. +$splitter: !pw.xpacks.llm.splitters.TokenCountSplitter + max_tokens: 400 + +# The LLM used to answer questions over the indexed video descriptions. +# The list of available Pathway LLM wrappers is available here: +# https://pathway.com/developers/api-docs/pathway-xpacks-llm/llms +$llm: !pw.xpacks.llm.llms.OpenAIChat + model: "gpt-4.1-mini" + retry_strategy: !pw.udfs.ExponentialBackoffRetryStrategy + max_retries: 6 + cache_strategy: !pw.udfs.DefaultCache {} + temperature: 0 + capacity: 8 + async_mode: "fully_async" + +# Builds the in-memory vector index over the Marengo embeddings. +$retriever_factory: !pw.indexing.UsearchKnnFactory + reserved_space: 1000 + embedder: $embedder + metric: !pw.indexing.USearchMetricKind.COS + +# Stores and retrieves the indexed video descriptions. +$document_store: !pw.xpacks.llm.document_store.DocumentStore + docs: $sources + parser: $parser + splitter: $splitter + retriever_factory: $retriever_factory + +# The RAG question-answering component served over HTTP. +question_answerer: !pw.xpacks.llm.question_answering.BaseRAGQuestionAnswerer + llm: $llm + indexer: $document_store + # You can set the number of documents to be included as the context of the query + # search_topk: 6 + # You can use your own prompt for querying. + # prompt_template: "Given these documents: {context}, please answer the question: {query}" + +# Change host and port of the webserver by uncommenting these lines +# host: "0.0.0.0" +# port: 8000 + +# By default, caching is enabled for UDFs with cache_strategy set. +# You can disable it by uncommenting the following line. +# persistence_mode: null diff --git a/templates/video_rag_twelvelabs/data/.gitkeep b/templates/video_rag_twelvelabs/data/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/templates/video_rag_twelvelabs/pathway_twelvelabs/__init__.py b/templates/video_rag_twelvelabs/pathway_twelvelabs/__init__.py new file mode 100644 index 0000000..e763ebd --- /dev/null +++ b/templates/video_rag_twelvelabs/pathway_twelvelabs/__init__.py @@ -0,0 +1,309 @@ +# Copyright © 2026 Pathway + +"""TwelveLabs components for the Pathway Live Data Framework. + +This module provides two opt-in building blocks that let Pathway pipelines work +directly with video: + +* :class:`TwelveLabsVideoParser` - a Pathway parser (``pw.UDF``) that turns raw + video bytes into text using TwelveLabs' `Pegasus + `_ video-understanding + model. The resulting text can be chunked, embedded and indexed by any of the + standard Pathway RAG components, exactly like the output of the built-in PDF + parsers. +* :class:`MarengoEmbedder` - a Pathway embedder (``BaseEmbedder``) backed by + TwelveLabs' `Marengo `_ + multimodal embedding model. It returns 512-dimensional vectors that live in the + same embedding space for text, image, audio and video, which makes it a natural + choice when indexing the text produced by ``TwelveLabsVideoParser``. + +Both components require the official ``twelvelabs`` Python SDK (``>=1.2.8``) and a +TwelveLabs API key. The key is read from the ``TWELVELABS_API_KEY`` environment +variable unless it is passed explicitly to the constructor. +""" + +import asyncio +import logging +import os +import time + +import numpy as np +import pathway as pw +from pathway import udfs +from pathway.xpacks.llm.embedders import BaseEmbedder + +DEFAULT_PEGASUS_MODEL = "pegasus1.5" +DEFAULT_MARENGO_MODEL = "marengo3.0" +DEFAULT_PROMPT = ( + "Describe this video in detail. Summarize what happens, who and what appears, " + "the setting, any spoken or on-screen text, and the overall topic. " + "Write the description so it can be used to answer questions about the video." +) + +logger = logging.getLogger(__name__) + + +def _resolve_api_key(api_key: str | None) -> str: + key = api_key or os.environ.get("TWELVELABS_API_KEY") + if not key: + raise ValueError( + "TwelveLabs API key is missing. Pass `api_key=...` or set the " + "`TWELVELABS_API_KEY` environment variable." + ) + return key + + +def _build_client(api_key: str | None): + try: + from twelvelabs import TwelveLabs + except ImportError as e: + raise ImportError( + "The `twelvelabs` package is required to use the TwelveLabs components. " + "Install it with `pip install twelvelabs>=1.2.8`." + ) from e + return TwelveLabs(api_key=_resolve_api_key(api_key)) + + +def _build_async_client(api_key: str | None): + try: + from twelvelabs import AsyncTwelveLabs + except ImportError as e: + raise ImportError( + "The `twelvelabs` package is required to use the TwelveLabs components. " + "Install it with `pip install twelvelabs>=1.2.8`." + ) from e + return AsyncTwelveLabs(api_key=_resolve_api_key(api_key)) + + +class TwelveLabsVideoParser(pw.UDF): + """Parse videos into text using the TwelveLabs Pegasus model. + + The parser uploads the incoming video bytes to TwelveLabs as an asset, waits + for the asset to be ready, and then asks Pegasus to produce a textual + description of the video using ``prompt``. The returned text is suitable for + chunking, embedding and indexing by the standard Pathway RAG components. + + By default the uploaded asset is deleted once the analysis finishes (even if + the analysis fails), so repeated runs do not flood the TwelveLabs asset list. + Set ``delete_assets=False`` to keep the assets around for reuse or + inspection; in that case the emitted ``twelvelabs_asset_id`` metadata refers + to a live, retrievable asset. + + Args: + prompt: Instruction sent to Pegasus describing what to extract from the + video. Defaults to a generic, RAG-oriented description prompt. + model: Pegasus model name. Defaults to ``"pegasus1.5"``. + api_key: TwelveLabs API key. If ``None``, the SDK reads it from the + ``TWELVELABS_API_KEY`` environment variable. + max_tokens: Maximum number of tokens Pegasus may generate. Defaults to 2048. + temperature: Sampling temperature for Pegasus. Defaults to ``None`` (SDK default). + asset_poll_interval: Seconds between asset-readiness checks. Defaults to 5. + asset_timeout: Maximum number of seconds to wait for an uploaded asset to + become ready before raising. Defaults to 600. + delete_assets: If ``True`` (the default), the uploaded asset is deleted + after the analysis completes, so repeated runs do not accumulate + assets in your TwelveLabs account. When ``True``, the emitted + ``twelvelabs_asset_id`` metadata is omitted because the asset no + longer exists. Set to ``False`` to keep assets (e.g. for reuse or + debugging), in which case the id is included in the metadata. + cache_strategy: Pathway caching strategy. To enable caching, pass a valid + :py:class:`~pathway.udfs.CacheStrategy`. Defaults to ``None``. + + Example: + + >>> import pathway as pw # doctest: +SKIP + >>> from pathway_twelvelabs import TwelveLabsVideoParser # doctest: +SKIP + >>> parser = TwelveLabsVideoParser() # doctest: +SKIP + """ + + def __init__( + self, + prompt: str = DEFAULT_PROMPT, + model: str = DEFAULT_PEGASUS_MODEL, + api_key: str | None = None, + max_tokens: int = 2048, + temperature: float | None = None, + asset_poll_interval: float = 5.0, + asset_timeout: float = 600.0, + delete_assets: bool = True, + cache_strategy: udfs.CacheStrategy | None = None, + ): + super().__init__(cache_strategy=cache_strategy) + self.prompt = prompt + self.model = model + self.max_tokens = max_tokens + self.temperature = temperature + self.asset_poll_interval = asset_poll_interval + self.asset_timeout = asset_timeout + self.delete_assets = delete_assets + self._api_key = api_key + self._client = None + + @property + def client(self): + if self._client is None: + self._client = _build_client(self._api_key) + return self._client + + def _upload_asset(self, contents: bytes) -> str: + """Upload video bytes and return the asset id once it is ready.""" + asset = self.client.assets.create( + method="direct", file=("video.mp4", contents), filename="video.mp4" + ) + deadline = time.monotonic() + self.asset_timeout + while asset.status not in ("ready", "failed"): + if time.monotonic() > deadline: + raise TimeoutError( + f"TwelveLabs asset {asset.id} was not ready after " + f"{self.asset_timeout}s (last status: {asset.status})." + ) + time.sleep(self.asset_poll_interval) + asset = self.client.assets.retrieve(asset.id) + if asset.status == "failed": + raise RuntimeError(f"TwelveLabs asset {asset.id} failed to process.") + return asset.id + + def __wrapped__(self, contents: bytes, **kwargs) -> list[tuple[str, dict]]: + from twelvelabs.types.video_context import VideoContext_AssetId + + asset_id = self._upload_asset(contents) + try: + logger.info("Analyzing TwelveLabs asset %s with Pegasus...", asset_id) + analyze_kwargs: dict = dict( + model_name=self.model, + video=VideoContext_AssetId(asset_id=asset_id), + prompt=self.prompt, + max_tokens=self.max_tokens, + ) + if self.temperature is not None: + analyze_kwargs["temperature"] = self.temperature + response = self.client.analyze(**analyze_kwargs) + text = response.data or "" + finally: + if self.delete_assets: + # Remove the per-run asset so repeated runs do not flood the + # TwelveLabs asset list. Best-effort: a cleanup failure must not + # mask the analysis result (or an analysis error above). + try: + self.client.assets.delete(asset_id) + except Exception: # noqa: BLE001 + logger.warning("Failed to delete TwelveLabs asset %s.", asset_id) + # When the asset has been deleted the id no longer resolves, so only + # surface it in the metadata when the asset is kept around. + metadata = {} if self.delete_assets else {"twelvelabs_asset_id": asset_id} + return [(text, metadata)] + + def __call__(self, contents: pw.ColumnExpression, **kwargs) -> pw.ColumnExpression: + """Parse the video document. + + Args: + contents: Column with the raw bytes of each video. + + Returns: + A column with a list of ``(text, metadata)`` pairs for each video. + When ``delete_assets=False`` the metadata records the TwelveLabs + ``twelvelabs_asset_id`` used for the analysis; with the default + ``delete_assets=True`` the asset is removed afterwards and the id is + omitted (it would no longer resolve). + """ + return super().__call__(contents, **kwargs) + + +class MarengoEmbedder(BaseEmbedder): + """Embed text using the TwelveLabs Marengo multimodal embedding model. + + Marengo returns 512-dimensional embeddings in a shared multimodal space, so the + text it produces is directly comparable with image, audio and video embeddings + from the same model. This makes it a natural retriever embedder for pipelines + that index video with :class:`TwelveLabsVideoParser`. + + Args: + model: Marengo model name. Defaults to ``"marengo3.0"``. + api_key: TwelveLabs API key. If ``None``, the SDK reads it from the + ``TWELVELABS_API_KEY`` environment variable. + capacity: Maximum number of concurrent operations. Defaults to ``None`` + (no specific limit). + retry_strategy: Strategy for handling retries. Defaults to + :py:class:`~pathway.udfs.ExponentialBackoffRetryStrategy`. + cache_strategy: Pathway caching strategy. Defaults to ``None``. + + Example: + + >>> import pathway as pw # doctest: +SKIP + >>> from pathway_twelvelabs import MarengoEmbedder # doctest: +SKIP + >>> embedder = MarengoEmbedder() # doctest: +SKIP + """ + + def __init__( + self, + *, + model: str = DEFAULT_MARENGO_MODEL, + api_key: str | None = None, + capacity: int | None = None, + retry_strategy: ( + udfs.AsyncRetryStrategy | None + ) = pw.udfs.ExponentialBackoffRetryStrategy(), + cache_strategy: udfs.CacheStrategy | None = None, + ): + executor = udfs.async_executor(capacity=capacity, retry_strategy=retry_strategy) + # Marengo embeds one text per request, so keep batches at size 1. + super().__init__( + executor=executor, cache_strategy=cache_strategy, max_batch_size=1 + ) + self.model = model + self._api_key = api_key + self._client = None + self._aclient = None + + @property + def client(self): + if self._client is None: + self._client = _build_client(self._api_key) + return self._client + + @property + def aclient(self): + if self._aclient is None: + self._aclient = _build_async_client(self._api_key) + return self._aclient + + def get_embedding_dimension(self, **kwargs) -> int: + """Return the embedding dimension (512 for Marengo). + + This is a one-time, setup-time probe: Pathway calls it once while + building the index, not on the per-document hot path. The single + synchronous request issued here is therefore intentional and acceptable + (the actual embedding hot path runs asynchronously via ``__wrapped__``). + + The base implementation probes ``__wrapped__`` with a single string and + takes ``len`` of the result; since this embedder always returns a list of + vectors, probe with a one-element list and measure the first vector + instead (mirroring Pathway's ``SentenceTransformerEmbedder``). + """ + return len(self._embed_one(".")) + + def _embed_one(self, text: str) -> np.ndarray: + """Synchronous single-text embed, used only for the setup-time probe.""" + response = self.client.embed.create(model_name=self.model, text=text) + vector = response.text_embedding.segments[0].float_ + return np.array(vector, dtype=np.float32) + + async def _aembed_one(self, text: str) -> np.ndarray: + resp = await self.aclient.embed.create(model_name=self.model, text=text) + vector = resp.text_embedding.segments[0].float_ + return np.array(vector, dtype=np.float32) + + async def __wrapped__(self, inputs: list[str], **kwargs) -> list[np.ndarray]: + """Embed the given texts with Marengo. + + Marengo embeds one text per request, so the requests are issued + concurrently on the async TwelveLabs client (``AsyncTwelveLabs``) rather + than serially, keeping the embedding hot path non-blocking. + + Args: + inputs: the strings to embed. + + Returns: + A list of 512-dimensional ``numpy`` arrays, one per input string. + """ + return list(await asyncio.gather(*[self._aembed_one(t) for t in inputs])) diff --git a/templates/video_rag_twelvelabs/requirements.txt b/templates/video_rag_twelvelabs/requirements.txt new file mode 100644 index 0000000..7f2ca56 --- /dev/null +++ b/templates/video_rag_twelvelabs/requirements.txt @@ -0,0 +1,3 @@ +pathway[all] +python-dotenv~=1.0 +twelvelabs>=1.2.8 diff --git a/templates/video_rag_twelvelabs/test_twelvelabs.py b/templates/video_rag_twelvelabs/test_twelvelabs.py new file mode 100644 index 0000000..f06e965 --- /dev/null +++ b/templates/video_rag_twelvelabs/test_twelvelabs.py @@ -0,0 +1,249 @@ +# Copyright © 2026 Pathway + +"""Tests for the TwelveLabs video-RAG components. + +The no-network tests stub the TwelveLabs SDK and run without any credentials. +The live test is skipped unless ``TWELVELABS_API_KEY`` is set in the environment. + +Run with:: + + pytest templates/video_rag_twelvelabs/test_twelvelabs.py +""" + +import os +import sys +import types + +import numpy as np +import pytest + +# Importing `pathway.xpacks.llm.embedders` runs the xpacks package __init__, +# which eagerly imports sibling submodules (parsers, document_store, ...) that +# in turn pull in heavy, optional document-parsing dependencies (docling, +# unstructured, ...). Those are present in the `pathwaycom/pathway` Docker image +# used to run this template but are not needed to exercise the TwelveLabs +# components. If they are unavailable, stub the siblings so the import chain +# succeeds in a lightweight test environment. +try: # pragma: no cover - only triggers when optional deps are missing + import pathway.xpacks.llm.embedders # noqa: F401 +except ImportError: + for _name in ( + "parsers", + "document_store", + "question_answering", + "rerankers", + "servers", + "splitters", + "vector_store", + "llms", + "prompts", + ): + _full = f"pathway.xpacks.llm.{_name}" + sys.modules.setdefault(_full, types.ModuleType(_full)) + +from pathway_twelvelabs import ( # noqa: E402 + DEFAULT_MARENGO_MODEL, + DEFAULT_PEGASUS_MODEL, + MarengoEmbedder, + TwelveLabsVideoParser, +) + + +class _FakeSegment: + def __init__(self, vector): + self.float_ = vector + + +class _FakeTextEmbedding: + def __init__(self, vector): + self.segments = [_FakeSegment(vector)] + + +class _FakeEmbeddingResponse: + def __init__(self, vector): + self.text_embedding = _FakeTextEmbedding(vector) + + +class _FakeEmbed: + def __init__(self, vector): + self._vector = vector + self.calls = [] + + def create(self, *, model_name, text): + self.calls.append((model_name, text)) + return _FakeEmbeddingResponse(self._vector) + + +class _FakeAsset: + def __init__(self, id, status): + self.id = id + self.status = status + + +class _FakeAssets: + def __init__(self): + self.uploaded = None + self.deleted = [] + + def create(self, *, method, file, filename): + self.uploaded = (method, filename) + return _FakeAsset("asset-123", "ready") + + def retrieve(self, asset_id): + return _FakeAsset(asset_id, "ready") + + def delete(self, asset_id): + self.deleted.append(asset_id) + + +class _FakeAnalyzeResponse: + def __init__(self, data): + self.data = data + + +class _FakeClient: + def __init__(self, *, vector=None, analyze_text="a description"): + self.embed = _FakeEmbed(vector or [0.0] * 512) + self.assets = _FakeAssets() + self._analyze_text = analyze_text + self.analyze_calls = [] + + def analyze(self, **kwargs): + self.analyze_calls.append(kwargs) + return _FakeAnalyzeResponse(self._analyze_text) + + +# --- No-network unit tests ------------------------------------------------- + + +class _FakeAsyncEmbed: + def __init__(self, vector): + self._vector = vector + self.calls = [] + + async def create(self, *, model_name, text): + self.calls.append((model_name, text)) + return _FakeEmbeddingResponse(self._vector) + + +class _FakeAsyncClient: + def __init__(self, *, vector=None): + self.embed = _FakeAsyncEmbed(vector or [0.0] * 512) + + +def test_embedder_returns_vector_array(): + embedder = MarengoEmbedder() + embedder._client = _FakeClient(vector=list(range(512))) + + out = embedder._embed_one("a red car") + + assert isinstance(out, np.ndarray) + assert out.shape == (512,) + assert out.dtype == np.float32 + assert embedder._client.embed.calls == [(DEFAULT_MARENGO_MODEL, "a red car")] + + +def test_embedder_defaults(): + embedder = MarengoEmbedder() + assert embedder.model == DEFAULT_MARENGO_MODEL + # Marengo embeds a single text per request. + assert embedder.max_batch_size == 1 + + +def test_embedder_wrapped_is_async_and_concurrent(): + # The hot path runs on the async client and returns one array per input. + import asyncio + + embedder = MarengoEmbedder() + embedder._aclient = _FakeAsyncClient(vector=list(range(512))) + + out = asyncio.run(embedder.__wrapped__(["a red car", "a blue boat"])) + + assert len(out) == 2 + for arr in out: + assert isinstance(arr, np.ndarray) + assert arr.shape == (512,) + assert arr.dtype == np.float32 + assert embedder._aclient.embed.calls == [ + (DEFAULT_MARENGO_MODEL, "a red car"), + (DEFAULT_MARENGO_MODEL, "a blue boat"), + ] + + +def test_embedding_dimension_probe_returns_512(): + # `BaseEmbedder.get_embedding_dimension` probes `__wrapped__` with a single + # string (not a list); the index factory relies on this returning the true + # vector size, so `__wrapped__` must handle a bare string input. + embedder = MarengoEmbedder() + embedder._client = _FakeClient(vector=[0.0] * 512) + assert embedder.get_embedding_dimension() == 512 + + +def test_video_parser_uploads_then_analyzes_and_deletes_asset(): + # Default: delete_assets=True -> asset is removed and id omitted from metadata. + parser = TwelveLabsVideoParser(prompt="What happens?") + parser._client = _FakeClient(analyze_text="A red car drives on a highway.") + + out = parser.__wrapped__(b"fake-video-bytes") + + assert out == [("A red car drives on a highway.", {})] + # Asset was uploaded via the direct method... + assert parser._client.assets.uploaded[0] == "direct" + # ...deleted afterwards so runs don't flood the asset list... + assert parser._client.assets.deleted == ["asset-123"] + # ...and Pegasus was called with the right model, prompt and asset. + (call,) = parser._client.analyze_calls + assert call["model_name"] == DEFAULT_PEGASUS_MODEL + assert call["prompt"] == "What happens?" + assert call["video"].asset_id == "asset-123" + + +def test_video_parser_keeps_asset_when_disabled(): + parser = TwelveLabsVideoParser(delete_assets=False) + parser._client = _FakeClient(analyze_text="desc") + + out = parser.__wrapped__(b"bytes") + + assert out == [("desc", {"twelvelabs_asset_id": "asset-123"})] + assert parser._client.assets.deleted == [] + + +def test_video_parser_failed_asset_raises(): + parser = TwelveLabsVideoParser() + client = _FakeClient() + client.assets.create = lambda **kw: _FakeAsset("a", "failed") + parser._client = client + + with pytest.raises(RuntimeError): + parser.__wrapped__(b"bytes") + # Failure happens during upload (before analyze); nothing was deleted. + assert client.assets.deleted == [] + + +def test_video_parser_deletes_asset_even_when_analyze_raises(): + parser = TwelveLabsVideoParser() + client = _FakeClient() + + def _boom(**kwargs): + raise RuntimeError("pegasus exploded") + + client.analyze = _boom + parser._client = client + + with pytest.raises(RuntimeError): + parser.__wrapped__(b"bytes") + # try/finally still cleaned up the uploaded asset. + assert client.assets.deleted == ["asset-123"] + + +# --- Live smoke test (requires TWELVELABS_API_KEY) ------------------------- + + +@pytest.mark.skipif( + not os.environ.get("TWELVELABS_API_KEY"), + reason="TWELVELABS_API_KEY not set; skipping live TwelveLabs call", +) +def test_marengo_live_embedding_is_512_dim(): + embedder = MarengoEmbedder() + vector = embedder._embed_one("a red car driving on a highway") + assert vector.shape == (512,)