From b28b59a49f7ccbe3fdb83dfd36c47720e7a28da2 Mon Sep 17 00:00:00 2001
From: Kian Jones <11655409+kianjones9@users.noreply.github.com>
Date: Fri, 29 Aug 2025 16:44:15 -0700
Subject: [PATCH] chore: add tbpuf secrets and testing (#4292)

* add tbpuf secret

* deploy tbpuf key to prod deploy for core

* add secrets in staging and ci

* rn turbopuff test and add to integration test suite

* add tbpuf

* rename all occurences to tpuf

* remove debug prints

* buffer-puf

* add letta prefix

* add debug

* fix tests
---
 letta/helpers/tpuf_client.py                  | 24 +++++++++----------
 ...ion.py => integration_test_turbopuffer.py} | 17 ++++++++++---
 2 files changed, 26 insertions(+), 15 deletions(-)
 rename tests/{test_turbopuffer_integration.py => integration_test_turbopuffer.py} (96%)

diff --git a/letta/helpers/tpuf_client.py b/letta/helpers/tpuf_client.py
index 69884579..7ed1e915 100644
--- a/letta/helpers/tpuf_client.py
+++ b/letta/helpers/tpuf_client.py
@@ -187,21 +187,21 @@ class TurbopufferClient:
         from turbopuffer import AsyncTurbopuffer
         from turbopuffer.types import QueryParam
 
+        # validate inputs based on search mode first
+        if search_mode == "vector" and query_embedding is None:
+            raise ValueError("query_embedding is required for vector search mode")
+        if search_mode == "fts" and query_text is None:
+            raise ValueError("query_text is required for FTS search mode")
+        if search_mode == "hybrid":
+            if query_embedding is None or query_text is None:
+                raise ValueError("Both query_embedding and query_text are required for hybrid search mode")
+        if search_mode not in ["vector", "fts", "hybrid", "timestamp"]:
+            raise ValueError(f"Invalid search_mode: {search_mode}. Must be 'vector', 'fts', 'hybrid', or 'timestamp'")
+
         # Check if we should fallback to timestamp-based retrieval
-        if query_embedding is None and query_text is None:
+        if query_embedding is None and query_text is None and search_mode not in ["timestamp"]:
             # Fallback to retrieving most recent passages when no search query is provided
             search_mode = "timestamp"
-        else:
-            # validate inputs based on search mode
-            if search_mode == "vector" and query_embedding is None:
-                raise ValueError("query_embedding is required for vector search mode")
-            if search_mode == "fts" and query_text is None:
-                raise ValueError("query_text is required for FTS search mode")
-            if search_mode == "hybrid":
-                if query_embedding is None or query_text is None:
-                    raise ValueError("Both query_embedding and query_text are required for hybrid search mode")
-            if search_mode not in ["vector", "fts", "hybrid"]:
-                raise ValueError(f"Invalid search_mode: {search_mode}. Must be 'vector', 'fts', or 'hybrid'")
 
         namespace_name = self._get_namespace_name(archive_id)
 
diff --git a/tests/test_turbopuffer_integration.py b/tests/integration_test_turbopuffer.py
similarity index 96%
rename from tests/test_turbopuffer_integration.py
rename to tests/integration_test_turbopuffer.py
index f95e986e..0bc8555d 100644
--- a/tests/test_turbopuffer_integration.py
+++ b/tests/integration_test_turbopuffer.py
@@ -7,6 +7,7 @@ from letta.config import LettaConfig
 from letta.helpers.tpuf_client import TurbopufferClient, should_use_tpuf
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import TagMatchMode, VectorDBProvider
+from letta.schemas.passage import Passage
 from letta.server.server import SyncServer
 from letta.settings import settings
 
@@ -374,17 +375,27 @@ class TestTurbopufferIntegration:
             assert all(isinstance(score, float) for _, score in vector_heavy_results)
 
             # Test error handling - missing embedding for vector mode
-            with pytest.raises(ValueError, match="query_embedding is required"):
+            with pytest.raises(ValueError, match="query_embedding is required for vector search mode"):
                 await client.query_passages(archive_id=archive_id, search_mode="vector", top_k=3)
 
             # Test error handling - missing text for FTS mode
-            with pytest.raises(ValueError, match="query_text is required"):
+            with pytest.raises(ValueError, match="query_text is required for FTS search mode"):
                 await client.query_passages(archive_id=archive_id, search_mode="fts", top_k=3)
 
-            # Test error handling - missing both for hybrid mode
+            # Test error handling - missing text for hybrid mode (embedding provided but text missing)
             with pytest.raises(ValueError, match="Both query_embedding and query_text are required"):
                 await client.query_passages(archive_id=archive_id, query_embedding=[1.0, 2.0, 3.0], search_mode="hybrid", top_k=3)
 
+            # Test error handling - missing embedding for hybrid mode (text provided but embedding missing)
+            with pytest.raises(ValueError, match="Both query_embedding and query_text are required"):
+                await client.query_passages(archive_id=archive_id, query_text="test", search_mode="hybrid", top_k=3)
+
+            # Test explicit timestamp mode
+            timestamp_results = await client.query_passages(archive_id=archive_id, search_mode="timestamp", top_k=3)
+            assert len(timestamp_results) <= 3
+            # Should return passages ordered by timestamp (most recent first)
+            assert all(isinstance(passage, Passage) for passage, _ in timestamp_results)
+
         finally:
             # Clean up
             try: