Merge branch 'main' of https://github.com/meta-llama/llama-stack into add_nemo_customizer

tests/unit/providers/agents/test_persistence_access_control.py

@@ -0,0 +1,175 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import shutil
+import tempfile
+import uuid
+from datetime import datetime
+from unittest.mock import patch
+
+import pytest
+
+from llama_stack.apis.agents import Turn
+from llama_stack.apis.inference import CompletionMessage, StopReason
+from llama_stack.distribution.datatypes import AccessAttributes
+from llama_stack.providers.inline.agents.meta_reference.persistence import AgentPersistence, AgentSessionInfo
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
+
+
+@pytest.fixture
+async def test_setup():
+    temp_dir = tempfile.mkdtemp()
+    db_path = os.path.join(temp_dir, "test_persistence_access_control.db")
+    kvstore_config = SqliteKVStoreConfig(db_path=db_path)
+    kvstore = SqliteKVStoreImpl(kvstore_config)
+    await kvstore.initialize()
+    agent_persistence = AgentPersistence(agent_id="test_agent", kvstore=kvstore)
+    yield agent_persistence
+    shutil.rmtree(temp_dir)
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.providers.inline.agents.meta_reference.persistence.get_auth_attributes")
+async def test_session_creation_with_access_attributes(mock_get_auth_attributes, test_setup):
+    agent_persistence = test_setup
+
+    # Set creator's attributes for the session
+    creator_attributes = {"roles": ["researcher"], "teams": ["ai-team"]}
+    mock_get_auth_attributes.return_value = creator_attributes
+
+    # Create a session
+    session_id = await agent_persistence.create_session("Test Session")
+
+    # Get the session and verify access attributes were set
+    session_info = await agent_persistence.get_session_info(session_id)
+    assert session_info is not None
+    assert session_info.access_attributes is not None
+    assert session_info.access_attributes.roles == ["researcher"]
+    assert session_info.access_attributes.teams == ["ai-team"]
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.providers.inline.agents.meta_reference.persistence.get_auth_attributes")
+async def test_session_access_control(mock_get_auth_attributes, test_setup):
+    agent_persistence = test_setup
+
+    # Create a session with specific access attributes
+    session_id = str(uuid.uuid4())
+    session_info = AgentSessionInfo(
+        session_id=session_id,
+        session_name="Restricted Session",
+        started_at=datetime.now(),
+        access_attributes=AccessAttributes(roles=["admin"], teams=["security-team"]),
+    )
+
+    await agent_persistence.kvstore.set(
+        key=f"session:{agent_persistence.agent_id}:{session_id}",
+        value=session_info.model_dump_json(),
+    )
+
+    # User with matching attributes can access
+    mock_get_auth_attributes.return_value = {"roles": ["admin", "user"], "teams": ["security-team", "other-team"]}
+    retrieved_session = await agent_persistence.get_session_info(session_id)
+    assert retrieved_session is not None
+    assert retrieved_session.session_id == session_id
+
+    # User without matching attributes cannot access
+    mock_get_auth_attributes.return_value = {"roles": ["user"], "teams": ["other-team"]}
+    retrieved_session = await agent_persistence.get_session_info(session_id)
+    assert retrieved_session is None
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.providers.inline.agents.meta_reference.persistence.get_auth_attributes")
+async def test_turn_access_control(mock_get_auth_attributes, test_setup):
+    agent_persistence = test_setup
+
+    # Create a session with restricted access
+    session_id = str(uuid.uuid4())
+    session_info = AgentSessionInfo(
+        session_id=session_id,
+        session_name="Restricted Session",
+        started_at=datetime.now(),
+        access_attributes=AccessAttributes(roles=["admin"]),
+    )
+
+    await agent_persistence.kvstore.set(
+        key=f"session:{agent_persistence.agent_id}:{session_id}",
+        value=session_info.model_dump_json(),
+    )
+
+    # Create a turn for this session
+    turn_id = str(uuid.uuid4())
+    turn = Turn(
+        session_id=session_id,
+        turn_id=turn_id,
+        steps=[],
+        started_at=datetime.now(),
+        input_messages=[],
+        output_message=CompletionMessage(
+            content="Hello",
+            stop_reason=StopReason.end_of_turn,
+        ),
+    )
+
+    # Admin can add turn
+    mock_get_auth_attributes.return_value = {"roles": ["admin"]}
+    await agent_persistence.add_turn_to_session(session_id, turn)
+
+    # Admin can get turn
+    retrieved_turn = await agent_persistence.get_session_turn(session_id, turn_id)
+    assert retrieved_turn is not None
+    assert retrieved_turn.turn_id == turn_id
+
+    # Regular user cannot get turn
+    mock_get_auth_attributes.return_value = {"roles": ["user"]}
+    with pytest.raises(ValueError):
+        await agent_persistence.get_session_turn(session_id, turn_id)
+
+    # Regular user cannot get turns for session
+    with pytest.raises(ValueError):
+        await agent_persistence.get_session_turns(session_id)
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.providers.inline.agents.meta_reference.persistence.get_auth_attributes")
+async def test_tool_call_and_infer_iters_access_control(mock_get_auth_attributes, test_setup):
+    agent_persistence = test_setup
+
+    # Create a session with restricted access
+    session_id = str(uuid.uuid4())
+    session_info = AgentSessionInfo(
+        session_id=session_id,
+        session_name="Restricted Session",
+        started_at=datetime.now(),
+        access_attributes=AccessAttributes(roles=["admin"]),
+    )
+
+    await agent_persistence.kvstore.set(
+        key=f"session:{agent_persistence.agent_id}:{session_id}",
+        value=session_info.model_dump_json(),
+    )
+
+    turn_id = str(uuid.uuid4())
+
+    # Admin user can set inference iterations
+    mock_get_auth_attributes.return_value = {"roles": ["admin"]}
+    await agent_persistence.set_num_infer_iters_in_turn(session_id, turn_id, 5)
+
+    # Admin user can get inference iterations
+    infer_iters = await agent_persistence.get_num_infer_iters_in_turn(session_id, turn_id)
+    assert infer_iters == 5
+
+    # Regular user cannot get inference iterations
+    mock_get_auth_attributes.return_value = {"roles": ["user"]}
+    infer_iters = await agent_persistence.get_num_infer_iters_in_turn(session_id, turn_id)
+    assert infer_iters is None
+
+    # Regular user cannot set inference iterations (should raise ValueError)
+    with pytest.raises(ValueError):
+        await agent_persistence.set_num_infer_iters_in_turn(session_id, turn_id, 10)

