fix: Removed answer_length and mode from answer (#65)

* fix: Removed answer_length and mode from answer

* test: Fixed test

Author: asafgardin

ai21/clients/common/answer_base.py

@@ -1,7 +1,7 @@
 from abc import ABC
-from typing import Optional, Any, Dict
+from typing import Any, Dict
 
-from ai21.models import Mode, AnswerLength, AnswerResponse
+from ai21.models import AnswerResponse
 
 
 class Answer(ABC):
@@ -11,17 +11,12 @@ def create(
         self,
         context: str,
         question: str,
-        *,
-        answer_length: Optional[AnswerLength] = None,
-        mode: Optional[Mode] = None,
         **kwargs,
     ) -> AnswerResponse:
         """
 
         :param context: A string containing the document context for which the question will be answered
         :param question: A string containing the question to be answered based on the provided context.
-        :param answer_length: Approximate length of the answer in words.
-        :param mode:
         :param kwargs:
         :return:
         """
@@ -34,7 +29,5 @@ def _create_body(
         self,
         context: str,
         question: str,
-        answer_length: Optional[AnswerLength],
-        mode: Optional[str],
     ) -> Dict[str, Any]:
-        return {"context": context, "question": question, "answerLength": answer_length, "mode": mode}
+        return {"context": context, "question": question}

ai21/clients/sagemaker/resources/sagemaker_answer.py

@@ -1,21 +1,16 @@
-from typing import Optional
-
 from ai21.clients.common.answer_base import Answer
 from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource
-from ai21.models import AnswerResponse, AnswerLength, Mode
+from ai21.models import AnswerResponse
 
 
 class SageMakerAnswer(SageMakerResource, Answer):
     def create(
         self,
         context: str,
         question: str,
-        *,
-        answer_length: Optional[AnswerLength] = None,
-        mode: Optional[Mode] = None,
         **kwargs,
     ) -> AnswerResponse:
-        body = self._create_body(context=context, question=question, answer_length=answer_length, mode=mode)
+        body = self._create_body(context=context, question=question)
         response = self._invoke(body)
 
         return self._json_to_response(response)

ai21/clients/studio/resources/studio_answer.py

@@ -1,23 +1,18 @@
-from typing import Optional
-
 from ai21.clients.common.answer_base import Answer
 from ai21.clients.studio.resources.studio_resource import StudioResource
-from ai21.models import AnswerLength, Mode, AnswerResponse
+from ai21.models import AnswerResponse
 
 
 class StudioAnswer(StudioResource, Answer):
     def create(
         self,
         context: str,
         question: str,
-        *,
-        answer_length: Optional[AnswerLength] = None,
-        mode: Optional[Mode] = None,
         **kwargs,
     ) -> AnswerResponse:
         url = f"{self._client.get_base_url()}/{self._module_name}"
 
-        body = self._create_body(context=context, question=question, answer_length=answer_length, mode=mode)
+        body = self._create_body(context=context, question=question)
 
         response = self._post(url=url, body=body)
 

ai21/clients/studio/resources/studio_library.py

@@ -2,7 +2,7 @@
 
 from ai21.ai21_http_client import AI21HTTPClient
 from ai21.clients.studio.resources.studio_resource import StudioResource
-from ai21.models import Mode, AnswerLength, FileResponse, LibraryAnswerResponse, LibrarySearchResponse
+from ai21.models import FileResponse, LibraryAnswerResponse, LibrarySearchResponse
 
 
 class StudioLibrary(StudioResource):
@@ -107,8 +107,6 @@ def create(
         path: Optional[str] = None,
         field_ids: Optional[List[str]] = None,
         labels: Optional[List[str]] = None,
-        answer_length: Optional[AnswerLength] = None,
-        mode: Optional[Mode] = None,
         **kwargs,
     ) -> LibraryAnswerResponse:
         url = f"{self._client.get_base_url()}/{self._module_name}"
@@ -117,8 +115,6 @@ def create(
             "path": path,
             "fieldIds": field_ids,
             "labels": labels,
-            "answerLength": answer_length,
-            "mode": mode,
         }
         raw_response = self._post(url=url, body=body)
         return LibraryAnswerResponse.from_dict(raw_response)

ai21/models/__init__.py

@@ -1,9 +1,7 @@
-from ai21.models.answer_length import AnswerLength
 from ai21.models.chat_message import ChatMessage
 from ai21.models.document_type import DocumentType
 from ai21.models.embed_type import EmbedType
 from ai21.models.improvement_type import ImprovementType
-from ai21.models.mode import Mode
 from ai21.models.paraphrase_style_type import ParaphraseStyleType
 from ai21.models.penalty import Penalty
 from ai21.models.responses.answer_response import AnswerResponse
@@ -32,8 +30,6 @@
 
 
 __all__ = [
-    "AnswerLength",
-    "Mode",
     "ChatMessage",
     "RoleType",
     "Penalty",

ai21/models/answer_length.py

@@ -1,6 +1,5 @@
 import pytest
 from ai21 import AI21Client
-from ai21.models import AnswerLength, Mode
 
 _CONTEXT = (
     "Holland is a geographical region[2] and former province on the western coast of"
@@ -27,8 +26,6 @@ def test_answer(question: str, is_answer_in_context: bool, expected_answer_type:
     response = client.answer.create(
         context=_CONTEXT,
         question=question,
-        answer_length=AnswerLength.LONG,
-        mode=Mode.FLEXIBLE,
     )
 
     assert response.answer_in_context == is_answer_in_context

ai21/models/mode.py

@@ -60,9 +60,7 @@ def get_studio_answer():
         {"context": _DUMMY_CONTEXT, "question": _DUMMY_QUESTION},
         "answer",
         {
-            "answerLength": None,
             "context": _DUMMY_CONTEXT,
-            "mode": None,
             "question": _DUMMY_QUESTION,
         },
         AnswerResponse(id="some-id", answer_in_context=True, answer="42"),

tests/integration_tests/clients/studio/test_answer.py

@@ -96,9 +96,7 @@ def test__create__when_pass_kwargs__should_not_pass_to_request(self, mock_ai21_s
             method="POST",
             url=_BASE_URL + "/answer",
             params={
-                "answerLength": None,
                 "context": _DUMMY_CONTEXT,
-                "mode": None,
                 "question": _DUMMY_QUESTION,
             },
             files=None,