API Reference 📚

This page provides the technical reference for the core Auto-Labeler classes.

auto_labeler.core.AutoLabeler

Main entry point for the Auto-Labeler library. Orchestrates discovery and labeling tasks using various strategies.

Source code in src/auto_labeler/core.py

class AutoLabeler:
    """
    Main entry point for the Auto-Labeler library.
    Orchestrates discovery and labeling tasks using various strategies.
    """
    def __init__(
        self, 
        model_name: str = "gemini/gemini-2.5-flash", 
        api_key: Optional[str] = None,
        use_cache: bool = True,
        cache_dir: str = ".auto_labeler_cache",
        log_level: str = "INFO"
    ):
        """
        Initialize the AutoLabeler with model configuration and preferences.
        """
        # Set log level
        setup_logger(level=log_level)

        self.config = AutoLabelerConfig(
            model_name=model_name,
            api_key=api_key,
            use_cache=use_cache,
            cache_dir=cache_dir,
            log_level=log_level
        )

        self.llm = LLMAdapter(
            model_name=self.config.model_name, 
            api_key=self.config.api_key,
            use_cache=self.config.use_cache,
            cache_dir=self.config.cache_dir
        )
        self.prompts_dir = pathlib.Path(__file__).parent / "prompts"

    def get_usage(self) -> dict:
        """
        Returns a summary of token usage and estimated cost for the current session.
        """
        return self.llm.tracker.get_summary()

    def _validate_df(self, df: pd.DataFrame, column: Optional[str] = None):
        """Internal validation for dataframes."""
        if df.empty:
            raise ValueError("Input DataFrame is empty.")
        if column and column not in df.columns:
            raise ValueError(f"Column '{column}' not found in DataFrame.")

    def suggest_labels(
        self, 
        df: pd.DataFrame, 
        context: str, 
        n_labels: int = 5,
        column: Optional[str] = None,
        strategy: Optional[Any] = None
    ) -> List[str]:
        """
        Suggests a list of labels based on the dataset content and context.
        """
        # Validate inputs via Pydantic
        config = DiscoveryConfig(context=context, n_labels=n_labels)
        self._validate_df(df, column)

        logger.info(f"Starting label discovery (asking for {config.n_labels} labels)...")

        # Default to Simple Strategy if none provided
        if not strategy:
            from .strategies import SimpleDiscoveryStrategy
            strategy = SimpleDiscoveryStrategy(self.llm)

        return strategy.suggest_labels(
            df=df,
            context=config.context,
            prompts_dir=self.prompts_dir,
            n_labels=config.n_labels
        )

    def label_dataset(
        self, 
        df: pd.DataFrame, 
        labels: List[str], 
        context: str, 
        target_column: str = "text",
        multi_label: bool = False,
        batch_size: int = 1,
        strategy: Optional[Any] = None,
        examples: Optional[List[dict]] = None
    ) -> pd.DataFrame:
        """
        Labels the dataset using the provided labels and context.
        """
        # Validate inputs via Pydantic
        config = LabelingConfig(
            context=context, 
            labels=labels, 
            target_column=target_column, 
            multi_label=multi_label,
            batch_size=batch_size
        )
        self._validate_df(df, config.target_column)

        logger.info(f"Starting labeling task on {len(df)} records (Batch Size: {config.batch_size})...")

        # Default to Simple Strategy if none provided
        if not strategy:
            from .strategies import SimpleLabelingStrategy
            strategy = SimpleLabelingStrategy(self.llm, batch_size=config.batch_size)
        elif hasattr(strategy, 'batch_size'):
            # Override strategy batch size if provided in call
            strategy.batch_size = config.batch_size

        return strategy.label(
            df=df, 
            labels=config.labels, 
            context=config.context, 
            prompts_dir=self.prompts_dir,
            target_column=config.target_column, 
            multi_label=config.multi_label,
            examples=examples
        )

Functions

