storage

Class info

Classes

Name	Children	Inherits
BaseStorageProviderConfig llmling_agent.models.storage	SQLStorageConfig TextLogConfig FileStorageConfig MemoryStorageConfig Mem0Config	BaseModel
FileStorageConfig llmling_agent.models.storage File storage configuration.		BaseStorageProviderConfig
Mem0Config llmling_agent.models.storage Configuration for mem0 storage.		BaseStorageProviderConfig
MemoryStorageConfig llmling_agent.models.storage In-memory storage configuration for testing.		BaseStorageProviderConfig
SQLStorageConfig llmling_agent.models.storage SQL database storage configuration.		BaseStorageProviderConfig
StorageManager llmling_agent.storage Manages multiple storage providers.		TaskManagerMixin
TaskManagerMixin llmling_agent.utils.tasks Mixin for managing async tasks.	MessageEmitter Agent StorageManager StorageProvider PromptManager	object
TextLogConfig llmling_agent.models.storage Text log configuration.		BaseStorageProviderConfig

🛈 DocStrings

Storage manager for handling multiple providers.

StorageManager

Bases: TaskManagerMixin

Manages multiple storage providers.

Handles: - Provider initialization and cleanup - Message distribution to providers - History loading from capable providers - Global logging filters

Source code in src/llmling_agent/storage.py

