From a127b4982c1406c1fce7bb70abc960605613bf17 Mon Sep 17 00:00:00 2001
From: Lina Tang <linatanglily@gmail.com>
Date: Fri, 9 Feb 2024 11:17:27 +0800
Subject: [PATCH] [Executor] Update prompt info in otel traces (#1969)

# Description
This pull request primarily focuses on enhancing the traceability of the
`promptflow` application. The changes involve modifying the function
parameters, adding new functions, and updating tests to validate the new
changes. The most significant changes include the addition of a new
`enrich_span_with_prompt_info` function in the `tracer.py` file, the
modification of the `render_template_jinja2` function in the
`template_rendering.py` file, and the addition of new tests in the
`test_traces.py` file.

Changes to function parameters:

*
[`src/promptflow-tools/promptflow/tools/template_rendering.py`](diffhunk://#diff-cd8fa4e9bc7912dced2d024c72d70fb9c3853cd69df5543ce19686bfb0ae196aR4-R9):
The `render_template_jinja2` function's `template` parameter type has
been changed from `str` to `PromptTemplate`. This change ensures that
the function only accepts `PromptTemplate` type inputs, providing a more
defined contract for the function usage.

New function additions:

*
[`src/promptflow/promptflow/_core/tracer.py`](diffhunk://#diff-8f8c2ae53e5ffd37a14e8a899119fbb2742486db8faab6df3fcf506e1b720ad8R267-R282):
A new function `enrich_span_with_prompt_info` has been added. This
function enriches a span with prompt information, assuming there is only
one prompt template parameter in the function. It is used in both the
`wrapped` functions within the same file.
[[1]](diffhunk://#diff-8f8c2ae53e5ffd37a14e8a899119fbb2742486db8faab6df3fcf506e1b720ad8R267-R282)
[[2]](diffhunk://#diff-8f8c2ae53e5ffd37a14e8a899119fbb2742486db8faab6df3fcf506e1b720ad8R370)
[[3]](diffhunk://#diff-8f8c2ae53e5ffd37a14e8a899119fbb2742486db8faab6df3fcf506e1b720ad8R421)

Updated tests:

*
[`src/promptflow/tests/executor/e2etests/test_traces.py`](diffhunk://#diff-a36689a949893c227689d9dadb4c1e7008a0551206871e371658f4cc70f16f07R390-R430):
New tests have been added to validate the changes made in the
`tracer.py` file. The `test_otel_trace_with_prompt` test checks if the
OpenTelemetry traces include the correct prompt template and variables.

Example:
After this PR, if the function has input with type `PromptTemplate`, we
will add prompt information in the otel traces:
```
"prompt.template": "<template>"
"prompt.variables": "<dumped_parameter_name_and_value_pairs>"
```

# All Promptflow Contribution checklist:
- [x] **The pull request does not introduce [breaking changes].**
- [ ] **CHANGELOG is updated for new features, bug fixes or other
significant changes.**
- [x] **I have read the [contribution guidelines](../CONTRIBUTING.md).**
- [ ] **Create an issue and link to the pull request to get dedicated
review from promptflow team. Learn more: [suggested
workflow](../CONTRIBUTING.md#suggested-workflow).**

## General Guidelines and Best Practices
- [x] Title of the pull request is clear and informative.
- [x] There are a small number of commits, each of which have an
informative message. This means that previously merged commits do not
appear in the history of the PR. For more information on cleaning up the
commits in your PR, [see this
page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).

### Testing Guidelines
- [x] Pull request includes test coverage for the included changes.

---------

Co-authored-by: Lina Tang <linatang@microsoft.com>
---
 src/promptflow/promptflow/_core/tracer.py     | 19 +++++++
 .../tests/executor/e2etests/test_traces.py    | 50 ++++++++++++++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/src/promptflow/promptflow/_core/tracer.py b/src/promptflow/promptflow/_core/tracer.py
index 5322cb6617f..51ae4533a1e 100644
--- a/src/promptflow/promptflow/_core/tracer.py
+++ b/src/promptflow/promptflow/_core/tracer.py
@@ -20,6 +20,7 @@
 from promptflow._core.operation_context import OperationContext
 from promptflow._utils.dataclass_serializer import serialize
 from promptflow._utils.multimedia_utils import default_json_encoder
+from promptflow._utils.tool_utils import get_inputs_for_prompt_template, get_prompt_param_name_from_func
 from promptflow.contracts.tool import ConnectionType
 from promptflow.contracts.trace import Trace, TraceType
 
@@ -263,6 +264,22 @@ def enrich_span_with_trace(span, trace):
         logging.warning(f"Failed to enrich span with trace: {e}")
 
 
+def enrich_span_with_prompt_info(span, func, kwargs):
+    try:
+        # Assume there is only one prompt template parameter in the function,
+        # we use the first one by default if there are multiple.
+        prompt_tpl_param_name = get_prompt_param_name_from_func(func)
+        if prompt_tpl_param_name is not None:
+            prompt_tpl = kwargs.get(prompt_tpl_param_name)
+            prompt_vars = {
+                key: kwargs.get(key) for key in get_inputs_for_prompt_template(prompt_tpl) if key in kwargs
+            }
+            prompt_info = {"prompt.template": prompt_tpl, "prompt.variables": serialize_attribute(prompt_vars)}
+            span.set_attributes(prompt_info)
+    except Exception as e:
+        logging.warning(f"Failed to enrich span with prompt info: {e}")
+
+
 def enrich_span_with_input(span, input):
     try:
         serialized_input = serialize_attribute(input)
@@ -350,6 +367,7 @@ async def wrapped(*args, **kwargs):
         span_name = get_node_name_from_context() if trace_type == TraceType.TOOL else trace.name
         with open_telemetry_tracer.start_as_current_span(span_name) as span:
             enrich_span_with_trace(span, trace)
+            enrich_span_with_prompt_info(span, func, kwargs)
 
             # Should not extract these codes to a separate function here.
             # We directly call func instead of calling Tracer.invoke,
@@ -400,6 +418,7 @@ def wrapped(*args, **kwargs):
         span_name = get_node_name_from_context() if trace_type == TraceType.TOOL else trace.name
         with open_telemetry_tracer.start_as_current_span(span_name) as span:
             enrich_span_with_trace(span, trace)
+            enrich_span_with_prompt_info(span, func, kwargs)
 
             # Should not extract these codes to a separate function here.
             # We directly call func instead of calling Tracer.invoke,
diff --git a/src/promptflow/tests/executor/e2etests/test_traces.py b/src/promptflow/tests/executor/e2etests/test_traces.py
index d8dc891b06c..30b47dcc1ee 100644
--- a/src/promptflow/tests/executor/e2etests/test_traces.py
+++ b/src/promptflow/tests/executor/e2etests/test_traces.py
@@ -8,12 +8,13 @@
 
 from promptflow._core.tracer import TraceType, trace
 from promptflow._utils.dataclass_serializer import serialize
+from promptflow._utils.tool_utils import get_inputs_for_prompt_template
 from promptflow.contracts.run_info import Status
 from promptflow.executor import FlowExecutor
 from promptflow.executor._result import LineResult
 
 from ..process_utils import execute_function_in_subprocess
-from ..utils import get_flow_sample_inputs, get_yaml_file, prepare_memory_exporter
+from ..utils import get_flow_folder, get_flow_sample_inputs, get_yaml_file, prepare_memory_exporter, load_content
 
 OPEN_AI_FUNCTION_NAMES = [
     "openai.resources.chat.completions.Completions.create",
@@ -28,6 +29,11 @@
     "__computed__.cumulative_token_count.total",
 ]
 
+SHOULD_INCLUDE_PROMPT_FUNCTION_NAMES = [
+    "render_template_jinja2",
+    "AzureOpenAI.chat",
+]
+
 
 def get_chat_input(stream):
     return {
@@ -380,6 +386,48 @@ def validate_openai_tokens(self, span_list):
                 for token_name in TOKEN_NAMES:
                     assert span.attributes[token_name] == token_dict[span_id][token_name]
 
+    @pytest.mark.parametrize(
+        "flow_file, inputs, prompt_tpl_file",
+        [
+            ("llm_tool", {"topic": "Hello", "stream": False}, "joke.jinja2"),
+            # Add back this test case after changing the interface of render_template_jinja2
+            # ("prompt_tools", {"text": "test"}, "summarize_text_content_prompt.jinja2"),
+        ]
+    )
+    def test_otel_trace_with_prompt(
+        self,
+        dev_connections,
+        flow_file,
+        inputs,
+        prompt_tpl_file,
+    ):
+        execute_function_in_subprocess(
+            self.assert_otel_traces_with_prompt, dev_connections, flow_file, inputs, prompt_tpl_file
+        )
+
+    def assert_otel_traces_with_prompt(self, dev_connections, flow_file, inputs, prompt_tpl_file):
+        memory_exporter = prepare_memory_exporter()
+
+        executor = FlowExecutor.create(get_yaml_file(flow_file), dev_connections)
+        line_run_id = str(uuid.uuid4())
+        resp = executor.exec_line(inputs, run_id=line_run_id)
+        assert isinstance(resp, LineResult)
+        assert isinstance(resp.output, dict)
+
+        prompt_tpl = load_content(get_flow_folder(flow_file) / prompt_tpl_file)
+        prompt_vars = list(get_inputs_for_prompt_template(prompt_tpl).keys())
+        span_list = memory_exporter.get_finished_spans()
+        for span in span_list:
+            assert span.status.status_code == StatusCode.OK
+            assert isinstance(span.name, str)
+            if span.attributes.get("function", "") in SHOULD_INCLUDE_PROMPT_FUNCTION_NAMES:
+                assert "prompt.template" in span.attributes
+                assert span.attributes["prompt.template"] == prompt_tpl
+                assert "prompt.variables" in span.attributes
+                for var in prompt_vars:
+                    if var in inputs:
+                        assert var in span.attributes["prompt.variables"]
+
     def test_flow_with_traced_function(self):
         execute_function_in_subprocess(self.assert_otel_traces_run_flow_then_traced_function)