vision-rag 0.1.0__tar.gz

vision_rag-0.1.0/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: vision-rag
+Version: 0.1.0
+Summary: Video RAG — a Python library for retrieval-augmented generation over video
+Requires-Python: >=3.9
+Requires-Dist: pymediainfo

vision_rag-0.1.0/README.md

@@ -0,0 +1,277 @@
+# vision-rag
+
+**Video RAG** — A Python library for Retrieval-Augmented Generation over video.
+
+Ask questions about any video and get answers using the transcript and visual frames.
+
+---
+
+## Install
+
+```bash
+pip install vision-rag
+```
+
+---
+
+## How it works
+
+```
+Video → Chunks → Embeddings → Vector Store → Retrieval → Answer
+```
+
+1. **Ingest** — reads video metadata
+2. **Chunk** — splits video into time-based overlapping chunks with frames and transcript
+3. **Embed** — converts text and frames into vectors (your choice of model)
+4. **Index** — stores vectors in FAISS or Chroma
+5. **Retrieve** — searches both text and image indexes for a query
+6. **Generate** — passes retrieved chunks to a VLM to generate the answer
+
+---
+
+## Quick Start
+
+```python
+from vision-rag.video_ingestion import VideoLoader
+from vision-rag.video_chunker import Chunker, WhisperLocalASR
+from vision-rag.embedding import EmbeddingBuilder, BaseTextEmbedder, BaseImageEmbedder
+from vision-rag.vectorstores import FAISS
+from vision-rag.retriever import Retriever
+from vision-rag.generator import Generator, OllamaGenerator
+import requests, base64
+
+# --- your choice of embedder (example: Jina v4) ---
+class JinaTextEmbedder(BaseTextEmbedder):
+    def __init__(self, api_key):
+        self.api_key = api_key
+    def embed(self, text):
+        r = requests.post("https://api.jina.ai/v1/embeddings",
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            json={"model": "jina-embeddings-v4", "input": [{"text": text}], "task": "retrieval.passage"})
+        return r.json()["data"][0]["embedding"]
+
+class JinaImageEmbedder(BaseImageEmbedder):
+    def __init__(self, api_key):
+        self.api_key = api_key
+    def embed(self, image_path):
+        with open(image_path, "rb") as f:
+            b64 = base64.b64encode(f.read()).decode()
+        r = requests.post("https://api.jina.ai/v1/embeddings",
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            json={"model": "jina-embeddings-v4", "input": [{"image": b64}], "task": "retrieval.passage"})
+        return r.json()["data"][0]["embedding"]
+
+# Stage 1 — Ingest
+video_doc = VideoLoader().load("video.mp4")
+
+# Stage 2 — Chunk
+chunks = Chunker(
+    asr=WhisperLocalASR(model_size="base"),
+    use_asr=True,
+    use_frames=True,
+    chunk_size=5.0,
+    chunk_overlap=1.0,
+).chunk("video.mp4")
+
+# Stage 3 — Embed
+text_embedder  = JinaTextEmbedder(api_key="your_jina_key")
+image_embedder = JinaImageEmbedder(api_key="your_jina_key")
+embedded_chunks = EmbeddingBuilder(
+    text_embedding=text_embedder,
+    image_embedding=image_embedder,
+).embed(chunks)
+
+# Stage 4 — Index
+store = FAISS()
+store.index(embedded_chunks)
+
+# Stage 5 + 6 — Retrieve and Generate
+query   = input("Ask a question: ")
+results = Retriever(store=store, text_embedder=text_embedder).retrieve(query)
+answer  = Generator(llm=OllamaGenerator(model="llava:7b")).generate(query=query, results=results)
+
+print(answer.text)
+```
+
+---
+
+## Chunker
+
+```python
+from vision-rag.video_chunker import Chunker, WhisperLocalASR
+
+chunker = Chunker(
+    asr=WhisperLocalASR(model_size="medium"),  # or OpenAIASR, DeepgramASR, or your own
+    use_asr=True,
+    use_frames=True,
+    chunk_size=5.0,      # seconds per chunk
+    chunk_overlap=1.0,   # overlap between chunks
+)
+chunks = chunker.chunk("video.mp4")
+```
+
+Each chunk contains:
+
+| Field | Description |
+|---|---|
+| `chunk.chunk_id` | chunk index |
+| `chunk.start` | start time in seconds |
+| `chunk.end` | end time in seconds |
+| `chunk.duration` | duration in seconds |
+| `chunk.text` | transcript for this chunk |
+| `chunk.frame_path` | path to keyframe image |
+| `chunk.metadata` | source info, asr provider, etc. |
+
+---
+
+## ASR — Bring Your Own
+
+vision-rag ships with built-in ASR providers but you can plug in anything:
+
+```python
+from vision-rag.video_chunker import BaseASR
+
+# built-in
+from vision-rag.video_chunker import WhisperLocalASR, OpenAIASR, DeepgramASR
+
+# your own — any model, any API
+class MyASR(BaseASR):
+    def transcribe(self, audio_path: str) -> list[dict]:
+        return [{"start": 0.0, "end": 5.0, "text": "..."}]
+
+chunker = Chunker(asr=MyASR(), use_asr=True)
+```
+
+---
+
+## Embedding — Bring Your Own
+
+```python
+from vision-rag.embedding import BaseTextEmbedder, BaseImageEmbedder
+
+# your own text embedder
+class MyTextEmbedder(BaseTextEmbedder):
+    def embed(self, text: str) -> list[float]:
+        return [...]  # your model or API
+
+# your own image embedder
+class MyImageEmbedder(BaseImageEmbedder):
+    def embed(self, image_path: str) -> list[float]:
+        return [...]  # your model or API
+
+embedder = EmbeddingBuilder(
+    text_embedding=MyTextEmbedder(),
+    image_embedding=MyImageEmbedder(),
+)
+```
+
+Built-in providers: `OpenAITextEmbedder`, `SentenceTransformerTextEmbedder`, `CLIPImageEmbedder`, `OpenAIImageEmbedder`
+
+---
+
+## Vector Stores
+
+```python
+from vision-rag.vectorstores import FAISS, Chroma
+
+# FAISS — fast local search
+store = FAISS()
+store.index(embedded_chunks)
+store.save("my_index")
+store.load("my_index")
+
+# Chroma — persistent local DB
+store = Chroma(path="my_chroma_db")
+store.index(embedded_chunks)
+```
+
+Plug in your own:
+
+```python
+from vision-rag.vectorstores import BaseVectorStore
+
+class MyVectorStore(BaseVectorStore):
+    def index(self, embedded_chunks): ...
+    def search_text(self, vector, top_k): ...
+    def search_image(self, vector, top_k): ...
+```
+
+---
+
+## Retrieval
+
+```python
+from vision-rag.retriever import Retriever
+
+retriever = Retriever(
+    store=store,
+    text_embedder=text_embedder,
+    top_k_text=5,
+    top_k_image=5,
+)
+
+# semantic search
+results = retriever.retrieve("What did they say about frozen yogurt?")
+results.text_results    # top text matches
+results.image_results   # top image matches
+results.all             # combined, ranked by score
+
+# time-based search
+chunks = retriever.retrieve_by_time(start=10.0, end=20.0)
+```
+
+---
+
+## Generation — Bring Your Own VLM
+
+```python
+from vision-rag.generator import Generator, OpenAIGenerator, AnthropicGenerator, GeminiGenerator, OllamaGenerator
+
+# GPT-4o
+generator = Generator(llm=OpenAIGenerator(api_key="sk-..."))
+
+# Claude
+generator = Generator(llm=AnthropicGenerator(api_key="sk-ant-..."))
+
+# Gemini
+generator = Generator(llm=GeminiGenerator(api_key="..."))
+
+# Ollama (local)
+generator = Generator(llm=OllamaGenerator(model="llava:7b"))
+
+# your own
+from vision-rag.generator import BaseGenerator
+
+class MyGenerator(BaseGenerator):
+    def generate(self, query: str, chunks) -> str:
+        return "answer..."
+
+generator = Generator(llm=MyGenerator())
+```
+
+---
+
+## Dependencies
+
+vision-rag ships with only one hard dependency — `pymediainfo`. Everything else is installed based on what you use:
+
+| Feature | Install |
+|---|---|
+| ASR (local Whisper) | `pip install faster-whisper` |
+| ASR (OpenAI) | `pip install openai` |
+| ASR (Deepgram) | `pip install deepgram-sdk` |
+| Frames + Audio | `brew install ffmpeg` |
+| FAISS vector store | `pip install faiss-cpu` |
+| Chroma vector store | `pip install chromadb` |
+| OpenAI embedding | `pip install openai` |
+| Sentence Transformers | `pip install sentence-transformers` |
+| CLIP image embedding | `pip install git+https://github.com/openai/CLIP.git torch Pillow` |
+| Ollama generation | `pip install ollama` |
+| Anthropic generation | `pip install anthropic` |
+| Gemini generation | `pip install google-genai` |
+
+---
+
+## License
+
+MIT