class StorageManager(TaskManagerMixin):
    """Manages multiple storage providers.

    Handles:
    - Provider initialization and cleanup
    - Message distribution to providers
    - History loading from capable providers
    - Global logging filters
    """

    def __init__(self, config: StorageConfig):
        """Initialize storage manager.

        Args:
            config: Storage configuration including providers and filters
        """
        self.config = config
        self.providers = [
            self._create_provider(cfg) for cfg in self.config.effective_providers
        ]

    def cleanup(self):
        """Clean up all providers."""
        for provider in self.providers:
            try:
                provider.cleanup()
            except Exception:
                logger.exception("Error cleaning up provider: %r", provider)
        self.providers.clear()

    def _create_provider(self, config: BaseStorageProviderConfig) -> StorageProvider:
        """Create provider instance from configuration."""
        # Extract common settings from BaseStorageProviderConfig
        match self.config.filter_mode:
            case "and" if self.config.agents and config.agents:
                logged_agents = self.config.agents & config.agents
            case "and":
                logged_agents = self.config.agents or config.agents or set()
            case "override":
                logged_agents = (
                    config.agents
                    if config.agents is not None
                    else self.config.agents or set()
                )

        provider_config = config.model_copy(
            update={
                "log_messages": config.log_messages and self.config.log_messages,
                "log_conversations": config.log_conversations
                and self.config.log_conversations,
                "log_tool_calls": config.log_tool_calls and self.config.log_tool_calls,
                "log_commands": config.log_commands and self.config.log_commands,
                "log_context": config.log_context and self.config.log_context,
                "agents": logged_agents,
            }
        )

        match provider_config:
            case SQLStorageConfig():
                from sqlmodel import create_engine

                from llmling_agent_storage.sql_provider import SQLModelProvider

                engine = create_engine(
                    provider_config.url, pool_size=provider_config.pool_size
                )
                return SQLModelProvider(provider_config, engine)
            case FileStorageConfig():
                from llmling_agent_storage.file_provider import FileProvider

                return FileProvider(provider_config)
            case TextLogConfig():
                from llmling_agent_storage.text_log_provider import TextLogProvider

                return TextLogProvider(provider_config)

            case Mem0Config():
                from llmling_agent_storage.mem0 import Mem0StorageProvider

                return Mem0StorageProvider(provider_config)

            case MemoryStorageConfig():
                from llmling_agent_storage.memory_provider import MemoryStorageProvider

                return MemoryStorageProvider(provider_config)
            case _:
                msg = f"Unknown provider type: {provider_config}"
                raise ValueError(msg)

    def get_history_provider(self, preferred: str | None = None) -> StorageProvider:
        """Get provider for loading history.

        Args:
            preferred: Optional preferred provider name

        Returns:
            First capable provider based on priority:
            1. Preferred provider if specified and capable
            2. Default provider if specified and capable
            3. First capable provider
            4. Raises error if no capable provider found
        """

        # Function to find capable provider by name
        def find_provider(name: str) -> StorageProvider | None:
            for p in self.providers:
                if (
                    not getattr(p, "write_only", False)
                    and p.can_load_history
                    and p.__class__.__name__.lower() == name.lower()
                ):
                    return p
            return None

        # Try preferred provider
        if preferred and (provider := find_provider(preferred)):
            return provider

        # Try default provider
        if self.config.default_provider:
            if provider := find_provider(self.config.default_provider):
                return provider
            msg = "Default provider %s not found or not capable of loading history"
            logger.warning(msg, self.config.default_provider)

        # Find first capable provider
        for provider in self.providers:
            if not getattr(provider, "write_only", False) and provider.can_load_history:
                return provider

        msg = "No capable provider found for loading history"
        raise RuntimeError(msg)

    async def filter_messages(
        self,
        query: SessionQuery,
        preferred_provider: str | None = None,
    ) -> list[ChatMessage[str]]:
        """Get messages matching query.

        Args:
            query: Filter criteria
            preferred_provider: Optional preferred provider to use
        """
        provider = self.get_history_provider(preferred_provider)
        return await provider.filter_messages(query)

    async def log_message(
        self,
        *,
        conversation_id: str,
        content: str,
        role: str,
        name: str | None = None,
        cost_info: TokenCost | None = None,
        model: str | None = None,
        response_time: float | None = None,
        forwarded_from: list[str] | None = None,
    ):
        """Log message to all providers."""
        if not self.config.log_messages:
            return
        for provider in self.providers:
            try:
                if not provider.should_log_agent(name or "no name"):
                    continue
                await provider.log_message(
                    conversation_id=conversation_id,
                    content=content,
                    role=role,
                    name=name,
                    cost_info=cost_info,
                    model=model,
                    response_time=response_time,
                    forwarded_from=forwarded_from,
                )
            except Exception:
                logger.exception("Error logging message to provider: %r", provider)

    async def log_conversation(
        self,
        *,
        conversation_id: str,
        node_name: str,
        start_time: datetime | None = None,
    ):
        """Log conversation to all providers."""
        if not self.config.log_conversations:
            return

        for provider in self.providers:
            try:
                await provider.log_conversation(
                    conversation_id=conversation_id,
                    node_name=node_name,
                    start_time=start_time,
                )
            except Exception:
                logger.exception("Error logging conversation to provider: %r", provider)

    async def log_tool_call(
        self,
        *,
        conversation_id: str,
        message_id: str,
        tool_call: ToolCallInfo,
    ):
        """Log tool call to all providers."""
        if not self.config.log_tool_calls:
            return

        for provider in self.providers:
            try:
                await provider.log_tool_call(
                    conversation_id=conversation_id,
                    message_id=message_id,
                    tool_call=tool_call,
                )
            except Exception:
                logger.exception("Error logging tool call to provider: %r", provider)

    async def log_command(
        self,
        *,
        agent_name: str,
        session_id: str,
        command: str,
        context_type: type | None = None,
        metadata: dict[str, JsonValue] | None = None,
    ):
        """Log command to all providers."""
        if not self.config.log_commands:
            return

        for provider in self.providers:
            try:
                await provider.log_command(
                    agent_name=agent_name,
                    session_id=session_id,
                    command=command,
                    context_type=context_type,
                    metadata=metadata,
                )
            except Exception:
                logger.exception("Error logging command to provider: %r", provider)

    async def get_commands(
        self,
        agent_name: str,
        session_id: str,
        *,
        limit: int | None = None,
        current_session_only: bool = False,
        preferred_provider: str | None = None,
    ) -> list[str]:
        """Get command history."""
        if not self.config.log_commands:
            return []

        provider = self.get_history_provider(preferred_provider)
        return await provider.get_commands(
            agent_name=agent_name,
            session_id=session_id,
            limit=limit,
            current_session_only=current_session_only,
        )

    async def log_context_message(
        self,
        *,
        conversation_id: str,
        content: str,
        role: str,
        name: str | None = None,
        model: str | None = None,
    ):
        """Log context message to all providers."""
        for provider in self.providers:
            try:
                await provider.log_context_message(
                    conversation_id=conversation_id,
                    content=content,
                    role=role,
                    name=name,
                    model=model,
                )
            except Exception:
                msg = "Error logging context message to provider: %r"
                logger.exception(msg, provider)

    async def reset(
        self,
        *,
        agent_name: str | None = None,
        hard: bool = False,
    ) -> tuple[int, int]:
        """Reset storage in all providers.

        Returns counts from primary provider.
        """
        counts = (0, 0)
        for provider in self.providers:
            try:
                counts = await provider.reset(agent_name=agent_name, hard=hard)
            except Exception:
                msg = "Error resetting provider: %r"
                logger.exception(msg, provider.__class__.__name__)
        return counts

    async def get_conversation_counts(
        self,
        *,
        agent_name: str | None = None,
    ) -> tuple[int, int]:
        """Get counts from primary provider."""
        provider = self.get_history_provider()
        return await provider.get_conversation_counts(agent_name=agent_name)

    # Sync wrappers
    def reset_sync(self, *args, **kwargs) -> tuple[int, int]:
        """Sync wrapper for reset."""
        return self.run_task_sync(self.reset(*args, **kwargs))

    def get_conversation_counts_sync(self, *args, **kwargs) -> tuple[int, int]:
        """Sync wrapper for get_conversation_counts."""
        return self.run_task_sync(self.get_conversation_counts(*args, **kwargs))

    def log_conversation_sync(self, *args, **kwargs):
        """Sync wrapper for log_conversation."""
        for provider in self.providers:
            provider.log_conversation_sync(*args, **kwargs)

    def log_message_sync(self, *args, **kwargs):
        """Sync wrapper for log_message."""
        for provider in self.providers:
            provider.log_message_sync(*args, **kwargs)

    def log_tool_call_sync(self, *args, **kwargs):
        """Sync wrapper for log_tool_call."""
        for provider in self.providers:
            provider.log_tool_call_sync(*args, **kwargs)

    def log_command_sync(self, *args, **kwargs):
        """Sync wrapper for log_command."""
        for provider in self.providers:
            provider.log_command_sync(*args, **kwargs)

    def get_commands_sync(self, *args, **kwargs) -> list[str]:
        """Sync wrapper for get_commands."""
        provider = self.get_history_provider()
        return provider.get_commands_sync(*args, **kwargs)

    def filter_messages_sync(self, *args, **kwargs) -> list[ChatMessage[str]]:
        """Sync wrapper for filter_messages."""
        provider = self.get_history_provider()
        return provider.filter_messages_sync(*args, **kwargs)

    def log_context_message_sync(self, *args, **kwargs):
        """Sync wrapper for log_context_message."""
        for provider in self.providers:
            provider.log_context_message_sync(*args, **kwargs)

