feat: adds metadata and batch_index to batch_run (#6318)

* Update batch_run.py

* updates to test component and fixes formatting

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: anovazzi1 <otavio2204@gmail.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
This commit is contained in:
Edwin Jose 2025-02-14 16:04:50 -05:00 • committed by GitHub
commit a1967bc472
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 287 additions and 47 deletions

View file

@ -21,6 +21,7 @@ class TestBatchRunComponent(ComponentTestBaseWithoutClient):
"model": MockLanguageModel(),
"df": DataFrame({"text": ["Hello"]}),
"column_name": "text",
"enable_metadata": True,
}
@pytest.fixture
@ -33,7 +34,11 @@ class TestBatchRunComponent(ComponentTestBaseWithoutClient):
test_df = DataFrame({"text": ["Hello", "World", "Test"]})
component = BatchRunComponent(
model=MockLanguageModel(), system_message="You are a helpful assistant", df=test_df, column_name="text"
model=MockLanguageModel(),
system_message="You are a helpful assistant",
df=test_df,
column_name="text",
enable_metadata=True,
)
# Run the batch process
@ -43,46 +48,188 @@ class TestBatchRunComponent(ComponentTestBaseWithoutClient):
assert isinstance(result, DataFrame)
assert "text_input" in result.columns
assert "model_response" in result.columns
assert "metadata" in result.columns
assert len(result) == 3
assert all(isinstance(resp, str) for resp in result["model_response"])
# Convert DataFrame to list of dicts for easier testing
result_dicts = result.to_dict("records")
# Verify metadata
assert all(row["metadata"]["has_system_message"] for row in result_dicts)
assert all(row["metadata"]["processing_status"] == "success" for row in result_dicts)
async def test_batch_run_without_system_message(self):
async def test_batch_run_without_metadata(self):
test_df = DataFrame({"text": ["Hello", "World"]})
component = BatchRunComponent(model=MockLanguageModel(), df=test_df, column_name="text")
component = BatchRunComponent(
model=MockLanguageModel(),
df=test_df,
column_name="text",
enable_metadata=False,
)
result = await component.run_batch()
assert isinstance(result, DataFrame)
assert len(result) == 2
assert "metadata" not in result.columns
assert all(isinstance(resp, str) for resp in result["model_response"])
async def test_batch_run_error_with_metadata(self):
component = BatchRunComponent(
model=MockLanguageModel(),
df="not_a_dataframe", # This will cause a TypeError
column_name="text",
enable_metadata=True,
)
with pytest.raises(TypeError, match=re.escape("Expected DataFrame input, got <class 'str'>")):
await component.run_batch()
async def test_batch_run_error_without_metadata(self):
component = BatchRunComponent(
model=MockLanguageModel(),
df="not_a_dataframe", # This will cause a TypeError
column_name="text",
enable_metadata=False,
)
with pytest.raises(TypeError, match=re.escape("Expected DataFrame input, got <class 'str'>")):
await component.run_batch()
async def test_operational_error_with_metadata(self):
# Create a mock model that raises an AttributeError during processing
class ErrorModel:
def with_config(self, *_, **__):
return self
async def abatch(self, *_):
msg = "Mock error during batch processing"
raise AttributeError(msg)
component = BatchRunComponent(
model=ErrorModel(),
df=DataFrame({"text": ["test1", "test2"]}),
column_name="text",
enable_metadata=True,
)
result = await component.run_batch()
assert isinstance(result, DataFrame)
assert len(result) == 1 # Component returns a single error row
error_row = result.iloc[0]
# Verify error metadata
assert error_row["metadata"]["processing_status"] == "failed"
assert "Mock error during batch processing" in error_row["metadata"]["error"]
# Verify base row structure
assert error_row["text_input"] == ""
assert error_row["model_response"] == ""
assert error_row["batch_index"] == -1
async def test_operational_error_without_metadata(self):
# Create a mock model that raises an AttributeError during processing
class ErrorModel:
def with_config(self, *_, **__):
return self
async def abatch(self, *_):
msg = "Mock error during batch processing"
raise AttributeError(msg)
component = BatchRunComponent(
model=ErrorModel(),
df=DataFrame({"text": ["test1", "test2"]}),
column_name="text",
enable_metadata=False,
)
result = await component.run_batch()
assert isinstance(result, DataFrame)
assert len(result) == 1 # Component returns a single error row
error_row = result.iloc[0]
# Verify no metadata
assert "metadata" not in error_row
# Verify base row structure
assert error_row["text_input"] == ""
assert error_row["model_response"] == ""
assert error_row["batch_index"] == -1
def test_create_base_row(self):
component = BatchRunComponent()
row = component._create_base_row(text_input="test_input", model_response="test_response", batch_index=1)
assert row == {
"text_input": "test_input",
"model_response": "test_response",
"batch_index": 1,
}
def test_add_metadata_success(self):
component = BatchRunComponent(enable_metadata=True)
row = component._create_base_row(text_input="test_input", model_response="test_response", batch_index=1)
component._add_metadata(row, success=True, system_msg="test_system")
assert "metadata" in row
assert row["metadata"]["has_system_message"] is True
assert row["metadata"]["processing_status"] == "success"
assert row["metadata"]["input_length"] == len("test_input")
assert row["metadata"]["response_length"] == len("test_response")
def test_add_metadata_failure(self):
component = BatchRunComponent(enable_metadata=True)
row = component._create_base_row()
component._add_metadata(row, success=False, error="test_error")
assert "metadata" in row
assert row["metadata"]["processing_status"] == "failed"
assert row["metadata"]["error"] == "test_error"
def test_metadata_disabled(self):
component = BatchRunComponent(enable_metadata=False)
row = component._create_base_row(text_input="test")
component._add_metadata(row, success=True)
assert "metadata" not in row
async def test_invalid_column_name(self):
component = BatchRunComponent(
model=MockLanguageModel(), df=DataFrame({"text": ["Hello"]}), column_name="nonexistent_column"
model=MockLanguageModel(),
df=DataFrame({"text": ["Hello"]}),
column_name="nonexistent_column",
enable_metadata=True,
)
with pytest.raises(ValueError, match=re.escape("Column 'nonexistent_column' not found in the DataFrame.")):
with pytest.raises(
ValueError,
match=re.escape("Column 'nonexistent_column' not found in the DataFrame. Available columns: text"),
):
await component.run_batch()
async def test_empty_dataframe(self):
component = BatchRunComponent(model=MockLanguageModel(), df=DataFrame({"text": []}), column_name="text")
component = BatchRunComponent(
model=MockLanguageModel(),
df=DataFrame({"text": []}),
column_name="text",
enable_metadata=True,
)
result = await component.run_batch()
assert isinstance(result, DataFrame)
assert len(result) == 0
async def test_non_string_column_conversion(self):
test_df = DataFrame(
{
"text": [123, 456, 789] # Numeric values
}
)
test_df = DataFrame({"text": [123, 456, 789]}) # Numeric values
component = BatchRunComponent(model=MockLanguageModel(), df=test_df, column_name="text")
component = BatchRunComponent(
model=MockLanguageModel(),
df=test_df,
column_name="text",
enable_metadata=True,
)
result = await component.run_batch()
assert isinstance(result, DataFrame)
assert all(isinstance(text, str) for text in result["text_input"])
assert all(str(num) in text for num, text in zip(test_df["text"], result["text_input"], strict=False))
result_dicts = result.to_dict("records")
assert all(row["metadata"]["processing_status"] == "success" for row in result_dicts)