Skip to content

fix: implement lazy loading for tools imports to improve package startup time #1172

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MervinPraison merged 3 commits into from
Mar 30, 2026
Merged

fix: implement lazy loading for tools imports to improve package startup time #1172

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d9cb153
fix: implement lazy loading for tools imports to improve package star…
github-actions[bot] Mar 30, 2026
309ba9d
fix: address lazy loading implementation issues
github-actions[bot] Mar 30, 2026
42cc31c
fix: address lazy loading implementation issues
github-actions[bot] Mar 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 74 additions & 31 deletions src/praisonai-agents/praisonaiagents/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,18 +59,7 @@
# Import configuration (lightweight, no heavy deps)
from . import _config

# Lightweight imports that don't trigger heavy dependency chains
from .tools.tools import Tools
from .tools.base import BaseTool, ToolResult, ToolValidationError, validate_tool
from .tools.decorator import tool, FunctionTool
from .tools.registry import get_registry, register_tool, get_tool, ToolRegistry
# db and obs are lazy-loaded via __getattr__ for performance

# Sub-packages for organized imports (pa.config, pa.tools, etc.)
# These enable: import praisonaiagents as pa; pa.config.MemoryConfig
from . import config
from . import tools
# Note: db, obs, knowledge and mcp are lazy-loaded via __getattr__ due to heavy deps
# Note: tools, config, memory, workflows, db, obs, knowledge and mcp are lazy-loaded via __getattr__ due to heavy deps

Copy link

Copilot AI Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tools is still listed in __all__/__dir__, but the eager from . import tools was removed and __getattr__ currently doesn’t return the .tools submodule when pa.tools is accessed (fallback checks attributes inside submodules, not the submodule itself). This makes import praisonaiagents as pa; pa.tools raise AttributeError. Consider adding an explicit handler (like memory/workflows) for name == 'tools' to return importlib.import_module('.tools', 'praisonaiagents'), or re-introduce a lightweight submodule import if acceptable.

Suggested change
# -----------------------------------------------------------------------------
# Lazy top-level access for `tools` submodule
# -----------------------------------------------------------------------------
# CodeQL noted that `tools` remains in __all__/__dir__ but is no longer eagerly
# imported. To ensure `import praisonaiagents as pa; pa.tools` works while
# preserving lazy loading, we expose a lightweight proxy that imports the real
# submodule on first attribute access.
import importlib as _importlib
class _LazyToolsModule:
_module = None
def _load(self):
if self._module is None:
# Import the actual tools submodule only when it is first used
self._module = _importlib.import_module(".tools", "praisonaiagents")
return self._module
def __getattr__(self, name):
return getattr(self._load(), name)
def __dir__(self):
return dir(self._load())
# Expose a lazy `tools` attribute at the package level
tools = _LazyToolsModule()

Copilot uses AI. Check for mistakes.
# Embedding API - LAZY LOADED via __getattr__ for performance
# Supports: embedding, embeddings, aembedding, aembeddings, EmbeddingResult, get_dimensions
Expand Down Expand Up @@ -125,6 +114,19 @@ def _get_lazy_cache():
# ============================================================================

