SqliteLogger

def start(self) -> str:
    try:
        query = """
            CREATE TABLE IF NOT EXISTS chat_completions(
                id INTEGER PRIMARY KEY,
                invocation_id TEXT,
                client_id INTEGER,
                wrapper_id INTEGER,
                session_id TEXT,
                source_name TEXT,
                request TEXT,
                response TEXT,
                is_cached INEGER,
                cost REAL,
                start_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                end_time DATETIME DEFAULT CURRENT_TIMESTAMP)
        """
        self._run_query(query=query)

        query = """
            CREATE TABLE IF NOT EXISTS agents (
                id INTEGER PRIMARY KEY,                             -- Key assigned by the database
                agent_id INTEGER,                                   -- result of python id(agent)
                wrapper_id INTEGER,                                 -- result of python id(agent.client)
                session_id TEXT,
                name TEXT,                                          -- agent.name
                class TEXT,                                         -- type or class name of agent
                init_args TEXT,                                     -- JSON serialization of constructor
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(agent_id, session_id))
        """
        self._run_query(query=query)

        query = """
            CREATE TABLE IF NOT EXISTS oai_wrappers (
                id INTEGER PRIMARY KEY,                             -- Key assigned by the database
                wrapper_id INTEGER,                                 -- result of python id(wrapper)
                session_id TEXT,
                init_args TEXT,                                     -- JSON serialization of constructor
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(wrapper_id, session_id))
        """
        self._run_query(query=query)

        query = """
            CREATE TABLE IF NOT EXISTS oai_clients (
                id INTEGER PRIMARY KEY,                             -- Key assigned by the database
                client_id INTEGER,                                  -- result of python id(client)
                wrapper_id INTEGER,                                 -- result of python id(wrapper)
                session_id TEXT,
                class TEXT,                                         -- type or class name of client
                init_args TEXT,                                     -- JSON serialization of constructor
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(client_id, session_id))
        """
        self._run_query(query=query)

        query = """
        CREATE TABLE IF NOT EXISTS version (
            id INTEGER PRIMARY KEY CHECK (id = 1),                  -- id of the logging database
            version_number INTEGER NOT NULL                         -- version of the logging database
        );
        """
        self._run_query(query=query)

        query = """
        CREATE TABLE IF NOT EXISTS events (
            event_name TEXT,
            source_id INTEGER,
            source_name TEXT,
            agent_module TEXT DEFAULT NULL,
            agent_class_name TEXT DEFAULT NULL,
            id INTEGER PRIMARY KEY,
            json_state TEXT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
        );
        """
        self._run_query(query=query)

        query = """
                    CREATE TABLE IF NOT EXISTS function_calls (
                        source_id INTEGER,
                        source_name TEXT,
                        function_name TEXT,
                        args TEXT DEFAULT NULL,
                        returns TEXT DEFAULT NULL,
                        timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
                    );
                    """
        self._run_query(query=query)

        current_version = self._get_current_db_version()
        if current_version is None:
            self._run_query(
                query="INSERT INTO version (id, version_number) VALUES (1, ?);", args=(SqliteLogger.schema_version,)
            )
        self._apply_migration()

    except sqlite3.Error as e:
        logger.error(f"[SqliteLogger] start logging error: {e}")
    finally:
        return self.session_id

def log_chat_completion(
    self,
    invocation_id: uuid.UUID,
    client_id: int,
    wrapper_id: int,
    source: str | Agent,
    request: dict[str, float | str | list[dict[str, str]]],
    response: str | ChatCompletion,
    is_cached: int,
    cost: float,
    start_time: str,
) -> None:
    """Log chat completion.

    Args:
        invocation_id (uuid.UUID): Invocation ID.
        client_id (int): Client ID.
        wrapper_id (int): Wrapper ID.
        source (str | Agent): Source of the chat completion.
        request (dict[str, float | str | list[dict[str, str]]]): Request for the chat completion.
        response (str | ChatCompletion): Response for the chat completion.
        is_cached (int): Whether the response is cached.
        cost (float): Cost of the chat completion.
        start_time (str): Start time of the chat completion.
    """
    if self.con is None:
        return

    end_time = get_current_ts()

    if response is None or isinstance(response, str):
        response_messages = json.dumps({"response": response})
    else:
        response_messages = json.dumps(to_dict(response), indent=4)

    source_name = (
        source
        if isinstance(source, str)
        else source.name
        if hasattr(source, "name") and source.name is not None
        else ""
    )

    query = """
        INSERT INTO chat_completions (
            invocation_id, client_id, wrapper_id, session_id, request, response, is_cached, cost, start_time, end_time, source_name
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    """
    args = (
        invocation_id,
        client_id,
        wrapper_id,
        self.session_id,
        json.dumps(request),
        response_messages,
        is_cached,
        cost,
        start_time,
        end_time,
        source_name,
    )

    self._run_query(query=query, args=args)