__init__

__init__(config: StorageConfig)

Initialize storage manager.

Parameters:

Name	Type	Description	Default
config	StorageConfig	Storage configuration including providers and filters	required

Source code in src/llmling_agent/storage.py

def __init__(self, config: StorageConfig):
    """Initialize storage manager.

    Args:
        config: Storage configuration including providers and filters
    """
    self.config = config
    self.providers = [
        self._create_provider(cfg) for cfg in self.config.effective_providers
    ]

_create_provider

_create_provider(config: BaseStorageProviderConfig) -> StorageProvider

Create provider instance from configuration.

Source code in src/llmling_agent/storage.py

def _create_provider(self, config: BaseStorageProviderConfig) -> StorageProvider:
    """Create provider instance from configuration."""
    # Extract common settings from BaseStorageProviderConfig
    match self.config.filter_mode:
        case "and" if self.config.agents and config.agents:
            logged_agents = self.config.agents & config.agents
        case "and":
            logged_agents = self.config.agents or config.agents or set()
        case "override":
            logged_agents = (
                config.agents
                if config.agents is not None
                else self.config.agents or set()
            )

    provider_config = config.model_copy(
        update={
            "log_messages": config.log_messages and self.config.log_messages,
            "log_conversations": config.log_conversations
            and self.config.log_conversations,
            "log_tool_calls": config.log_tool_calls and self.config.log_tool_calls,
            "log_commands": config.log_commands and self.config.log_commands,
            "log_context": config.log_context and self.config.log_context,
            "agents": logged_agents,
        }
    )

    match provider_config:
        case SQLStorageConfig():
            from sqlmodel import create_engine

            from llmling_agent_storage.sql_provider import SQLModelProvider

            engine = create_engine(
                provider_config.url, pool_size=provider_config.pool_size
            )
            return SQLModelProvider(provider_config, engine)
        case FileStorageConfig():
            from llmling_agent_storage.file_provider import FileProvider

            return FileProvider(provider_config)
        case TextLogConfig():
            from llmling_agent_storage.text_log_provider import TextLogProvider

            return TextLogProvider(provider_config)

        case Mem0Config():
            from llmling_agent_storage.mem0 import Mem0StorageProvider

            return Mem0StorageProvider(provider_config)

        case MemoryStorageConfig():
            from llmling_agent_storage.memory_provider import MemoryStorageProvider

            return MemoryStorageProvider(provider_config)
        case _:
            msg = f"Unknown provider type: {provider_config}"
            raise ValueError(msg)

