cleanup

Signed-off-by: Mandana Vaziri <[email protected]>

Author: vazirim

src/pdl/pdl_interpreter.py

@@ -1088,7 +1088,7 @@ def generate_client_response(  # pylint: disable=too-many-arguments
     state: InterpreterState,
     block: BamModelBlock | LitellmModelBlock,
     model_input: Messages,
-) -> Generator[YieldMessage, Any, Message]:
+) -> Generator[YieldMessage, Any, tuple[Message, Any]]:
     raw_result = None
     match state.batch:
         case 0:
@@ -1111,7 +1111,7 @@ def generate_client_response_streaming(
     block: BamModelBlock | LitellmModelBlock,
     model_input: Messages,
 ) -> Generator[YieldMessage, Any, tuple[Message, Any]]:
-    msg_stream: Generator[Message, Any, None]
+    msg_stream: Generator[Message, Any, Any]
     model_input_str = messages_to_str(block.model, model_input)
     match block:
         case BamModelBlock():
@@ -1150,7 +1150,7 @@ def generate_client_response_streaming(
     if block.modelResponse is not None:
         raw_result = wrapped_gen.value
     if complete_msg is None:
-        return Message(role=state.role, content="")
+        return Message(role=state.role, content=""), raw_result
     return complete_msg, raw_result
 
 
@@ -1169,7 +1169,7 @@ def generate_client_response_single(
     state: InterpreterState,
     block: BamModelBlock | LitellmModelBlock,
     model_input: Messages,
-) -> Generator[YieldMessage, Any, Message]:
+) -> Generator[YieldMessage, Any, tuple[Message, Any]]:
     msg: Message
     model_input_str = messages_to_str(block.model, model_input)
     match block:

src/pdl/pdl_llms.py

@@ -51,7 +51,7 @@ def generate_text(  # pylint: disable=too-many-arguments,too-many-positional-arg
         parameters: Optional[dict | BamTextGenerationParameters],
         moderations: Optional[BamModerationParameters],
         data: Optional[BamPromptTemplateData],
-    ) -> tuple[Message, list[Any]]:
+    ) -> tuple[Message, Any]:
         client = BamModel.get_model()
         params = set_default_model_params(parameters)
         text = ""
@@ -79,7 +79,7 @@ def generate_text_stream(  # pylint: disable=too-many-arguments,too-many-positio
         parameters: Optional[dict | BamTextGenerationParameters],
         moderations: Optional[BamModerationParameters],
         data: Optional[BamPromptTemplateData],
-    ) -> Generator[Message, Any, list[Any]]:
+    ) -> Generator[Message, Any, Any]:
         client = BamModel.get_model()
         params = set_default_model_params(parameters)
         responses = []
@@ -149,7 +149,7 @@ def generate_text(
         model_id: str,
         messages: list[Message],
         parameters: dict[str, Any],
-    ) -> Message:
+    ) -> tuple[Message, Any]:
         if "granite" in model_id and "granite-20b-code-instruct-r1.1" not in model_id:
             parameters = set_default_granite_model_parameters(model_id, parameters)
         if parameters.get("mock_response") is not None:
@@ -167,7 +167,7 @@ def generate_text_stream(
         model_id: str,
         messages: list[Message],
         parameters: dict[str, Any],
-    ) -> Generator[Message, Any, list[Any]]:
+    ) -> Generator[Message, Any, Any]:
         if "granite" in model_id and "granite-20b-code-instruct-r1.1" not in model_id:
             parameters = set_default_granite_model_parameters(model_id, parameters)
         response = completion(