__init__(model_name='gemini/gemini-2.5-flash', api_key=None, use_cache=True, cache_dir='.auto_labeler_cache', log_level='INFO')

Initialize the AutoLabeler with model configuration and preferences.

Source code in src/auto_labeler/core.py

def __init__(
    self, 
    model_name: str = "gemini/gemini-2.5-flash", 
    api_key: Optional[str] = None,
    use_cache: bool = True,
    cache_dir: str = ".auto_labeler_cache",
    log_level: str = "INFO"
):
    """
    Initialize the AutoLabeler with model configuration and preferences.
    """
    # Set log level
    setup_logger(level=log_level)

    self.config = AutoLabelerConfig(
        model_name=model_name,
        api_key=api_key,
        use_cache=use_cache,
        cache_dir=cache_dir,
        log_level=log_level
    )

    self.llm = LLMAdapter(
        model_name=self.config.model_name, 
        api_key=self.config.api_key,
        use_cache=self.config.use_cache,
        cache_dir=self.config.cache_dir
    )
    self.prompts_dir = pathlib.Path(__file__).parent / "prompts"

get_usage()

Returns a summary of token usage and estimated cost for the current session.

Source code in src/auto_labeler/core.py

def get_usage(self) -> dict:
    """
    Returns a summary of token usage and estimated cost for the current session.
    """
    return self.llm.tracker.get_summary()

label_dataset(df, labels, context, target_column='text', multi_label=False, batch_size=1, strategy=None, examples=None)

Labels the dataset using the provided labels and context.

Source code in src/auto_labeler/core.py

def label_dataset(
    self, 
    df: pd.DataFrame, 
    labels: List[str], 
    context: str, 
    target_column: str = "text",
    multi_label: bool = False,
    batch_size: int = 1,
    strategy: Optional[Any] = None,
    examples: Optional[List[dict]] = None
) -> pd.DataFrame:
    """
    Labels the dataset using the provided labels and context.
    """
    # Validate inputs via Pydantic
    config = LabelingConfig(
        context=context, 
        labels=labels, 
        target_column=target_column, 
        multi_label=multi_label,
        batch_size=batch_size
    )
    self._validate_df(df, config.target_column)

    logger.info(f"Starting labeling task on {len(df)} records (Batch Size: {config.batch_size})...")

    # Default to Simple Strategy if none provided
    if not strategy:
        from .strategies import SimpleLabelingStrategy
        strategy = SimpleLabelingStrategy(self.llm, batch_size=config.batch_size)
    elif hasattr(strategy, 'batch_size'):
        # Override strategy batch size if provided in call
        strategy.batch_size = config.batch_size

    return strategy.label(
        df=df, 
        labels=config.labels, 
        context=config.context, 
        prompts_dir=self.prompts_dir,
        target_column=config.target_column, 
        multi_label=config.multi_label,
        examples=examples
    )

suggest_labels(df, context, n_labels=5, column=None, strategy=None)

Suggests a list of labels based on the dataset content and context.

Source code in src/auto_labeler/core.py

def suggest_labels(
    self, 
    df: pd.DataFrame, 
    context: str, 
    n_labels: int = 5,
    column: Optional[str] = None,
    strategy: Optional[Any] = None
) -> List[str]:
    """
    Suggests a list of labels based on the dataset content and context.
    """
    # Validate inputs via Pydantic
    config = DiscoveryConfig(context=context, n_labels=n_labels)
    self._validate_df(df, column)

    logger.info(f"Starting label discovery (asking for {config.n_labels} labels)...")

    # Default to Simple Strategy if none provided
    if not strategy:
        from .strategies import SimpleDiscoveryStrategy
        strategy = SimpleDiscoveryStrategy(self.llm)

    return strategy.suggest_labels(
        df=df,
        context=config.context,
        prompts_dir=self.prompts_dir,
        n_labels=config.n_labels
    )

auto_labeler.llm.LLMAdapter

