added --litellm-tracing parameter

Author: tututzi

nerve/cli/run.py

@@ -31,6 +31,7 @@ def run(
         args.log_path,
         level="DEBUG" if args.debug else "SUCCESS" if args.quiet else "INFO",
         litellm_debug=args.litellm_debug,
+        litellm_tracing=args.litellm_tracing,
     )
     logger.info(f"🧠 nerve v{nerve.__version__}")
 

nerve/cli/utils.py

@@ -70,6 +70,10 @@ def _get_run_args(
         bool,
         typer.Option("--litellm-debug", help="Enable litellm debug logging"),
     ] = False,
+    litellm_tracing: t.Annotated[
+        str | None,
+        typer.Option("--litellm-tracing", help="Set litellm callbacks for tracing"),
+    ] = None,
     quiet: t.Annotated[
         bool,
         typer.Option("--quiet", "-q", help="Quiet mode"),
@@ -112,6 +116,7 @@ def _get_run_args(
         interactive=interactive,
         debug=debug,
         litellm_debug=litellm_debug,
+        litellm_tracing=litellm_tracing,
         quiet=quiet,
         max_steps=max_steps,
         max_cost=max_cost,

nerve/runtime/logging.py

@@ -14,6 +14,7 @@ def init(
     log_path: pathlib.Path | None = None,
     level: str = "INFO",
     litellm_debug: bool = False,
+    litellm_tracing: str | None = None,
     target: t.TextIO = sys.stderr,
 ) -> None:
     """
@@ -23,6 +24,7 @@ def init(
         log_path: The path to the log file.
         level: The log level to use.
         litellm_debug: Whether to enable litellm debug logging.
+        litellm_tracing: The callback to set for litellm tracing.        
     """
 
     if level != "DEBUG":
@@ -40,6 +42,9 @@ def init(
         # https://github.com/BerriAI/litellm/issues/4825
         litellm.suppress_debug_info = True
 
+    if litellm_tracing:
+        litellm.callbacks = [litellm_tracing]
+
     if log_path:
         logger.add(log_path, format=format, level=level)
 

nerve/runtime/runner.py

@@ -23,6 +23,7 @@ class Arguments(BaseModel):
     interactive: bool
     debug: bool
     litellm_debug: bool
+    litellm_tracing: str | None
     quiet: bool
     max_steps: int
     max_cost: float
@@ -72,6 +73,10 @@ def _create_command_line(
     # if run_args.litellm_debug:
     #     command_line.append("--litellm-debug")
 
+    if run_args.litellm_tracing:
+        command_line.append("--litellm-tracing")
+        command_line.append(run_args.litellm_tracing)
+
     # if the task is set, add it to the command line
     if "task" in input_state:
         command_line.append("--task")