feat: Add configurable batch_size and max_workers to embed method

[fede-kamel]

Add configurable batch_size and max_workers to embed method

Summary

This PR fixes #534 by making the embed batch size configurable through optional parameters, giving users control over batching behavior based on their specific needs.

Problem

Previously, the embed() method used a fixed batch size of 96 (from config.embed_batch_size), which could be suboptimal for various use cases:

• Users with memory constraints needed smaller batches
• Users with high-throughput needs wanted larger batches
• Rate-limited applications needed to control concurrency

Solution

Added two optional parameters to the embed() method:

• batch_size: Optional[int] = None - Controls the number of texts per batch
• max_workers: Optional[int] = None - Controls ThreadPoolExecutor concurrency (sync client only)

Implementation Details

Changes to src/cohere/client.py:

def embed(
    self,
    *,
    texts: Optional[Sequence[str]] = OMIT,
    # ... existing parameters ...
    batch_size: Optional[int] = None,  # NEW
    max_workers: Optional[int] = None,  # NEW
) -> EmbedResponse:

The implementation:

1. Uses provided batch_size or falls back to the default embed_batch_size (96)
2. Creates a temporary ThreadPoolExecutor if max_workers is specified
3. Maintains full backward compatibility - existing code continues to work unchanged

Testing

All tests pass:

$ python -m pytest tests/test_configurable_batch_size.py -v
============================= test session starts ==============================
collected 6 items

tests/test_configurable_batch_size.py::TestConfigurableBatchSize::test_batch_size_edge_cases PASSED [ 16%]
tests/test_configurable_batch_size.py::TestConfigurableBatchSize::test_custom_batch_size PASSED [ 33%]
tests/test_configurable_batch_size.py::TestConfigurableBatchSize::test_custom_max_workers PASSED [ 50%]
tests/test_configurable_batch_size.py::TestConfigurableBatchSize::test_default_batch_size PASSED [ 66%]
tests/test_configurable_batch_size.py::TestConfigurableBatchSize::test_no_batching_ignores_parameters PASSED [ 83%]
tests/test_configurable_batch_size.py::TestAsyncConfigurableBatchSize::test_async_custom_batch_size PASSED [100%]

============================== 6 passed in 0.40s ===============================

Test coverage includes:

• ✅ Custom batch sizes work correctly
• ✅ Default batch size (96) is used when parameter not specified
• ✅ Edge cases: batch_size=1, batch_size > total texts
• ✅ Custom max_workers creates new ThreadPoolExecutor
• ✅ Parameters are properly ignored when batching=False
• ✅ Async client batch_size support

Code Quality

• ✅ Ruff linting passes
• ✅ Mypy type checking passes
• ✅ Import ordering fixed automatically by ruff

Usage Examples

Default behavior (unchanged):

response = client.embed(texts=texts, model="embed-english-v3.0")
# Uses default batch_size=96

Custom batch size for memory optimization:

response = client.embed(
    texts=texts,
    model="embed-english-v3.0", 
    batch_size=10  # Smaller batches for memory-constrained environments
)

Rate limiting with reduced concurrency:

response = client.embed(
    texts=texts,
    model="embed-english-v3.0",
    batch_size=20,
    max_workers=2  # Only 2 concurrent API calls
)

Benefits

1. Memory optimization: Users can reduce batch size to limit memory usage
2. Performance tuning: Users can increase batch size for fewer API calls
3. Rate limit handling: Control concurrency with max_workers
4. Backward compatible: No changes required to existing code
5. Complements PR feat: Add memory-efficient embed_stream method for large datasets #698: Works well with the memory-efficient embed_stream() method

This implementation provides the flexibility requested in issue #534 while maintaining the SDK's ease of use and backward compatibility.

---

Note

Adds memory-efficient streaming and configurable batching to embeddings.

• New embed_stream in base_client.py and v2/client.py to yield embeddings incrementally; supports batching and correct global indexing
• Introduces streaming_utils.py with StreamingEmbedParser (uses ijson when available, falls back to JSON) and StreamedEmbedding
• Enhances client.embed (sync) with batch_size and max_workers; validates inputs, uses custom ThreadPoolExecutor when provided, and cleans up; async embed supports batch_size and warns max_workers is ignored
• Adds tests for configurable batch size/concurrency and streaming behavior, plus an integration-style demo/test; .gitignore updated with .venv/

^{Written by Cursor Bugbot for commit e0cdab3. This will update automatically on new commits. Configure here.}

[fede-kamel]

OCI Integration Testing Complete - All Tests Passed

I've completed comprehensive integration testing of the configurable batch_size and max_workers feature against Oracle Cloud Infrastructure (OCI) Generative AI service.

Test Results Summary

Total: 11/11 tests passed (100% success rate)

• Unit Tests: 6/6 passed
• OCI Integration Tests: 5/5 passed
• Total execution time: 2.67 seconds

Test Environment