vision_rag-0.1.0/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "vision-rag"
+version = "0.1.0"
+description = "Video RAG — a Python library for retrieval-augmented generation over video"
+requires-python = ">=3.9"
+dependencies = [
+    "pymediainfo",
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["vision-rag*"]

vision_rag-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

vision_rag-0.1.0/vision-rag/__init__.py

@@ -0,0 +1,38 @@
+from vision-rag.video_ingestion import VideoLoader, VideoDocument
+from vision-rag.video_chunker import Chunker, Chunk, BaseASR, WhisperLocalASR, OpenAIASR, DeepgramASR
+from vision-rag.embedding import (
+    EmbeddingBuilder, EmbeddedChunk,
+    BaseTextEmbedder, BaseImageEmbedder,
+    OpenAITextEmbedder, SentenceTransformerTextEmbedder,
+    CLIPImageEmbedder, OpenAIImageEmbedder,
+)
+from vision-rag.vectorstores import BaseVectorStore, SearchResult, FAISS, Chroma
+from vision-rag.retriever import Retriever, RetrievalResult
+from vision-rag.generator import (
+    Generator, GeneratorAnswer,
+    BaseGenerator,
+    OpenAIGenerator, AnthropicGenerator, GeminiGenerator, OllamaGenerator,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    # Stage 1
+    "VideoLoader", "VideoDocument",
+    # Stage 2
+    "Chunker", "Chunk",
+    "BaseASR", "WhisperLocalASR", "OpenAIASR", "DeepgramASR",
+    # Stage 3
+    "EmbeddingBuilder", "EmbeddedChunk",
+    "BaseTextEmbedder", "BaseImageEmbedder",
+    "OpenAITextEmbedder", "SentenceTransformerTextEmbedder",
+    "CLIPImageEmbedder", "OpenAIImageEmbedder",
+    # Stage 4
+    "BaseVectorStore", "SearchResult",
+    "FAISS", "Chroma",
+    # Stage 5
+    "Retriever", "RetrievalResult",
+    # Stage 6
+    "Generator", "GeneratorAnswer",
+    "BaseGenerator",
+    "OpenAIGenerator", "AnthropicGenerator", "GeminiGenerator", "OllamaGenerator",
+]

vision_rag-0.1.0/vision-rag/embedding.py

@@ -0,0 +1,328 @@
+"""
+vision-rag/embedding.py
+
+Stage 3 of the vision-rag pipeline — Embedding.
+Converts chunk.text and chunk.frame_path into vectors.
+
+Usage:
+    from vision-rag.embedding import EmbeddingBuilder, CLIPImageEmbedder, OpenAITextEmbedder
+
+    embedder = EmbeddingBuilder(
+        text_embedding=OpenAITextEmbedder(),
+        image_embedding=CLIPImageEmbedder()
+    )
+    embedded_chunks = embedder.embed(chunks)
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+from vision-rag.video_chunker import Chunk
+
+
+# ──────────────────────────────────────────────────────────────
+# EmbeddedChunk — Chunk + vectors
+# ──────────────────────────────────────────────────────────────
+
+@dataclass
+class EmbeddedChunk:
+    """
+    A Chunk with text and image vectors attached.
+    This is what flows into the indexing stage.
+    """
+    # -- original chunk fields (unchanged) --
+    chunk_id:     int
+    video_path:   str
+    start:        float
+    end:          float
+    duration:     float
+    frame_path:   Optional[str]
+    text:         Optional[str]
+    metadata:     dict = field(default_factory=dict)
+
+    # -- new: vectors --
+    text_vector:  Optional[list[float]] = None   # None if text=None or text_embedding not set
+    image_vector: Optional[list[float]] = None   # None if frame_path=None or image_embedding not set
+
+    @classmethod
+    def from_chunk(cls, chunk: Chunk) -> "EmbeddedChunk":
+        """Create an EmbeddedChunk from a Chunk (vectors are None until embed() is called)."""
+        return cls(
+            chunk_id   = chunk.chunk_id,
+            video_path = chunk.video_path,
+            start      = chunk.start,
+            end        = chunk.end,
+            duration   = chunk.duration,
+            frame_path = chunk.frame_path,
+            text       = chunk.text,
+            metadata   = chunk.metadata.copy(),
+        )
+
+    def __repr__(self) -> str:
+        tv = f"dim={len(self.text_vector)}"  if self.text_vector  else "None"
+        iv = f"dim={len(self.image_vector)}" if self.image_vector else "None"
+        return (
+            f"EmbeddedChunk("
+            f"id={self.chunk_id}, "
+            f"start={self.start:.2f}s, "
+            f"end={self.end:.2f}s, "
+            f"text_vector={tv}, "
+            f"image_vector={iv}"
+            f")"
+        )
+
+
+# ──────────────────────────────────────────────────────────────
+# Base classes — plug in anything
+# ──────────────────────────────────────────────────────────────
+
+class BaseTextEmbedder:
+    """
+    Base class for all text embedding providers.
+    Subclass and implement embed() to use any model or API.
+
+    Example:
+        class MyTextEmbedder(BaseTextEmbedder):
+            def embed(self, text: str) -> list[float]:
+                # your model or API here
+                return [...]
+    """
+
+    def embed(self, text: str) -> list[float]:
+        raise NotImplementedError("Implement embed() in your TextEmbedder subclass.")
+
+    @property
+    def provider_name(self) -> str:
+        return self.__class__.__name__
+
+
+class BaseImageEmbedder:
+    """
+    Base class for all image embedding providers.
+    Subclass and implement embed() to use any model or API.
+
+    Example:
+        class MyImageEmbedder(BaseImageEmbedder):
+            def embed(self, image_path: str) -> list[float]:
+                # your model or API here
+                return [...]
+    """
+
+    def embed(self, image_path: str) -> list[float]:
+        raise NotImplementedError("Implement embed() in your ImageEmbedder subclass.")
+
+    @property
+    def provider_name(self) -> str:
+        return self.__class__.__name__
+
+
+# ──────────────────────────────────────────────────────────────
+# Built-in Text Embedders
+# ──────────────────────────────────────────────────────────────
+
+class OpenAITextEmbedder(BaseTextEmbedder):
+    """
+    OpenAI text embeddings (text-embedding-3-small by default).
+
+    Usage:
+        embedder = OpenAITextEmbedder(api_key="sk-...")
+        # or set OPENAI_API_KEY env var and just do OpenAITextEmbedder()
+    """
+
+    def __init__(self, api_key: Optional[str] = None, model: str = "text-embedding-3-small"):
+        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "OpenAI API key required. Pass api_key= or set OPENAI_API_KEY env var."
+            )
+        self.model = model
+        self._client = None
+
+    def embed(self, text: str) -> list[float]:
+        try:
+            from openai import OpenAI
+        except ImportError:
+            raise RuntimeError("pip install openai")
+
+        if self._client is None:
+            self._client = OpenAI(api_key=self.api_key)
+
+        response = self._client.embeddings.create(input=text, model=self.model)
+        return response.data[0].embedding
+
+
+class SentenceTransformerTextEmbedder(BaseTextEmbedder):
+    """
+    Local text embeddings using SentenceTransformers. No API key needed.
+
+    Usage:
+        embedder = SentenceTransformerTextEmbedder()
+        embedder = SentenceTransformerTextEmbedder(model="all-mpnet-base-v2")
+
+    pip install sentence-transformers
+    """
+
+    def __init__(self, model: str = "all-MiniLM-L6-v2"):
+        self.model_name = model
+        self._model = None
+
+    def embed(self, text: str) -> list[float]:
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError:
+            raise RuntimeError("pip install sentence-transformers")
+
+        if self._model is None:
+            self._model = SentenceTransformer(self.model_name)
+
+        return self._model.encode(text).tolist()
+
+
+# ──────────────────────────────────────────────────────────────
+# Built-in Image Embedders
+# ──────────────────────────────────────────────────────────────
+
+class CLIPImageEmbedder(BaseImageEmbedder):
+    """
+    Local image embeddings using OpenAI CLIP. No API key needed.
+    Industry standard for image embeddings.
+
+    Usage:
+        embedder = CLIPImageEmbedder()
+        embedder = CLIPImageEmbedder(model="ViT-B/32")
+
+    pip install git+https://github.com/openai/CLIP.git Pillow torch
+    """
+
+    def __init__(self, model: str = "ViT-B/32", device: str = "cpu"):
+        self.model_name = model
+        self.device = device
+        self._model = None
+        self._preprocess = None
+
+    def embed(self, image_path: str) -> list[float]:
+        try:
+            import clip
+            import torch
+            from PIL import Image
+        except ImportError:
+            raise RuntimeError(
+                "pip install git+https://github.com/openai/CLIP.git Pillow torch"
+            )
+
+        if self._model is None:
+            self._model, self._preprocess = clip.load(self.model_name, device=self.device)
+
+        image = self._preprocess(Image.open(image_path)).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            vector = self._model.encode_image(image)
+        return vector.squeeze().tolist()
+
+
+class OpenAIImageEmbedder(BaseImageEmbedder):
+    """
+    Image embeddings via OpenAI vision model.
+    Encodes the image as base64 and gets an embedding via the API.
+
+    Usage:
+        embedder = OpenAIImageEmbedder(api_key="sk-...")
+        # or set OPENAI_API_KEY env var
+
+    pip install openai
+    """
+
+    def __init__(self, api_key: Optional[str] = None, model: str = "text-embedding-3-small"):
+        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "OpenAI API key required. Pass api_key= or set OPENAI_API_KEY env var."
+            )
+        self.model = model
+        self._client = None
+
+    def embed(self, image_path: str) -> list[float]:
+        import base64
+        try:
+            from openai import OpenAI
+        except ImportError:
+            raise RuntimeError("pip install openai")
+
+        if self._client is None:
+            self._client = OpenAI(api_key=self.api_key)
+
+        with open(image_path, "rb") as f:
+            b64 = base64.b64encode(f.read()).decode("utf-8")
+
+        response = self._client.embeddings.create(
+            input=f"data:image/jpeg;base64,{b64}",
+            model=self.model,
+        )
+        return response.data[0].embedding
+
+
+# ──────────────────────────────────────────────────────────────
+# EmbeddingBuilder — main entry point
+# ──────────────────────────────────────────────────────────────
+
+class EmbeddingBuilder:
+    """
+    Embeds a list of Chunk objects into EmbeddedChunk objects.
+
+    Parameters
+    ----------
+    text_embedding : BaseTextEmbedder | None
+        Any text embedding provider. None = skip text embedding.
+    image_embedding : BaseImageEmbedder | None
+        Any image embedding provider. None = skip image embedding.
+
+    Usage:
+        embedder = EmbeddingBuilder(
+            text_embedding=OpenAITextEmbedder(),
+            image_embedding=CLIPImageEmbedder()
+        )
+        embedded_chunks = embedder.embed(chunks)
+    """
+
+    def __init__(
+        self,
+        text_embedding:  Optional[BaseTextEmbedder]  = None,
+        image_embedding: Optional[BaseImageEmbedder] = None,
+    ):
+        if text_embedding is None and image_embedding is None:
+            raise ValueError(
+                "At least one of text_embedding or image_embedding must be provided."
+            )
+        self.text_embedding  = text_embedding
+        self.image_embedding = image_embedding
+
+    def embed(self, chunks: list[Chunk]) -> list[EmbeddedChunk]:
+        """
+        Embed a list of chunks. Returns a list of EmbeddedChunk objects.
+
+        Parameters
+        ----------
+        chunks : list[Chunk]
+            Output from Chunker.chunk()
+        """
+        embedded = []
+        for chunk in chunks:
+            ec = EmbeddedChunk.from_chunk(chunk)
+
+            # -- text vector --
+            if self.text_embedding and chunk.text:
+                ec.text_vector = self.text_embedding.embed(chunk.text)
+                ec.metadata["text_embedder"] = self.text_embedding.provider_name
+
+            # -- image vector --
+            if self.image_embedding and chunk.frame_path:
+                if Path(chunk.frame_path).exists():
+                    ec.image_vector = self.image_embedding.embed(chunk.frame_path)
+                    ec.metadata["image_embedder"] = self.image_embedding.provider_name
+
+            embedded.append(ec)
+
+        return embedded