cleanup

cleanup()

Clean up all providers.

Source code in src/llmling_agent/storage.py

def cleanup(self):
    """Clean up all providers."""
    for provider in self.providers:
        try:
            provider.cleanup()
        except Exception:
            logger.exception("Error cleaning up provider: %r", provider)
    self.providers.clear()

filter_messages async

filter_messages(
    query: SessionQuery, preferred_provider: str | None = None
) -> list[ChatMessage[str]]

Get messages matching query.

Parameters:

Name	Type	Description	Default
query	SessionQuery	Filter criteria	required
preferred_provider	str | None	Optional preferred provider to use	None

Source code in src/llmling_agent/storage.py

async def filter_messages(
    self,
    query: SessionQuery,
    preferred_provider: str | None = None,
) -> list[ChatMessage[str]]:
    """Get messages matching query.

    Args:
        query: Filter criteria
        preferred_provider: Optional preferred provider to use
    """
    provider = self.get_history_provider(preferred_provider)
    return await provider.filter_messages(query)

filter_messages_sync

filter_messages_sync(*args, **kwargs) -> list[ChatMessage[str]]

Sync wrapper for filter_messages.

Source code in src/llmling_agent/storage.py

def filter_messages_sync(self, *args, **kwargs) -> list[ChatMessage[str]]:
    """Sync wrapper for filter_messages."""
    provider = self.get_history_provider()
    return provider.filter_messages_sync(*args, **kwargs)

get_commands async

get_commands(
    agent_name: str,
    session_id: str,
    *,
    limit: int | None = None,
    current_session_only: bool = False,
    preferred_provider: str | None = None,
) -> list[str]

Get command history.

Source code in src/llmling_agent/storage.py

async def get_commands(
    self,
    agent_name: str,
    session_id: str,
    *,
    limit: int | None = None,
    current_session_only: bool = False,
    preferred_provider: str | None = None,
) -> list[str]:
    """Get command history."""
    if not self.config.log_commands:
        return []

    provider = self.get_history_provider(preferred_provider)
    return await provider.get_commands(
        agent_name=agent_name,
        session_id=session_id,
        limit=limit,
        current_session_only=current_session_only,
    )

get_commands_sync

get_commands_sync(*args, **kwargs) -> list[str]

Sync wrapper for get_commands.

Source code in src/llmling_agent/storage.py

def get_commands_sync(self, *args, **kwargs) -> list[str]:
    """Sync wrapper for get_commands."""
    provider = self.get_history_provider()
    return provider.get_commands_sync(*args, **kwargs)

get_conversation_counts async

get_conversation_counts(*, agent_name: str | None = None) -> tuple[int, int]

Get counts from primary provider.

Source code in src/llmling_agent/storage.py

async def get_conversation_counts(
    self,
    *,
    agent_name: str | None = None,
) -> tuple[int, int]:
    """Get counts from primary provider."""
    provider = self.get_history_provider()
    return await provider.get_conversation_counts(agent_name=agent_name)

get_conversation_counts_sync

get_conversation_counts_sync(*args, **kwargs) -> tuple[int, int]

Sync wrapper for get_conversation_counts.