A unified wrapper around LiteLLM to handle different providers (OpenAI, Gemini, etc.) and standard interactions like text generation and structured outputs. Supports persistent disk caching of responses.

Source code in src/auto_labeler/llm.py

class LLMAdapter:
    """
    A unified wrapper around LiteLLM to handle different providers (OpenAI, Gemini, etc.)
    and standard interactions like text generation and structured outputs.
    Supports persistent disk caching of responses.
    """
    def __init__(
        self, 
        model_name: str, 
        api_key: Optional[str] = None, 
        tracker: Optional[UsageTracker] = None,
        use_cache: bool = True,
        cache_dir: str = ".auto_labeler_cache"
    ):
        """
        Args:
            model_name: The LiteLLM model identifier (e.g. 'gpt-3.5-turbo').
            api_key: Optional API key override.
            tracker: Optional UsageTracker to aggregate token usage.
            use_cache: Whether to use disk caching for responses.
            cache_dir: Directory to store the cache.
        """
        self.model_name = model_name
        self.api_key = api_key
        self.tracker = tracker if tracker else UsageTracker()
        self.use_cache = use_cache
        self.cache = None

        if use_cache and Cache:
            self.cache = Cache(cache_dir)
            logger.debug(f"Caching enabled. Storage: {cache_dir}")
        elif use_cache and not Cache:
            logger.warning("diskcache not installed. Caching will be disabled.")

    def _get_cache_key(self, prompt: str, system_prompt: Optional[str] = None, **kwargs) -> str:
        """Generates a stable cache key for a request."""
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "system_prompt": system_prompt,
            **kwargs
        }
        encoded = json.dumps(payload, sort_keys=True).encode()
        key = hashlib.md5(encoded).hexdigest()
        logger.debug(f"Generated cache key for {self.model_name}: {key}")
        return key

    def generate(self, prompt: str, system_prompt: Optional[str] = None) -> str:
        """
        Generates a simple text response from the LLM.
        """
        cache_key = self._get_cache_key(prompt, system_prompt)
        if self.cache is not None:
            cached_val = self.cache.get(cache_key)
            if cached_val is not None:
                logger.debug("Cache hit for simple generation.")
                return cached_val

        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        response = litellm.completion(
            model=self.model_name,
            api_key=self.api_key,
            messages=messages
        )
        self.tracker.update(response)
        result = response.choices[0].message.content

        if self.cache is not None:
            self.cache[cache_key] = result
        return result

    async def agenerate(self, prompt: str, system_prompt: Optional[str] = None) -> str:
        """
        Asynchronous version of generate.
        """
        cache_key = self._get_cache_key(prompt, system_prompt, async_call=True)
        if self.cache is not None:
            cached_val = self.cache.get(cache_key)
            if cached_val is not None:
                return cached_val

        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        response = await litellm.acompletion(
            model=self.model_name,
            api_key=self.api_key,
            messages=messages
        )
        self.tracker.update(response)
        result = response.choices[0].message.content

        if self.cache is not None:
            self.cache[cache_key] = result
        return result

    def generate_structured(
        self, 
        prompt: str, 
        response_schema: Dict[str, Any], 
        system_prompt: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Generates a structured JSON response based on a schema.
        Uses JSON mode where available or prompts the model to output JSON.
        """
        cache_key = self._get_cache_key(prompt, system_prompt, schema=response_schema)
        if self.cache is not None:
            cached_val = self.cache.get(cache_key)
            if cached_val is not None:
                logger.debug("Cache hit for structured generation.")
                return cached_val

        # Append instruction to output JSON if not present
        if "json" not in prompt.lower():
            prompt += "\n\nPlease output the result as a raw JSON object."

        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        # Litellm supports response_format with schema for many providers
        response = litellm.completion(
            model=self.model_name,
            api_key=self.api_key,
            messages=messages,
            response_format=response_schema if response_schema else {"type": "json_object"},
            num_retries=3
        )
        self.tracker.update(response)

        content = response.choices[0].message.content
        result = self._parse_json_content(content)

        if self.cache is not None:
            self.cache[cache_key] = result
        return result

    async def agenerate_structured(
        self, 
        prompt: str, 
        response_schema: Dict[str, Any], 
        system_prompt: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Asynchronous version of generate_structured.
        """
        cache_key = self._get_cache_key(prompt, system_prompt, schema=response_schema, async_call=True)
        if self.cache is not None:
            cached_val = self.cache.get(cache_key)
            if cached_val is not None:
                return cached_val

        if "json" not in prompt.lower():
            prompt += "\n\nPlease output the result as a raw JSON object."

        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        response = await litellm.acompletion(
            model=self.model_name,
            api_key=self.api_key,
            messages=messages,
            response_format=response_schema if response_schema else {"type": "json_object"},
            num_retries=3
        )
        self.tracker.update(response)

        content = response.choices[0].message.content
        result = self._parse_json_content(content)

        if self.cache is not None:
            self.cache[cache_key] = result
        return result

    def _parse_json_content(self, content: str) -> Dict[str, Any]:
        # Clean markdown code blocks if present
        clean_content = content.strip()
        if clean_content.startswith("```"):
            # Remove ```json or just ```
            lines = clean_content.splitlines()
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].startswith("```"):
                lines = lines[:-1]
            clean_content = "\n".join(lines).strip()

        try:
            return json.loads(clean_content)
        except json.JSONDecodeError:
            raise ValueError(f"Failed to parse LLM response as JSON: {clean_content}")

    def _default_embedding_model(self) -> str:
        """Derives a sensible embedding model from the adapter's LLM model name."""
        name = self.model_name.lower()
        if name.startswith("gemini/") or name.startswith("google/"):
            return "gemini/text-embedding-004"
        if name.startswith("anthropic/") or name.startswith("claude"):
            # Anthropic has no native embedding API; fall back to OpenAI default
            return "text-embedding-3-small"
        # OpenAI and everything else
        return "text-embedding-3-small"

    def get_embedding(self, text: Union[str, List[str]], model: Optional[str] = None) -> Union[List[float], List[List[float]]]:
        """
        Generates embeddings for a string or list of strings.

        Args:
            text: String or list of strings to embed.
            model: Embedding model to use. Defaults to a provider-appropriate model
                   derived from the adapter's LLM (e.g. ``gemini/text-embedding-004``
                   for Gemini users, ``text-embedding-3-small`` for OpenAI users).
        """
        if model is None:
            model = self._default_embedding_model()
        cache_key = self._get_cache_key(str(text), model=model, type="split_embedding")
        if self.cache and cache_key in self.cache:
            return self.cache[cache_key]

        response = litellm.embedding(
            model=model,
            input=text,
            api_key=self.api_key
        )
        self.tracker.update(response)

        # Handle both single string and list inputs
        data = response.data
        if isinstance(text, str):
            result = data[0]["embedding"]
        else:
            result = [item["embedding"] for item in data]

        if self.cache is not None:
            self.cache[cache_key] = result
        return result

