Merge pull request #844 from parea-ai/PAI-1093-suppress-logging-durin…

…g-ds-py-optimization

feat: enable disabling logging

parea/experiment/experiment.py

@@ -14,9 +14,9 @@
 from tqdm.asyncio import tqdm_asyncio
 
 from parea import Parea
-from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID, TURN_OFF_PAREA_LOGGING
+from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID
 from parea.experiment.dvc import save_results_to_dvc_if_init
-from parea.helpers import duplicate_dicts, gen_random_name
+from parea.helpers import duplicate_dicts, gen_random_name, is_logging_disabled
 from parea.schemas import EvaluationResult
 from parea.schemas.models import CreateExperimentRequest, ExperimentSchema, ExperimentStatsSchema, FinishExperimentRequestSchema
 from parea.utils.trace_utils import thread_ids_running_evals, trace_data
@@ -212,7 +212,7 @@ def run(self, run_name: Optional[str] = None) -> None:
         param run_name: The run name of the experiment. This name must be unique across experiment runs.
         If no run name is provided a memorable name will be generated automatically.
         """
-        if TURN_OFF_PAREA_LOGGING:
+        if is_logging_disabled():
             print("Parea logging is turned off. Experiment can't be run without logging. Set env var TURN_OFF_PAREA_LOGGING to False to enable.")
             return
 
@@ -232,7 +232,7 @@ async def arun(self, run_name: Optional[str] = None) -> None:
         param run_name: The run name of the experiment. This name must be unique across experiment runs.
         If no run name is provided a memorable name will be generated automatically.
         """
-        if TURN_OFF_PAREA_LOGGING:
+        if is_logging_disabled():
             print("Parea logging is turned off. Experiment can't be run without logging. Set env var TURN_OFF_PAREA_LOGGING to False to enable.")
             return
 

parea/helpers.py

@@ -10,7 +10,7 @@
 from attr import asdict, fields_dict
 from cattrs import GenConverter
 
-from parea.constants import ADJECTIVES, NOUNS
+from parea.constants import ADJECTIVES, NOUNS, TURN_OFF_PAREA_LOGGING
 from parea.schemas.models import Completion, TraceLog, UpdateLog
 from parea.utils.universal_encoder import json_dumps
 
@@ -97,3 +97,29 @@ def structure_union_type(obj: Any, cl: type) -> Any:
 
 def structure_trace_logs_from_api(data: List[dict]) -> List[TraceLog]:
     return [structure_trace_log_from_api(d) for d in data]
+
+
+PAREA_LOGGING_DISABLED = False
+
+
+def disable_parea_logging():
+    global PAREA_LOGGING_DISABLED
+    PAREA_LOGGING_DISABLED = True
+
+
+def enable_logging():
+    global PAREA_LOGGING_DISABLED
+    PAREA_LOGGING_DISABLED = False
+
+
+def is_logging_disabled() -> bool:
+    global PAREA_LOGGING_DISABLED
+    return PAREA_LOGGING_DISABLED or TURN_OFF_PAREA_LOGGING
+
+
+class TurnOffPareaLogging:
+    def __enter__(self):
+        disable_parea_logging()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        enable_logging()

parea/utils/trace_integrations/langchain.py

@@ -5,7 +5,7 @@
 from langchain_core.tracers import BaseTracer
 from langchain_core.tracers.schemas import ChainRun, LLMRun, Run, ToolRun
 
-from parea.constants import TURN_OFF_PAREA_LOGGING
+from parea.helpers import is_logging_disabled
 from parea.parea_logger import parea_logger
 from parea.schemas import UpdateTraceScenario
 from parea.schemas.log import TraceIntegrations
@@ -18,7 +18,7 @@ class PareaAILangchainTracer(BaseTracer):
     _parea_parent_trace_id: str = None
 
     def _persist_run(self, run: Union[Run, LLMRun, ChainRun, ToolRun]) -> None:
