ag2 0.10.2__py3-none-any.whl

autogen/retrieve_utils.py

@@ -0,0 +1,490 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+import glob
+import hashlib
+import logging
+import os
+import re
+from collections.abc import Callable
+from typing import Any
+from urllib.parse import urlparse
+
+import requests
+
+from .import_utils import optional_import_block, require_optional_import
+from .token_count_utils import count_token
+
+with optional_import_block():
+    import chromadb
+    import markdownify
+    from bs4 import BeautifulSoup
+
+    if chromadb.__version__ < "0.4.15":
+        from chromadb.api import API
+    else:
+        from chromadb.api import ClientAPI as API  # noqa: N814
+    import chromadb.utils.embedding_functions as ef
+    import pypdf
+    from chromadb.api.types import QueryResult
+
+
+with optional_import_block() as result:
+    from unstructured.partition.auto import partition
+
+HAS_UNSTRUCTURED = result.is_successful
+
+logger = logging.getLogger(__name__)
+TEXT_FORMATS = [
+    "txt",
+    "json",
+    "csv",
+    "tsv",
+    "md",
+    "html",
+    "htm",
+    "rtf",
+    "rst",
+    "jsonl",
+    "log",
+    "xml",
+    "yaml",
+    "yml",
+    "pdf",
+    "mdx",
+]
+UNSTRUCTURED_FORMATS = [
+    "doc",
+    "docx",
+    "epub",
+    "msg",
+    "odt",
+    "org",
+    "pdf",
+    "ppt",
+    "pptx",
+    "rtf",
+    "rst",
+    "xlsx",
+]  # These formats will be parsed by the 'unstructured' library, if installed.
+if HAS_UNSTRUCTURED:
+    TEXT_FORMATS += UNSTRUCTURED_FORMATS
+    TEXT_FORMATS = list(set(TEXT_FORMATS))
+VALID_CHUNK_MODES = frozenset({"one_line", "multi_lines"})
+RAG_MINIMUM_MESSAGE_LENGTH = int(os.environ.get("RAG_MINIMUM_MESSAGE_LENGTH", 5))
+
+
+def split_text_to_chunks(
+    text: str,
+    max_tokens: int = 4000,
+    chunk_mode: str = "multi_lines",
+    must_break_at_empty_line: bool = True,
+    overlap: int = 0,  # number of overlapping lines
+):
+    """Split a long text into chunks of max_tokens."""
+    if chunk_mode not in VALID_CHUNK_MODES:
+        raise AssertionError
+    if chunk_mode == "one_line":
+        must_break_at_empty_line = False
+        overlap = 0
+    chunks = []
+    lines = text.split("\n")
+    num_lines = len(lines)
+    if num_lines < 3 and must_break_at_empty_line:
+        logger.warning("The input text has less than 3 lines. Set `must_break_at_empty_line` to `False`")
+        must_break_at_empty_line = False
+    lines_tokens = [count_token(line) for line in lines]
+    sum_tokens = sum(lines_tokens)
+    while sum_tokens > max_tokens:
+        estimated_line_cut = 2 if chunk_mode == "one_line" else max(int(max_tokens / sum_tokens * len(lines)), 2)
+        cnt = 0
+        prev = ""
+        for cnt in reversed(range(estimated_line_cut)):
+            if must_break_at_empty_line and lines[cnt].strip() != "":
+                continue
+            if sum(lines_tokens[:cnt]) <= max_tokens:
+                prev = "\n".join(lines[:cnt])
+                break
+        if cnt == 0:
+            logger.warning(
+                f"max_tokens is too small to fit a single line of text. Breaking this line:\n\t{lines[0][:100]} ..."
+            )
+            if not must_break_at_empty_line:
+                split_len = max(
+                    int(max_tokens / (lines_tokens[0] * 0.9 * len(lines[0]) + 0.1)), RAG_MINIMUM_MESSAGE_LENGTH
+                )
+                prev = lines[0][:split_len]
+                lines[0] = lines[0][split_len:]
+                lines_tokens[0] = count_token(lines[0])
+            else:
+                logger.warning("Failed to split docs with must_break_at_empty_line being True, set to False.")
+                must_break_at_empty_line = False
+        (
+            chunks.append(prev) if len(prev) >= RAG_MINIMUM_MESSAGE_LENGTH else None
+        )  # don't add chunks less than RAG_MINIMUM_MESSAGE_LENGTH characters
+        lines = lines[cnt - overlap if cnt > overlap else cnt :]
+        lines_tokens = lines_tokens[cnt - overlap if cnt > overlap else cnt :]
+        sum_tokens = sum(lines_tokens)
+    text_to_chunk = "\n".join(lines).strip()
+    (
+        chunks.append(text_to_chunk) if len(text_to_chunk) >= RAG_MINIMUM_MESSAGE_LENGTH else None
+    )  # don't add chunks less than RAG_MINIMUM_MESSAGE_LENGTH characters
+    return chunks
+
+
+@require_optional_import("pypdf", "retrievechat")
+def extract_text_from_pdf(file: str) -> str:
+    """Extract text from PDF files"""
+    text = ""
+    with open(file, "rb") as f:
+        reader = pypdf.PdfReader(f)
+        if reader.is_encrypted:  # Check if the PDF is encrypted
+            try:
+                reader.decrypt("")
+            except pypdf.errors.FileNotDecryptedError as e:
+                logger.warning(f"Could not decrypt PDF {file}, {e}")
+                return text  # Return empty text if PDF could not be decrypted
+
+        for page_num in range(len(reader.pages)):
+            page = reader.pages[page_num]
+            text += page.extract_text()
+
+    if not text.strip():  # Debugging line to check if text is empty
+        logger.warning(f"Could not decrypt PDF {file}")
+
+    return text
+
+
+def split_files_to_chunks(
+    files: list[tuple[str, str] | str],
+    max_tokens: int = 4000,
+    chunk_mode: str = "multi_lines",
+    must_break_at_empty_line: bool = True,
+    custom_text_split_function: Callable[[str], list[str]] | None = None,
+) -> tuple[list[str], list[dict[str, Any]]]:
+    """Split a list of files into chunks of max_tokens."""
+    chunks = []
+    sources = []
+
+    for file in files:
+        if isinstance(file, tuple):
+            url = file[1]
+            file = file[0]
+        else:
+            url = None
+        _, file_extension = os.path.splitext(file)
+        file_extension = file_extension.lower()
+
+        if HAS_UNSTRUCTURED and file_extension[1:] in UNSTRUCTURED_FORMATS:
+            text = partition(file)
+            text = "\n".join([t.text for t in text]) if len(text) > 0 else ""
+        elif file_extension == ".pdf":
+            text = extract_text_from_pdf(file)
+        else:  # For non-PDF text-based files
+            with open(file, encoding="utf-8", errors="ignore") as f:
+                text = f.read()
+
+        if not text.strip():  # Debugging line to check if text is empty after reading
+            logger.warning(f"No text available in file: {file}")
+            continue  # Skip to the next file if no text is available
+
+        if custom_text_split_function is not None:
+            tmp_chunks = custom_text_split_function(text)
+        else:
+            tmp_chunks = split_text_to_chunks(text, max_tokens, chunk_mode, must_break_at_empty_line)
+        chunks += tmp_chunks
+        sources += [{"source": url if url else file}] * len(tmp_chunks)
+
+    return chunks, sources
+
+
+def get_files_from_dir(dir_path: str | list[str], types: list[str] = TEXT_FORMATS, recursive: bool = True) -> list[Any]:
+    """Return a list of all the files in a given directory, a url, a file path or a list of them."""
+    if len(types) == 0:
+        raise ValueError("types cannot be empty.")
+    types = [t[1:].lower() if t.startswith(".") else t.lower() for t in set(types)]
+    types += [t.upper() for t in types]
+
+    files = []
+    # If the path is a list of files or urls, process and return them
+    if isinstance(dir_path, list):
+        for item in dir_path:
+            if os.path.isfile(item):
+                files.append(item)
+            elif is_url(item):
+                filepath = get_file_from_url(item)
+                if filepath:
+                    files.append(filepath)
+            elif os.path.exists(item):
+                try:
+                    files.extend(get_files_from_dir(item, types, recursive))
+                except ValueError:
+                    logger.warning(f"Directory {item} does not exist. Skipping.")
+            else:
+                logger.warning(f"File {item} does not exist. Skipping.")
+        return files
+
+    # If the path is a file, return it
+    if os.path.isfile(dir_path):
+        return [dir_path]
+
+    # If the path is a url, download it and return the downloaded file
+    if is_url(dir_path):
+        filepath = get_file_from_url(dir_path)
+        if filepath:
+            return [filepath]
+        else:
+            return []
+
+    if os.path.exists(dir_path):
+        for type in types:
+            if recursive:
+                files += glob.glob(os.path.join(dir_path, f"**/*.{type}"), recursive=True)
+            else:
+                files += glob.glob(os.path.join(dir_path, f"*.{type}"), recursive=False)
+    else:
+        logger.error(f"Directory {dir_path} does not exist.")
+        raise ValueError(f"Directory {dir_path} does not exist.")
+    return files
+
+
+@require_optional_import(["markdownify", "bs4"], "retrievechat")
+def parse_html_to_markdown(html: str, url: str = None) -> str:
+    """Parse HTML to markdown."""
+    soup = BeautifulSoup(html, "html.parser")
+    title = soup.title.string
+    # Remove javascript and style blocks
+    for script in soup(["script", "style"]):
+        script.extract()
+
+    # Convert to markdown -- Wikipedia gets special attention to get a clean version of the page
+    if isinstance(url, str) and url.startswith("https://en.wikipedia.org/"):
+        body_elm = soup.find("div", {"id": "mw-content-text"})
+        title_elm = soup.find("span", {"class": "mw-page-title-main"})
+
+        if body_elm:
+            # What's the title
+            main_title = soup.title.string
+            if title_elm and len(title_elm) > 0:
+                main_title = title_elm.string
+            webpage_text = "# " + main_title + "\n\n" + markdownify.MarkdownConverter().convert_soup(body_elm)
+        else:
+            webpage_text = markdownify.MarkdownConverter().convert_soup(soup)
+    else:
+        webpage_text = markdownify.MarkdownConverter().convert_soup(soup)
+
+    # Convert newlines
+    webpage_text = re.sub(r"\r\n", "\n", webpage_text)
+    webpage_text = re.sub(r"\n{2,}", "\n\n", webpage_text).strip()
+    webpage_text = "# " + title + "\n\n" + webpage_text
+    return webpage_text
+
+
+def _generate_file_name_from_url(url: str, max_length=255) -> str:
+    url_bytes = url.encode("utf-8")
+    hash = hashlib.blake2b(url_bytes).hexdigest()
+    parsed_url = urlparse(url)
+    file_name = os.path.basename(url)
+    file_name = (
+        f"{parsed_url.netloc}_{file_name}_{hash[: min(8, max_length - len(parsed_url.netloc) - len(file_name) - 1)]}"
+    )
+    return file_name
+
+
+def get_file_from_url(url: str, save_path: str = None) -> tuple[str, str]:
+    """Download a file from a URL."""
+    if save_path is None:
+        save_path = "tmp/chromadb"
+        os.makedirs(save_path, exist_ok=True)
+    if os.path.isdir(save_path):
+        filename = _generate_file_name_from_url(url)
+        save_path = os.path.join(save_path, filename)
+    else:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+    custom_headers = {
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36",
+    }
+    try:
+        response = requests.get(url, stream=True, headers=custom_headers, timeout=30)
+        response.raise_for_status()
+    except requests.exceptions.RequestException as e:
+        logger.warning(f"Failed to download {url}, {e}")
+        return None
+
+    content_type = response.headers.get("content-type", "")
+    if "text/html" in content_type:
+        # Get the content of the response
+        html = ""
+        for chunk in response.iter_content(chunk_size=8192, decode_unicode=True):
+            html += chunk
+        text = parse_html_to_markdown(html, url)
+        with open(save_path, "w", encoding="utf-8") as f:
+            f.write(text)
+    else:
+        with open(save_path, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+    return save_path, url
+
+
+def is_url(string: str):
+    """Return True if the string is a valid URL."""
+    try:
+        result = urlparse(string)
+        return all([result.scheme, result.netloc])
+    except ValueError:
+        return False
+
+
+@require_optional_import("chromadb", "retrievechat")
+def create_vector_db_from_dir(
+    dir_path: str | list[str],
+    max_tokens: int = 4000,
+    client: "API" = None,
+    db_path: str = "tmp/chromadb.db",
+    collection_name: str = "all-my-documents",
+    get_or_create: bool = False,
+    chunk_mode: str = "multi_lines",
+    must_break_at_empty_line: bool = True,
+    embedding_model: str = "all-MiniLM-L6-v2",
+    embedding_function: Callable = None,
+    custom_text_split_function: Callable = None,
+    custom_text_types: list[str] = TEXT_FORMATS,
+    recursive: bool = True,
+    extra_docs: bool = False,
+) -> "API":
+    """Create a vector db from all the files in a given directory, the directory can also be a single file or a url to
+        a single file. We support chromadb compatible APIs to create the vector db, this function is not required if
+        you prepared your own vector db.
+
+    Args:
+        dir_path (Union[str, List[str]]): the path to the directory, file, url or a list of them.
+        max_tokens (Optional, int): the maximum number of tokens per chunk. Default is 4000.
+        client (Optional, API): the chromadb client. Default is None.
+        db_path (Optional, str): the path to the chromadb. Default is "tmp/chromadb.db". The default was `/tmp/chromadb.db` for version `<=0.2.24`.
+        collection_name (Optional, str): the name of the collection. Default is "all-my-documents".
+        get_or_create (Optional, bool): Whether to get or create the collection. Default is False. If True, the collection
+            will be returned if it already exists. Will raise ValueError if the collection already exists and get_or_create is False.
+        chunk_mode (Optional, str): the chunk mode. Default is "multi_lines".
+        must_break_at_empty_line (Optional, bool): Whether to break at empty line. Default is True.
+        embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if
+            embedding_function is not None.
+        embedding_function (Optional, Callable): the embedding function to use. Default is None, SentenceTransformer with
+            the given `embedding_model` will be used. If you want to use OpenAI, Cohere, HuggingFace or other embedding
+            functions, you can pass it here, follow the examples in `https://docs.trychroma.com/embeddings`.
+        custom_text_split_function (Optional, Callable): a custom function to split a string into a list of strings.
+            Default is None, will use the default function in `autogen.retrieve_utils.split_text_to_chunks`.
+        custom_text_types (Optional, List[str]): a list of file types to be processed. Default is TEXT_FORMATS.
+        recursive (Optional, bool): whether to search documents recursively in the dir_path. Default is True.
+        extra_docs (Optional, bool): whether to add more documents in the collection. Default is False
+
+    Returns:
+    The chromadb client.
+    """
+    if client is None:
+        client = chromadb.PersistentClient(path=db_path)
+    try:
+        embedding_function = (
+            ef.SentenceTransformerEmbeddingFunction(embedding_model)
+            if embedding_function is None
+            else embedding_function
+        )
+        collection = client.create_collection(
+            collection_name,
+            get_or_create=get_or_create,
+            embedding_function=embedding_function,
+            # https://github.com/nmslib/hnswlib#supported-distances
+            # https://github.com/chroma-core/chroma/blob/566bc80f6c8ee29f7d99b6322654f32183c368c4/chromadb/segment/impl/vector/local_hnsw.py#L184
+            # https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md
+            metadata={"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32},  # ip, l2, cosine
+        )
+
+        length = 0
+        if extra_docs:
+            length = len(collection.get()["ids"])
+
+        if custom_text_split_function is not None:
+            chunks, sources = split_files_to_chunks(
+                get_files_from_dir(dir_path, custom_text_types, recursive),
+                custom_text_split_function=custom_text_split_function,
+            )
+        else:
+            chunks, sources = split_files_to_chunks(
+                get_files_from_dir(dir_path, custom_text_types, recursive),
+                max_tokens,
+                chunk_mode,
+                must_break_at_empty_line,
+            )
+        logger.info(f"Found {len(chunks)} chunks.")
+        # Upsert in batch of 40000 or less if the total number of chunks is less than 40000
+        for i in range(0, len(chunks), min(40000, len(chunks))):
+            end_idx = i + min(40000, len(chunks) - i)
+            collection.upsert(
+                documents=chunks[i:end_idx],
+                ids=[f"doc_{j + length}" for j in range(i, end_idx)],  # unique for each doc
+                metadatas=sources[i:end_idx],
+            )
+    except ValueError as e:
+        logger.warning(f"{e}")
+    return client
+
+
+@require_optional_import("chromadb", "retrievechat")
+def query_vector_db(
+    query_texts: list[str],
+    n_results: int = 10,
+    client: "API" = None,
+    db_path: str = "tmp/chromadb.db",
+    collection_name: str = "all-my-documents",
+    search_string: str = "",
+    embedding_model: str = "all-MiniLM-L6-v2",
+    embedding_function: Callable = None,
+) -> "QueryResult":
+    """Query a vector db. We support chromadb compatible APIs, it's not required if you prepared your own vector db
+        and query function.
+
+    Args:
+        query_texts (List[str]): the list of strings which will be used to query the vector db.
+        n_results (Optional, int): the number of results to return. Default is 10.
+        client (Optional, API): the chromadb compatible client. Default is None, a chromadb client will be used.
+        db_path (Optional, str): the path to the vector db. Default is "tmp/chromadb.db". The default was `/tmp/chromadb.db` for version `<=0.2.24`.
+        collection_name (Optional, str): the name of the collection. Default is "all-my-documents".
+        search_string (Optional, str): the search string. Only docs that contain an exact match of this string will be retrieved. Default is "".
+        embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if
+            embedding_function is not None.
+        embedding_function (Optional, Callable): the embedding function to use. Default is None, SentenceTransformer with
+            the given `embedding_model` will be used. If you want to use OpenAI, Cohere, HuggingFace or other embedding
+            functions, you can pass it here, follow the examples in `https://docs.trychroma.com/embeddings`.
+
+    Returns:
+        The query result. The format is:
+
+    ```python
+    class QueryResult(TypedDict):
+        ids: List[IDs]
+        embeddings: Optional[List[List[Embedding]]]
+        documents: Optional[List[List[Document]]]
+        metadatas: Optional[List[List[Metadata]]]
+        distances: Optional[List[List[float]]]
+    ```
+    """
+    if client is None:
+        client = chromadb.PersistentClient(path=db_path)
+    # the collection's embedding function is always the default one, but we want to use the one we used to create the
+    # collection. So we compute the embeddings ourselves and pass it to the query function.
+    collection = client.get_collection(collection_name)
+    embedding_function = (
+        ef.SentenceTransformerEmbeddingFunction(embedding_model) if embedding_function is None else embedding_function
+    )
+    query_embeddings = embedding_function(query_texts)
+    # Query/search n most similar results. You can also .get by id
+    results = collection.query(
+        query_embeddings=query_embeddings,
+        n_results=n_results,
+        where_document={"$contains": search_string} if search_string else None,  # optional filter
+    )
+    return results

autogen/runtime_logging.py

@@ -0,0 +1,161 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+from __future__ import annotations
+
+import logging
+import sqlite3
+import uuid
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
+
+from .logger.base_logger import BaseLogger, LLMConfig
+from .logger.logger_factory import LoggerFactory
+
+if TYPE_CHECKING:
+    from openai import AzureOpenAI, OpenAI
+    from openai.types.chat import ChatCompletion
+
+    from . import Agent, ConversableAgent, OpenAIWrapper
+    from .oai.anthropic import AnthropicClient
+    from .oai.bedrock import BedrockClient
+    from .oai.cerebras import CerebrasClient
+    from .oai.cohere import CohereClient
+    from .oai.gemini import GeminiClient
+    from .oai.groq import GroqClient
+    from .oai.mistral import MistralAIClient
+    from .oai.ollama import OllamaClient
+    from .oai.together import TogetherClient
+
+logger = logging.getLogger(__name__)
+
+autogen_logger = None
+is_logging = False
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def start(
+    logger: BaseLogger | None = None,
+    logger_type: Literal["sqlite", "file"] = "sqlite",
+    config: dict[str, Any] | None = None,
+) -> str:
+    """Start logging for the runtime.
+
+    Args:
+        logger (BaseLogger):    A logger instance
+        logger_type (str):      The type of logger to use (default: sqlite)
+        config (dict):          Configuration for the logger
+    Returns:
+        session_id (str(uuid.uuid4)):       a unique id for the logging session
+    """
+    global autogen_logger
+    global is_logging
+
+    autogen_logger = logger or LoggerFactory.get_logger(logger_type=logger_type, config=config)
+
+    try:
+        session_id = autogen_logger.start()
+        is_logging = True
+    except Exception as e:
+        logger.error(f"[runtime logging] Failed to start logging: {e}")
+    finally:
+        return session_id
+
+
+def log_chat_completion(
+    invocation_id: uuid.UUID,
+    client_id: int,
+    wrapper_id: int,
+    agent: str | Agent,
+    request: dict[str, float | str | list[dict[str, str]]],
+    response: str | ChatCompletion,
+    is_cached: int,
+    cost: float,
+    start_time: str,
+) -> None:
+    if autogen_logger is None:
+        logger.error("[runtime logging] log_chat_completion: autogen logger is None")
+        return
+
+    autogen_logger.log_chat_completion(
+        invocation_id, client_id, wrapper_id, agent, request, response, is_cached, cost, start_time
+    )
+
+
+def log_new_agent(agent: ConversableAgent, init_args: dict[str, Any]) -> None:
+    if autogen_logger is None:
+        logger.error("[runtime logging] log_new_agent: autogen logger is None")
+        return
+
+    autogen_logger.log_new_agent(agent, init_args)
+
+
+def log_event(source: str | Agent, name: str, **kwargs: dict[str, Any]) -> None:
+    if autogen_logger is None:
+        logger.error("[runtime logging] log_event: autogen logger is None")
+        return
+
+    autogen_logger.log_event(source, name, **kwargs)
+
+
+def log_function_use(agent: str | Agent, function: F, args: dict[str, Any], returns: any):
+    if autogen_logger is None:
+        logger.error("[runtime logging] log_function_use: autogen logger is None")
+        return
+
+    autogen_logger.log_function_use(agent, function, args, returns)
+
+
+def log_new_wrapper(wrapper: OpenAIWrapper, init_args: dict[str, LLMConfig | list[LLMConfig]]) -> None:
+    if autogen_logger is None:
+        logger.error("[runtime logging] log_new_wrapper: autogen logger is None")
+        return
+
+    autogen_logger.log_new_wrapper(wrapper, init_args)
+
+
+def log_new_client(
+    client: (
+        AzureOpenAI
+        | OpenAI
+        | CerebrasClient
+        | GeminiClient
+        | AnthropicClient
+        | MistralAIClient
+        | TogetherClient
+        | GroqClient
+        | CohereClient
+        | OllamaClient
+        | BedrockClient
+    ),
+    wrapper: OpenAIWrapper,
+    init_args: dict[str, Any],
+) -> None:
+    if autogen_logger is None:
+        logger.error("[runtime logging] log_new_client: autogen logger is None")
+        return
+
+    autogen_logger.log_new_client(client, wrapper, init_args)
+
+
+def stop() -> None:
+    global is_logging
+    if autogen_logger:
+        autogen_logger.stop()
+    is_logging = False
+
+
+def get_connection() -> None | sqlite3.Connection:
+    if autogen_logger is None:
+        logger.error("[runtime logging] get_connection: autogen logger is None")
+        return None
+
+    return autogen_logger.get_connection()
+
+
+def logging_enabled() -> bool:
+    return is_logging

autogen/testing/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .messages import ToolCall, tools_message
+from .test_agent import TestAgent
+
+__all__ = (
+    "TestAgent",
+    "ToolCall",
+    "tools_message",
+)