Functions

__init__(model_name, api_key=None, tracker=None, use_cache=True, cache_dir='.auto_labeler_cache')

Parameters:

Name	Type	Description	Default
model_name	str	The LiteLLM model identifier (e.g. 'gpt-3.5-turbo').	required
api_key	Optional[str]	Optional API key override.	None
tracker	Optional[UsageTracker]	Optional UsageTracker to aggregate token usage.	None
use_cache	bool	Whether to use disk caching for responses.	True
cache_dir	str	Directory to store the cache.	'.auto_labeler_cache'

Source code in src/auto_labeler/llm.py

def __init__(
    self, 
    model_name: str, 
    api_key: Optional[str] = None, 
    tracker: Optional[UsageTracker] = None,
    use_cache: bool = True,
    cache_dir: str = ".auto_labeler_cache"
):
    """
    Args:
        model_name: The LiteLLM model identifier (e.g. 'gpt-3.5-turbo').
        api_key: Optional API key override.
        tracker: Optional UsageTracker to aggregate token usage.
        use_cache: Whether to use disk caching for responses.
        cache_dir: Directory to store the cache.
    """
    self.model_name = model_name
    self.api_key = api_key
    self.tracker = tracker if tracker else UsageTracker()
    self.use_cache = use_cache
    self.cache = None

    if use_cache and Cache:
        self.cache = Cache(cache_dir)
        logger.debug(f"Caching enabled. Storage: {cache_dir}")
    elif use_cache and not Cache:
        logger.warning("diskcache not installed. Caching will be disabled.")