Source code in src/llmling_agent/storage.py

def get_conversation_counts_sync(self, *args, **kwargs) -> tuple[int, int]:
    """Sync wrapper for get_conversation_counts."""
    return self.run_task_sync(self.get_conversation_counts(*args, **kwargs))

get_history_provider

get_history_provider(preferred: str | None = None) -> StorageProvider

Get provider for loading history.

Parameters:

Name	Type	Description	Default
preferred	str | None	Optional preferred provider name	None

Returns:

Type	Description
StorageProvider	First capable provider based on priority:
StorageProvider	Preferred provider if specified and capable
StorageProvider	Default provider if specified and capable
StorageProvider	First capable provider
StorageProvider	Raises error if no capable provider found

Source code in src/llmling_agent/storage.py

def get_history_provider(self, preferred: str | None = None) -> StorageProvider:
    """Get provider for loading history.

    Args:
        preferred: Optional preferred provider name

    Returns:
        First capable provider based on priority:
        1. Preferred provider if specified and capable
        2. Default provider if specified and capable
        3. First capable provider
        4. Raises error if no capable provider found
    """

    # Function to find capable provider by name
    def find_provider(name: str) -> StorageProvider | None:
        for p in self.providers:
            if (
                not getattr(p, "write_only", False)
                and p.can_load_history
                and p.__class__.__name__.lower() == name.lower()
            ):
                return p
        return None

    # Try preferred provider
    if preferred and (provider := find_provider(preferred)):
        return provider

    # Try default provider
    if self.config.default_provider:
        if provider := find_provider(self.config.default_provider):
            return provider
        msg = "Default provider %s not found or not capable of loading history"
        logger.warning(msg, self.config.default_provider)

    # Find first capable provider
    for provider in self.providers:
        if not getattr(provider, "write_only", False) and provider.can_load_history:
            return provider

    msg = "No capable provider found for loading history"
    raise RuntimeError(msg)

log_command async

log_command(
    *,
    agent_name: str,
    session_id: str,
    command: str,
    context_type: type | None = None,
    metadata: dict[str, JsonValue] | None = None,
)

Log command to all providers.

Source code in src/llmling_agent/storage.py

async def log_command(
    self,
    *,
    agent_name: str,
    session_id: str,
    command: str,
    context_type: type | None = None,
    metadata: dict[str, JsonValue] | None = None,
):
    """Log command to all providers."""
    if not self.config.log_commands:
        return

    for provider in self.providers:
        try:
            await provider.log_command(
                agent_name=agent_name,
                session_id=session_id,
                command=command,
                context_type=context_type,
                metadata=metadata,
            )
        except Exception:
            logger.exception("Error logging command to provider: %r", provider)

log_command_sync

log_command_sync(*args, **kwargs)

Sync wrapper for log_command.

Source code in src/llmling_agent/storage.py

def log_command_sync(self, *args, **kwargs):
    """Sync wrapper for log_command."""
    for provider in self.providers:
        provider.log_command_sync(*args, **kwargs)

log_context_message async

log_context_message(
    *,
    conversation_id: str,
    content: str,
    role: str,
    name: str | None = None,
    model: str | None = None,
)

Log context message to all providers.

Source code in src/llmling_agent/storage.py

async def log_context_message(
    self,
    *,
    conversation_id: str,
    content: str,
    role: str,
    name: str | None = None,
    model: str | None = None,
):
    """Log context message to all providers."""
    for provider in self.providers:
        try:
            await provider.log_context_message(
                conversation_id=conversation_id,
                content=content,
                role=role,
                name=name,
                model=model,
            )
        except Exception:
            msg = "Error logging context message to provider: %r"
            logger.exception(msg, provider)

log_context_message_sync

log_context_message_sync(*args, **kwargs)

Sync wrapper for log_context_message.

Source code in src/llmling_agent/storage.py

def log_context_message_sync(self, *args, **kwargs):
    """Sync wrapper for log_context_message."""
    for provider in self.providers:
        provider.log_context_message_sync(*args, **kwargs)

log_conversation async

log_conversation(
    *, conversation_id: str, node_name: str, start_time: datetime | None = None
)

Log conversation to all providers.

Source code in src/llmling_agent/storage.py

