remove default logger handlers when using libcli with notebook

This commit is contained in:
Dinesh Yeduguru 2024-12-30 15:55:24 -08:00
parent 8ba29b19f2
commit b2137ab705

View file

@ -7,6 +7,7 @@
import asyncio import asyncio
import inspect import inspect
import json import json
import logging
import os import os
import queue import queue
import threading import threading
@ -16,7 +17,6 @@ from pathlib import Path
from typing import Any, Generator, get_args, get_origin, Optional, TypeVar from typing import Any, Generator, get_args, get_origin, Optional, TypeVar
import httpx import httpx
import yaml import yaml
from llama_stack_client import ( from llama_stack_client import (
APIResponse, APIResponse,
@ -28,7 +28,6 @@ from llama_stack_client import (
) )
from pydantic import BaseModel, TypeAdapter from pydantic import BaseModel, TypeAdapter
from rich.console import Console from rich.console import Console
from termcolor import cprint from termcolor import cprint
from llama_stack.distribution.build import print_pip_install_help from llama_stack.distribution.build import print_pip_install_help
@ -41,7 +40,6 @@ from llama_stack.distribution.stack import (
get_stack_run_config_from_template, get_stack_run_config_from_template,
replace_env_vars, replace_env_vars,
) )
from llama_stack.providers.utils.telemetry.tracing import ( from llama_stack.providers.utils.telemetry.tracing import (
end_trace, end_trace,
setup_logger, setup_logger,
@ -186,9 +184,20 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
import nest_asyncio import nest_asyncio
nest_asyncio.apply() nest_asyncio.apply()
self._remove_root_logger_handlers()
return asyncio.run(self.async_client.initialize()) return asyncio.run(self.async_client.initialize())
def _remove_root_logger_handlers(self):
"""
Remove all handlers from the root logger. Needed to avoid polluting the console with logs.
"""
root_logger = logging.getLogger()
for handler in root_logger.handlers[:]:
root_logger.removeHandler(handler)
print(f"Removed handler {handler.__class__.__name__} from root logger")
def _get_path( def _get_path(
self, self,
cast_to: Any, cast_to: Any,