update

retrieve.py

@@ -41,19 +41,14 @@ def search(
             input_type="search_query",
         )
 
-        # In Cohere V2, embeddings are returned as a list-like object directly
-        # The response.embeddings is iterable, and we want the first element
-        try:
-            # First try: if embeddings is directly a list of embeddings
-            if isinstance(embedding_response.embeddings, list):
-                query_embedding = list(embedding_response.embeddings[0])
-            else:
-                # Convert to list if it's an iterable object
-                embeddings_list = [e for e in embedding_response.embeddings]
-                query_embedding = list(embeddings_list[0])
-        except (TypeError, IndexError) as e:
-            logger.error(f"Failed to extract embedding: {type(e).__name__}: {e}")
-            raise
+        # Extract embedding from Cohere V2 response
+        # The response.embeddings is EmbedByTypeResponseEmbeddings object
+        # Access first embedding with direct indexing
+        query_embedding = embedding_response.embeddings[0]
+
+        # Ensure it's a plain list of floats
+        if not isinstance(query_embedding, list):
+            query_embedding = list(query_embedding)
 
         logger.debug(f"Generated embedding dimension: {len(query_embedding)}")
 
@@ -85,4 +80,4 @@ def search(
             f"Search failed for query '{query_text[:100]}...': {type(e).__name__}: {e}",
             exc_info=True,
         )
-        raise
+        raise