Merge pull request #761 from parea-ai/PAI-1056-capture-templated-messages-from-openai-wrapper

feat: capture inputs to auto-traced LLM calls

Author: joschkabraun

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

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 <[email protected]>"]