_LAZY_IMPORTS = {
# Tools (moved from eager imports for lazy loading)
'Tools': ('praisonaiagents.tools.tools', 'Tools'),
'BaseTool': ('praisonaiagents.tools.base', 'BaseTool'),
'ToolResult': ('praisonaiagents.tools.base', 'ToolResult'),
'ToolValidationError': ('praisonaiagents.tools.base', 'ToolValidationError'),
'validate_tool': ('praisonaiagents.tools.base', 'validate_tool'),
'tool': ('praisonaiagents.tools.decorator', 'tool'),
'FunctionTool': ('praisonaiagents.tools.decorator', 'FunctionTool'),
'get_registry': ('praisonaiagents.tools.registry', 'get_registry'),
'register_tool': ('praisonaiagents.tools.registry', 'register_tool'),
'get_tool': ('praisonaiagents.tools.registry', 'get_tool'),
'ToolRegistry': ('praisonaiagents.tools.registry', 'ToolRegistry'),
Comment on lines +117 to +128
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Restore the full top-level tools surface before disabling fallback.

Only a subset of praisonaiagents.tools.__all__ is mapped here. With fallback_modules=[], root-level imports that used to resolve through the tools fallback—add_tool, list_tools, RetryPolicy, ToolProfile, and the tool functions exported from praisonaiagents.tools—now fail with AttributeError. Either lazy-map the remaining public tool exports or keep a bounded fallback keyed off a known allowlist instead of removing fallback entirely.

Based on learnings: "Public API changes require a deprecation cycle: emit DeprecationWarning for one release before breaking change"

Also applies to: 594-595


# Main display utilities (imports rich)
'TaskOutput': ('praisonaiagents.main', 'TaskOutput'),
'ReflectionOutput': ('praisonaiagents.main', 'ReflectionOutput'),
Expand Down Expand Up @@ -179,9 +181,7 @@ def _get_lazy_cache():
'HandoffDepthError': ('praisonaiagents.agent.handoff', 'HandoffDepthError'),
'HandoffTimeoutError': ('praisonaiagents.agent.handoff', 'HandoffTimeoutError'),

# Embedding API
'embedding': ('praisonaiagents.embedding.embed', 'embedding'),
'embeddings': ('praisonaiagents.embedding.embed', 'embedding'),
# Embedding API (Note: embedding/embeddings handled in custom_handler to override subpackage)
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The comment here states that embedding/embeddings are handled in custom_handler. However, these are now handled directly by the _EmbeddingProxy instances assigned at the module level, making the custom_handler logic for these names redundant. Please update the comment to reflect the actual implementation.

Suggested change
# Embedding API (Note: embedding/embeddings handled in custom_handler to override subpackage)
# Embedding API (Note: embedding/embeddings handled by _EmbeddingProxy to override subpackage)

'aembedding': ('praisonaiagents.embedding.embed', 'aembedding'),
'aembeddings': ('praisonaiagents.embedding.embed', 'aembedding'),
'embed': ('praisonaiagents.embedding.embed', 'embed'),
Expand Down Expand Up @@ -506,21 +506,27 @@ def _custom_handler(name, cache):
cache['Agents'] = value
return value

# Task removed in v4.0.0 - use Task instead
if name == "Task":
raise ImportError(
"Task has been removed in v4.0.0. Use Task instead.\n"
"Migration: Replace 'from praisonaiagents import Task' with 'from praisonaiagents import Task'\n"
"Task supports all Task features including action, handler, loop_over, etc."
)

# Module imports (return the module itself)
if name == 'tools':
import importlib
mod = importlib.import_module('.tools', 'praisonaiagents')
cache['tools'] = mod
return mod
if name == 'config':
import importlib
mod = importlib.import_module('.config', 'praisonaiagents')
cache['config'] = mod
return mod
if name == 'memory':
import importlib
return importlib.import_module('.memory', 'praisonaiagents')
mod = importlib.import_module('.memory', 'praisonaiagents')
cache['memory'] = mod
return mod
if name == 'workflows':
import importlib
return importlib.import_module('.workflows', 'praisonaiagents')
mod = importlib.import_module('.workflows', 'praisonaiagents')
cache['workflows'] = mod
return mod

raise AttributeError(f"Not handled by custom_handler: {name}")

Expand All @@ -532,20 +538,57 @@ def _custom_handler(name, cache):
# Override them here to return the function instead of the module.
# ============================================================================

# Override 'embedding' to return the function, not the subpackage
from .embedding.embed import embedding as _embedding_func
embedding = _embedding_func
# Override 'embedding' and 'embeddings' at module level to prevent subpackage import
# These need to be set after _LAZY_IMPORTS is defined but before __getattr__ is created
def _get_embedding_func():
"""Lazy getter for embedding function."""
# Import with alias to avoid overwriting the module proxy
from .embedding.embed import embedding as _embedding_func
return _embedding_func
Comment on lines +543 to +547
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify whether the embedding proxy survives after first invocation

python3 - <<'EOF'
import sys

# Import the package
import praisonaiagents as pa

print(f"Before call: type(pa.embedding) = {type(pa.embedding).__name__}")
print(f"Before call: pa.embedding is _EmbeddingProxy? = {type(pa.embedding).__name__ == '_EmbeddingProxy'}")

# Check if embedding is callable (proxy should be callable)
print(f"Before call: callable(pa.embedding) = {callable(pa.embedding)}")

# Now trigger the lazy load by accessing __signature__ (doesn't require actual API call)
try:
    import inspect
    sig = inspect.signature(pa.embedding)
    print(f"After signature access: type(pa.embedding) = {type(pa.embedding).__name__}")
    print(f"After signature access: callable(pa.embedding) = {callable(pa.embedding)}")
except Exception as e:
    print(f"Error getting signature: {e}")

# Check what's in sys.modules
if 'praisonaiagents.embedding' in sys.modules:
    print(f"sys.modules['praisonaiagents.embedding'] exists")
    submod = sys.modules['praisonaiagents.embedding']
    print(f"  type: {type(submod)}")
    
# Final check
print(f"Final: type(pa.embedding) = {type(pa.embedding).__name__}")
EOF

Repository: MervinPraison/PraisonAI

Length of output: 194

🏁 Script executed:

fd -t f "__init__.py" "praisonaiagents" | grep -E "praisonaiagents.*__init__\.py$"

Repository: MervinPraison/PraisonAI

Length of output: 169

🏁 Script executed:

git ls-files "src/praisonai-agents/praisonaiagents/__init__.py"

Repository: MervinPraison/PraisonAI

Length of output: 115

🏁 Script executed:

cat -n src/praisonai-agents/praisonaiagents/__init__.py | head -80

Repository: MervinPraison/PraisonAI

Length of output: 4636

🏁 Script executed:

wc -l src/praisonai-agents/praisonaiagents/__init__.py

Repository: MervinPraison/PraisonAI

Length of output: 119

🏁 Script executed:

sed -n '530,580p' src/praisonai-agents/praisonaiagents/__init__.py

Repository: MervinPraison/PraisonAI

Length of output: 1905

🏁 Script executed:

sed -n '500,570p' src/praisonai-agents/praisonaiagents/__init__.py | cat -n

Repository: MervinPraison/PraisonAI

Length of output: 3155

🏁 Script executed:

sed -n '570,620p' src/praisonai-agents/praisonaiagents/__init__.py

Repository: MervinPraison/PraisonAI

Length of output: 1768

🏁 Script executed:

sed -n '580,650p' src/praisonai-agents/praisonaiagents/__init__.py | cat -n

Repository: MervinPraison/PraisonAI

Length of output: 3065

🏁 Script executed:

rg -A 10 "def _get_embedding_func" src/praisonai-agents/praisonaiagents/__init__.py

Repository: MervinPraison/PraisonAI

Length of output: 487

Apply sys.modules restoration to prevent embedding proxy from being overwritten by submodule import.

When _get_embedding_func() executes from .embedding.embed import embedding as _embedding_func, Python's import machinery binds the .embedding subpackage to sys.modules['praisonaiagents'].embedding, overwriting the _EmbeddingProxy() instance set at module load. The alias as _embedding_func only avoids local variable conflict but does not prevent submodule binding. After the first invocation, praisonaiagents.embedding will refer to the submodule rather than the callable proxy.

Restore the proxy after import to maintain the intended lazy-load behavior:

def _get_embedding_func():
    """Lazy getter for embedding function."""
    from .embedding.embed import embedding as _embedding_func
    # Restore proxy after import to prevent submodule from overwriting it
    sys.modules[__name__].embedding = embedding
    sys.modules[__name__].embeddings = embeddings
    return _embedding_func
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/praisonai-agents/praisonaiagents/__init__.py` around lines 543 - 547, The
import in _get_embedding_func() binds the .embedding subpackage into sys.modules
and overwrites the module-level embedding proxy; after importing embedding from
.embedding.embed you must restore the proxy names on the package module so the
lazy proxy remains in place. Modify _get_embedding_func() to import the callable
(embedding) as before, then reassign sys.modules[__name__].embedding and
sys.modules[__name__].embeddings back to the original proxy objects (the
module-level embedding and embeddings proxies) before returning the imported
callable; reference the function name _get_embedding_func and the module-level
symbols embedding and embeddings when making the change.


# Create lazy properties that override the submodule
class _EmbeddingProxy:
"""Proxy object that loads embedding function on first access."""
def __init__(self):
self._func = None

def _load(self):
"""Load the actual embedding function if not already loaded."""
if self._func is None:
self._func = _get_embedding_func()
return self._func

def __call__(self, *args, **kwargs):
return self._load()(*args, **kwargs)

def __getattr__(self, name):
return getattr(self._load(), name)

@property
def __wrapped__(self):
"""Support for inspect.signature() and functools.wraps."""
return self._load()

@property
def __signature__(self):
"""Support for inspect.signature()."""
import inspect
return inspect.signature(self._load())

def __repr__(self):
return f"<lazy proxy for {self._load()!r}>"

# Also provide embeddings alias
embeddings = _embedding_func
# Override the submodule with our function proxy
embedding = _EmbeddingProxy()
embeddings = embedding # embeddings is an alias
Comment on lines +543 to +583
Copy link

Copilot AI Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new _EmbeddingProxy changes the public embedding/embeddings export from a real Python function to a proxy instance. While it stays callable, this is an observable API change (e.g., inspect.signature, help(), inspect.isfunction, pickling, and some type-checking behaviors will differ), which conflicts with the “zero breaking changes” goal. A safer approach is to keep exporting an actual function object (even if it’s a thin lazy wrapper) and, if needed, swap the module global to the real function after first load so subsequent imports/introspection see the true function.

Copilot uses AI. Check for mistakes.
Comment thread
coderabbitai[bot] marked this conversation as resolved.


# Create the __getattr__ function using centralized utility
__getattr__ = create_lazy_getattr_with_fallback(
mapping=_LAZY_IMPORTS,
module_name=__name__,
cache=_lazy_cache,
fallback_modules=['tools', 'memory', 'config', 'workflows'],
fallback_modules=[], # Note: 'embedding' excluded to avoid conflict with embedding() function
custom_handler=_custom_handler
)

Expand Down