Merge pull request #1047 from parea-ai/PAI-1461-add-log-sampling-to-completion-lc-logging-endpoint

feat: add log sampling

Author: joschkabraun

cookbook/openai/tracing_with_openai_with_functions.py

@@ -17,6 +17,7 @@
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 p = Parea(api_key=os.getenv("PAREA_API_KEY"))
+p.wrap_openai_client(openai)
 
 functions = [
     {

parea/schemas/models.py

@@ -27,6 +27,7 @@ class Completion:
     log_omit_inputs: bool = False
     log_omit_outputs: bool = False
     log_omit: bool = False
+    log_sample_rate: float = 1.0
     experiment_uuid: Optional[str] = None
     project_uuid: str = "default"
 
@@ -123,6 +124,7 @@ class TraceLog(EvaluatedLog):
     output_for_eval_metrics: Optional[str] = None
     evaluation_metric_names: Optional[List[str]] = field(factory=list)
     apply_eval_frac: float = 1.0
+    log_sample_rate: float = 1.0
     feedback_score: Optional[float] = None
 
     # info filled from decorator

parea/utils/trace_integrations/langchain.py

@@ -24,6 +24,7 @@ class PareaAILangchainTracer(BaseTracer):
     _metadata: Dict[str, Any] = {}
     _end_user_identifier: Optional[str] = None
     _deployment_id: Optional[str] = None
+    _log_sample_rate: Optional[float] = 1.0
 
     def __init__(
         self,
@@ -32,12 +33,14 @@ def __init__(
         metadata: Optional[Dict[str, Any]] = None,
         end_user_identifier: Optional[str] = None,
         deployment_id: Optional[str] = None,
+        log_sample_rate: Optional[float] = 1.0,
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
         self._session_id = session_id
         self._end_user_identifier = end_user_identifier
         self._deployment_id = deployment_id
+        self._log_sample_rate = log_sample_rate
         if tags:
             self._tags = tags
         if metadata:
@@ -56,6 +59,7 @@ def _persist_run(self, run: Union[Run, LLMRun, ChainRun, ToolRun]) -> None:
             data["_metadata"] = self._metadata
             data["_end_user_identifier"] = self._end_user_identifier
             data["_deployment_id"] = self._deployment_id
+            data["_log_sample_rate"] = self._log_sample_rate
             # check if run has an attribute execution order
             if (hasattr(run, "execution_order") and run.execution_order == 1) or run.parent_run_id is None:
                 data["_parea_parent_trace_id"] = self._parea_parent_trace_id or None

parea/utils/trace_utils.py

@@ -130,8 +130,11 @@ def fill_trace_data(trace_id: str, data: Dict[str, Any], scenario: UpdateTraceSc
             trace_data.get()[trace_id].error = data["error"]
             trace_data.get()[trace_id].status = "error"
         elif scenario == UpdateTraceScenario.CHAIN:
-            trace_data.get()[trace_id].parent_trace_id = data["parent_trace_id"]
-            trace_data.get()[data["parent_trace_id"]].children.append(trace_id)
+            parent_trace_id = data["parent_trace_id"]
+            trace_data.get()[trace_id].parent_trace_id = parent_trace_id
+            trace_data.get()[parent_trace_id].children.append(trace_id)
+            if trace_data.get()[parent_trace_id].log_sample_rate != trace_data.get()[trace_id].log_sample_rate:
+                trace_data.get()[trace_id].log_sample_rate = trace_data.get()[parent_trace_id].log_sample_rate
         elif scenario == UpdateTraceScenario.LANGCHAIN_CHILD:
             trace_data.get()[data["parent_trace_id"]].children.append(trace_id)
         elif scenario == UpdateTraceScenario.OPENAICONFIG:
@@ -162,6 +165,7 @@ def trace(
     log_omit_outputs: Optional[bool] = False,
     overwrite_trace_id: Optional[str] = None,
     overwrite_inputs: Optional[Dict[str, Any]] = None,
+    log_sample_rate: Optional[float] = 1.0,
 ):
     def init_trace(func_name, _parea_target_field, args, kwargs, func) -> Tuple[str, datetime, contextvars.Token]:
         start_time = timezone_aware_now()
@@ -224,6 +228,7 @@ def init_trace(func_name, _parea_target_field, args, kwargs, func) -> Tuple[str,
                 deployment_id=deployment_id,
                 depth=depth,
                 execution_order=execution_order,
+                log_sample_rate=log_sample_rate,
             )
 
             if parent_trace_id:

pyproject.toml

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