From 1f5fcfd1a0e6659eff70a810c9c6e7294b7c5b4b Mon Sep 17 00:00:00 2001
From: Joschka Braun <joschka.braun@gmail.com>
Date: Sun, 21 Apr 2024 13:17:13 -0400
Subject: [PATCH] feat: capture inputs to auto-traced LLM calls

---
 parea/wrapper/wrapper.py | 13 +++++++++----
 pyproject.toml           |  2 +-
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/parea/wrapper/wrapper.py b/parea/wrapper/wrapper.py
index cdce9676..ba308d08 100644
--- a/parea/wrapper/wrapper.py
+++ b/parea/wrapper/wrapper.py
@@ -72,13 +72,18 @@ def _get_decorator(self, unwrapped_func: Callable, original_func: Callable):
         else:
             return self.sync_decorator(original_func)
 
-    def _init_trace(self) -> Tuple[str, datetime, contextvars.Token]:
+    def _init_trace(self, kwargs) -> Tuple[str, datetime, contextvars.Token]:
         start_time = timezone_aware_now()
         trace_id = str(uuid4())
 
         new_trace_context = trace_context.get() + [trace_id]
         token = trace_context.set(new_trace_context)
 
+        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)
+
         if TURN_OFF_PAREA_LOGGING:
             return trace_id, start_time, token
         try:
@@ -93,7 +98,7 @@ def _init_trace(self) -> Tuple[str, datetime, contextvars.Token]:
                 metadata=None,
                 target=None,
                 tags=None,
-                inputs={},
+                inputs=template_inputs,
                 experiment_uuid=os.getenv(PAREA_OS_ENV_EXPERIMENT_UUID, None),
             )
 
@@ -109,7 +114,7 @@ def _init_trace(self) -> Tuple[str, datetime, contextvars.Token]:
     def async_decorator(self, orig_func: Callable) -> Callable:
         @functools.wraps(orig_func)
         async def wrapper(*args, **kwargs):
-            trace_id, start_time, context_token = self._init_trace()
+            trace_id, start_time, context_token = self._init_trace(kwargs)
             response = None
             exception = None
             error = None
@@ -141,7 +146,7 @@ async def wrapper(*args, **kwargs):
     def sync_decorator(self, orig_func: Callable) -> Callable:
         @functools.wraps(orig_func)
         def wrapper(*args, **kwargs):
-            trace_id, start_time, context_token = self._init_trace()
+            trace_id, start_time, context_token = self._init_trace(kwargs)
             response = None
             error = None
             cache_hit = False
diff --git a/pyproject.toml b/pyproject.toml
index a446b8ce..0c7ef1f4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "parea-ai"
 packages = [{ include = "parea" }]
-version = "0.2.130"
+version = "0.2.131"
 description = "Parea python sdk"
 readme = "README.md"
 authors = ["joel-parea-ai <joel@parea.ai>"]