single trace log prep for ch

Author: jalexanderII

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
 
@@ -336,3 +337,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()
 
@@ -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"] = safe_format_template_to_prompt(m["content"], **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.131"
+version = "0.2.132"
 description = "Parea python sdk"
 readme = "README.md"
 authors = ["joel-parea-ai <[email protected]>"]