dataflow.registry.importers.litellm_importer

LiteLLM model importer for the Haive registry system.

This module provides functionality to import LLM and embedding models from LiteLLM’s model data repository into the Haive registry system. It fetches model information including pricing, context windows, and capabilities, then registers them in the registry database.

The importer handles various model providers including OpenAI, Anthropic, Google, Mistral, and others. It extracts model metadata, capabilities, and pricing information, ensuring comprehensive model registration.

Key functions: - Fetching model data from LiteLLM’s GitHub repository - Creating provider types and model entries in the registry - Extracting model capabilities and specifications - Tracking import operations and results

Typical usage example:

from haive.dataflow.registry.importers.litellm_importer import import_litellm_models

# Import all models from LiteLLM session_id, import_count = import_litellm_models() print(f”Imported {import_count} models in session {session_id}”)

# Import models from a specific provider session_id, import_count = import_litellm_models(provider_filter=”openai”) print(f”Imported {import_count} OpenAI models”)

Functions

add_import_log(entity_name, entity_type, status, message)	Add an import log entry to the audit.import_logs table.
get_or_create_provider(provider_name, provider_type)	Get or create a provider and return its data.
get_or_create_provider_type(type_name, display_name)	Get or create a provider type in the registry.
import_embedding_models()	Import embedding models.
import_from_env()	Extract embedding models from environment variables.
import_llm_models()	Import LLM models from LiteLLM.
main()	Main function to run the import.

Module Contents

dataflow.registry.importers.litellm_importer.add_import_log(entity_name, entity_type, status, message)

Add an import log entry to the audit.import_logs table.

Parameters:

• entity_name (str)

• entity_type (str)

• status (str)

• message (str)

Return type:

None

dataflow.registry.importers.litellm_importer.get_or_create_provider(provider_name, provider_type)

Get or create a provider and return its data.

Parameters:

• provider_name (str)

• provider_type (str)

Return type:

dict[str, Any] | None

dataflow.registry.importers.litellm_importer.get_or_create_provider_type(type_name, display_name)

Get or create a provider type in the registry.

This function checks if a provider type exists in the registry database, and creates it if it doesn’t exist. It’s used to ensure that all required provider types are available before importing models.

Parameters:

• type_name (str) – The internal name/ID of the provider type

• display_name (str) – The human-readable name of the provider

Returns:

The ID of the provider type, or None if creation failed

Return type:

Optional[str]

Examples

>>> provider_id = get_or_create_provider_type("openai", "OpenAI")
>>> if provider_id:
...     print(f"Provider type ID: {provider_id}")

dataflow.registry.importers.litellm_importer.import_embedding_models()

Import embedding models.

Returns the number of models imported.

Return type:

int

dataflow.registry.importers.litellm_importer.import_from_env()

Extract embedding models from environment variables.

Look for vars like OPENAI_EMBEDDING_MODEL, AZURE_EMBEDDING_MODEL, etc.

Return type:

list[dict[str, Any]]

dataflow.registry.importers.litellm_importer.import_llm_models()

Import LLM models from LiteLLM.

Returns the number of models imported.

Return type:

int

dataflow.registry.importers.litellm_importer.main()

Main function to run the import.