fix: Address review feedback for embed_stream

Fixes for issues identified by Cursor bugbot:

1. Multiple embedding types IndexError (High):
   - Track text index separately per embedding type
   - Use type_indices dict to correctly map embeddings to texts

2. Image embeddings IndexError (Medium):
   - Remove images parameter from v2 embed_stream (text-only)
   - Document that images should use regular embed()

3. Fallback fails after ijson consumes stream (Medium):
   - Buffer response content before attempting ijson parsing
   - Fallback can now use buffered content if ijson fails

4. OMIT default causes TypeError (Low):
   - Check explicitly for None or OMIT sentinel
   - Handle ellipsis default value correctly

5. Zero/negative batch_size crashes (Low):
   - Add validation: raise ValueError if batch_size < 1

Author: fede-kamel

src/cohere/base_client.py

@@ -1193,19 +1193,26 @@ def embed_stream(
             print(f"Embedding {embedding.index}: {embedding.embedding[:5]}...")
             # Process/save embedding immediately
         """
-        if not texts:
+        # Validate inputs
+        if texts is None or texts is OMIT:
             return
-            
+        if batch_size < 1:
+            raise ValueError("batch_size must be at least 1")
+
         from .streaming_utils import StreamingEmbedParser
-        
+
         # Process texts in batches
-        texts_list = list(texts) if texts else []
-        total_embeddings_yielded = 0
-        
+        texts_list = list(texts)
+        if not texts_list:
+            return
+
+        # Track text index separately from embedding index (for multiple embedding types)
+        global_text_index = 0
+
         for batch_start in range(0, len(texts_list), batch_size):
             batch_end = min(batch_start + batch_size, len(texts_list))
             batch_texts = texts_list[batch_start:batch_end]
-            
+
             # Get response for this batch
             response = self._raw_client.embed(
                 texts=batch_texts,
@@ -1215,15 +1222,15 @@ def embed_stream(
                 truncate=truncate,
                 request_options=request_options,
             )
-            
+
             # 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 tracks text index per embedding type
+                # Adjust text reference to use batch_texts mapping
+                text_index_in_batch = batch_texts.index(embedding.text) if embedding.text in batch_texts else 0
+                embedding.index = batch_start + text_index_in_batch
                 yield embedding
-            total_embeddings_yielded += len(batch_texts)
 
     def rerank(
         self,

src/cohere/streaming_utils.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import io
+import json
 from dataclasses import dataclass
 from typing import Iterator, List, Optional, Union
 
@@ -21,84 +23,89 @@ class StreamedEmbedding:
     embedding: Union[List[float], List[int], str]  # float, int8, uint8, binary, ubinary, base64
     embedding_type: str
     text: Optional[str] = None
-    
+
 
 class StreamingEmbedParser:
     """
     Parses embed responses incrementally using ijson for memory efficiency.
     Falls back to regular JSON parsing if ijson is not available.
     """
-    
+
     def __init__(self, response: httpx.Response, batch_texts: Optional[List[str]] = None):
         """
         Initialize the streaming parser.
-        
+
         Args:
             response: The httpx response object
             batch_texts: The original texts for this batch (for correlation)
         """
         self.response = response
         self.batch_texts = batch_texts or []
         self.embeddings_yielded = 0
-        
+        self._response_content: Optional[bytes] = None
+
     def iter_embeddings(self) -> Iterator[StreamedEmbedding]:
         """
         Iterate over embeddings one at a time without loading all into memory.
-        
+
         Yields:
             StreamedEmbedding objects as they are parsed from the response
         """
-        if not IJSON_AVAILABLE:
-            # Fallback to regular parsing if ijson not available
+        # Try to buffer the response content first to allow fallback if ijson fails
+        # This trades some memory for reliability
+        if self._response_content is None:
+            try:
+                content = self.response.content
+                if isinstance(content, bytes):
+                    self._response_content = content
+            except Exception:
+                # Content not available as bytes, will use json() method
+                pass
+
+        if not IJSON_AVAILABLE or self._response_content is None:
+            # Fallback to regular parsing if ijson not available or no bytes content
             yield from self._iter_embeddings_fallback()
             return
-            
+
         try:
             # Use ijson for memory-efficient parsing
-            parser = ijson.parse(self.response.iter_bytes(chunk_size=65536))
+            parser = ijson.parse(io.BytesIO(self._response_content))
             yield from self._parse_with_ijson(parser)
         except Exception:
             # If ijson parsing fails, fallback to regular parsing
             yield from self._iter_embeddings_fallback()
 
     def _parse_with_ijson(self, parser) -> Iterator[StreamedEmbedding]:
         """Parse embeddings using ijson incremental parser."""
-        current_path: List[str] = []
-        current_embedding = []
-        embedding_index = 0
-        embedding_type = "float"
+        current_embedding: List[Union[float, int]] = []
         response_type = None
-        in_embeddings = False
-        
+        # Track index per embedding type to properly map to texts
+        type_indices: dict[str, int] = {}
+
         for prefix, event, value in parser:
-            # Track current path
-            if event == 'map_key':
-                if current_path and current_path[-1] == 'embeddings':
-                    # This is an embedding type key (float_, int8, etc.)
-                    embedding_type = value.rstrip('_')
-            
             # Detect response type
             if prefix == 'response_type':
                 response_type = value
-                
+
             # Handle embeddings based on response type
             if response_type == 'embeddings_floats':
                 # Simple float array format
                 if prefix.startswith('embeddings.item.item'):
                     current_embedding.append(value)
                 elif prefix.startswith('embeddings.item') and event == 'end_array':
                     # Complete embedding
-                    text = self.batch_texts[embedding_index] if embedding_index < len(self.batch_texts) else None
+                    text_index = type_indices.get('float', 0)
+                    text = self.batch_texts[text_index] if text_index < len(self.batch_texts) else None
                     yield StreamedEmbedding(
                         index=self.embeddings_yielded,
-                        embedding=current_embedding,
+                        embedding=list(current_embedding),
                         embedding_type='float',
                         text=text
                     )
                     self.embeddings_yielded += 1
-                    embedding_index += 1
+                    type_indices['float'] = text_index + 1
                     current_embedding = []
-                    
+
             elif response_type == 'embeddings_by_type':
                 # Complex format with multiple embedding types
                 # Pattern: embeddings.<type>.item.item
@@ -108,66 +115,73 @@ def _parse_with_ijson(self, parser) -> Iterator[StreamedEmbedding]:
                         current_embedding.append(value)
                     elif prefix.startswith(f'embeddings.{emb_type}.item') and event == 'end_array':
                         # Complete embedding of this type
-                        text = self.batch_texts[embedding_index] if embedding_index < len(self.batch_texts) else None
+                        # Use separate index per type to correctly map to texts
+                        text_index = type_indices.get(type_name, 0)
+                        text = self.batch_texts[text_index] if text_index < len(self.batch_texts) else None
                         yield StreamedEmbedding(
                             index=self.embeddings_yielded,
-                            embedding=current_embedding,
+                            embedding=list(current_embedding),
                             embedding_type=type_name,
                             text=text
                         )
                         self.embeddings_yielded += 1
-                        embedding_index += 1
+                        type_indices[type_name] = text_index + 1
                         current_embedding = []
-                        
+
                 # Handle base64 embeddings (string format)
                 if prefix.startswith('embeddings.base64.item') and event == 'string':
-                    text = self.batch_texts[embedding_index] if embedding_index < len(self.batch_texts) else None
+                    text_index = type_indices.get('base64', 0)
+                    text = self.batch_texts[text_index] if text_index < len(self.batch_texts) else None
                     yield StreamedEmbedding(
                         index=self.embeddings_yielded,
                         embedding=value,  # base64 string
                         embedding_type='base64',
                         text=text
                     )
                     self.embeddings_yielded += 1
-                    embedding_index += 1
+                    type_indices['base64'] = text_index + 1
 
     def _iter_embeddings_fallback(self) -> Iterator[StreamedEmbedding]:
         """Fallback method using regular JSON parsing."""
-        # This still loads the full response but at least provides the same interface
-        if hasattr(self.response, 'json'):
+        # Use buffered content if available, otherwise read from response
+        if self._response_content is not None and isinstance(self._response_content, bytes):
+            data = json.loads(self._response_content)
+        elif hasattr(self.response, 'json') and callable(self.response.json):
             data = self.response.json()
         elif hasattr(self.response, '_response'):
             data = self.response._response.json()  # type: ignore
         else:
             raise ValueError("Response object does not have a json() method")
+
         response_type = data.get('response_type', '')
-        
+        texts = data.get('texts', self.batch_texts)
+
         if response_type == 'embeddings_floats':
             embeddings = data.get('embeddings', [])
-            texts = data.get('texts', [])
             for i, embedding in enumerate(embeddings):
                 yield StreamedEmbedding(
-                    index=i,
+                    index=self.embeddings_yielded,
                     embedding=embedding,
                     embedding_type='float',
                     text=texts[i] if i < len(texts) else None
                 )
-                
+                self.embeddings_yielded += 1
+
         elif response_type == 'embeddings_by_type':
             embeddings_obj = data.get('embeddings', {})
-            texts = data.get('texts', [])
-            
+
             # Iterate through each embedding type
             for emb_type, embeddings_list in embeddings_obj.items():
                 type_name = emb_type.rstrip('_')
                 if isinstance(embeddings_list, list):
                     for i, embedding in enumerate(embeddings_list):
                         yield StreamedEmbedding(
-                            index=i,
+                            index=self.embeddings_yielded,
                             embedding=embedding,
                             embedding_type=type_name,
                             text=texts[i] if i < len(texts) else None
                         )
+                        self.embeddings_yielded += 1
 
 
 def stream_embed_response(response: httpx.Response, texts: List[str]) -> Iterator[StreamedEmbedding]:

src/cohere/v2/client.py

@@ -495,7 +495,6 @@ def embed_stream(
         model: str,
         input_type: EmbedInputType,
         texts: typing.Optional[typing.Sequence[str]] = OMIT,
-        images: typing.Optional[typing.Sequence[str]] = OMIT,
         max_tokens: typing.Optional[int] = OMIT,
         output_dimension: typing.Optional[int] = OMIT,
         embedding_types: typing.Optional[typing.Sequence[EmbeddingType]] = OMIT,
@@ -505,11 +504,14 @@ def embed_stream(
     ) -> typing.Iterator[typing.Any]:  # Returns Iterator[StreamedEmbedding]
         """
         Memory-efficient streaming version of embed that yields embeddings one at a time.
-        
+
         This method processes texts in batches and yields individual embeddings as they are
         parsed from the response, without loading all embeddings into memory at once.
         Ideal for processing large datasets where memory usage is a concern.
 
+        Note: This method only supports text embeddings. For image embeddings, use the
+        regular embed() method.
+
         Parameters
         ----------
         model : str
@@ -521,9 +523,6 @@ def embed_stream(
         texts : typing.Optional[typing.Sequence[str]]
             An array of strings for the model to embed. Will be processed in batches.
 
-        images : typing.Optional[typing.Sequence[str]]
-            An array of image data URIs for the model to embed.
-
         max_tokens : typing.Optional[int]
             The maximum number of tokens to embed per input.
 
@@ -556,7 +555,7 @@ def embed_stream(
             client_name="YOUR_CLIENT_NAME",
             token="YOUR_TOKEN",
         )
-        
+
         # Process embeddings one at a time without loading all into memory
         for embedding in client.v2.embed_stream(
             model="embed-v4.0",
@@ -567,40 +566,43 @@ def embed_stream(
             print(f"Embedding {embedding.index}: {embedding.embedding[:5]}...")
             # Process/save embedding immediately
         """
-        if not texts:
+        # Validate inputs
+        if texts is None or texts is OMIT:
             return
-            
+        if batch_size < 1:
+            raise ValueError("batch_size must be at least 1")
+
         from ..streaming_utils import StreamingEmbedParser
-        
+
         # Process texts in batches
-        texts_list = list(texts) if texts else []
-        total_embeddings_yielded = 0
-        
+        texts_list = list(texts)
+        if not texts_list:
+            return
+
         for batch_start in range(0, len(texts_list), batch_size):
             batch_end = min(batch_start + batch_size, len(texts_list))
             batch_texts = texts_list[batch_start:batch_end]
-            
+
             # Get response for this batch
             response = self._raw_client.embed(
                 model=model,
                 input_type=input_type,
                 texts=batch_texts,
-                images=images if batch_start == 0 else None,  # Only include images in first batch
                 max_tokens=max_tokens,
                 output_dimension=output_dimension,
                 embedding_types=embedding_types,
                 truncate=truncate,
                 request_options=request_options,
             )
-            
+
             # 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 tracks text index per embedding type
+                # Adjust text reference to use batch_texts mapping
+                text_index_in_batch = batch_texts.index(embedding.text) if embedding.text in batch_texts else 0
+                embedding.index = batch_start + text_index_in_batch
                 yield embedding
-            total_embeddings_yielded += len(batch_texts)
 
     def rerank(
         self,