import json from unittest.mock import MagicMock, patch import numpy as np import pytest from langflow.schema.data import Data from langflow.schema.dataframe import DataFrame from langflow.schema.message import Message from lfx.base.knowledge_bases import get_knowledge_bases from lfx.components.files_and_knowledge import KnowledgeIngestionComponent from tests.base import ComponentTestBaseWithClient class TestKnowledgeIngestionComponent(ComponentTestBaseWithClient): @pytest.fixture def component_class(self): """Return the component class to test.""" return KnowledgeIngestionComponent @pytest.fixture(autouse=True) def mock_knowledge_base_path(self, tmp_path): """Mock the knowledge base root path directly.""" with patch("lfx.components.files_and_knowledge.ingestion._KNOWLEDGE_BASES_ROOT_PATH", tmp_path): yield @pytest.fixture def default_kwargs(self, tmp_path, active_user): """Return default kwargs for component instantiation.""" # Create a sample DataFrame data_df = DataFrame( {"text": ["Sample text 1", "Sample text 2"], "title": ["Title 1", "Title 2"], "category": ["cat1", "cat2"]} ) # Create column configuration column_config = [ {"column_name": "text", "vectorize": True, "identifier": False}, {"column_name": "title", "vectorize": False, "identifier": False}, {"column_name": "category", "vectorize": False, "identifier": True}, ] # Create knowledge base directory kb_name = "test_kb" kb_path = tmp_path / active_user.username / kb_name kb_path.mkdir(parents=True, exist_ok=True) # Create embedding metadata file (new format with model_selection) metadata = { "embedding_provider": "HuggingFace", "embedding_model": "sentence-transformers/all-MiniLM-L6-v2", "model_selection": { "name": "sentence-transformers/all-MiniLM-L6-v2", "provider": "HuggingFace", "metadata": { "embedding_class": "HuggingFaceEmbeddings", "param_mapping": {"model": "model_name"}, }, }, "api_key": None, "api_key_used": False, "chunk_size": 1000, "created_at": "2024-01-01T00:00:00Z", } (kb_path / "embedding_metadata.json").write_text(json.dumps(metadata)) return { "knowledge_base": kb_name, "input_df": data_df, "column_config": column_config, "chunk_size": 1000, "kb_root_path": str(tmp_path), "api_key": None, "allow_duplicates": False, "silent_errors": False, "_user_id": active_user.id, } @pytest.fixture def file_names_mapping(self): """Return file names mapping for version testing.""" # This is a new component, so it doesn't exist in older versions return [] def test_validate_column_config_valid(self, component_class, default_kwargs): """Test column configuration validation with valid config.""" component = component_class(**default_kwargs) data_df = default_kwargs["input_df"] config_list = component._validate_column_config(data_df) assert len(config_list) == 3 assert config_list[0]["column_name"] == "text" assert config_list[0]["vectorize"] is True def test_validate_column_config_invalid_column(self, component_class, default_kwargs): """Test column configuration validation with invalid column name.""" # Modify column config to include non-existent column invalid_config = [{"column_name": "nonexistent", "vectorize": True, "identifier": False}] default_kwargs["column_config"] = invalid_config # Instantiate the component with the modified config component = component_class(**default_kwargs) data_df = default_kwargs["input_df"] # Should raise ValueError since column does not exist in DataFrame with pytest.raises(ValueError, match="Column 'nonexistent' not found in DataFrame"): component._validate_column_config(data_df) def test_new_knowledge_dialog_uses_provider_credentials(self, component_class, default_kwargs): """Test the create-knowledge dialog no longer exposes a redundant API key override.""" component = component_class(**default_kwargs) dialog_inputs = component.inputs[0].dialog_inputs["fields"]["data"]["node"] embedding_model_input = dialog_inputs["template"]["02_embedding_model"] assert dialog_inputs["field_order"] == ["01_new_kb_name", "02_embedding_model"] assert "03_api_key" not in dialog_inputs["template"] assert "configured credentials" in embedding_model_input.info @patch("lfx.components.files_and_knowledge.ingestion.get_settings_service") @patch("lfx.components.files_and_knowledge.ingestion.encrypt_api_key") def test_build_embedding_metadata(self, mock_encrypt, mock_get_settings, component_class, default_kwargs): """Test building embedding metadata.""" component = component_class(**default_kwargs) mock_settings = MagicMock() mock_get_settings.return_value = mock_settings mock_encrypt.return_value = "encrypted_key" model_selection = [ {"name": "sentence-transformers/all-MiniLM-L6-v2", "provider": "HuggingFace", "metadata": {}} ] metadata = component._build_embedding_metadata(model_selection, "test-key") assert metadata["embedding_provider"] == "HuggingFace" assert metadata["embedding_model"] == "sentence-transformers/all-MiniLM-L6-v2" assert metadata["api_key"] == "encrypted_key" # pragma:allowlist secret assert metadata["api_key_used"] is True assert metadata["chunk_size"] == 1000 assert "created_at" in metadata def test_update_metadata_metrics_persists_chunk_stats(self, component_class, default_kwargs, tmp_path, active_user): """Test updating the persisted KB metrics from the Chroma collection.""" component = component_class(**default_kwargs) kb_path = tmp_path / active_user.username / "test_kb" mock_collection = MagicMock() mock_collection.count.return_value = 2 mock_collection.get.return_value = {"documents": ["hello world", "foo"]} mock_chroma = MagicMock() mock_chroma._collection = mock_collection component._update_metadata_metrics(kb_path, mock_chroma) stored_metadata = json.loads((kb_path / "embedding_metadata.json").read_text()) assert stored_metadata["chunks"] == 2 assert stored_metadata["words"] == 3 assert stored_metadata["characters"] == 14 assert stored_metadata["avg_chunk_size"] == 7.0 assert stored_metadata["size"] > 0 def test_update_metadata_metrics_no_metadata_file(self, component_class, default_kwargs, tmp_path, active_user): """Test _update_metadata_metrics silently returns when embedding_metadata.json does not exist.""" component = component_class(**default_kwargs) kb_path = tmp_path / active_user.username / "no_metadata_kb" kb_path.mkdir(parents=True, exist_ok=True) # Do NOT create embedding_metadata.json mock_chroma = MagicMock() # Should return without error component._update_metadata_metrics(kb_path, mock_chroma) # Metadata file should still not exist assert not (kb_path / "embedding_metadata.json").exists() def test_update_metadata_metrics_empty_collection(self, component_class, default_kwargs, tmp_path, active_user): """Test _update_metadata_metrics when Chroma collection has 0 chunks.""" component = component_class(**default_kwargs) kb_path = tmp_path / active_user.username / "test_kb" mock_collection = MagicMock() mock_collection.count.return_value = 0 mock_chroma = MagicMock() mock_chroma._collection = mock_collection component._update_metadata_metrics(kb_path, mock_chroma) stored_metadata = json.loads((kb_path / "embedding_metadata.json").read_text()) assert stored_metadata["chunks"] == 0 # update_text_metrics does not write words/characters/avg_chunk_size when chunks == 0 # unless they were already present — verify no crash occurred mock_collection.get.assert_not_called() def test_build_column_metadata(self, component_class, default_kwargs): """Test building column metadata.""" component = component_class(**default_kwargs) data_df = default_kwargs["input_df"] config_list = default_kwargs["column_config"] metadata = component._build_column_metadata(config_list, data_df) assert metadata["total_columns"] == 3 assert metadata["mapped_columns"] == 3 assert metadata["unmapped_columns"] == 0 assert len(metadata["columns"]) == 3 assert "text" in metadata["summary"]["vectorized_columns"] assert "category" in metadata["summary"]["identifier_columns"] async def test_convert_df_to_data_objects(self, component_class, default_kwargs): """Test converting DataFrame to Data objects.""" component = component_class(**default_kwargs) data_df = default_kwargs["input_df"] config_list = default_kwargs["column_config"] # Mock Chroma to avoid actual vector store operations with patch("lfx.components.files_and_knowledge.ingestion.Chroma") as mock_chroma: mock_chroma_instance = MagicMock() mock_chroma_instance.get.return_value = {"metadatas": []} mock_chroma.return_value = mock_chroma_instance data_objects = await component._convert_df_to_data_objects(data_df, config_list) assert len(data_objects) == 2 assert all(isinstance(obj, Data) for obj in data_objects) # Check first data object first_obj = data_objects[0] assert "text" in first_obj.data assert "title" in first_obj.data assert "category" in first_obj.data assert "_id" in first_obj.data async def test_convert_df_to_data_objects_no_duplicates(self, component_class, default_kwargs): """Test converting DataFrame to Data objects with duplicate prevention.""" default_kwargs["allow_duplicates"] = False component = component_class(**default_kwargs) data_df = default_kwargs["input_df"] config_list = default_kwargs["column_config"] # Mock Chroma with existing hash with patch("lfx.components.files_and_knowledge.ingestion.Chroma") as mock_chroma: # Simulate existing document with same hash existing_hash = "some_existing_hash" mock_chroma_instance = MagicMock() mock_chroma_instance.get.return_value = {"metadatas": [{"_id": existing_hash}]} mock_chroma.return_value = mock_chroma_instance # Mock hashlib to return the existing hash for first row with patch("lfx.components.files_and_knowledge.ingestion.hashlib.sha256") as mock_hash: mock_hash_obj = MagicMock() mock_hash_obj.hexdigest.side_effect = [existing_hash, "different_hash"] mock_hash.return_value = mock_hash_obj data_objects = await component._convert_df_to_data_objects(data_df, config_list) # Should only return one object (second row) since first is duplicate assert len(data_objects) == 1 def test_is_valid_collection_name(self, component_class, default_kwargs): """Test collection name validation.""" component = component_class(**default_kwargs) # Valid names assert component.is_valid_collection_name("valid_name") is True assert component.is_valid_collection_name("valid-name") is True assert component.is_valid_collection_name("ValidName123") is True # Invalid names assert component.is_valid_collection_name("ab") is False # Too short assert component.is_valid_collection_name("a" * 64) is False # Too long assert component.is_valid_collection_name("_invalid") is False # Starts with underscore assert component.is_valid_collection_name("invalid_") is False # Ends with underscore assert component.is_valid_collection_name("invalid@name") is False # Invalid character @patch("lfx.components.files_and_knowledge.ingestion.get_embeddings") async def test_build_kb_info_success(self, mock_get_embeddings, component_class, default_kwargs): """Test successful KB info building.""" component = component_class(**default_kwargs) mock_embedding_fn = MagicMock() mock_get_embeddings.return_value = mock_embedding_fn # Mock vector store creation with patch.object(component, "_create_vector_store"), patch.object(component, "_save_kb_files"): result = await component.build_kb_info() assert isinstance(result, Data) assert "kb_id" in result.data assert "kb_name" in result.data assert "rows" in result.data assert result.data["rows"] == 2 async def test_get_knowledge_bases(self, tmp_path, active_user): """Test getting list of knowledge bases.""" # Create additional test directories (tmp_path / active_user.username / "kb1").mkdir(parents=True, exist_ok=True) (tmp_path / active_user.username / "kb2").mkdir(parents=True, exist_ok=True) (tmp_path / active_user.username / ".hidden").mkdir(parents=True, exist_ok=True) # Should be ignored kb_list = await get_knowledge_bases(tmp_path, user_id=active_user.id) assert "test_kb" in kb_list assert "kb1" in kb_list assert "kb2" in kb_list assert ".hidden" not in kb_list @patch("lfx.components.files_and_knowledge.ingestion.get_embeddings") async def test_update_build_config_new_kb(self, mock_get_embeddings, component_class, default_kwargs): """Test updating build config for new knowledge base creation.""" component = component_class(**default_kwargs) build_config = {"knowledge_base": {"value": None, "options": [], "dialog_inputs": {}}} model_selection = [ {"name": "sentence-transformers/all-MiniLM-L6-v2", "provider": "HuggingFace", "metadata": {}} ] field_value = { "01_new_kb_name": "new_test_kb", "02_embedding_model": model_selection, } # Mock embedding validation mock_embeddings = MagicMock() mock_embeddings.embed_query.return_value = [0.1, 0.2, 0.3] mock_get_embeddings.return_value = mock_embeddings with patch.object(component, "_save_embedding_metadata") as mock_save_metadata: result = await component.update_build_config(build_config, field_value, "knowledge_base") assert result["knowledge_base"]["value"] == "new_test_kb" assert "new_test_kb" in result["knowledge_base"]["options"] assert "api_key" not in mock_get_embeddings.call_args.kwargs assert "api_key" not in mock_save_metadata.call_args.kwargs @patch("lfx.components.files_and_knowledge.ingestion.get_embeddings") async def test_build_kb_info_with_message_input(self, mock_get_embeddings, component_class, default_kwargs): """Test that Message input is accepted and converted to DataFrame.""" # Replace the DataFrame input with a Message default_kwargs["input_df"] = Message(text="Sample text 1") default_kwargs["column_config"] = [ {"column_name": "text", "vectorize": True, "identifier": True}, ] component = component_class(**default_kwargs) mock_embedding_fn = MagicMock() mock_get_embeddings.return_value = mock_embedding_fn with patch.object(component, "_create_vector_store"), patch.object(component, "_save_kb_files"): result = await component.build_kb_info() assert isinstance(result, Data) assert result.data["rows"] == 1 assert result.data["kb_name"] == "test_kb" async def test_update_build_config_invalid_kb_name(self, component_class, default_kwargs): """Test updating build config with invalid KB name.""" component = component_class(**default_kwargs) build_config = {"knowledge_base": {"value": None, "options": []}} field_value = { "01_new_kb_name": "invalid@name", # Invalid character "02_embedding_model": "sentence-transformers/all-MiniLM-L6-v2", } with pytest.raises(ValueError, match="Invalid knowledge base name"): await component.update_build_config(build_config, field_value, "knowledge_base") @patch("lfx.components.files_and_knowledge.ingestion.get_embeddings") async def test_build_kb_info_with_new_format_metadata( self, mock_get_embeddings, component_class, default_kwargs, tmp_path, active_user ): """Test that build_kb_info uses model_selection directly from new-format metadata.""" # Overwrite the default metadata file to use the new format (includes model_selection key). # The old format only had embedding_model/embedding_provider strings; the new format # stores the full model_selection dict so get_embeddings() can reconstruct the client. kb_path = tmp_path / active_user.username / "test_kb" new_format_metadata = { "embedding_provider": "HuggingFace", "embedding_model": "sentence-transformers/all-MiniLM-L6-v2", "model_selection": { "name": "sentence-transformers/all-MiniLM-L6-v2", "provider": "HuggingFace", "metadata": {}, }, "api_key": None, "api_key_used": False, "chunk_size": 1000, "created_at": "2024-01-01T00:00:00Z", } (kb_path / "embedding_metadata.json").write_text(json.dumps(new_format_metadata)) component = component_class(**default_kwargs) mock_get_embeddings.return_value = MagicMock() with patch.object(component, "_create_vector_store"), patch.object(component, "_save_kb_files"): result = await component.build_kb_info() assert isinstance(result, Data) assert result.data["rows"] == 2 # Verify get_embeddings was called with the full model_selection from the new-format metadata, # not a minimal reconstructed dict from the backward-compat path. call_kwargs = mock_get_embeddings.call_args passed_model = call_kwargs.kwargs.get("model") or call_kwargs.args[0] assert isinstance(passed_model, list) assert passed_model[0]["name"] == "sentence-transformers/all-MiniLM-L6-v2" assert passed_model[0]["provider"] == "HuggingFace" async def test_convert_df_to_data_objects_allow_duplicates(self, component_class, default_kwargs): """Test that allow_duplicates=True returns all rows even when their hashes already exist.""" default_kwargs["allow_duplicates"] = True component = component_class(**default_kwargs) data_df = default_kwargs["input_df"] config_list = default_kwargs["column_config"] with patch("lfx.components.files_and_knowledge.ingestion.Chroma") as mock_chroma: mock_chroma_instance = MagicMock() # Simulate all rows as already-existing duplicates in the store mock_chroma_instance.get.return_value = {"metadatas": [{"_id": "hash_1"}, {"_id": "hash_2"}]} mock_chroma.return_value = mock_chroma_instance with patch("lfx.components.files_and_knowledge.ingestion.hashlib.sha256") as mock_hash: mock_hash_obj = MagicMock() # Return hashes that match the existing IDs above mock_hash_obj.hexdigest.side_effect = ["hash_1", "hash_2"] mock_hash.return_value = mock_hash_obj data_objects = await component._convert_df_to_data_objects(data_df, config_list) # All rows should be included — duplicates are allowed assert len(data_objects) == 2 async def test_build_kb_info_no_metadata_file_raises_error( self, component_class, default_kwargs, tmp_path, active_user ): """Test that build_kb_info raises RuntimeError when no embedding metadata file exists.""" # Remove the metadata file so model_selection cannot be determined kb_path = tmp_path / active_user.username / "test_kb" (kb_path / "embedding_metadata.json").unlink() component = component_class(**default_kwargs) with pytest.raises(RuntimeError, match="No embedding model configuration found"): await component.build_kb_info() @patch("lfx.components.files_and_knowledge.ingestion.get_embedding_model_options") @patch("lfx.components.files_and_knowledge.ingestion.get_embeddings") async def test_build_kb_info_old_format_unrecognized_model( self, mock_get_embeddings, # noqa: ARG002 mock_get_options, component_class, default_kwargs, tmp_path, active_user, ): """Test that old-format metadata with an unrecognized model name raises a clear error.""" # Overwrite metadata to use old format (no model_selection key) with a model name # that is not in the current registry. kb_path = tmp_path / active_user.username / "test_kb" old_format_metadata = { "embedding_provider": "SomeOldProvider", "embedding_model": "old-model-that-no-longer-exists", "api_key": None, "api_key_used": False, "chunk_size": 1000, "created_at": "2024-01-01T00:00:00Z", } (kb_path / "embedding_metadata.json").write_text(json.dumps(old_format_metadata)) # Registry returns models that do NOT include the old model name mock_get_options.return_value = [ {"name": "text-embedding-3-small", "provider": "OpenAI", "metadata": {}}, ] component = component_class(**default_kwargs) # Should raise a RuntimeError wrapping a ValueError with a clear message with pytest.raises(RuntimeError, match="no longer recognized"): await component.build_kb_info() def test_scalar_notna_with_scalar_values(self, component_class, default_kwargs): """Test _scalar_notna returns correct results for scalar values.""" component = component_class(**default_kwargs) assert component._scalar_notna("hello") is True assert component._scalar_notna(42) is True assert component._scalar_notna(0) is True assert component._scalar_notna("") is True assert component._scalar_notna(None) is False assert component._scalar_notna(float("nan")) is False def test_scalar_notna_with_numpy_arrays(self, component_class, default_kwargs): """Test _scalar_notna handles numpy arrays without raising ambiguous truth value errors.""" component = component_class(**default_kwargs) # Empty array — should be falsy (no valid data) assert not component._scalar_notna(np.array([])) # Array with valid values — should be truthy assert component._scalar_notna(np.array([1, 2, 3])) # Array containing NaN — should be falsy (not all values are non-NA) assert not component._scalar_notna(np.array([1, float("nan"), 3])) # Array of strings — should be truthy assert component._scalar_notna(np.array(["a", "b"])) def test_scalar_notna_with_lists(self, component_class, default_kwargs): """Test _scalar_notna handles plain lists safely.""" component = component_class(**default_kwargs) assert not component._scalar_notna([]) assert component._scalar_notna([1, 2]) async def test_convert_df_to_data_objects_with_array_cells(self, component_class, default_kwargs): """Test that _convert_df_to_data_objects handles DataFrame rows containing numpy arrays. This reproduces the bug where Split Text output contains metadata columns with array values, causing 'truth value of an empty array is ambiguous' errors. """ # Build a DataFrame with an array-valued metadata column (mimics Split Text output) data_df = DataFrame( { "text": ["chunk 1", "chunk 2"], "source": ["file.txt", "file.txt"], "tags": [np.array([]), np.array(["important"])], } ) default_kwargs["input_df"] = data_df default_kwargs["column_config"] = [ {"column_name": "text", "vectorize": True, "identifier": False}, {"column_name": "source", "vectorize": False, "identifier": True}, {"column_name": "tags", "vectorize": False, "identifier": False}, ] component = component_class(**default_kwargs) config_list = default_kwargs["column_config"] with patch("lfx.components.files_and_knowledge.ingestion.Chroma") as mock_chroma: mock_chroma_instance = MagicMock() mock_chroma_instance.get.return_value = {"metadatas": []} mock_chroma.return_value = mock_chroma_instance # This should NOT raise "truth value of an empty array is ambiguous" data_objects = await component._convert_df_to_data_objects(data_df, config_list) assert len(data_objects) == 2 assert all(isinstance(obj, Data) for obj in data_objects) def test_build_embedding_metadata_without_api_key(self, component_class, default_kwargs): """Test _build_embedding_metadata with no API key stores model_selection for later use.""" component = component_class(**default_kwargs) model_selection = [ {"name": "sentence-transformers/all-MiniLM-L6-v2", "provider": "HuggingFace", "metadata": {}} ] metadata = component._build_embedding_metadata(model_selection, api_key=None) assert metadata["embedding_provider"] == "HuggingFace" assert metadata["embedding_model"] == "sentence-transformers/all-MiniLM-L6-v2" assert metadata["api_key"] is None assert metadata["api_key_used"] is False # New in this PR: full model_selection is stored alongside the string fields so # build_kb_info() can reconstruct the embedding client without hitting the model registry. assert "model_selection" in metadata assert metadata["model_selection"]["name"] == "sentence-transformers/all-MiniLM-L6-v2" assert metadata["model_selection"]["provider"] == "HuggingFace"