tests/unit/providers/vector_io/test_sqlite_vec.py

@@ -5,17 +5,16 @@
 # the root directory of this source tree.
 
 import asyncio
-import sqlite3
 
 import numpy as np
 import pytest
 import pytest_asyncio
-import sqlite_vec
 
 from llama_stack.apis.vector_io import Chunk, QueryChunksResponse
 from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import (
     SQLiteVecIndex,
     SQLiteVecVectorIOAdapter,
+    _create_sqlite_connection,
     generate_chunk_id,
 )
 
@@ -36,29 +35,25 @@ def loop():
     return asyncio.new_event_loop()
 
 
-@pytest.fixture(scope="session", autouse=True)
-def sqlite_connection(loop):
-    conn = sqlite3.connect(":memory:")
-    try:
-        conn.enable_load_extension(True)
-        sqlite_vec.load(conn)
-        yield conn
-    finally:
-        conn.close()
-
-
 @pytest_asyncio.fixture(scope="session", autouse=True)
-async def sqlite_vec_index(sqlite_connection, embedding_dimension):
-    return await SQLiteVecIndex.create(dimension=embedding_dimension, connection=sqlite_connection, bank_id="test_bank")
+async def sqlite_vec_index(embedding_dimension, tmp_path_factory):
+    temp_dir = tmp_path_factory.getbasetemp()
+    db_path = str(temp_dir / "test_sqlite.db")
+    index = await SQLiteVecIndex.create(dimension=embedding_dimension, db_path=db_path, bank_id="test_bank")
+    yield index
+    await index.delete()
 
 
 @pytest.mark.asyncio
 async def test_add_chunks(sqlite_vec_index, sample_chunks, sample_embeddings):
     await sqlite_vec_index.add_chunks(sample_chunks, sample_embeddings, batch_size=2)
-    cur = sqlite_vec_index.connection.cursor()
+    connection = _create_sqlite_connection(sqlite_vec_index.db_path)
+    cur = connection.cursor()
     cur.execute(f"SELECT COUNT(*) FROM {sqlite_vec_index.metadata_table}")
     count = cur.fetchone()[0]
     assert count == len(sample_chunks)
