turbopuffer cleanup

Author: tschellenbach

agents-core/vision_agents/core/rag/__init__.py

@@ -0,0 +1,3 @@
+from .rag import Document, RAG
+
+__all__ = ["Document", "RAG"]

agents-core/vision_agents/core/rag/rag.py

@@ -0,0 +1,99 @@
+import abc
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class Document:
+    """A document to be indexed in the RAG system."""
+
+    text: str
+    source: str
+    metadata: dict | None = None
+
+
+class RAG(abc.ABC):
+    """
+    Abstract base class for RAG (Retrieval Augmented Generation) implementations.
+
+    The full complexities of RAG are beyond the scope of this project.
+    We ship with examples including TurboPuffer RAG with hybrid search.
+
+    The documentation explains in greater detail how to build RAG.
+    """
+
+    @abc.abstractmethod
+    async def add_documents(self, documents: list[Document]) -> int:
+        """
+        Add documents to the RAG index.
+
+        Args:
+            documents: List of documents to index.
+
+        Returns:
+            Number of chunks indexed.
+        """
+
+    async def add_directory(
+        self,
+        path: str | Path,
+        extensions: list[str] | None = None,
+    ) -> int:
+        """
+        Add all files from a directory to the RAG index.
+
+        Args:
+            path: Path to directory containing files.
+            extensions: File extensions to include (e.g., ['.md', '.txt']).
+                       Defaults to ['.md', '.txt'].
+
+        Returns:
+            Total number of chunks indexed.
+        """
+        directory = Path(path)
+        if not directory.is_dir():
+            raise NotADirectoryError(f"Not a directory: {directory}")
+
+        if extensions is None:
+            extensions = [".md", ".txt"]
+
+        # Normalize extensions
+        extensions = [
+            ext.lower() if ext.startswith(".") else f".{ext.lower()}"
+            for ext in extensions
+        ]
+
+        files = [
+            f
+            for f in directory.iterdir()
+            if f.is_file() and f.suffix.lower() in extensions
+        ]
+
+        if not files:
+            return 0
+
+        documents = [
+            Document(text=f.read_text(), source=f.name) for f in files
+        ]
+
+        return await self.add_documents(documents)
+
+    @abc.abstractmethod
+    async def search(self, query: str, top_k: int = 3) -> str:
+        """
+        Search the knowledge base.
+
+        Args:
+            query: Search query.
+            top_k: Number of results to return.
+
+        Returns:
+            Formatted string with search results.
+        """
+
+    @abc.abstractmethod
+    async def clear(self) -> None:
+        """Clear all indexed documents."""
+
+    async def close(self) -> None:
+        """Close any open resources. Override if needed."""

examples/03_phone_and_rag_example/inbound_phone_and_rag_example.py

@@ -28,7 +28,7 @@
 from fastapi.responses import JSONResponse
 
 from vision_agents.core import User, Agent
-from vision_agents.plugins import getstream, gemini, twilio, elevenlabs, deepgram
+from vision_agents.plugins import getstream, gemini, twilio, elevenlabs, deepgram, turbopuffer
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -111,7 +111,7 @@ async def media_stream(websocket: WebSocket, call_id: str, token: str):
         call_registry.remove(call_id)
 
 
-async def create_rag_knowledge():
+async def create_rag_from_directory():
     """Initialize the RAG backend based on RAG_BACKEND environment variable."""
     global file_search_store, rag
 
@@ -120,74 +120,38 @@ async def create_rag_knowledge():
         return
 
     if RAG_BACKEND == "turbopuffer":
-        await _init_turbopuffer_rag()
+        logger.info(f"📚 Initializing TurboPuffer RAG from {KNOWLEDGE_DIR}")
+        rag = await turbopuffer.create_rag(
+            namespace="stream-product-knowledge-gemini",
+            knowledge_dir=KNOWLEDGE_DIR,
+            extensions=[".md"],
+        )
+        logger.info(f"✅ TurboPuffer RAG ready with {len(rag._indexed_files)} documents indexed")
     else:
-        await _init_gemini_rag()
+        logger.info(f"📚 Initializing Gemini File Search from {KNOWLEDGE_DIR}")
+        file_search_store = await gemini.create_file_search_store(
+            name="stream-product-knowledge",
+            knowledge_dir=KNOWLEDGE_DIR,
+            extensions=[".md"],
+        )
+        logger.info(f"✅ Gemini RAG ready with {len(file_search_store._uploaded_files)} documents")
 
 
-async def _init_gemini_rag():
-    """Initialize Gemini File Search RAG."""
-    global file_search_store
 
