bug fixes

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

Author: vazirim

examples/hello/hello_model_raw.pdl

@@ -71,13 +71,13 @@
 from .pdl_parser import PDLParseError, parse_file
 from .pdl_scheduler import (
     CodeYieldResultMessage,
+    GeneratorWrapper,
     ModelCallMessage,
     ModelYieldResultMessage,
     YieldBackgroundMessage,
     YieldMessage,
     YieldResultMessage,
     schedule,
-    GeneratorWrapper,
 )
 from .pdl_schema_validator import type_check_args, type_check_spec
 from .pdl_utils import messages_concat, messages_to_str, stringify
@@ -1059,7 +1059,9 @@ def get_transformed_inputs(kwargs):
 
         litellm.input_callback = [get_transformed_inputs]
         # append_log(state, "Model Input", messages_to_str(model_input))
-        msg, raw_result = yield from generate_client_response(state, concrete_block, model_input)
+        msg, raw_result = yield from generate_client_response(
+            state, concrete_block, model_input
+        )
         if "input" in litellm_params:
             append_log(state, "Model Input", litellm_params["input"])
         else:
@@ -1087,13 +1089,14 @@ def generate_client_response(  # pylint: disable=too-many-arguments
     block: BamModelBlock | LitellmModelBlock,
     model_input: Messages,
 ) -> Generator[YieldMessage, Any, Message]:
+    raw_result = None
     match state.batch:
         case 0:
             model_output, raw_result = yield from generate_client_response_streaming(
                 state, block, model_input
             )
         case 1:
-            model_output = yield from generate_client_response_single(
+            model_output, raw_result = yield from generate_client_response_single(
                 state, block, model_input
             )
         case _:
@@ -1171,7 +1174,7 @@ def generate_client_response_single(
     model_input_str = messages_to_str(block.model, model_input)
     match block:
         case BamModelBlock():
-            msg = BamModel.generate_text(
+            msg, raw_result = BamModel.generate_text(
                 model_id=block.model,
                 prompt_id=block.prompt_id,
                 model_input=model_input_str,
@@ -1180,7 +1183,7 @@ def generate_client_response_single(
                 data=block.data,
             )
         case LitellmModelBlock():
-            msg = LitellmModel.generate_text(
+            msg, raw_result = LitellmModel.generate_text(
                 model_id=block.model,
                 messages=model_input,
                 parameters=litellm_parameters_to_dict(block.parameters),
@@ -1189,7 +1192,7 @@ def generate_client_response_single(
         yield YieldResultMessage(msg["content"])
     if state.yield_background:
         yield YieldBackgroundMessage([msg])
-    return msg
+    return msg, raw_result
 
 
 def generate_client_response_batching(  # pylint: disable=too-many-arguments

src/pdl/pdl_interpreter.py

@@ -1,7 +1,7 @@
+import json
 from typing import Any, Generator, Optional
 
 import litellm
-import json
 from dotenv import load_dotenv
 from genai.client import Client as BamClient
 from genai.credentials import Credentials as BamCredentials
@@ -51,10 +51,11 @@ def generate_text(  # pylint: disable=too-many-arguments,too-many-positional-arg
         parameters: Optional[dict | BamTextGenerationParameters],
         moderations: Optional[BamModerationParameters],
         data: Optional[BamPromptTemplateData],
-    ) -> Message:
+    ) -> tuple[Message, list[Any]]:
         client = BamModel.get_model()
         params = set_default_model_params(parameters)
         text = ""
+        responses = []
         for response in client.text.generation.create(
             model_id=model_id,
             prompt_id=prompt_id,
@@ -64,10 +65,11 @@ def generate_text(  # pylint: disable=too-many-arguments,too-many-positional-arg
             data=data,
         ):
             # XXX TODO: moderation
+            responses.append(response)
             for result in response.results:
                 if result.generated_text:
                     text += result.generated_text
-        return {"role": None, "content": text}
+        return {"role": None, "content": text}, responses
 
     @staticmethod
     def generate_text_stream(  # pylint: disable=too-many-arguments,too-many-positional-arguments
@@ -101,6 +103,7 @@ def generate_text_stream(  # pylint: disable=too-many-arguments,too-many-positio
                 if result.generated_text:
                     yield {"role": None, "content": result.generated_text}
         return responses
+
     # @staticmethod
     # def generate_text_lazy(  # pylint: disable=too-many-arguments
     #     model_id: str,
@@ -157,14 +160,14 @@ def generate_text(
         msg = response.choices[0].message  # pyright: ignore
         if msg.content is None:
             assert False, "TODO"  # XXX TODO XXX
-        return {"role": msg.role, "content": msg.content}
+        return {"role": msg.role, "content": msg.content}, response.json()
 
     @staticmethod
     def generate_text_stream(
         model_id: str,
         messages: list[Message],
         parameters: dict[str, Any],
-    ) -> Generator[Message, Any, None]:
+    ) -> Generator[Message, Any, list[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(
@@ -173,8 +176,11 @@ def generate_text_stream(
             stream=True,
             **parameters,
         )
+        result = []
         for chunk in response:
+            result.append(chunk.json())
             msg = chunk.choices[0].delta  # pyright: ignore
             if msg.content is None:
                 break
             yield {"role": msg.role, "content": msg.content}
+        return result