agenerate(prompt, system_prompt=None) async

Asynchronous version of generate.

Source code in src/auto_labeler/llm.py

async def agenerate(self, prompt: str, system_prompt: Optional[str] = None) -> str:
    """
    Asynchronous version of generate.
    """
    cache_key = self._get_cache_key(prompt, system_prompt, async_call=True)
    if self.cache is not None:
        cached_val = self.cache.get(cache_key)
        if cached_val is not None:
            return cached_val

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    response = await litellm.acompletion(
        model=self.model_name,
        api_key=self.api_key,
        messages=messages
    )
    self.tracker.update(response)
    result = response.choices[0].message.content

    if self.cache is not None:
        self.cache[cache_key] = result
    return result

agenerate_structured(prompt, response_schema, system_prompt=None) async

Asynchronous version of generate_structured.

Source code in src/auto_labeler/llm.py

async def agenerate_structured(
    self, 
    prompt: str, 
    response_schema: Dict[str, Any], 
    system_prompt: Optional[str] = None
) -> Dict[str, Any]:
    """
    Asynchronous version of generate_structured.
    """
    cache_key = self._get_cache_key(prompt, system_prompt, schema=response_schema, async_call=True)
    if self.cache is not None:
        cached_val = self.cache.get(cache_key)
        if cached_val is not None:
            return cached_val

    if "json" not in prompt.lower():
        prompt += "\n\nPlease output the result as a raw JSON object."

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    response = await litellm.acompletion(
        model=self.model_name,
        api_key=self.api_key,
        messages=messages,
        response_format=response_schema if response_schema else {"type": "json_object"},
        num_retries=3
    )
    self.tracker.update(response)

    content = response.choices[0].message.content
    result = self._parse_json_content(content)

    if self.cache is not None:
        self.cache[cache_key] = result
    return result

generate(prompt, system_prompt=None)

Generates a simple text response from the LLM.

Source code in src/auto_labeler/llm.py

def generate(self, prompt: str, system_prompt: Optional[str] = None) -> str:
    """
    Generates a simple text response from the LLM.
    """
    cache_key = self._get_cache_key(prompt, system_prompt)
    if self.cache is not None:
        cached_val = self.cache.get(cache_key)
        if cached_val is not None:
            logger.debug("Cache hit for simple generation.")
            return cached_val

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    response = litellm.completion(
        model=self.model_name,
        api_key=self.api_key,
        messages=messages
    )
    self.tracker.update(response)
    result = response.choices[0].message.content

    if self.cache is not None:
        self.cache[cache_key] = result
    return result

generate_structured(prompt, response_schema, system_prompt=None)

Generates a structured JSON response based on a schema. Uses JSON mode where available or prompts the model to output JSON.

Source code in src/auto_labeler/llm.py

def generate_structured(
    self, 
    prompt: str, 
    response_schema: Dict[str, Any], 
    system_prompt: Optional[str] = None
) -> Dict[str, Any]:
    """
    Generates a structured JSON response based on a schema.
    Uses JSON mode where available or prompts the model to output JSON.
    """
    cache_key = self._get_cache_key(prompt, system_prompt, schema=response_schema)
    if self.cache is not None:
        cached_val = self.cache.get(cache_key)
        if cached_val is not None:
            logger.debug("Cache hit for structured generation.")
            return cached_val

    # Append instruction to output JSON if not present
    if "json" not in prompt.lower():
        prompt += "\n\nPlease output the result as a raw JSON object."

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    # Litellm supports response_format with schema for many providers
    response = litellm.completion(
        model=self.model_name,
        api_key=self.api_key,
        messages=messages,
        response_format=response_schema if response_schema else {"type": "json_object"},
        num_retries=3
    )
    self.tracker.update(response)

    content = response.choices[0].message.content
    result = self._parse_json_content(content)

    if self.cache is not None:
        self.cache[cache_key] = result
    return result