def log_new_agent(self, agent: ConversableAgent, init_args: dict[str, Any]) -> None:
    """Log new agent.

    Args:
        agent (ConversableAgent): Agent to log.
        init_args (dict[str, Any]): Initialization arguments of the agent
    """
    from .. import Agent

    if self.con is None:
        return

    args = to_dict(
        init_args,
        exclude=(
            "self",
            "__class__",
            "api_key",
            "organization",
            "base_url",
            "azure_endpoint",
            "azure_ad_token",
            "azure_ad_token_provider",
        ),
        no_recursive=(Agent,),
    )

    # We do an upsert since both the superclass and subclass may call this method (in that order)
    query = """
    INSERT INTO agents (agent_id, wrapper_id, session_id, name, class, init_args, timestamp) VALUES (?, ?, ?, ?, ?, ?, ?)
    ON CONFLICT (agent_id, session_id) DO UPDATE SET
        wrapper_id = excluded.wrapper_id,
        name = excluded.name,
        class = excluded.class,
        init_args = excluded.init_args,
        timestamp = excluded.timestamp
    """
    args = (
        id(agent),
        agent.client.wrapper_id if hasattr(agent, "client") and agent.client is not None else "",
        self.session_id,
        agent.name if hasattr(agent, "name") and agent.name is not None else "",
        type(agent).__name__,
        json.dumps(args),
        get_current_ts(),
    )
    self._run_query(query=query, args=args)

def log_event(self, source: str | Agent, name: str, **kwargs: dict[str, Any]) -> None:
    """Log event.

    Args:
        source (str | Agent): Source of the event.
        name (str): Name of the event.
        **kwargs (dict[str, Any]): Additional arguments for the event.
    """
    from autogen import Agent

    if self.con is None:
        return

    json_args = json.dumps(kwargs, default=lambda o: f"<<non-serializable: {type(o).__qualname__}>>")

    if isinstance(source, Agent):
        query = """
        INSERT INTO events (source_id, source_name, event_name, agent_module, agent_class_name, json_state, timestamp) VALUES (?, ?, ?, ?, ?, ?, ?)
        """
        args = (
            id(source),
            source.name if hasattr(source, "name") else source,
            name,
            source.__module__,
            source.__class__.__name__,
            json_args,
            get_current_ts(),
        )
        self._run_query(query=query, args=args)
    else:
        query = """
        INSERT INTO events (source_id, source_name, event_name, json_state, timestamp) VALUES (?, ?, ?, ?, ?)
        """
        args_str_based = (
            id(source),
            source.name if hasattr(source, "name") else source,
            name,
            json_args,
            get_current_ts(),
        )
        self._run_query(query=query, args=args_str_based)

def log_new_wrapper(self, wrapper: OpenAIWrapper, init_args: dict[str, LLMConfig | list[LLMConfig]]) -> None:
    """Log new wrapper.

    Args:
        wrapper (OpenAIWrapper): Wrapper to log.
        init_args (dict[str, LLMConfig | list[LLMConfig]]): Initialization arguments of the wrapper
    """
    if self.con is None:
        return

    args = to_dict(
        init_args,
        exclude=(
            "self",
            "__class__",
            "api_key",
            "organization",
            "base_url",
            "azure_endpoint",
            "azure_ad_token",
            "azure_ad_token_provider",
        ),
    )

    query = """
    INSERT INTO oai_wrappers (wrapper_id, session_id, init_args, timestamp) VALUES (?, ?, ?, ?)
    ON CONFLICT (wrapper_id, session_id) DO NOTHING;
    """
    args = (
        id(wrapper),
        self.session_id,
        json.dumps(args),
        get_current_ts(),
    )
    self._run_query(query=query, args=args)

def log_function_use(self, source: str | Agent, function: F, args: dict[str, Any], returns: Any) -> None:
    """Log function use.

    Args:
        source (str | Agent): Source of the function use.
        function (F): Function to log.
        args (dict[str, Any]): Arguments of the function.
        returns (Any): Returns of the function.
    """
    if self.con is None:
        return

    query = """
    INSERT INTO function_calls (source_id, source_name, function_name, args, returns, timestamp) VALUES (?, ?, ?, ?, ?, ?)
    """
    query_args: tuple[Any, ...] = (
        id(source),
        source.name if hasattr(source, "name") else source,
        function.__name__,
        safe_serialize(args),
        safe_serialize(returns),
        get_current_ts(),
    )
    self._run_query(query=query, args=query_args)

def log_new_client(
    self,
    client: (
        AzureOpenAI
        | OpenAI
        | CerebrasClient
        | GeminiClient
        | AnthropicClient
        | MistralAIClient
        | TogetherClient
        | GroqClient
        | CohereClient
        | OllamaClient
        | BedrockClient
    ),
    wrapper: OpenAIWrapper,
    init_args: dict[str, Any],
) -> None:
    """Log new client.

    Args:
        client (AzureOpenAI | OpenAI | CerebrasClient | GeminiClient | AnthropicClient | MistralAIClient | TogetherClient | GroqClient | CohereClient | OllamaClient | BedrockClient): Client to log.
        wrapper (OpenAIWrapper): Wrapper of the client.
        init_args (dict[str, Any]): Initialization arguments of the client.
    """
    if self.con is None:
        return

    args = to_dict(
        init_args,
        exclude=(
            "self",
            "__class__",
            "api_key",
            "organization",
            "base_url",
            "azure_endpoint",
            "azure_ad_token",
            "azure_ad_token_provider",
        ),
    )

    query = """
    INSERT INTO oai_clients (client_id, wrapper_id, session_id, class, init_args, timestamp) VALUES (?, ?, ?, ?, ?, ?)
    ON CONFLICT (client_id, session_id) DO NOTHING;
    """
    args = (
        id(client),
        id(wrapper),
        self.session_id,
        type(client).__name__,
        json.dumps(args),
        get_current_ts(),
    )
    self._run_query(query=query, args=args)

def stop(self) -> None:
    """Stop the logger"""
    if self.con:
        self.con.close()

def get_connection(self) -> None | sqlite3.Connection:
    """Get connection."""
    if self.con:
        return self.con
    return None