fix: Address review feedback for configurable batch_size

Fixes for issues identified by Cursor bugbot:

1. Missing batch_size validation in embed method (Medium):
   - Added validation to raise ValueError if batch_size < 1
   - Applied to both sync and async embed methods

2. IndexError when using multiple embedding types with embed_stream (High):
   - Fixed index calculation to use text position from parser
   - Parser correctly tracks text index per embedding type

3. Fallback causes duplicate embeddings after partial ijson failure (Low):
   - Collect all ijson embeddings into list before yielding
   - Reset embeddings_yielded counter before fallback
   - Only yield after successful complete parsing

Author: fede-kamel

src/cohere/base_client.py

@@ -1222,12 +1222,13 @@ def embed_stream(
 
             # Parse embeddings from response incrementally
             parser = StreamingEmbedParser(response._response, batch_texts)
-            for i, embedding in enumerate(parser.iter_embeddings()):
-                # Adjust index for global position
-                embedding.index = batch_start + i
-                embedding.text = texts_list[embedding.index]
+            for embedding in parser.iter_embeddings():
+                # The parser sets embedding.text correctly for multiple embedding types
+                # Adjust the global index based on text position in batch
+                if embedding.text and embedding.text in batch_texts:
+                    text_idx_in_batch = batch_texts.index(embedding.text)
+                    embedding.index = batch_start + text_idx_in_batch
                 yield embedding
-            total_embeddings_yielded += len(batch_texts)
 
     def rerank(
         self,

src/cohere/client.py

@@ -203,6 +203,10 @@ def embed(
                 request_options=request_options,
             )
 
+        # Validate batch_size
+        if batch_size is not None and batch_size < 1:
+            raise ValueError("batch_size must be at least 1")
+
         textsarr: typing.Sequence[str]  = texts if texts is not OMIT and texts is not None else []
         effective_batch_size = batch_size if batch_size is not None else embed_batch_size
         texts_batches = [textsarr[i : i + effective_batch_size] for i in range(0, len(textsarr), effective_batch_size)]
@@ -408,6 +412,10 @@ async def embed(
                 request_options=request_options,
             )
 
+        # Validate batch_size
+        if batch_size is not None and batch_size < 1:
+            raise ValueError("batch_size must be at least 1")
+
         textsarr: typing.Sequence[str]  = texts if texts is not OMIT and texts is not None else []
         effective_batch_size = batch_size if batch_size is not None else embed_batch_size
         texts_batches = [textsarr[i : i + effective_batch_size] for i in range(0, len(textsarr), effective_batch_size)]

src/cohere/streaming_utils.py

@@ -50,21 +50,31 @@ def iter_embeddings(self) -> Iterator[StreamedEmbedding]:
         Yields:
             StreamedEmbedding objects as they are parsed from the response
         """
-        if not IJSON_AVAILABLE:
-            # Fallback to regular parsing if ijson not available
+        # Try to get response content as bytes for ijson
+        response_content: Optional[bytes] = None
+        try:
+            content = self.response.content
+            if isinstance(content, bytes):
+                response_content = content
+        except Exception:
+            pass
+
+        if not IJSON_AVAILABLE or response_content is None:
+            # Fallback to regular parsing if ijson not available or no bytes content
             yield from self._iter_embeddings_fallback()
             return
 
-        # Buffer response content first to allow fallback if ijson fails
-        # This prevents partial parsing issues where ijson yields some embeddings then fails
-        response_content = self.response.content
-
         try:
             # Use ijson for memory-efficient parsing
+            # Collect all embeddings first to avoid partial yields before failure
             parser = ijson.parse(io.BytesIO(response_content))
-            yield from self._parse_with_ijson(parser)
+            embeddings = list(self._parse_with_ijson(parser))
+            # Only yield after successful complete parsing
+            yield from embeddings
         except Exception:
             # If ijson parsing fails, fallback to regular parsing using buffered content
+            # Reset embeddings_yielded since we collected but didn't yield
+            self.embeddings_yielded = 0
             data = json.loads(response_content)
             yield from self._iter_embeddings_fallback_from_dict(data)
 

src/cohere/v2/client.py

@@ -603,12 +603,13 @@ def embed_stream(
 
             # Parse embeddings from response incrementally
             parser = StreamingEmbedParser(response._response, batch_texts)
-            for i, embedding in enumerate(parser.iter_embeddings()):
-                # Adjust index for global position
-                embedding.index = batch_start + i
-                embedding.text = texts_list[embedding.index]
+            for embedding in parser.iter_embeddings():
+                # The parser sets embedding.text correctly for multiple embedding types
+                # Adjust the global index based on text position in batch
+                if embedding.text and embedding.text in batch_texts:
+                    text_idx_in_batch = batch_texts.index(embedding.text)
+                    embedding.index = batch_start + text_idx_in_batch
                 yield embedding
-            total_embeddings_yielded += len(batch_texts)
 
     def rerank(
         self,