+    cur.close()
+    connection.close()
 
 
 @pytest.mark.asyncio
@@ -79,13 +74,14 @@ async def test_chunk_id_conflict(sqlite_vec_index, sample_chunks, embedding_dime
     sample_embeddings = np.random.rand(len(sample_chunks), embedding_dimension).astype(np.float32)
 
     await sqlite_vec_index.add_chunks(sample_chunks, sample_embeddings, batch_size=batch_size)
-
-    cur = sqlite_vec_index.connection.cursor()
+    connection = _create_sqlite_connection(sqlite_vec_index.db_path)
+    cur = connection.cursor()
 
     # Retrieve all chunk IDs to check for duplicates
     cur.execute(f"SELECT id FROM {sqlite_vec_index.metadata_table}")
     chunk_ids = [row[0] for row in cur.fetchall()]
     cur.close()
+    connection.close()
 
     # Ensure all chunk IDs are unique
     assert len(chunk_ids) == len(set(chunk_ids)), "Duplicate chunk IDs detected across batches!"

tests/unit/registry/test_registry.py

@@ -12,6 +12,7 @@ import pytest_asyncio
 from llama_stack.apis.inference import Model
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.distribution.store.registry import (
+    KEY_FORMAT,
     CachedDiskDistributionRegistry,
     DiskDistributionRegistry,
 )
@@ -197,3 +198,72 @@ async def test_get_all_objects(config):
         assert stored_vector_db.embedding_model == original_vector_db.embedding_model
         assert stored_vector_db.provider_id == original_vector_db.provider_id
         assert stored_vector_db.embedding_dimension == original_vector_db.embedding_dimension
+
+
+@pytest.mark.asyncio
+async def test_parse_registry_values_error_handling(config):
+    kvstore = await kvstore_impl(config)
+
+    valid_db = VectorDB(
+        identifier="valid_vector_db",
+        embedding_model="all-MiniLM-L6-v2",
+        embedding_dimension=384,
+        provider_resource_id="valid_vector_db",
+        provider_id="test-provider",
+    )
+
+    await kvstore.set(KEY_FORMAT.format(type="vector_db", identifier="valid_vector_db"), valid_db.model_dump_json())
+
+    await kvstore.set(KEY_FORMAT.format(type="vector_db", identifier="corrupted_json"), "{not valid json")
+
+    await kvstore.set(
+        KEY_FORMAT.format(type="vector_db", identifier="missing_fields"),
+        '{"type": "vector_db", "identifier": "missing_fields"}',
+    )
+
+    test_registry = DiskDistributionRegistry(kvstore)
+    await test_registry.initialize()
+
+    # Get all objects, which should only return the valid one
+    all_objects = await test_registry.get_all()
+
+    # Should have filtered out the invalid entries
+    assert len(all_objects) == 1
+    assert all_objects[0].identifier == "valid_vector_db"
+
+    # Check that the get method also handles errors correctly
+    invalid_obj = await test_registry.get("vector_db", "corrupted_json")
+    assert invalid_obj is None
+
+    invalid_obj = await test_registry.get("vector_db", "missing_fields")
+    assert invalid_obj is None
+
+
+@pytest.mark.asyncio
+async def test_cached_registry_error_handling(config):
+    kvstore = await kvstore_impl(config)
+
+    valid_db = VectorDB(
+        identifier="valid_cached_db",
+        embedding_model="all-MiniLM-L6-v2",
+        embedding_dimension=384,
+        provider_resource_id="valid_cached_db",
+        provider_id="test-provider",
+    )
+
+    await kvstore.set(KEY_FORMAT.format(type="vector_db", identifier="valid_cached_db"), valid_db.model_dump_json())
+
+    await kvstore.set(
+        KEY_FORMAT.format(type="vector_db", identifier="invalid_cached_db"),
+        '{"type": "vector_db", "identifier": "invalid_cached_db", "embedding_model": 12345}',  # Should be string
+    )
+
+    cached_registry = CachedDiskDistributionRegistry(kvstore)
+    await cached_registry.initialize()
+
+    all_objects = await cached_registry.get_all()
+    assert len(all_objects) == 1
+    assert all_objects[0].identifier == "valid_cached_db"
+
+    invalid_obj = await cached_registry.get("vector_db", "invalid_cached_db")
+    assert invalid_obj is None