src/second_brain_online/opik_utils.py

#!/usr/bin/env python3
"""
OPIK utilities for Second Brain AI Assistant.
"""

import os
from loguru import logger
import opik
from opik import OpikConfigurator

from second_brain_online.config import settings


def configure() -> None:
    if settings.COMET_API_KEY:
        try:
            client = OpikConfigurator(api_key=settings.COMET_API_KEY)
            default_workspace = client._get_default_workspace()
        except Exception:
            logger.warning(
                "Default workspace not found. Setting workspace to None and enabling interactive mode."
            )
            default_workspace = None

        os.environ["OPIK_PROJECT_NAME"] = settings.COMET_PROJECT or "second_brain_course"

        opik.configure(
            api_key=settings.COMET_API_KEY,
            workspace=default_workspace,
            use_local=False,
            force=True,
        )
        logger.info(
            f"Opik configured successfully using workspace '{default_workspace}'"
        )
    else:
        logger.warning(
            "COMET_API_KEY is not set. Set it to enable prompt monitoring with Opik (powered by Comet ML)."
        )


def get_or_create_dataset(name: str, prompts: list[str]) -> opik.Dataset | None:
    client = opik.Opik()
    try:
        dataset = client.get_dataset(name)
        if dataset is None:
            dataset = client.create_dataset(name, prompts)
        return dataset
    except Exception as e:
        logger.error(f"Error getting or creating dataset: {e}")
        return None