From 01bc4e98aff4f0799189b3354288d47c28b86b2b Mon Sep 17 00:00:00 2001 From: innainu Date: Wed, 23 Aug 2023 15:07:19 -0400 Subject: [PATCH] add chat in mapping so e2e tests don't fail --- cohere/responses/custom_model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cohere/responses/custom_model.py b/cohere/responses/custom_model.py index b24c4f8c6..2481e0e2d 100644 --- a/cohere/responses/custom_model.py +++ b/cohere/responses/custom_model.py @@ -28,12 +28,13 @@ "QUEUED", ] -INTERNAL_CUSTOM_MODEL_TYPE = Literal["GENERATIVE", "CLASSIFICATION", "RERANK"] -CUSTOM_MODEL_TYPE = Literal["GENERATIVE", "CLASSIFY", "RERANK"] +INTERNAL_CUSTOM_MODEL_TYPE = Literal["GENERATIVE", "CLASSIFICATION", "RERANK", "CHAT"] +CUSTOM_MODEL_TYPE = Literal["GENERATIVE", "CLASSIFY", "RERANK", "CHAT"] CUSTOM_MODEL_PRODUCT_MAPPING: Dict[CUSTOM_MODEL_TYPE, INTERNAL_CUSTOM_MODEL_TYPE] = { "GENERATIVE": "GENERATIVE", "CLASSIFY": "CLASSIFICATION", "RERANK": "RERANK", + "CHAT": "CHAT", } REVERSE_CUSTOM_MODEL_PRODUCT_MAPPING: Dict[INTERNAL_CUSTOM_MODEL_TYPE, CUSTOM_MODEL_TYPE] = { v: k for k, v in CUSTOM_MODEL_PRODUCT_MAPPING.items()