feat: Add support for vectors in SQLite (#3505)

letta/__init__.py

@@ -10,10 +10,12 @@ except PackageNotFoundError:
 if os.environ.get("LETTA_VERSION"):
     __version__ = os.environ["LETTA_VERSION"]
 
-
 # import clients
 from letta.client.client import RESTClient
 
+# Import sqlite_functions early to ensure event handlers are registered
+from letta.orm import sqlite_functions
+
 # # imports for easier access
 from letta.schemas.agent import AgentState
 from letta.schemas.block import Block

letta/helpers/converters.py

@@ -1,4 +1,3 @@
-import base64
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
@@ -43,7 +42,10 @@ from letta.schemas.tool_rule import (
     TerminalToolRule,
     ToolRule,
 )
+from letta.settings import DatabaseChoice, settings
 
+if settings.database_engine == DatabaseChoice.SQLITE:
+    import sqlite_vec
 # --------------------------
 # LLMConfig Serialization
 # --------------------------
@@ -272,22 +274,28 @@ def deserialize_message_content(data: Optional[List[Dict]]) -> List[MessageConte
 
 
 def serialize_vector(vector: Optional[Union[List[float], np.ndarray]]) -> Optional[bytes]:
-    """Convert a NumPy array or list into a base64-encoded byte string."""
+    """Convert a NumPy array or list into serialized format using sqlite-vec."""
     if vector is None:
         return None
     if isinstance(vector, list):
         vector = np.array(vector, dtype=np.float32)
+    else:
+        vector = vector.astype(np.float32)
 
-    return base64.b64encode(vector.tobytes())
+    return sqlite_vec.serialize_float32(vector.tolist())
 
 
 def deserialize_vector(data: Optional[bytes], dialect: Dialect) -> Optional[np.ndarray]:
-    """Convert a base64-encoded byte string back into a NumPy array."""
+    """Convert serialized data back into a NumPy array using sqlite-vec format."""
     if not data:
         return None
 
     if dialect.name == "sqlite":
-        data = base64.b64decode(data)
+        # Use sqlite-vec format
+        if len(data) % 4 == 0:  # Must be divisible by 4 for float32
+            return np.frombuffer(data, dtype=np.float32)
+        else:
+            raise ValueError(f"Invalid sqlite-vec binary data length: {len(data)}")
 
     return np.frombuffer(data, dtype=np.float32)
 

letta/orm/sqlite_functions.py

@@ -1,4 +1,3 @@
-import base64
 import sqlite3
 from typing import Optional, Union
 
@@ -7,11 +6,15 @@ from sqlalchemy import event
 from sqlalchemy.engine import Engine
 
 from letta.constants import MAX_EMBEDDING_DIM
+from letta.settings import DatabaseChoice, settings
+
+if settings.database_engine == DatabaseChoice.SQLITE:
+    import sqlite_vec
 
 
 def adapt_array(arr):
     """
-    Converts numpy array to binary for SQLite storage
+    Converts numpy array to binary for SQLite storage using sqlite-vec
     """
     if arr is None:
         return None
@@ -21,15 +24,14 @@ def adapt_array(arr):
     elif not isinstance(arr, np.ndarray):
         raise ValueError(f"Unsupported type: {type(arr)}")
 
-    # Convert to bytes and then base64 encode
-    bytes_data = arr.tobytes()
-    base64_data = base64.b64encode(bytes_data)
-    return sqlite3.Binary(base64_data)
+    # Ensure float32 for compatibility
+    arr = arr.astype(np.float32)
+    return sqlite_vec.serialize_float32(arr.tolist())
 
 
 def convert_array(text):
     """
-    Converts binary back to numpy array
+    Converts binary back to numpy array using sqlite-vec format
     """
     if text is None:
         return None
@@ -41,13 +43,11 @@ def convert_array(text):
     # Handle both bytes and sqlite3.Binary
     binary_data = bytes(text) if isinstance(text, sqlite3.Binary) else text
 
-    try:
-        # First decode base64
-        decoded_data = base64.b64decode(binary_data)
-        # Then convert to numpy array
-        return np.frombuffer(decoded_data, dtype=np.float32)
-    except Exception:
-        return None
+    # Use sqlite-vec native format
+    if len(binary_data) % 4 == 0:  # Must be divisible by 4 for float32
+        return np.frombuffer(binary_data, dtype=np.float32)
+    else:
+        raise ValueError(f"Invalid sqlite-vec binary data length: {len(binary_data)}")
 
 
 def verify_embedding_dimension(embedding: np.ndarray, expected_dim: int = MAX_EMBEDDING_DIM) -> bool:
@@ -131,11 +131,55 @@ def cosine_distance(embedding1, embedding2, expected_dim=MAX_EMBEDDING_DIM):
     return distance
 
 
+# Note: sqlite-vec provides native SQL functions for vector operations
+# We don't need custom Python distance functions since sqlite-vec handles this at the SQL level
+
+
 @event.listens_for(Engine, "connect")
 def register_functions(dbapi_connection, connection_record):
-    """Register SQLite functions"""
-    if isinstance(dbapi_connection, sqlite3.Connection):
-        dbapi_connection.create_function("cosine_distance", 2, cosine_distance)
+    """Register SQLite functions and enable sqlite-vec extension"""
+    # Check for both sync SQLite connections and async aiosqlite connections
+    is_sqlite_connection = isinstance(dbapi_connection, sqlite3.Connection)
+    is_aiosqlite_connection = hasattr(dbapi_connection, "_connection") and str(type(dbapi_connection)).find("aiosqlite") != -1
+
+    if is_sqlite_connection or is_aiosqlite_connection:
+        # Get the actual SQLite connection for async connections
+        actual_connection = dbapi_connection._connection if is_aiosqlite_connection else dbapi_connection
+
+        # Enable sqlite-vec extension
+        try:
+            if is_aiosqlite_connection:
+                # For aiosqlite connections, we cannot use async operations in sync event handlers
+                # The extension will need to be loaded per-connection when actually used
+                print("Detected aiosqlite connection - sqlite-vec will be loaded per-query")
+            else:
+                # For sync connections
+                dbapi_connection.enable_load_extension(True)
+                sqlite_vec.load(dbapi_connection)
+                dbapi_connection.enable_load_extension(False)
+                print("Successfully loaded sqlite-vec extension (sync)")
+        except Exception as e:
+            raise RuntimeError(f"Failed to load sqlite-vec extension: {e}")
+
+        # Register custom cosine_distance function for backward compatibility
+        try:
+            if is_aiosqlite_connection:
+                # Try to register function on the actual connection, even though it might be async
+                # This may require the function to be registered per-connection
+                print("Attempting function registration for aiosqlite connection")
+                # For async connections, we need to register the function differently
+                # We'll use the sync-style registration on the underlying connection
+                raw_conn = getattr(actual_connection, "_connection", actual_connection)
+                if hasattr(raw_conn, "create_function"):
+                    raw_conn.create_function("cosine_distance", 2, cosine_distance)
+                    print("Successfully registered cosine_distance for aiosqlite")
+            else:
+                dbapi_connection.create_function("cosine_distance", 2, cosine_distance)
+                print("Successfully registered cosine_distance for sync connection")
+        except Exception as e:
+            raise RuntimeError(f"Failed to register cosine_distance function: {e}")
+    else:
+        print(f"Warning: Not a SQLite connection, but instead {type(dbapi_connection)}: skipping function registration")
 
 
 # Register adapters and converters for numpy arrays

pyproject.toml

@@ -130,6 +130,8 @@ pytest-json-report = "^1.5.0"
 
 [tool.poetry.group.sqlite.dependencies]
 aiosqlite = "^0.21.0"
+# https://github.com/asg017/sqlite-vec/issues/148
+sqlite-vec = "^0.1.7a2"
 
 [tool.black]
 line-length = 140

tests/test_managers.py

@@ -2118,12 +2118,6 @@ async def test_list_agents_by_tags_pagination(server: SyncServer, default_user,
 
 
 @pytest.mark.asyncio
-@pytest.mark.skipif(
-    not hasattr(__import__("letta.settings"), "settings")
-    or not getattr(__import__("letta.settings").settings, "letta_pg_uri_no_default", None)
-    or USING_SQLITE,
-    reason="Skipping vector-related tests when using SQLite (vector search requires PostgreSQL)",
-)
 async def test_list_agents_query_text_pagination(server: SyncServer, default_user, default_organization, event_loop):
     """Test listing agents with query text filtering and pagination."""
     # Create test agents with specific names and descriptions
@@ -2463,7 +2457,7 @@ async def test_attach_block(server: SyncServer, sarah_agent, default_block, defa
     assert agent.memory.blocks[0].label == default_block.label
 
 
-@pytest.mark.skipif(USING_SQLITE, reason="Test not applicable when using SQLite.")
+# Test should work with both SQLite and PostgreSQL
 def test_attach_block_duplicate_label(server: SyncServer, sarah_agent, default_block, other_block, default_user):
     """Test attempting to attach a block with a duplicate label."""
     # Set up both blocks with same label
@@ -2716,11 +2710,6 @@ async def test_agent_list_passages_filtering(server, default_user, sarah_agent,
 
 
 @pytest.mark.asyncio
-@pytest.mark.skipif(
-    not hasattr(__import__("letta.settings"), "settings")
-    or not getattr(__import__("letta.settings").settings, "letta_pg_uri_no_default", None),
-    reason="Skipping vector-related tests when using SQLite (vector search requires PostgreSQL)",
-)
 async def test_agent_list_passages_vector_search(server, default_user, sarah_agent, default_source, default_file, event_loop):
     """Test vector search functionality of agent passages"""
     embed_model = embedding_model(DEFAULT_EMBEDDING_CONFIG)
@@ -3482,7 +3471,7 @@ def test_create_mcp_tool(server: SyncServer, mcp_tool, default_user, default_org
     assert mcp_tool.metadata_[MCP_TOOL_TAG_NAME_PREFIX]["server_name"] == "test"
 
 
-@pytest.mark.skipif(USING_SQLITE, reason="Test not applicable when using SQLite.")
+# Test should work with both SQLite and PostgreSQL
 def test_create_tool_duplicate_name(server: SyncServer, print_tool, default_user, default_organization):
     data = print_tool.model_dump(exclude=["id"])
     tool = PydanticTool(**data)

tests/test_sources.py

@@ -477,11 +477,6 @@ def test_agent_uses_open_close_file_correctly(disable_pinecone, client: LettaSDK
     print("✓ File successfully opened with different range - content differs as expected")
 
 
-@pytest.mark.skipif(
-    not hasattr(__import__("letta.settings"), "settings")
-    or not getattr(__import__("letta.settings").settings, "letta_pg_uri_no_default", None),
-    reason="Skipping vector-related tests when using SQLite (vector search requires PostgreSQL)",
-)
 def test_agent_uses_search_files_correctly(disable_pinecone, client: LettaSDKClient, agent_state: AgentState):
     # Create a new source
     source = client.sources.create(name="test_source", embedding="openai/text-embedding-3-small")
@@ -853,11 +848,6 @@ def test_open_files_schema_descriptions(disable_pinecone, client: LettaSDKClient
 # --- Pinecone Tests ---
 
 
-@pytest.mark.skipif(
-    not hasattr(__import__("letta.settings"), "settings")
-    or not getattr(__import__("letta.settings").settings, "letta_pg_uri_no_default", None),
-    reason="Skipping vector-related tests when using SQLite (vector search requires PostgreSQL)",
-)
 def test_pinecone_search_files_tool(client: LettaSDKClient):
     """Test that search_files tool uses Pinecone when enabled"""
     from letta.helpers.pinecone_utils import should_use_pinecone
@@ -909,11 +899,6 @@ def test_pinecone_search_files_tool(client: LettaSDKClient):
     ), f"Search results should contain relevant content: {search_results}"
 
 
-@pytest.mark.skipif(
-    not hasattr(__import__("letta.settings"), "settings")
-    or not getattr(__import__("letta.settings").settings, "letta_pg_uri_no_default", None),
-    reason="Skipping vector-related tests when using SQLite (vector search requires PostgreSQL)",
-)
 def test_pinecone_lifecycle_file_and_source_deletion(client: LettaSDKClient):
     """Test that file and source deletion removes records from Pinecone"""
     import asyncio

tests/test_vector_embeddings.py

@@ -1,39 +0,0 @@
-import numpy as np
-
-from letta.orm.sqlalchemy_base import adapt_array
-from letta.orm.sqlite_functions import convert_array, verify_embedding_dimension
-
-
-def test_vector_conversions():
-    """Test the vector conversion functions"""
-    # Create test data
-    original = np.random.random(4096).astype(np.float32)
-    print(f"Original shape: {original.shape}")
-
-    # Test full conversion cycle
-    encoded = adapt_array(original)
-    print(f"Encoded type: {type(encoded)}")
-    print(f"Encoded length: {len(encoded)}")
-
-    decoded = convert_array(encoded)
-    print(f"Decoded shape: {decoded.shape}")
-    print(f"Dimension verification: {verify_embedding_dimension(decoded)}")
-
-    # Verify data integrity
-    np.testing.assert_array_almost_equal(original, decoded)
-    print("✓ Data integrity verified")
-
-    # Test with a list
-    list_data = original.tolist()
-    encoded_list = adapt_array(list_data)
-    decoded_list = convert_array(encoded_list)
-    np.testing.assert_array_almost_equal(original, decoded_list)
-    print("✓ List conversion verified")
-
-    # Test None handling
-    assert adapt_array(None) is None
-    assert convert_array(None) is None
-    print("✓ None handling verified")
-
-
-# Run the tests