chore: organize the code

Author: zakahan

veadk/models/ark_llm.py

@@ -60,20 +60,77 @@ def __init__(self):
     async def acompletion(
         self, model, messages, tools, **kwargs
     ) -> Union[ModelResponse, CustomStreamWrapper]:
-        # 1.1. Get optional_params using get_optional_params function
+        # 1 Modify messages
+        # Keep the header system-prompt and the user's messages
+        messages = messages[:1] + messages[-1:]
+
+        # 2 Get request params
+        (
+            request_data,
+            optional_params,
+            litellm_params,
+            logging_obj,
+            custom_llm_provider,
+        ) = self._get_request_data(model, messages, tools, **kwargs)
+
+        # 3. Call litellm.aresponses with the transformed request data
+        raw_response = await aresponses(
+            **request_data,
+        )
+        # 4. Transform ResponsesAPIResponse
+        # 4.1 Create model_response object
+        model_response = ModelResponse()
+        setattr(model_response, "usage", litellm.Usage())
+
+        # 4.2 Transform ResponsesAPIResponse to ModelResponses
+        if isinstance(raw_response, ResponsesAPIResponse):
+            response = self.transformation_handler.transform_response(
+                model=model,
+                raw_response=raw_response,
+                model_response=model_response,
+                logging_obj=logging_obj,
+                request_data=request_data,
+                messages=messages,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                encoding=kwargs.get("encoding"),
+                api_key=kwargs.get("api_key"),
+                json_mode=kwargs.get("json_mode"),
+            )
+            # 4.2.1 Modify ModelResponse id
+            if raw_response and hasattr(raw_response, "id"):
+                response.id = raw_response.id
+            return response
+
+        else:
+            completion_stream = self.transformation_handler.get_model_response_iterator(
+                streaming_response=raw_response,  # type: ignore
+                sync_stream=True,
+                json_mode=kwargs.get("json_mode"),
+            )
+            streamwrapper = CustomStreamWrapper(
+                completion_stream=completion_stream,
+                model=model,
+                custom_llm_provider=custom_llm_provider,
+                logging_obj=logging_obj,
+            )
+            return streamwrapper
+
+    def _get_request_data(self, model, messages, tools, **kwargs) -> tuple:
+        # 1. Get optional_params using get_optional_params function
         optional_params = get_optional_params(model=model, tools=tools, **kwargs)
 
-        # 1.2. Get litellm_params using get_litellm_params function
+        # 2. Get litellm_params using get_litellm_params function
         litellm_params = get_litellm_params(**kwargs)
 
-        # 1.3. Get headers by merging kwargs headers and extra_headers
+        # 3. Get headers by merging kwargs headers and extra_headers
         headers = kwargs.get("headers", None) or kwargs.get("extra_headers", None)
         if headers is None:
             headers = {}
         if kwargs.get("extra_headers") is not None:
             headers.update(kwargs.get("extra_headers"))
 
-        # 1.4. Get logging_obj from kwargs or create new LiteLLMLoggingObj
+        # 4. Get logging_obj from kwargs or create new LiteLLMLoggingObj
         logging_obj = kwargs.get("litellm_logging_obj", None)
         if logging_obj is None:
             logging_obj = Logging(
@@ -86,7 +143,7 @@ async def acompletion(
                 start_time=datetime.now(),
                 kwargs=kwargs,
             )
-        # 1.5. Convert Message to `llm_provider` format
+        # 4. Convert Message to `llm_provider` format
         _, custom_llm_provider, _, _ = get_llm_provider(model=model)
         if custom_llm_provider is not None and custom_llm_provider in [
             provider.value for provider in LlmProviders
@@ -98,10 +155,8 @@ async def acompletion(
                 messages = provider_config.translate_developer_role_to_system_role(
                     messages=messages
                 )
-        # 1.6 Add response_id to llm_response
-        # Keep the header system-prompt and the user's messages
-        messages = messages[:1] + messages[-1:]
-        # 1.7 Transform request to responses api format
+
+        # 5 Transform request to responses api format
         request_data = self.transformation_handler.transform_request(
             model=model,
             messages=messages,
@@ -112,49 +167,22 @@ async def acompletion(
             client=kwargs.get("client"),
         )
 
-        # 2. Call litellm.aresponses with the transformed request data
-        raw_response = await aresponses(
-            **request_data,
+        # 6 handler Missing field supply
+        if "extra_body" not in request_data and kwargs.get("extra_body"):
+            request_data["extra_body"] = kwargs.get("extra_body")
+        if "extra_query" not in request_data and kwargs.get("extra_query"):
+            request_data["extra_query"] = kwargs.get("extra_query")
+        if "extra_headers" not in request_data and kwargs.get("extra_headers"):
+            request_data["extra_headers"] = kwargs.get("extra_headers")
+
+        return (
+            request_data,
+            optional_params,
+            litellm_params,
+            logging_obj,
+            custom_llm_provider,
         )
 
-        # 3.1 Create model_response object
-        model_response = ModelResponse()
-        setattr(model_response, "usage", litellm.Usage())
-
-        # 3.2 Transform ResponsesAPIResponse to ModelResponses
-        if isinstance(raw_response, ResponsesAPIResponse):
-            response = self.transformation_handler.transform_response(
-                model=model,
-                raw_response=raw_response,
-                model_response=model_response,
-                logging_obj=logging_obj,
-                request_data=request_data,
-                messages=messages,
-                optional_params=optional_params,
-                litellm_params=litellm_params,
-                encoding=kwargs.get("encoding"),
-                api_key=kwargs.get("api_key"),
-                json_mode=kwargs.get("json_mode"),
-            )
-            # 3.2.1 Modify ModelResponse id
-            if raw_response and hasattr(raw_response, "id"):
-                response.id = raw_response.id
-            return response
-
-        else:
-            completion_stream = self.transformation_handler.get_model_response_iterator(
-                streaming_response=raw_response,  # type: ignore
-                sync_stream=True,
-                json_mode=kwargs.get("json_mode"),
-            )
-            streamwrapper = CustomStreamWrapper(
-                completion_stream=completion_stream,
-                model=model,
-                custom_llm_provider=custom_llm_provider,
-                logging_obj=logging_obj,
-            )
-            return streamwrapper
-
 
 class ArkLlm(LiteLlm):
     llm_client: ArkLlmClient = Field(default_factory=ArkLlmClient)