Skip to content

Commit 8838856

Browse files
QuinnGTbigadsoleiman
authored andcommitted
Update task with enum
1 parent a6e6eaa commit 8838856

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

lib/shared/layers/python-sdk/python/genai_core/embeddings.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ def _generate_embeddings_amazon(model: EmbeddingsModel, input: List[str], bedroc
111111

112112

113113
def _generate_embeddings_cohere(model: EmbeddingsModel, input: List[str], task: Task, bedrock):
114-
input_type = "search_query" if task == Task.RETRIEVE else "search_document"
114+
input_type = Task.SEARCH_QUERY if task == Task.RETRIEVE else Task.SEARCH_DOCUMENT
115115
body = json.dumps({"texts": input, "input_type": input_type})
116116
response = bedrock.invoke_model(
117117
body=body,

0 commit comments

Comments
 (0)