-        if TURN_OFF_PAREA_LOGGING:
+        if is_logging_disabled():
             return
         self.parent_trace_id = run.id
         # using .dict() since langchain Run class currently set to Pydantic v1

parea/utils/trace_utils.py

@@ -11,8 +11,8 @@
 from functools import wraps
 from random import random
 
-from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID, TURN_OFF_PAREA_LOGGING
-from parea.helpers import gen_trace_id, timezone_aware_now
+from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID
+from parea.helpers import gen_trace_id, is_logging_disabled, timezone_aware_now
 from parea.parea_logger import parea_logger
 from parea.schemas import EvaluationResult
 from parea.schemas.models import TraceLog, UpdateTraceScenario
@@ -153,7 +153,7 @@ def init_trace(func_name, _parea_target_field, args, kwargs, func) -> Tuple[str,
         new_trace_context = trace_context.get() + [trace_id]
         token = trace_context.set(new_trace_context)
 
-        if TURN_OFF_PAREA_LOGGING:
+        if is_logging_disabled():
             return trace_id, start_time, token
 
         try:
@@ -226,7 +226,7 @@ async def async_wrapper(*args, **kwargs):
             output_as_list = check_multiple_return_values(func)
             try:
                 result = await func(*args, **kwargs)
-                if not TURN_OFF_PAREA_LOGGING and not log_omit_outputs:
+                if not is_logging_disabled() and not log_omit_outputs:
                     fill_trace_data(trace_id, {"result": result, "output_as_list": output_as_list, "eval_funcs_names": eval_funcs_names}, UpdateTraceScenario.RESULT)
                 return result
             except Exception as e:
@@ -246,7 +246,7 @@ def wrapper(*args, **kwargs):
             output_as_list = check_multiple_return_values(func)
             try:
                 result = func(*args, **kwargs)
-                if not TURN_OFF_PAREA_LOGGING and not log_omit_outputs:
+                if not is_logging_disabled() and not log_omit_outputs:
                     fill_trace_data(trace_id, {"result": result, "output_as_list": output_as_list, "eval_funcs_names": eval_funcs_names}, UpdateTraceScenario.RESULT)
                 return result
             except Exception as e:

parea/wrapper/openai_beta_wrapper.py

@@ -10,8 +10,8 @@
 from openai import AsyncOpenAI
 from openai.types.beta.threads import Run
 
-from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID, TURN_OFF_PAREA_LOGGING
-from parea.helpers import gen_trace_id, timezone_aware_now
+from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID
+from parea.helpers import gen_trace_id, is_logging_disabled, timezone_aware_now
 from parea.schemas import TraceLog, UpdateTraceScenario
 from parea.utils.trace_utils import check_multiple_return_values, logger_all_possible, make_output, trace_context, trace_data
 from parea.utils.universal_encoder import json_dumps
@@ -58,7 +58,7 @@ def init_trace(self, func_name, args, kwargs, func):
         new_trace_context = trace_context.get() + [trace_id]
         token = trace_context.set(new_trace_context)
 
-        if TURN_OFF_PAREA_LOGGING:
+        if is_logging_disabled():
             return trace_id, start_time, token
 
         try:

parea/wrapper/wrapper.py

@@ -9,9 +9,9 @@
 from uuid import uuid4
 
 from parea.cache.cache import Cache
-from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID, TURN_OFF_PAREA_LOGGING
+from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID
 from parea.evals.utils import _make_evaluations
-from parea.helpers import timezone_aware_now
+from parea.helpers import is_logging_disabled, 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 safe_format_template_to_prompt, skip_decorator_if_func_in_stack
@@ -84,7 +84,7 @@ def _init_trace(self, kwargs) -> Tuple[str, datetime, contextvars.Token]:
                 if isinstance(m, dict) and "content" in m:
                     m["content"] = safe_format_template_to_prompt(m["content"], **template_inputs)
 
-        if TURN_OFF_PAREA_LOGGING:
+        if is_logging_disabled():
             return trace_id, start_time, token
         try:
             trace_data.get()[trace_id] = TraceLog(

pyproject.toml

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