Merge pull request #762 from parea-ai/fix-json-templated-messages

fix: capture inputs to auto-traced LLM calls with json fields

Author: joschkabraun

parea/cookbook/tracing_templated_llm_calls.py

@@ -0,0 +1,21 @@
+import os
+
+from dotenv import load_dotenv
+from openai import OpenAI
+
+from parea import Parea
+
+load_dotenv()
+
+client = OpenAI()
+p = Parea(api_key=os.getenv("PAREA_API_KEY"))
+p.wrap_openai_client(client)
+
+response = client.chat.completions.create(
+    model="gpt-4",
+    messages=[
+        {"role": "user", "content": "Make up {{number}} people. Some {abc}: def"},
+    ],
+    template_inputs={"number": "three"},
+)
+print(response.choices[0].message.content)

parea/wrapper/utils.py

@@ -1,6 +1,7 @@
 from typing import Callable, Dict, List, Optional, Union
 
 import json
+import re
 import sys
 from functools import lru_cache, wraps
 
@@ -338,3 +339,16 @@ def convert_openai_raw_stream_to_log(content: list, tools: dict, data: dict, tra
 
 def convert_openai_raw_to_log(r: dict, data: dict):
     log_in_thread(_process_response, {"response": ChatCompletion(**r), "model_inputs": data, "trace_id": get_current_trace_id()})
+
+
+def safe_format_template_to_prompt(_template: str, **kwargs) -> str:
+    """Replaces langchain.prompts.PromptTemplate.format in a safe manner.
+
+    Only variables '{{...}}' will be replaced, not any in '{...}'
+    """
+
+    def replace(match):
+        var_name = match.group(1)
+        return str(kwargs.get(var_name, match.group(0)))
+
+    return re.sub(r"{{(\w+)}}", replace, _template)

parea/wrapper/wrapper.py

@@ -14,7 +14,7 @@
 from parea.helpers import timezone_aware_now
 from parea.schemas.models import TraceLog, UpdateTraceScenario
 from parea.utils.trace_utils import call_eval_funcs_then_log, fill_trace_data, trace_context, trace_data
-from parea.wrapper.utils import skip_decorator_if_func_in_stack
+from parea.wrapper.utils import safe_format_template_to_prompt, skip_decorator_if_func_in_stack
 
 logger = logging.getLogger()
 
@@ -82,7 +82,7 @@ def _init_trace(self, kwargs) -> Tuple[str, datetime, contextvars.Token]:
         if template_inputs := kwargs.pop("template_inputs", None):
             for m in kwargs["messages"] or []:
                 if isinstance(m, dict) and "content" in m:
-                    m["content"] = m["content"].format(**template_inputs)
+                    m["content"] = safe_format_template_to_prompt(m["content"], **template_inputs)
 
         if TURN_OFF_PAREA_LOGGING:
             return trace_id, start_time, token

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "parea-ai"
 packages = [{ include = "parea" }]
-version = "0.2.131"
+version = "0.2.132"
 description = "Parea python sdk"
 readme = "README.md"
 authors = ["joel-parea-ai <[email protected]>"]