get_embedding(text, model=None)

Generates embeddings for a string or list of strings.

Parameters:

Name	Type	Description	Default
text	Union[str, List[str]]	String or list of strings to embed.	required
model	Optional[str]	Embedding model to use. Defaults to a provider-appropriate model derived from the adapter's LLM (e.g. gemini/text-embedding-004 for Gemini users, text-embedding-3-small for OpenAI users).	None

Source code in src/auto_labeler/llm.py

def get_embedding(self, text: Union[str, List[str]], model: Optional[str] = None) -> Union[List[float], List[List[float]]]:
    """
    Generates embeddings for a string or list of strings.

    Args:
        text: String or list of strings to embed.
        model: Embedding model to use. Defaults to a provider-appropriate model
               derived from the adapter's LLM (e.g. ``gemini/text-embedding-004``
               for Gemini users, ``text-embedding-3-small`` for OpenAI users).
    """
    if model is None:
        model = self._default_embedding_model()
    cache_key = self._get_cache_key(str(text), model=model, type="split_embedding")
    if self.cache and cache_key in self.cache:
        return self.cache[cache_key]

    response = litellm.embedding(
        model=model,
        input=text,
        api_key=self.api_key
    )
    self.tracker.update(response)

    # Handle both single string and list inputs
    data = response.data
    if isinstance(text, str):
        result = data[0]["embedding"]
    else:
        result = [item["embedding"] for item in data]

    if self.cache is not None:
        self.cache[cache_key] = result
    return result

auto_labeler.llm.UsageTracker

Stays updated with the cumulative token usage and estimated cost for a session.

Source code in src/auto_labeler/llm.py

class UsageTracker:
    """
    Stays updated with the cumulative token usage and estimated cost for a session.
    """
    def __init__(self):
        self.prompt_tokens = 0
        self.completion_tokens = 0
        self.total_tokens = 0
        self.total_cost_usd = 0.0

    def update(self, response: Any):
        """
        Updates usage from a LiteLLM response object.
        """
        usage = getattr(response, "usage", None)
        if usage:
            self.prompt_tokens += getattr(usage, 'prompt_tokens', 0)
            self.completion_tokens += getattr(usage, 'completion_tokens', 0)
            self.total_tokens += getattr(usage, 'total_tokens', 0)

        # Calculate cost using litellm
        try:
            cost = litellm.completion_cost(completion_response=response)
            self.total_cost_usd += float(cost) if cost else 0.0
        except Exception:
            # Fallback if cost calculation fails
            pass

    def get_summary(self) -> Dict[str, Any]:
        return {
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_tokens": self.total_tokens,
            "total_cost_usd": round(self.total_cost_usd, 6)
        }

Functions

update(response)

Updates usage from a LiteLLM response object.

Source code in src/auto_labeler/llm.py

def update(self, response: Any):
    """
    Updates usage from a LiteLLM response object.
    """
    usage = getattr(response, "usage", None)
    if usage:
        self.prompt_tokens += getattr(usage, 'prompt_tokens', 0)
        self.completion_tokens += getattr(usage, 'completion_tokens', 0)
        self.total_tokens += getattr(usage, 'total_tokens', 0)

    # Calculate cost using litellm
    try:
        cost = litellm.completion_cost(completion_response=response)
        self.total_cost_usd += float(cost) if cost else 0.0
    except Exception:
        # Fallback if cost calculation fails
        pass