-    logger.info(f"📚 Initializing Gemini File Search from {KNOWLEDGE_DIR}")
-    file_search_store = await gemini.create_file_search_store(
-        name="stream-product-knowledge",
-        knowledge_dir=KNOWLEDGE_DIR,
-        extensions=[".md"],
-    )
-    logger.info(f"✅ Gemini RAG ready with {len(file_search_store._uploaded_files)} documents")
-
-
-async def _init_turbopuffer_rag():
-    """Initialize TurboPuffer + LangChain RAG."""
-    global rag
-
-    from rag_turbopuffer import create_rag
-
-    logger.info(f"📚 Initializing TurboPuffer RAG from {KNOWLEDGE_DIR}")
-    rag = await create_rag(
-        namespace="stream-product-knowledge-gemini",
-        knowledge_dir=KNOWLEDGE_DIR,
-        extensions=[".md"],
-    )
-    logger.info(f"✅ TurboPuffer RAG ready with {len(rag._indexed_files)} documents indexed")
-
-
-
-async def create_agent(**kwargs) -> Agent:
+async def create_agent() -> Agent:
     """Create an agent with RAG capabilities."""
-    if RAG_BACKEND == "turbopuffer":
-        return await _create_agent_turbopuffer()
-    else:
-        return await _create_agent_gemini()
-
-
-async def _create_agent_gemini() -> Agent:
-    """Create agent with Gemini File Search RAG."""
     instructions = """Read the instructions in @instructions.md"""
 
-    return Agent(
-        edge=getstream.Edge(),
-        agent_user=User(id="ai-agent", name="AI"),
-        instructions=instructions,
-        tts=elevenlabs.TTS(voice_id="FGY2WhTYpPnrIDTdsKH5"),
-        stt=deepgram.STT(eager_turn_detection=True),
-        llm=gemini.LLM("gemini-2.5-flash-lite", file_search_store=file_search_store),
-    )
-
-
-async def _create_agent_turbopuffer() -> Agent:
-    """Create agent with TurboPuffer RAG via function calling."""
-    instructions = """Read the instructions in @instructions.md"""
-
-    llm = gemini.LLM("gemini-2.5-flash-lite")
+    if RAG_BACKEND == "turbopuffer":
+        llm = gemini.LLM("gemini-2.5-flash-lite")
 
-    # Register RAG search as a callable function
-    @llm.register_function(
-        description="Search Stream's product knowledge base for detailed information about Chat, Video, Feeds, and Moderation APIs."
-    )
-    async def search_knowledge(query: str) -> str:
-        return await rag.search(query, top_k=3)
+        @llm.register_function(
+            description="Search Stream's product knowledge base for detailed information about Chat, Video, Feeds, and Moderation APIs."
+        )
+        async def search_knowledge(query: str) -> str:
+            return await rag.search(query, top_k=3)
+    else:
+        llm = gemini.LLM("gemini-2.5-flash-lite", file_search_store=file_search_store)
 
     return Agent(
         edge=getstream.Edge(),
@@ -200,6 +164,6 @@ async def search_knowledge(query: str) -> str:
 
 
 if __name__ == "__main__":
-    asyncio.run(create_rag_knowledge())
+    asyncio.run(create_rag_from_directory())
     logger.info(f"Starting with RAG_BACKEND={RAG_BACKEND}")
     uvicorn.run(app, host="localhost", port=8000)

plugins/gemini/vision_agents/plugins/gemini/__init__.py

@@ -1,13 +1,14 @@
 from .gemini_llm import GeminiLLM as LLM
 from .gemini_realtime import GeminiRealtime as Realtime
-from .file_search import FileSearchStore, create_file_search_store
+from .file_search import GeminiFilesearchRAG, FileSearchStore, create_file_search_store
 from google.genai.types import ThinkingLevel, MediaResolution
 
 __all__ = [
     "Realtime",
     "LLM",
     "ThinkingLevel",
     "MediaResolution",
-    "FileSearchStore",
+    "GeminiFilesearchRAG",
+    "FileSearchStore",  # Backwards compatibility alias
     "create_file_search_store",
 ]