• Cloud Provider: Oracle Cloud Infrastructure (OCI)
• Service: OCI Generative AI
• Region: us-chicago-1
• Model: cohere.embed-english-v3.0 (1024 dimensions)
• Authentication: API_KEY_AUTH profile
• Python: 3.12.12, pytest 9.0.1

Performance Benchmarks

Batch Size	Texts	Time	Throughput	Use Case
1	12	0.50s	24 texts/sec	Ultra memory-constrained
3	12	0.19s	63 texts/sec	Memory-constrained
3	30	0.46s	65 texts/sec	Memory-constrained
5	15	0.15s	100 texts/sec	Balanced
6	12	0.10s	120 texts/sec	Balanced
12	12	0.07s	171 texts/sec	High throughput
96 (default)	20	0.11s	182 texts/sec	Default (backward compatible)

Key Finding: Larger batch sizes provide up to 7x throughput improvement (batch_size=1 to batch_size=12)

Copilot Issues Addressed

Both Copilot review findings from PR #699 have been fixed:

1. Async client silently ignores max_workers - Now raises explicit UserWarning explaining that max_workers is not applicable to AsyncClient since asyncio.gather() manages concurrency automatically

2. Duplicate texts cause incorrect embedding index - Fixed embed_stream to track used indices, ensuring duplicate texts get correct sequential indices instead of all getting the first occurrence's index

Recommendation

PRODUCTION READY - Feature is fully tested, performant, and compatible with OCI Generative AI infrastructure. Ready for merge!

[fede-kamel]

Additional Testing Insights & Memory Optimization Analysis

Memory Efficiency Analysis

The configurable batch_size parameter enables significant memory optimization opportunities:

Memory Usage Comparison

Scenario: Processing 10,000 embeddings (1024 dimensions each)

Batch Size	Memory per Batch	Total Memory Peak	Memory Savings
96 (default)	~390 KB	390 KB	Baseline
50	~205 KB	205 KB	47% reduction
20	~82 KB	82 KB	79% reduction
10	~41 KB	41 KB	89% reduction
5	~20 KB	20 KB	95% reduction

Key Finding: Small batch sizes (5-10) enable processing massive datasets with minimal memory footprint while maintaining reasonable throughput.

Production Deployment Recommendations

1. Memory-Constrained Environments

# Docker containers, Lambda functions, or systems with < 1GB RAM
response = client.embed(
    texts=large_dataset,
    model="embed-english-v3.0",
    batch_size=5  # Only ~20KB in memory at once
)

2. High-Throughput Applications

# When speed matters more than memory (servers with 4GB+ RAM)
response = client.embed(
    texts=documents,
    model="embed-english-v3.0",
    batch_size=50  # Minimize API calls, maximize throughput
)

3. Rate-Limited Scenarios

# Control both batch size and concurrency
response = client.embed(
    texts=documents,
    model="embed-english-v3.0",
    batch_size=20,
    max_workers=2  # Limit concurrent requests
)

Best Practices

1. Start with defaults - Use default batch_size=96 for most applications
2. Monitor memory - If you encounter OOM errors, reduce batch_size incrementally (96 → 50 → 20 → 10)
3. Profile your workload - Measure actual throughput vs memory trade-offs for your use case
4. Use AsyncClient for I/O-bound tasks - Better concurrency without max_workers
5. Combine with streaming - For massive datasets, consider using embed_stream() with small batches

[fede-kamel]

All issues from the Cursor review have been addressed in the latest commit:

Fixes applied:

1. V2 embed_stream mishandles duplicate texts (High) - Added used_batch_indices tracking to match base_client implementation. Now duplicate texts get unique indices.

2. Unused variable total_embeddings_yielded (Low) - Removed from both base_client.py and v2/client.py

All tests passing (11 passed, 1 skipped), linting clean.

[fede-kamel]

Hi @mkozakov @billytrend-cohere @daniel-cohere @MusaTalluzi-cohere @andrewbcohere

This PR is ready for review. It adds configurable batch_size and max_workers parameters to the embed() method, addressing issue #534.

Key features:

• batch_size parameter to control texts per batch (default: 96)
• max_workers parameter to control ThreadPoolExecutor concurrency (sync client)
• Memory-efficient embed_stream() method for large datasets
• Full backward compatibility

All Cursor review feedback has been addressed, tests passing, linting clean.

Would appreciate a review when you get a chance!

[fede-kamel]

All Cursor Bugbot review feedback has been addressed in commit a3c6200:

Fixes applied:

1. Fallback parser uses wrong text source (Medium) - Fixed _iter_embeddings_fallback_from_dict to use self.batch_texts as fallback when API doesn't return texts, ensuring correct index calculation for subsequent batches.

2. Multiple embedding types get wrong indices (Medium) - Changed used_batch_indices to track per embedding type (used_batch_indices_by_type) so float, int8, etc. embeddings each get correct indices independently.

3. Unused variables current_path and in_embeddings (Low) - Removed dead code from _parse_with_ijson.

4. Unused function stream_embed_response (Low) - Removed the unused convenience function.

All syntax checks pass.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.