async def log_conversation(
    self,
    *,
    conversation_id: str,
    node_name: str,
    start_time: datetime | None = None,
):
    """Log conversation to all providers."""
    if not self.config.log_conversations:
        return

    for provider in self.providers:
        try:
            await provider.log_conversation(
                conversation_id=conversation_id,
                node_name=node_name,
                start_time=start_time,
            )
        except Exception:
            logger.exception("Error logging conversation to provider: %r", provider)

log_conversation_sync

log_conversation_sync(*args, **kwargs)

Sync wrapper for log_conversation.

Source code in src/llmling_agent/storage.py

def log_conversation_sync(self, *args, **kwargs):
    """Sync wrapper for log_conversation."""
    for provider in self.providers:
        provider.log_conversation_sync(*args, **kwargs)

log_message async

log_message(
    *,
    conversation_id: str,
    content: str,
    role: str,
    name: str | None = None,
    cost_info: TokenCost | None = None,
    model: str | None = None,
    response_time: float | None = None,
    forwarded_from: list[str] | None = None,
)

Log message to all providers.

Source code in src/llmling_agent/storage.py

async def log_message(
    self,
    *,
    conversation_id: str,
    content: str,
    role: str,
    name: str | None = None,
    cost_info: TokenCost | None = None,
    model: str | None = None,
    response_time: float | None = None,
    forwarded_from: list[str] | None = None,
):
    """Log message to all providers."""
    if not self.config.log_messages:
        return
    for provider in self.providers:
        try:
            if not provider.should_log_agent(name or "no name"):
                continue
            await provider.log_message(
                conversation_id=conversation_id,
                content=content,
                role=role,
                name=name,
                cost_info=cost_info,
                model=model,
                response_time=response_time,
                forwarded_from=forwarded_from,
            )
        except Exception:
            logger.exception("Error logging message to provider: %r", provider)

log_message_sync

log_message_sync(*args, **kwargs)

Sync wrapper for log_message.

Source code in src/llmling_agent/storage.py

def log_message_sync(self, *args, **kwargs):
    """Sync wrapper for log_message."""
    for provider in self.providers:
        provider.log_message_sync(*args, **kwargs)

log_tool_call async

log_tool_call(*, conversation_id: str, message_id: str, tool_call: ToolCallInfo)

Log tool call to all providers.

Source code in src/llmling_agent/storage.py

async def log_tool_call(
    self,
    *,
    conversation_id: str,
    message_id: str,
    tool_call: ToolCallInfo,
):
    """Log tool call to all providers."""
    if not self.config.log_tool_calls:
        return

    for provider in self.providers:
        try:
            await provider.log_tool_call(
                conversation_id=conversation_id,
                message_id=message_id,
                tool_call=tool_call,
            )
        except Exception:
            logger.exception("Error logging tool call to provider: %r", provider)

log_tool_call_sync

log_tool_call_sync(*args, **kwargs)

Sync wrapper for log_tool_call.

Source code in src/llmling_agent/storage.py

def log_tool_call_sync(self, *args, **kwargs):
    """Sync wrapper for log_tool_call."""
    for provider in self.providers:
        provider.log_tool_call_sync(*args, **kwargs)

reset async

reset(*, agent_name: str | None = None, hard: bool = False) -> tuple[int, int]

Reset storage in all providers.

Returns counts from primary provider.

Source code in src/llmling_agent/storage.py

async def reset(
    self,
    *,
    agent_name: str | None = None,
    hard: bool = False,
) -> tuple[int, int]:
    """Reset storage in all providers.

    Returns counts from primary provider.
    """
    counts = (0, 0)
    for provider in self.providers:
        try:
            counts = await provider.reset(agent_name=agent_name, hard=hard)
        except Exception:
            msg = "Error resetting provider: %r"
            logger.exception(msg, provider.__class__.__name__)
    return counts

reset_sync

reset_sync(*args, **kwargs) -> tuple[int, int]

Sync wrapper for reset.

Source code in src/llmling_agent/storage.py

def reset_sync(self, *args, **kwargs) -> tuple[int, int]:
    """Sync wrapper for reset."""
    return self.run_task_sync(self.reset(*args, **kwargs))