model raw output

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

Author: vazirim

examples/hello/hello_model_raw.pdl

@@ -0,0 +1,17 @@
+description: Hello world
+text:
+- text: "Hello\n"
+  contribute: [context]
+- model: ibm/granite-8b-code-instruct
+  platform: bam
+  def: output
+  modelResponse: raw_output
+  parameters:
+    decoding_method: greedy
+    stop_sequences: ["!"]
+    include_stop_sequence: true
+    return_options: 
+      generated_tokens: True  
+      token_logprobs: True  
+  contribute: []
+- ${ raw_output }

src/pdl/pdl-schema.json

@@ -1311,6 +1311,18 @@
           "default": null,
           "title": "Trace"
         },
+        "modelResponse": {
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Modelresponse"
+        },
         "platform": {
           "const": "bam",
           "enum": [
@@ -9132,6 +9144,18 @@
           "default": null,
           "title": "Trace"
         },
+        "modelResponse": {
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Modelresponse"
+        },
         "platform": {
           "const": "litellm",
           "default": "litellm",

src/pdl/pdl_ast.py

@@ -256,6 +256,7 @@ class ModelBlock(Block):
     model: str | ExpressionType
     input: Optional["BlocksType"] = None
     trace: Optional["BlockType"] = None
+    modelResponse: Optional[str] = None
 
 
 class BamModelBlock(ModelBlock):

src/pdl/pdl_dumper.py

@@ -104,6 +104,8 @@ def block_to_dict(block: pdl_ast.BlockType, json_compatible: bool) -> DumpedBloc
                 d["data"] = block.data
             if block.constraints is not None:
                 d["constraints"] = block.constraints
+            if block.modelResponse is not None:
+                d["modelResponse"] = block.modelResponse
         case LitellmModelBlock():
             d["platform"] = block.platform
             d["model"] = block.model
@@ -116,6 +118,8 @@ def block_to_dict(block: pdl_ast.BlockType, json_compatible: bool) -> DumpedBloc
                     )
                 else:
                     d["parameters"] = block.parameters
+            if block.modelResponse is not None:
+                d["modelResponse"] = block.modelResponse
         case CodeBlock():
             d["lang"] = block.lang
             d["code"] = blocks_to_dict(block.code, json_compatible)

src/pdl/pdl_interpreter.py

@@ -77,6 +77,7 @@
     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
@@ -1058,7 +1059,7 @@ def get_transformed_inputs(kwargs):
 
         litellm.input_callback = [get_transformed_inputs]
         # append_log(state, "Model Input", messages_to_str(model_input))
-        msg = 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:
@@ -1069,6 +1070,8 @@ def get_transformed_inputs(kwargs):
         result = msg["content"]
         append_log(state, "Model Output", result)
         trace = block.model_copy(update={"result": result, "trace": concrete_block})
+        if block.modelResponse is not None:
+            scope = scope | {block.modelResponse: raw_result}
         return result, background, scope, trace
     except Exception as exc:
         message = f"Error during model call: {repr(exc)}"
@@ -1086,7 +1089,7 @@ def generate_client_response(  # pylint: disable=too-many-arguments
 ) -> Generator[YieldMessage, Any, Message]:
     match state.batch:
         case 0:
-            model_output = yield from generate_client_response_streaming(
+            model_output, raw_result = yield from generate_client_response_streaming(
                 state, block, model_input
             )
         case 1:
@@ -1097,14 +1100,14 @@ def generate_client_response(  # pylint: disable=too-many-arguments
             model_output = yield from generate_client_response_batching(
                 state, block, model_input
             )
-    return model_output
+    return model_output, raw_result
 
 
 def generate_client_response_streaming(
     state: InterpreterState,
     block: BamModelBlock | LitellmModelBlock,
     model_input: Messages,
-) -> Generator[YieldMessage, Any, Message]:
+) -> Generator[YieldMessage, Any, tuple[Message, Any]]:
     msg_stream: Generator[Message, Any, None]
     model_input_str = messages_to_str(block.model, model_input)
     match block:
@@ -1127,7 +1130,8 @@ def generate_client_response_streaming(
             assert False
     complete_msg: Optional[Message] = None
     role = None
-    for chunk in msg_stream:
+    wrapped_gen = GeneratorWrapper(msg_stream)
+    for chunk in wrapped_gen:
         if state.yield_result:
             yield ModelYieldResultMessage(chunk["content"])
         if state.yield_background:
@@ -1139,9 +1143,12 @@ def generate_client_response_streaming(
             chunk_role = chunk["role"]
             if chunk_role is None or chunk_role == role:
                 complete_msg["content"] += chunk["content"]
+    raw_result = None
+    if block.modelResponse is not None:
+        raw_result = wrapped_gen.value
     if complete_msg is None:
         return Message(role=state.role, content="")
-    return complete_msg
+    return complete_msg, raw_result
 
 
 def litellm_parameters_to_dict(

src/pdl/pdl_llms.py

@@ -1,6 +1,7 @@
 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
@@ -76,9 +77,10 @@ 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, None]:
+    ) -> Generator[Message, Any, list[Any]]:
         client = BamModel.get_model()
         params = set_default_model_params(parameters)
+        responses = []
         for response in client.text.generation.create_stream(
             model_id=model_id,
             prompt_id=prompt_id,
@@ -87,6 +89,7 @@ def generate_text_stream(  # pylint: disable=too-many-arguments,too-many-positio
             moderations=moderations,
             data=data,
         ):
+            responses.append(json.loads(response.model_dump_json()))
             if response.results is None:
                 # append_log(
                 #     state,
@@ -97,7 +100,7 @@ def generate_text_stream(  # pylint: disable=too-many-arguments,too-many-positio
             for result in response.results:
                 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,