Compare commits

...
1 Commits
Author SHA1 Message Date
Mervin Praison 9cbab01785 Release v2.5.0 2025-12-26 02:20:59 +00:00
33 changed files with 2592 additions and 60 deletions

No files matched your search

+3 -1
View File
@@ -80,4 +80,6 @@ CopilotKit*
issue_details.json
# MCP Registry tokens
.mcpregistry_*
.mcpregistry_*
*.log
+4
View File
@@ -261,6 +261,10 @@ npm install praisonai
| ↳ 100+ LLM Support | [Example](examples/python/providers/openai/openai_gpt4_example.py) | [📖](https://docs.praison.ai/models) |
| ↳ Callback Agents | [Example](examples/python/general/advanced-callback-systems.py) | [📖](https://docs.praison.ai/features/callbacks) |
| ↳ Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
| ↳ Middleware System | [Example](examples/middleware/basic_middleware.py) | [📖](https://docs.praison.ai/features/middleware) |
| ↳ Configurable Model | [Example](examples/middleware/configurable_model.py) | [📖](https://docs.praison.ai/features/configurable-model) |
| ↳ Rate Limiter | [Example](examples/middleware/rate_limiter.py) | [📖](https://docs.praison.ai/features/rate-limiter) |
| ↳ Injected Tool State | [Example](examples/middleware/injected_state.py) | [📖](https://docs.praison.ai/features/injected-state) |
| ↳ Shadow Git Checkpoints | [Example](#10-shadow-git-checkpoints) | [📖](https://docs.praison.ai/features/checkpoints) |
| ↳ Background Tasks | [Example](examples/background/basic_background.py) | [📖](https://docs.praison.ai/features/background-tasks) |
| ↳ Policy Engine | [Example](examples/policy/basic_policy.py) | [📖](https://docs.praison.ai/features/policy-engine) |
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.4.4" \
"praisonai>=2.5.0" \
"praisonai[chat]" \
"embedchain[github,youtube]"
+1 -1
View File
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.4.4" \
"praisonai>=2.5.0" \
"praisonai[ui]" \
"praisonai[chat]" \
"praisonai[realtime]" \
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.4.4" \
"praisonai>=2.5.0" \
"praisonai[ui]" \
"praisonai[crewai]"
+55
View File
@@ -0,0 +1,55 @@
"""
Basic Middleware Example - PraisonAI Agents
Demonstrates before/after hooks and wrap decorators for model and tool calls.
"""
from praisonaiagents import Agent, tool
from praisonaiagents.hooks import (
before_model, after_model, wrap_model_call,
before_tool, after_tool, wrap_tool_call,
InvocationContext, ModelRequest, ModelResponse
)
# Simple tool for demonstration
@tool
def get_weather(city: str) -> str:
"""Get weather for a city."""
return f"Sunny, 22°C in {city}"
# Before model hook - adds context
@before_model
def add_context(request: ModelRequest) -> ModelRequest:
print(f"[before_model] Adding context to request")
return request
# After model hook - logs response
@after_model
def log_response(response: ModelResponse) -> ModelResponse:
print(f"[after_model] Response received")
return response
# Wrap tool call - retry on error
@wrap_tool_call
def retry_on_error(tool_call, call_next):
print(f"[wrap_tool_call] Executing tool")
try:
return call_next(tool_call)
except Exception as e:
print(f"[wrap_tool_call] Retrying after error: {e}")
return call_next(tool_call)
# Create agent with hooks
agent = Agent(
name="WeatherBot",
instructions="You help with weather queries.",
tools=[get_weather],
hooks=[add_context, log_response, retry_on_error]
)
if __name__ == "__main__":
# Test the tool directly
result = get_weather("London")
print(f"Weather: {result}")
print("\n✓ Middleware example complete")
+31
View File
@@ -0,0 +1,31 @@
"""
Configurable Model Example - PraisonAI Agents
Demonstrates runtime model switching without recreating the agent.
"""
from praisonaiagents import Agent
# Create agent with configurable model enabled
agent = Agent(
name="FlexBot",
instructions="You are a helpful assistant.",
llm="gpt-4o-mini",
llm_config={"configurable": True}
)
if __name__ == "__main__":
# Default model call
print("Using default model (gpt-4o-mini)...")
# response = agent.chat("Say hello in 5 words")
# Override model per-call
print("\nOverriding to different model...")
# response = agent.chat("Say hello in 5 words", config={"model": "gpt-4o"})
# Override temperature
print("\nOverriding temperature...")
# response = agent.chat("Say hello creatively", config={"temperature": 0.9})
print("\n✓ Configurable model example complete")
print("Note: Uncomment agent.chat() calls to run with API key")
+50
View File
@@ -0,0 +1,50 @@
"""
Injected State Example - PraisonAI Agents
Demonstrates injecting agent state into tools without exposing it in the schema.
"""
from praisonaiagents import Agent, tool
from praisonaiagents.tools import Injected
from praisonaiagents.tools.injected import AgentState, with_injection_context
# Tool with injected state - state param is NOT in the public schema
@tool
def show_context(query: str, state: Injected[dict]) -> str:
"""Show the current agent context."""
session_id = state.get('session_id', 'unknown')
agent_id = state.get('agent_id', 'unknown')
return f"Query: {query}, Session: {session_id}, Agent: {agent_id}"
# Create agent with the tool
agent = Agent(
name="ContextBot",
instructions="You help show context information.",
tools=[show_context],
session_id="my-session-123"
)
if __name__ == "__main__":
# Verify injected param is not in schema
schema = show_context.get_schema()
params = schema['function']['parameters']['properties']
print(f"Schema params: {list(params.keys())}")
assert 'state' not in params, "state should NOT be in schema"
print("✓ 'state' correctly excluded from schema")
# Test with manual injection context
mock_state = AgentState(
agent_id="test-agent",
run_id="run-1",
session_id="session-abc"
)
with with_injection_context(mock_state):
result = show_context(query="hello")
print(f"Result: {result}")
# Test via agent.execute_tool
result = agent.execute_tool("show_context", {"query": "test"})
print(f"Agent result: {result}")
print("\n✓ Injected state example complete")
+29
View File
@@ -0,0 +1,29 @@
"""
Rate Limiter Example - PraisonAI Agents
Demonstrates token bucket rate limiting for LLM API calls.
"""
from praisonaiagents import Agent
from praisonaiagents.llm import RateLimiter
# Create rate limiter: 60 requests per minute
limiter = RateLimiter(requests_per_minute=60, burst=5)
# Create agent with rate limiter
agent = Agent(
name="RateLimitedBot",
instructions="You are a helpful assistant.",
rate_limiter=limiter
)
if __name__ == "__main__":
print(f"Rate limiter: {limiter}")
print(f"Available tokens: {limiter.available_tokens}")
# Demonstrate rate limiting
for i in range(3):
limiter.acquire()
print(f"Request {i+1} acquired, tokens left: {limiter.available_tokens:.1f}")
print("\n✓ Rate limiter example complete")
+6 -11
View File
@@ -26,23 +26,18 @@ Expected Output:
Messages persisted: 2
"""
from praisonaiagents import Agent
from praisonai.db import PostgresDB
from praisonaiagents import Agent, db
# Create database adapter
db = PostgresDB(
host="localhost",
port=5432,
database="praisonai",
user="postgres",
password="praison123"
# Create database adapter (simplified import)
db_instance = db(
database_url="postgresql://postgres:praison123@localhost:5432/praisonai"
)
# Create agent with persistence enabled
agent = Agent(
name="Assistant",
instructions="You are a helpful assistant. Keep responses brief.",
db=db,
db=db_instance,
session_id="demo-session-001", # Same session_id = resume conversation
verbose=False
)
@@ -57,7 +52,7 @@ print(f"Session ID: {agent.session_id}")
print(f"Messages in history: {len(agent.chat_history)}")
# Close DB connection
db.close()
db_instance.close()
print("\n=== Session persisted! ===")
print("Run again with same session_id to resume.")
@@ -254,7 +254,10 @@ class Agent:
skills: Optional[List[str]] = None,
skills_dirs: Optional[List[str]] = None,
db: Optional[Any] = None,
session_id: Optional[str] = None
session_id: Optional[str] = None,
hooks: Optional[List[Any]] = None,
llm_config: Optional[Dict[str, Any]] = None,
rate_limiter: Optional[Any] = None
):
"""Initialize an Agent instance.
@@ -395,6 +398,17 @@ class Agent:
# Flag to track if final result has been displayed to prevent duplicates
self._final_display_shown = False
# Store hooks for middleware system (zero overhead when empty)
self._hooks = hooks or []
self._middleware_manager = None # Lazy init
# Store llm_config for configurable model switching
self._llm_config = llm_config or {}
self._llm_configurable = self._llm_config.get('configurable', False)
# Store rate limiter (optional, zero overhead when None)
self._rate_limiter = rate_limiter
# Store OpenAI client parameters for lazy initialization
self._openai_api_key = api_key
self._openai_base_url = base_url
@@ -1622,8 +1636,32 @@ Your Goal: {self.goal}"""
def execute_tool(self, function_name, arguments):
"""
Execute a tool dynamically based on the function name and arguments.
Injects agent state for tools with Injected[T] parameters.
"""
logging.debug(f"{self.name} executing tool {function_name} with arguments: {arguments}")
# Set up injection context for tools with Injected parameters
from ..tools.injected import AgentState, with_injection_context
state = AgentState(
agent_id=self.name,
run_id=getattr(self, '_current_run_id', 'unknown'),
session_id=getattr(self, '_session_id', None) or 'default',
last_user_message=self.chat_history[-1].get('content') if self.chat_history else None,
metadata={'agent_name': self.name}
)
# Execute within injection context
return self._execute_tool_with_context(function_name, arguments, state)
def _execute_tool_with_context(self, function_name, arguments, state):
"""Execute tool within injection context."""
from ..tools.injected import with_injection_context
with with_injection_context(state):
return self._execute_tool_impl(function_name, arguments)
def _execute_tool_impl(self, function_name, arguments):
"""Internal tool execution implementation."""
# Check if approval is required for this tool
from ..approval import is_approval_required, console_approval_callback, get_risk_level, mark_approved, ApprovalDecision, get_approval_callback
@@ -2058,7 +2096,11 @@ Your Goal: {self.goal}"""
"""Get the current session ID."""
return self._session_id
def chat(self, prompt, temperature=1.0, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, stream=None, task_name=None, task_description=None, task_id=None):
def chat(self, prompt, temperature=1.0, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, stream=None, task_name=None, task_description=None, task_id=None, config=None):
# Apply rate limiter if configured (before any LLM call)
if self._rate_limiter is not None:
self._rate_limiter.acquire()
# Initialize DB session on first chat (lazy)
self._init_db_session()
@@ -65,6 +65,23 @@ __all__ = [
"AfterAgentInput",
"SessionStartInput",
"SessionEndInput",
# Middleware types
"InvocationContext",
"ModelRequest",
"ModelResponse",
"ToolRequest",
"ToolResponse",
# Middleware decorators
"before_model",
"after_model",
"wrap_model_call",
"before_tool",
"after_tool",
"wrap_tool_call",
# Middleware utilities
"MiddlewareChain",
"AsyncMiddlewareChain",
"MiddlewareManager",
]
@@ -94,4 +111,27 @@ def __getattr__(name: str):
)
return locals()[name]
# Middleware types
if name in ("InvocationContext", "ModelRequest", "ModelResponse",
"ToolRequest", "ToolResponse"):
from .middleware import (
InvocationContext, ModelRequest, ModelResponse,
ToolRequest, ToolResponse
)
return locals()[name]
# Middleware decorators
if name in ("before_model", "after_model", "wrap_model_call",
"before_tool", "after_tool", "wrap_tool_call"):
from .middleware import (
before_model, after_model, wrap_model_call,
before_tool, after_tool, wrap_tool_call
)
return locals()[name]
# Middleware utilities
if name in ("MiddlewareChain", "AsyncMiddlewareChain", "MiddlewareManager"):
from .middleware import MiddlewareChain, AsyncMiddlewareChain, MiddlewareManager
return locals()[name]
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,439 @@
"""
Middleware System for PraisonAI Agents.
Provides wrap_model_call and wrap_tool_call patterns for intercepting
and modifying agent behavior at model/tool call boundaries.
Zero Performance Impact:
- All middleware is optional
- Empty middleware chain is O(1) fast path
- No overhead when middleware not registered
Usage:
from praisonaiagents.hooks import before_model, wrap_model_call, wrap_tool_call
@before_model
def add_context(request):
request.messages.append({"role": "system", "content": "Extra"})
return request
@wrap_model_call
def retry_on_error(request, call_next):
for _ in range(3):
try:
return call_next(request)
except Exception:
pass
raise RuntimeError("All retries failed")
agent = Agent(name="Test", hooks=[add_context, retry_on_error])
"""
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, TypeVar, Generic
from functools import wraps
@dataclass
class InvocationContext:
"""Context passed through middleware chain.
Contains identifiers for multi-agent safety and optional metadata.
"""
agent_id: str
run_id: str
session_id: str
tool_name: Optional[str] = None
model_name: Optional[str] = None
metadata: Dict[str, Any] = field(default_factory=dict)
@dataclass
class ModelRequest:
"""Request data for model calls."""
messages: List[Dict[str, Any]]
model: str
temperature: float = 1.0
context: Optional[InvocationContext] = None
tools: Optional[List[Any]] = None
extra: Dict[str, Any] = field(default_factory=dict)
@dataclass
class ModelResponse:
"""Response data from model calls."""
content: str
model: str
usage: Optional[Dict[str, int]] = None
context: Optional[InvocationContext] = None
tool_calls: Optional[List[Any]] = None
extra: Dict[str, Any] = field(default_factory=dict)
@dataclass
class ToolRequest:
"""Request data for tool calls."""
tool_name: str
arguments: Dict[str, Any]
context: Optional[InvocationContext] = None
extra: Dict[str, Any] = field(default_factory=dict)
@dataclass
class ToolResponse:
"""Response data from tool calls."""
tool_name: str
result: Any
error: Optional[str] = None
context: Optional[InvocationContext] = None
extra: Dict[str, Any] = field(default_factory=dict)
# Type aliases for middleware functions
T = TypeVar('T')
BeforeModelFn = Callable[[ModelRequest], ModelRequest]
AfterModelFn = Callable[[ModelResponse], ModelResponse]
WrapModelCallFn = Callable[[ModelRequest, Callable[[ModelRequest], ModelResponse]], ModelResponse]
BeforeToolFn = Callable[[ToolRequest], ToolRequest]
AfterToolFn = Callable[[ToolResponse], ToolResponse]
WrapToolCallFn = Callable[[ToolRequest, Callable[[ToolRequest], ToolResponse]], ToolResponse]
class MiddlewareChain:
"""Synchronous middleware chain executor.
Composes middleware functions in registration order.
First registered middleware wraps all subsequent ones.
"""
def __init__(self, middlewares: Optional[List[Callable]] = None):
self._middlewares = middlewares or []
def execute(self, request: Any, final_handler: Callable[[Any], Any]) -> Any:
"""Execute the middleware chain with a final handler.
Args:
request: The request object to pass through
final_handler: The actual operation to perform
Returns:
The result from the chain
"""
if not self._middlewares:
return final_handler(request)
# Build the chain from inside out
# Last middleware wraps final_handler
# First middleware is outermost
def build_chain(index: int) -> Callable[[Any], Any]:
if index >= len(self._middlewares):
return final_handler
middleware = self._middlewares[index]
next_handler = build_chain(index + 1)
def handler(req: Any) -> Any:
return middleware(req, next_handler)
return handler
chain = build_chain(0)
return chain(request)
def __len__(self) -> int:
return len(self._middlewares)
class AsyncMiddlewareChain:
"""Asynchronous middleware chain executor.
Same as MiddlewareChain but for async middleware and handlers.
"""
def __init__(self, middlewares: Optional[List[Callable]] = None):
self._middlewares = middlewares or []
async def execute(self, request: Any, final_handler: Callable) -> Any:
"""Execute the async middleware chain.
Args:
request: The request object to pass through
final_handler: The actual async operation to perform
Returns:
The result from the chain
"""
if not self._middlewares:
return await final_handler(request)
# Build the chain from inside out
def build_chain(index: int) -> Callable:
if index >= len(self._middlewares):
return final_handler
middleware = self._middlewares[index]
next_handler = build_chain(index + 1)
async def handler(req: Any) -> Any:
return await middleware(req, next_handler)
return handler
chain = build_chain(0)
return await chain(request)
def __len__(self) -> int:
return len(self._middlewares)
# Decorator functions for tagging middleware
def before_model(func: BeforeModelFn) -> BeforeModelFn:
"""Decorator to mark a function as a before_model hook.
The function receives a ModelRequest and should return a (possibly modified) ModelRequest.
Example:
@before_model
def add_context(request):
request.messages.append({"role": "system", "content": "Extra"})
return request
"""
func._hook_type = 'before_model'
return func
def after_model(func: AfterModelFn) -> AfterModelFn:
"""Decorator to mark a function as an after_model hook.
The function receives a ModelResponse and should return a (possibly modified) ModelResponse.
Example:
@after_model
def log_response(response):
print(f"Model returned: {response.content[:50]}")
return response
"""
func._hook_type = 'after_model'
return func
def wrap_model_call(func: WrapModelCallFn) -> WrapModelCallFn:
"""Decorator to mark a function as a wrap_model_call middleware.
The function receives (request, call_next) and should call call_next(request)
to continue the chain, or return early to short-circuit.
Example:
@wrap_model_call
def retry_on_error(request, call_next):
for _ in range(3):
try:
return call_next(request)
except Exception:
pass
raise RuntimeError("All retries failed")
"""
func._hook_type = 'wrap_model_call'
return func
def before_tool(func: BeforeToolFn) -> BeforeToolFn:
"""Decorator to mark a function as a before_tool hook.
The function receives a ToolRequest and should return a (possibly modified) ToolRequest.
Example:
@before_tool
def validate_args(request):
if 'dangerous' in request.arguments:
raise ValueError("Dangerous argument detected")
return request
"""
func._hook_type = 'before_tool'
return func
def after_tool(func: AfterToolFn) -> AfterToolFn:
"""Decorator to mark a function as an after_tool hook.
The function receives a ToolResponse and should return a (possibly modified) ToolResponse.
Example:
@after_tool
def log_result(response):
print(f"Tool {response.tool_name} returned: {response.result}")
return response
"""
func._hook_type = 'after_tool'
return func
def wrap_tool_call(func: WrapToolCallFn) -> WrapToolCallFn:
"""Decorator to mark a function as a wrap_tool_call middleware.
The function receives (request, call_next) and should call call_next(request)
to continue the chain, or return early to short-circuit.
Example:
@wrap_tool_call
def retry_flaky_tool(request, call_next):
last_error = None
for _ in range(3):
try:
return call_next(request)
except Exception as e:
last_error = e
raise last_error
"""
func._hook_type = 'wrap_tool_call'
return func
def get_hook_type(func: Callable) -> Optional[str]:
"""Get the hook type of a decorated function."""
return getattr(func, '_hook_type', None)
def is_middleware(func: Callable) -> bool:
"""Check if a function is a middleware (wrap_* type)."""
hook_type = get_hook_type(func)
return hook_type in ('wrap_model_call', 'wrap_tool_call')
def categorize_hooks(hooks: List[Callable]) -> Dict[str, List[Callable]]:
"""Categorize hooks by their type.
Returns:
Dict with keys: before_model, after_model, wrap_model_call,
before_tool, after_tool, wrap_tool_call
"""
categories = {
'before_model': [],
'after_model': [],
'wrap_model_call': [],
'before_tool': [],
'after_tool': [],
'wrap_tool_call': [],
}
for hook in hooks or []:
hook_type = get_hook_type(hook)
if hook_type and hook_type in categories:
categories[hook_type].append(hook)
return categories
class MiddlewareManager:
"""Manages middleware for an agent.
Provides methods to execute before/after hooks and middleware chains
for both model and tool calls.
"""
def __init__(self, hooks: Optional[List[Callable]] = None):
self._hooks = hooks or []
self._categorized = categorize_hooks(self._hooks)
# Pre-build chains for efficiency
self._model_chain = MiddlewareChain(self._categorized['wrap_model_call'])
self._tool_chain = MiddlewareChain(self._categorized['wrap_tool_call'])
@property
def has_model_hooks(self) -> bool:
"""Check if any model hooks are registered."""
return bool(
self._categorized['before_model'] or
self._categorized['after_model'] or
self._categorized['wrap_model_call']
)
@property
def has_tool_hooks(self) -> bool:
"""Check if any tool hooks are registered."""
return bool(
self._categorized['before_tool'] or
self._categorized['after_tool'] or
self._categorized['wrap_tool_call']
)
def run_before_model(self, request: ModelRequest) -> ModelRequest:
"""Run all before_model hooks."""
for hook in self._categorized['before_model']:
request = hook(request)
return request
def run_after_model(self, response: ModelResponse) -> ModelResponse:
"""Run all after_model hooks (in reverse order)."""
for hook in reversed(self._categorized['after_model']):
response = hook(response)
return response
def run_before_tool(self, request: ToolRequest) -> ToolRequest:
"""Run all before_tool hooks."""
for hook in self._categorized['before_tool']:
request = hook(request)
return request
def run_after_tool(self, response: ToolResponse) -> ToolResponse:
"""Run all after_tool hooks (in reverse order)."""
for hook in reversed(self._categorized['after_tool']):
response = hook(response)
return response
def execute_model_call(
self,
request: ModelRequest,
model_fn: Callable[[ModelRequest], ModelResponse]
) -> ModelResponse:
"""Execute a model call with all hooks and middleware.
Order:
1. before_model hooks (in order)
2. wrap_model_call middleware chain
3. actual model call
4. after_model hooks (in reverse order)
"""
# Fast path: no hooks
if not self.has_model_hooks:
return model_fn(request)
# Run before hooks
request = self.run_before_model(request)
# Build wrapped handler that includes after hooks
def wrapped_model_fn(req: ModelRequest) -> ModelResponse:
response = model_fn(req)
return self.run_after_model(response)
# Execute through middleware chain
return self._model_chain.execute(request, wrapped_model_fn)
def execute_tool_call(
self,
request: ToolRequest,
tool_fn: Callable[[ToolRequest], ToolResponse]
) -> ToolResponse:
"""Execute a tool call with all hooks and middleware.
Order:
1. before_tool hooks (in order)
2. wrap_tool_call middleware chain
3. actual tool call
4. after_tool hooks (in reverse order)
"""
# Fast path: no hooks
if not self.has_tool_hooks:
return tool_fn(request)
# Run before hooks
request = self.run_before_tool(request)
# Build wrapped handler that includes after hooks
def wrapped_tool_fn(req: ToolRequest) -> ToolResponse:
response = tool_fn(req)
return self.run_after_tool(response)
# Execute through middleware chain
return self._tool_chain.execute(request, wrapped_tool_fn)
@@ -27,6 +27,7 @@ from .model_router import (
TaskComplexity,
create_routing_agent
)
from .rate_limiter import RateLimiter
__all__ = [
"LLM",
@@ -46,5 +47,6 @@ __all__ = [
"ModelRouter",
"ModelProfile",
"TaskComplexity",
"create_routing_agent"
"create_routing_agent",
"RateLimiter"
]
@@ -0,0 +1,147 @@
"""
Rate Limiter for PraisonAI Agents.
Provides token bucket rate limiting for LLM API calls.
Zero Performance Impact:
- When rate_limiter is None, no overhead
- Uses monotonic clock for accuracy
- Injectable clock/sleep for deterministic testing
Usage:
from praisonaiagents.llm import RateLimiter
from praisonaiagents import Agent
limiter = RateLimiter(requests_per_minute=60)
agent = Agent(name="Test", rate_limiter=limiter)
"""
import time
import asyncio
from dataclasses import dataclass, field
from typing import Optional, Callable
@dataclass
class RateLimiter:
"""Token bucket rate limiter for API calls.
Uses the token bucket algorithm:
- Tokens are added at a fixed rate (requests_per_minute / 60 per second)
- Each request consumes one token
- If no tokens available, wait until one is available
- Burst allows multiple requests in quick succession
Args:
requests_per_minute: Maximum requests per minute
tokens_per_minute: Optional token-based rate limiting (for future use)
burst: Maximum burst size (default: 1)
Example:
limiter = RateLimiter(requests_per_minute=60) # 1 req/sec
limiter.acquire() # Blocks if rate exceeded
# Async usage
await limiter.acquire_async()
"""
requests_per_minute: int
tokens_per_minute: Optional[int] = None
burst: int = 1
# Internal state
_tokens: float = field(default=None, init=False, repr=False)
_last_update: float = field(default=None, init=False, repr=False)
_lock: asyncio.Lock = field(default=None, init=False, repr=False)
# Injectable functions for testing
_get_time: Callable[[], float] = field(default=None, init=False, repr=False)
_sleep: Callable[[float], None] = field(default=None, init=False, repr=False)
_async_sleep: Callable[[float], None] = field(default=None, init=False, repr=False)
def __post_init__(self):
"""Initialize internal state."""
self._tokens = float(self.burst)
self._last_update = time.monotonic()
self._lock = asyncio.Lock()
# Default implementations (can be overridden for testing)
if self._get_time is None:
self._get_time = time.monotonic
if self._sleep is None:
self._sleep = time.sleep
if self._async_sleep is None:
self._async_sleep = asyncio.sleep
@property
def _rate(self) -> float:
"""Tokens per second."""
return self.requests_per_minute / 60.0
def _refill(self) -> None:
"""Refill tokens based on elapsed time."""
now = self._get_time()
elapsed = now - self._last_update
self._tokens = min(self.burst, self._tokens + elapsed * self._rate)
self._last_update = now
def _wait_time(self) -> float:
"""Calculate time to wait for next token."""
if self._tokens >= 1.0:
return 0.0
return (1.0 - self._tokens) / self._rate
def acquire(self) -> None:
"""Acquire a token, blocking if necessary.
This is the synchronous version. Use acquire_async() for async code.
"""
self._refill()
wait = self._wait_time()
if wait > 0:
self._sleep(wait)
self._refill()
self._tokens -= 1.0
async def acquire_async(self) -> None:
"""Acquire a token asynchronously, waiting if necessary.
This is the async version. Use acquire() for sync code.
"""
async with self._lock:
self._refill()
wait = self._wait_time()
if wait > 0:
await self._async_sleep(wait)
self._refill()
self._tokens -= 1.0
def try_acquire(self) -> bool:
"""Try to acquire a token without blocking.
Returns:
True if token acquired, False if rate limit exceeded
"""
self._refill()
if self._tokens >= 1.0:
self._tokens -= 1.0
return True
return False
def reset(self) -> None:
"""Reset the rate limiter to initial state."""
self._tokens = float(self.burst)
self._last_update = self._get_time()
@property
def available_tokens(self) -> float:
"""Get current available tokens (for monitoring)."""
self._refill()
return self._tokens
def __repr__(self) -> str:
return f"RateLimiter(rpm={self.requests_per_minute}, burst={self.burst})"
@@ -2,6 +2,9 @@
from importlib import import_module
from typing import Any
# Export Injected type for tool state injection
from .injected import Injected
# Map of function names to their module and class (if any)
TOOL_MAPPINGS = {
# Direct functions
@@ -286,4 +289,4 @@ def __getattr__(name: str) -> Any:
method = getattr(_instances[class_name], name)
return method
__all__ = list(TOOL_MAPPINGS.keys())
__all__ = list(TOOL_MAPPINGS.keys()) + ['Injected']
@@ -14,6 +14,14 @@ Usage:
@tool(name="web_search", description="Search the internet")
def search(query: str, max_results: int = 5) -> list:
return [...]
# With injected state:
from praisonaiagents.tools import Injected
@tool
def my_tool(query: str, state: Injected[dict]) -> str:
'''Tool with injected state.'''
return f"session={state.get('session_id')}"
"""
import inspect
@@ -22,12 +30,19 @@ import logging
from typing import Any, Callable, Dict, Optional, Union, get_type_hints
from .base import BaseTool
from .injected import (
is_injected_type,
get_injected_params,
inject_state_into_kwargs,
filter_injected_from_schema
)
class FunctionTool(BaseTool):
"""A BaseTool wrapper for plain functions.
Created automatically by the @tool decorator.
Supports Injected[T] parameters for state injection.
"""
def __init__(
@@ -42,6 +57,9 @@ class FunctionTool(BaseTool):
self.description = description or func.__doc__ or f"Tool: {self.name}"
self.version = version
# Detect injected parameters
self._injected_params = get_injected_params(func)
# Generate schema from the original function, not run()
self.parameters = self._generate_schema_from_func(func)
@@ -55,8 +73,16 @@ class FunctionTool(BaseTool):
if not self.description:
self.description = self.__class__.__doc__ or f"Tool: {self.name}"
@property
def injected_params(self) -> Dict[str, Any]:
"""Get the injected parameters for this tool."""
return self._injected_params
def _generate_schema_from_func(self, func: Callable) -> Dict[str, Any]:
"""Generate JSON Schema from the wrapped function's signature."""
"""Generate JSON Schema from the wrapped function's signature.
Injected parameters are excluded from the schema.
"""
schema = {
"type": "object",
"properties": {},
@@ -71,8 +97,17 @@ class FunctionTool(BaseTool):
if param_name in ('self', 'cls'):
continue
# Skip injected parameters - they don't go in schema
if param_name in self._injected_params:
continue
# Get type hint
param_type = hints.get(param_name, Any)
# Double-check it's not an Injected type
if is_injected_type(param_type):
continue
json_type = BaseTool._python_type_to_json(param_type)
schema["properties"][param_name] = {"type": json_type}
@@ -86,11 +121,18 @@ class FunctionTool(BaseTool):
return schema
def run(self, **kwargs) -> Any:
"""Execute the wrapped function."""
"""Execute the wrapped function with injected state."""
# Inject state for any Injected parameters
kwargs = inject_state_into_kwargs(kwargs, self._injected_params)
return self._func(**kwargs)
def __call__(self, *args, **kwargs) -> Any:
"""Allow calling with positional args like the original function."""
"""Allow calling with positional args like the original function.
Injects state for Injected parameters.
"""
# Inject state for any Injected parameters
kwargs = inject_state_into_kwargs(kwargs, self._injected_params)
return self._func(*args, **kwargs)
@@ -0,0 +1,213 @@
"""
Injected State for PraisonAI Tools.
Provides a mechanism for tools to receive injected state (agent context)
without exposing it in the tool's public schema.
Zero Performance Impact:
- When no Injected params, no overhead
- Context is thread-local for safety
- Lazy evaluation of state
Usage:
from praisonaiagents import tool
from praisonaiagents.tools import Injected
@tool
def my_tool(query: str, state: Injected[dict]) -> str:
'''Tool with injected state.'''
session_id = state.get('session_id')
return f"Query: {query}, Session: {session_id}"
"""
from dataclasses import dataclass, field
from typing import Any, Dict, Generic, List, Optional, TypeVar, get_type_hints, get_origin, get_args
from contextvars import ContextVar
from contextlib import contextmanager
T = TypeVar('T')
class Injected(Generic[T]):
"""Type marker for injected parameters.
Parameters annotated with Injected[T] will:
- Not appear in the tool's public schema
- Be automatically injected at runtime with the current agent state
Supported types:
- Injected[dict] - Receives state as a dictionary
- Injected[AgentState] - Receives state as AgentState dataclass
Example:
@tool
def my_tool(query: str, state: Injected[dict]) -> str:
return f"session={state.get('session_id')}"
"""
pass
@dataclass
class AgentState:
"""State object injected into tools.
Contains context about the current agent execution.
"""
agent_id: str
run_id: str
session_id: str
last_user_message: Optional[str] = None
last_agent_message: Optional[str] = None
memory: Optional[Any] = None
previous_tool_results: List[Any] = field(default_factory=list)
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary for Injected[dict] usage."""
return {
'agent_id': self.agent_id,
'run_id': self.run_id,
'session_id': self.session_id,
'last_user_message': self.last_user_message,
'last_agent_message': self.last_agent_message,
'previous_tool_results': self.previous_tool_results,
'metadata': self.metadata,
}
# Thread-local context for current state
_current_state: ContextVar[Optional[AgentState]] = ContextVar('current_state', default=None)
def get_current_state() -> Optional[AgentState]:
"""Get the current injection state."""
return _current_state.get()
def set_current_state(state: Optional[AgentState]) -> None:
"""Set the current injection state."""
_current_state.set(state)
@contextmanager
def with_injection_context(state: AgentState):
"""Context manager for setting injection state.
Usage:
with with_injection_context(state):
result = my_tool(query="test")
"""
token = _current_state.set(state)
try:
yield
finally:
_current_state.reset(token)
def is_injected_type(annotation: Any) -> bool:
"""Check if a type annotation is Injected[T]."""
origin = get_origin(annotation)
if origin is Injected:
return True
# Handle string annotations
if isinstance(annotation, str) and 'Injected' in annotation:
return True
return False
def get_injected_type(annotation: Any) -> Optional[type]:
"""Get the inner type T from Injected[T]."""
origin = get_origin(annotation)
if origin is Injected:
args = get_args(annotation)
if args:
return args[0]
return None
def resolve_injected_value(annotation: Any, state: Optional[AgentState]) -> Any:
"""Resolve the value to inject based on annotation type.
Args:
annotation: The Injected[T] annotation
state: The current agent state
Returns:
The resolved value (dict or AgentState)
"""
if state is None:
# Return empty dict if no state available
return {}
inner_type = get_injected_type(annotation)
if inner_type is dict or inner_type is Dict:
return state.to_dict()
elif inner_type is AgentState:
return state
else:
# Default to dict
return state.to_dict()
def get_injected_params(func) -> Dict[str, Any]:
"""Get all injected parameters from a function.
Returns:
Dict mapping param name to its Injected[T] annotation
"""
try:
hints = get_type_hints(func)
except Exception:
# If we can't get hints, return empty
return {}
injected = {}
for name, hint in hints.items():
if name == 'return':
continue
if is_injected_type(hint):
injected[name] = hint
return injected
def filter_injected_from_schema(properties: Dict[str, Any], required: List[str],
injected_params: Dict[str, Any]) -> tuple:
"""Remove injected parameters from schema.
Args:
properties: Schema properties dict
required: List of required param names
injected_params: Dict of injected param names
Returns:
Tuple of (filtered_properties, filtered_required)
"""
filtered_props = {k: v for k, v in properties.items() if k not in injected_params}
filtered_required = [r for r in required if r not in injected_params]
return filtered_props, filtered_required
def inject_state_into_kwargs(kwargs: Dict[str, Any], injected_params: Dict[str, Any]) -> Dict[str, Any]:
"""Inject state values into kwargs for injected parameters.
Args:
kwargs: The original kwargs
injected_params: Dict of injected param names to annotations
Returns:
Updated kwargs with injected values
"""
if not injected_params:
return kwargs
state = get_current_state()
result = kwargs.copy()
for param_name, annotation in injected_params.items():
if param_name not in result:
result[param_name] = resolve_injected_value(annotation, state)
return result
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "praisonaiagents"
version = "0.2.2"
version = "0.3.0"
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
readme = "README.md"
requires-python = ">=3.10"
@@ -0,0 +1,234 @@
"""
TDD Tests for Configurable Model (runtime model switching).
These tests are written FIRST before implementation.
"""
import pytest
from unittest.mock import Mock, patch
class TestConfigurableModelBasic:
"""Test basic configurable model functionality."""
def test_agent_accepts_llm_config_configurable(self):
"""Agent should accept llm_config with configurable=True."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
assert agent is not None
assert getattr(agent, '_llm_configurable', False) is True
def test_agent_chat_accepts_config_parameter(self):
"""Agent.chat() should accept config parameter for model override."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
# Should not raise - config parameter should be accepted
# We'll mock the actual LLM call
with patch.object(agent, '_chat_completion', return_value="mocked"):
result = agent.chat(
"Hello",
config={"model": "claude-3-5-sonnet"}
)
def test_config_model_override(self):
"""config.model should override default model for that call."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
# Track which model was used
used_model = [None]
original_chat = agent._chat_completion
def mock_chat(*args, **kwargs):
# The model should be accessible somehow
used_model[0] = getattr(agent, '_current_model', agent.llm)
return "response"
with patch.object(agent, '_chat_completion', side_effect=mock_chat):
agent.chat("Hello", config={"model": "claude-3-5-sonnet"})
# Model should have been switched for this call
# Implementation will set _current_model or similar
def test_config_does_not_mutate_default(self):
"""Per-call config should not change agent's default model."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
original_llm = agent.llm
with patch.object(agent, '_chat_completion', return_value="mocked"):
agent.chat("Hello", config={"model": "claude-3-5-sonnet"})
# Default should be unchanged
assert agent.llm == original_llm
def test_config_temperature_override(self):
"""config.temperature should override for that call."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
# Should accept temperature in config
with patch.object(agent, '_chat_completion', return_value="mocked") as mock:
agent.chat("Hello", config={"temperature": 0.2})
# Temperature should be passed through
class TestConfigurableModelCaching:
"""Test model client caching for configurable models."""
def test_model_client_cached(self):
"""Model clients should be cached to avoid recreation."""
from praisonaiagents.llm import LLM
llm = LLM(model="gpt-4o", configurable=True)
# First call with model A
# Second call with model A should reuse client
# This is internal behavior - we test via call count or similar
def test_cache_size_limited(self):
"""Model client cache should have a size limit."""
from praisonaiagents.llm import LLM
llm = LLM(model="gpt-4o", configurable=True, cache_size=3)
# Should not grow unbounded
class TestConfigurableModelThreadSafety:
"""Test thread safety of configurable models."""
def test_concurrent_calls_different_models(self):
"""Concurrent calls with different models should not interfere."""
import threading
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
results = {}
errors = []
def call_with_model(model_name, thread_id):
try:
with patch.object(agent, '_chat_completion', return_value=f"response_{model_name}"):
result = agent.chat("Hello", config={"model": model_name})
results[thread_id] = result
except Exception as e:
errors.append(e)
threads = [
threading.Thread(target=call_with_model, args=("model_a", 1)),
threading.Thread(target=call_with_model, args=("model_b", 2)),
]
for t in threads:
t.start()
for t in threads:
t.join()
assert len(errors) == 0
class TestConfigurableModelValidation:
"""Test validation of configurable model settings."""
def test_config_without_configurable_flag_ignored(self):
"""If llm_config.configurable is False, config param should be ignored."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o"
# No llm_config or configurable=False
)
# Should still work, just ignore the config
with patch.object(agent, '_chat_completion', return_value="mocked"):
result = agent.chat("Hello", config={"model": "other"})
# Should use default model, not "other"
def test_invalid_config_keys_ignored(self):
"""Unknown keys in config should be ignored, not error."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
# Should not raise for unknown keys
with patch.object(agent, '_chat_completion', return_value="mocked"):
result = agent.chat("Hello", config={"unknown_key": "value"})
class TestConfigurableModelProvider:
"""Test provider switching in configurable models."""
def test_config_provider_override(self):
"""config.provider should allow switching providers."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
# Should accept provider in config
with patch.object(agent, '_chat_completion', return_value="mocked"):
agent.chat("Hello", config={"provider": "anthropic", "model": "claude-3-5-sonnet"})
def test_model_string_with_provider_prefix(self):
"""Model string like 'anthropic/claude-3' should work."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent",
llm="gpt-4o",
llm_config={"configurable": True}
)
with patch.object(agent, '_chat_completion', return_value="mocked"):
agent.chat("Hello", config={"model": "anthropic/claude-3-5-sonnet"})
@@ -0,0 +1,367 @@
"""
TDD Tests for InjectedState for Tools.
These tests are written FIRST before implementation.
"""
import pytest
from unittest.mock import Mock, patch
class TestInjectedTypeExists:
"""Test that Injected type marker exists."""
def test_injected_type_exists(self):
"""Injected type should exist in praisonaiagents.tools."""
from praisonaiagents.tools import Injected
# Should be a generic type marker
assert Injected is not None
def test_injected_generic_usage(self):
"""Injected should work as a generic type annotation."""
from praisonaiagents.tools import Injected
# Should be usable as Injected[dict], Injected[SomeType], etc.
def my_tool(query: str, state: Injected[dict]) -> str:
return query
# Should have annotations
annotations = my_tool.__annotations__
assert 'state' in annotations
def test_injected_with_agent_state(self):
"""Injected[AgentState] should work."""
from praisonaiagents.tools import Injected
from praisonaiagents.tools.injected import AgentState
def my_tool(query: str, state: Injected[AgentState]) -> str:
return query
annotations = my_tool.__annotations__
assert 'state' in annotations
class TestAgentStateType:
"""Test AgentState type for injection."""
def test_agent_state_exists(self):
"""AgentState type should exist."""
from praisonaiagents.tools.injected import AgentState
assert AgentState is not None
def test_agent_state_has_required_fields(self):
"""AgentState should have agent_id, run_id, session_id."""
from praisonaiagents.tools.injected import AgentState
state = AgentState(
agent_id="agent-123",
run_id="run-456",
session_id="session-789"
)
assert state.agent_id == "agent-123"
assert state.run_id == "run-456"
assert state.session_id == "session-789"
def test_agent_state_optional_fields(self):
"""AgentState should have optional fields."""
from praisonaiagents.tools.injected import AgentState
state = AgentState(
agent_id="a",
run_id="r",
session_id="s",
last_user_message="Hello",
last_agent_message="Hi there",
memory=Mock(),
previous_tool_results=["result1", "result2"]
)
assert state.last_user_message == "Hello"
assert state.last_agent_message == "Hi there"
assert state.memory is not None
assert len(state.previous_tool_results) == 2
class TestToolDecoratorInjection:
"""Test @tool decorator handles Injected parameters."""
def test_tool_decorator_detects_injected(self):
"""@tool should detect Injected parameters."""
from praisonaiagents import tool
from praisonaiagents.tools import Injected
@tool
def my_tool(query: str, state: Injected[dict]) -> str:
"""A tool with injected state."""
return f"query={query}, session={state.get('session_id')}"
# Should have metadata about injected params
assert hasattr(my_tool, '_injected_params') or hasattr(my_tool, 'injected_params')
def test_injected_params_not_in_schema(self):
"""Injected parameters should not appear in tool schema."""
from praisonaiagents import tool
from praisonaiagents.tools import Injected
@tool
def my_tool(query: str, state: Injected[dict]) -> str:
"""A tool with injected state."""
return query
schema = my_tool.get_schema()
params = schema['function']['parameters']['properties']
# 'query' should be in schema
assert 'query' in params
# 'state' should NOT be in schema (it's injected)
assert 'state' not in params
def test_injected_params_not_required(self):
"""Injected parameters should not be in required list."""
from praisonaiagents import tool
from praisonaiagents.tools import Injected
@tool
def my_tool(query: str, state: Injected[dict]) -> str:
"""A tool with injected state."""
return query
schema = my_tool.get_schema()
required = schema['function']['parameters'].get('required', [])
assert 'state' not in required
assert 'query' in required
class TestToolExecutionInjection:
"""Test that injected state is provided at runtime."""
def test_injected_state_provided_at_runtime(self):
"""When tool is called, injected state should be provided."""
from praisonaiagents import tool
from praisonaiagents.tools import Injected
from praisonaiagents.tools.injected import AgentState
received_state = [None]
@tool
def my_tool(query: str, state: Injected[AgentState]) -> str:
"""A tool with injected state."""
received_state[0] = state
return f"session={state.session_id}"
# Simulate execution with injection context
from praisonaiagents.tools.injected import with_injection_context
mock_state = AgentState(
agent_id="a",
run_id="r",
session_id="test-session"
)
with with_injection_context(mock_state):
result = my_tool(query="hello")
assert received_state[0] is not None
assert received_state[0].session_id == "test-session"
def test_injected_dict_state(self):
"""Injected[dict] should provide dict with state info."""
from praisonaiagents import tool
from praisonaiagents.tools import Injected
from praisonaiagents.tools.injected import with_injection_context, AgentState
received_state = [None]
@tool
def my_tool(query: str, state: Injected[dict]) -> str:
"""A tool with injected dict state."""
received_state[0] = state
return "ok"
mock_state = AgentState(
agent_id="a",
run_id="r",
session_id="s"
)
with with_injection_context(mock_state):
my_tool(query="test")
assert isinstance(received_state[0], dict)
assert 'session_id' in received_state[0]
class TestAgentToolExecution:
"""Test Agent executes tools with injection."""
def test_agent_injects_state_to_tools(self):
"""Agent should inject state when executing tools."""
from praisonaiagents import Agent, tool
from praisonaiagents.tools import Injected
received_session = [None]
@tool
def show_state(query: str, state: Injected[dict]) -> str:
"""Show the injected state."""
received_session[0] = state.get('session_id')
return f"session={state.get('session_id')}"
agent = Agent(
name="Test",
instructions="Test agent",
tools=[show_state],
session_id="my-session-123"
)
# Execute the tool directly through agent
result = agent.execute_tool("show_state", {"query": "test"})
# State should have been injected
assert received_session[0] == "my-session-123"
class TestInjectionContextManager:
"""Test injection context manager."""
def test_with_injection_context_exists(self):
"""with_injection_context should exist."""
from praisonaiagents.tools.injected import with_injection_context
assert with_injection_context is not None
def test_context_manager_sets_state(self):
"""Context manager should set current state."""
from praisonaiagents.tools.injected import (
with_injection_context,
get_current_state,
AgentState
)
state = AgentState(agent_id="a", run_id="r", session_id="s")
with with_injection_context(state):
current = get_current_state()
assert current is state
# After context, should be None
assert get_current_state() is None
def test_nested_contexts(self):
"""Nested contexts should work correctly."""
from praisonaiagents.tools.injected import (
with_injection_context,
get_current_state,
AgentState
)
state1 = AgentState(agent_id="a1", run_id="r1", session_id="s1")
state2 = AgentState(agent_id="a2", run_id="r2", session_id="s2")
with with_injection_context(state1):
assert get_current_state().agent_id == "a1"
with with_injection_context(state2):
assert get_current_state().agent_id == "a2"
# Back to state1
assert get_current_state().agent_id == "a1"
class TestAsyncInjection:
"""Test async tool injection."""
@pytest.mark.asyncio
async def test_async_tool_injection(self):
"""Async tools should also receive injected state."""
from praisonaiagents import tool
from praisonaiagents.tools import Injected
from praisonaiagents.tools.injected import with_injection_context, AgentState
received_state = [None]
@tool
async def async_tool(query: str, state: Injected[dict]) -> str:
"""Async tool with injected state."""
received_state[0] = state
return "ok"
state = AgentState(agent_id="a", run_id="r", session_id="async-session")
with with_injection_context(state):
await async_tool(query="test")
assert received_state[0] is not None
assert received_state[0].get('session_id') == "async-session"
class TestInjectionZeroOverhead:
"""Test that injection has zero overhead when not used."""
def test_no_injected_params_no_overhead(self):
"""Tools without Injected params should have no overhead."""
from praisonaiagents import tool
@tool
def simple_tool(query: str) -> str:
"""Simple tool without injection."""
return query
# Should work normally without any injection context
result = simple_tool(query="hello")
assert result == "hello"
def test_schema_generation_fast(self):
"""Schema generation should be fast even with Injected params."""
from praisonaiagents import tool
from praisonaiagents.tools import Injected
import time
@tool
def my_tool(a: str, b: int, c: Injected[dict], d: float) -> str:
"""Tool with mixed params."""
return a
start = time.monotonic()
for _ in range(100):
my_tool.get_schema()
elapsed = time.monotonic() - start
# Should be very fast
assert elapsed < 0.1 # 100 calls in < 100ms
class TestBackwardCompatibility:
"""Test backward compatibility with existing tools."""
def test_existing_tools_still_work(self):
"""Existing tools without Injected should continue to work."""
from praisonaiagents import tool
@tool
def legacy_tool(query: str, count: int = 5) -> str:
"""Legacy tool."""
return f"{query} x {count}"
result = legacy_tool(query="test", count=3)
assert result == "test x 3"
def test_existing_tool_schema_unchanged(self):
"""Existing tool schemas should be unchanged."""
from praisonaiagents import tool
@tool
def legacy_tool(query: str, count: int = 5) -> str:
"""Legacy tool."""
return f"{query} x {count}"
schema = legacy_tool.get_schema()
params = schema['function']['parameters']['properties']
assert 'query' in params
assert 'count' in params
assert len(params) == 2
@@ -0,0 +1,403 @@
"""
TDD Tests for Middleware System (wrap_model_call, wrap_tool_call).
These tests are written FIRST before implementation.
"""
import pytest
from unittest.mock import Mock, MagicMock, patch, AsyncMock
from dataclasses import dataclass
from typing import Any, Dict, Optional
class TestMiddlewareTypes:
"""Test middleware type definitions."""
def test_invocation_context_exists(self):
"""InvocationContext dataclass should exist with required fields."""
from praisonaiagents.hooks.middleware import InvocationContext
ctx = InvocationContext(
agent_id="agent-123",
run_id="run-456",
session_id="session-789"
)
assert ctx.agent_id == "agent-123"
assert ctx.run_id == "run-456"
assert ctx.session_id == "session-789"
def test_invocation_context_optional_fields(self):
"""InvocationContext should have optional tool_name, model_name, metadata."""
from praisonaiagents.hooks.middleware import InvocationContext
ctx = InvocationContext(
agent_id="a",
run_id="r",
session_id="s",
tool_name="my_tool",
model_name="gpt-4o",
metadata={"key": "value"}
)
assert ctx.tool_name == "my_tool"
assert ctx.model_name == "gpt-4o"
assert ctx.metadata == {"key": "value"}
def test_model_request_exists(self):
"""ModelRequest dataclass for before_model hooks."""
from praisonaiagents.hooks.middleware import ModelRequest
req = ModelRequest(
messages=[{"role": "user", "content": "Hello"}],
model="gpt-4o",
temperature=0.7,
context=Mock()
)
assert req.messages == [{"role": "user", "content": "Hello"}]
assert req.model == "gpt-4o"
assert req.temperature == 0.7
def test_model_response_exists(self):
"""ModelResponse dataclass for after_model hooks."""
from praisonaiagents.hooks.middleware import ModelResponse
resp = ModelResponse(
content="Hello back!",
model="gpt-4o",
usage={"prompt_tokens": 10, "completion_tokens": 5},
context=Mock()
)
assert resp.content == "Hello back!"
assert resp.model == "gpt-4o"
def test_tool_request_exists(self):
"""ToolRequest dataclass for before_tool hooks."""
from praisonaiagents.hooks.middleware import ToolRequest
req = ToolRequest(
tool_name="get_weather",
arguments={"city": "Paris"},
context=Mock()
)
assert req.tool_name == "get_weather"
assert req.arguments == {"city": "Paris"}
def test_tool_response_exists(self):
"""ToolResponse dataclass for after_tool hooks."""
from praisonaiagents.hooks.middleware import ToolResponse
resp = ToolResponse(
tool_name="get_weather",
result="Sunny, 22°C",
error=None,
context=Mock()
)
assert resp.result == "Sunny, 22°C"
assert resp.error is None
class TestMiddlewareDecorators:
"""Test middleware decorator functions."""
def test_before_model_decorator_exists(self):
"""before_model decorator should exist and tag functions."""
from praisonaiagents.hooks import before_model
@before_model
def my_hook(request):
return request
assert hasattr(my_hook, '_hook_type')
assert my_hook._hook_type == 'before_model'
def test_after_model_decorator_exists(self):
"""after_model decorator should exist and tag functions."""
from praisonaiagents.hooks import after_model
@after_model
def my_hook(response):
return response
assert hasattr(my_hook, '_hook_type')
assert my_hook._hook_type == 'after_model'
def test_wrap_model_call_decorator_exists(self):
"""wrap_model_call decorator should exist and tag functions."""
from praisonaiagents.hooks import wrap_model_call
@wrap_model_call
def my_middleware(request, call_next):
return call_next(request)
assert hasattr(my_middleware, '_hook_type')
assert my_middleware._hook_type == 'wrap_model_call'
def test_before_tool_decorator_exists(self):
"""before_tool decorator should exist and tag functions."""
from praisonaiagents.hooks import before_tool
@before_tool
def my_hook(request):
return request
assert hasattr(my_hook, '_hook_type')
assert my_hook._hook_type == 'before_tool'
def test_after_tool_decorator_exists(self):
"""after_tool decorator should exist and tag functions."""
from praisonaiagents.hooks import after_tool
@after_tool
def my_hook(response):
return response
assert hasattr(my_hook, '_hook_type')
assert my_hook._hook_type == 'after_tool'
def test_wrap_tool_call_decorator_exists(self):
"""wrap_tool_call decorator should exist and tag functions."""
from praisonaiagents.hooks import wrap_tool_call
@wrap_tool_call
def my_middleware(request, call_next):
return call_next(request)
assert hasattr(my_middleware, '_hook_type')
assert my_middleware._hook_type == 'wrap_tool_call'
class TestMiddlewareChaining:
"""Test middleware chain composition."""
def test_wrap_model_call_chain_order(self):
"""wrap_model_call hooks should compose in registration order."""
from praisonaiagents.hooks.middleware import MiddlewareChain
calls = []
def mw1(req, call_next):
calls.append("mw1_before")
result = call_next(req)
calls.append("mw1_after")
return result
def mw2(req, call_next):
calls.append("mw2_before")
result = call_next(req)
calls.append("mw2_after")
return result
def final_handler(req):
calls.append("final")
return "response"
chain = MiddlewareChain([mw1, mw2])
result = chain.execute({"test": True}, final_handler)
assert calls == ["mw1_before", "mw2_before", "final", "mw2_after", "mw1_after"]
assert result == "response"
def test_empty_middleware_chain(self):
"""Empty middleware chain should just call final handler."""
from praisonaiagents.hooks.middleware import MiddlewareChain
def final_handler(req):
return "direct"
chain = MiddlewareChain([])
result = chain.execute({"test": True}, final_handler)
assert result == "direct"
def test_middleware_can_modify_request(self):
"""Middleware should be able to modify request before passing on."""
from praisonaiagents.hooks.middleware import MiddlewareChain
def add_context(req, call_next):
req["extra"] = "added"
return call_next(req)
def final_handler(req):
return req.get("extra")
chain = MiddlewareChain([add_context])
result = chain.execute({"original": True}, final_handler)
assert result == "added"
def test_middleware_can_modify_response(self):
"""Middleware should be able to modify response after call_next."""
from praisonaiagents.hooks.middleware import MiddlewareChain
def uppercase_response(req, call_next):
result = call_next(req)
return result.upper()
def final_handler(req):
return "hello"
chain = MiddlewareChain([uppercase_response])
result = chain.execute({}, final_handler)
assert result == "HELLO"
def test_middleware_can_short_circuit(self):
"""Middleware can return early without calling call_next."""
from praisonaiagents.hooks.middleware import MiddlewareChain
def block_middleware(req, call_next):
if req.get("block"):
return "blocked"
return call_next(req)
def final_handler(req):
return "reached"
chain = MiddlewareChain([block_middleware])
# Should be blocked
result = chain.execute({"block": True}, final_handler)
assert result == "blocked"
# Should pass through
result = chain.execute({"block": False}, final_handler)
assert result == "reached"
class TestMiddlewareRetry:
"""Test retry logic in middleware."""
def test_wrap_tool_call_retry_on_error(self):
"""wrap_tool_call should support retry patterns."""
from praisonaiagents.hooks.middleware import MiddlewareChain
attempt_count = [0]
def retry_middleware(req, call_next):
last_error = None
for i in range(3):
try:
return call_next(req)
except Exception as e:
last_error = e
raise last_error
def flaky_handler(req):
attempt_count[0] += 1
if attempt_count[0] < 3:
raise RuntimeError("flaky")
return "success"
chain = MiddlewareChain([retry_middleware])
result = chain.execute({}, flaky_handler)
assert result == "success"
assert attempt_count[0] == 3
class TestAsyncMiddleware:
"""Test async middleware support."""
@pytest.mark.asyncio
async def test_async_wrap_model_call(self):
"""Async middleware should work with async handlers."""
from praisonaiagents.hooks.middleware import AsyncMiddlewareChain
calls = []
async def async_mw(req, call_next):
calls.append("before")
result = await call_next(req)
calls.append("after")
return result
async def async_handler(req):
calls.append("handler")
return "async_result"
chain = AsyncMiddlewareChain([async_mw])
result = await chain.execute({}, async_handler)
assert calls == ["before", "handler", "after"]
assert result == "async_result"
@pytest.mark.asyncio
async def test_async_middleware_chain_order(self):
"""Async middleware should maintain order."""
from praisonaiagents.hooks.middleware import AsyncMiddlewareChain
order = []
async def mw1(req, call_next):
order.append(1)
result = await call_next(req)
order.append(-1)
return result
async def mw2(req, call_next):
order.append(2)
result = await call_next(req)
order.append(-2)
return result
async def handler(req):
order.append(0)
return "done"
chain = AsyncMiddlewareChain([mw1, mw2])
await chain.execute({}, handler)
assert order == [1, 2, 0, -2, -1]
class TestAgentMiddlewareIntegration:
"""Test middleware integration with Agent class."""
def test_agent_accepts_hooks_parameter(self):
"""Agent should accept hooks parameter."""
from praisonaiagents import Agent
from praisonaiagents.hooks import before_model
@before_model
def my_hook(req):
return req
# Should not raise
agent = Agent(
name="Test",
instructions="Test agent",
hooks=[my_hook]
)
assert agent is not None
def test_agent_hooks_empty_by_default(self):
"""Agent should have empty hooks by default."""
from praisonaiagents import Agent
agent = Agent(name="Test", instructions="Test")
# Should have no hooks or empty hooks list
hooks = getattr(agent, '_hooks', None) or getattr(agent, 'hooks', [])
assert hooks is None or len(hooks) == 0
class TestMiddlewareZeroOverhead:
"""Test that middleware has zero overhead when not used."""
def test_no_hooks_fast_path(self):
"""When no hooks registered, should take fast path."""
from praisonaiagents.hooks.middleware import MiddlewareChain
# Empty chain should be essentially a no-op
chain = MiddlewareChain([])
call_count = [0]
def handler(req):
call_count[0] += 1
return "result"
result = chain.execute({}, handler)
assert result == "result"
assert call_count[0] == 1
def test_middleware_chain_is_none_check(self):
"""Should handle None middleware list gracefully."""
from praisonaiagents.hooks.middleware import MiddlewareChain
chain = MiddlewareChain(None)
result = chain.execute({}, lambda r: "ok")
assert result == "ok"
@@ -0,0 +1,260 @@
"""
TDD Tests for Rate Limiter.
These tests are written FIRST before implementation.
"""
import pytest
import time
from unittest.mock import Mock, patch
class TestRateLimiterBasic:
"""Test basic rate limiter functionality."""
def test_rate_limiter_exists(self):
"""RateLimiter class should exist in praisonaiagents.llm."""
from praisonaiagents.llm import RateLimiter
limiter = RateLimiter(requests_per_minute=60)
assert limiter is not None
def test_rate_limiter_rpm_parameter(self):
"""RateLimiter should accept requests_per_minute."""
from praisonaiagents.llm import RateLimiter
limiter = RateLimiter(requests_per_minute=60)
assert limiter.requests_per_minute == 60
def test_rate_limiter_acquire_sync(self):
"""RateLimiter.acquire() should block if rate exceeded."""
from praisonaiagents.llm import RateLimiter
# 1 request per second = 60 per minute
limiter = RateLimiter(requests_per_minute=60)
# First acquire should be instant
start = time.monotonic()
limiter.acquire()
elapsed = time.monotonic() - start
assert elapsed < 0.1 # Should be nearly instant
def test_rate_limiter_blocks_when_exceeded(self):
"""RateLimiter should block when rate is exceeded."""
from praisonaiagents.llm import RateLimiter
# Very low rate: 6 per minute = 1 per 10 seconds
# But we'll use a mock clock to test without waiting
limiter = RateLimiter(requests_per_minute=6)
# Inject a mock clock for testing
mock_time = [0.0]
limiter._get_time = lambda: mock_time[0]
limiter._sleep = lambda s: mock_time.__setitem__(0, mock_time[0] + s)
# First request should pass
limiter.acquire()
# Second request should need to wait ~10 seconds
# (but our mock sleep just advances time)
limiter.acquire()
# Time should have advanced
assert mock_time[0] >= 10.0
def test_rate_limiter_token_bucket(self):
"""RateLimiter should use token bucket algorithm."""
from praisonaiagents.llm import RateLimiter
limiter = RateLimiter(requests_per_minute=60, burst=5)
# Reset internal state with mock time
mock_time = [0.0]
limiter._get_time = lambda: mock_time[0]
limiter._sleep = lambda s: mock_time.__setitem__(0, mock_time[0] + s)
limiter._tokens = 5.0 # Reset tokens to burst size
limiter._last_update = 0.0 # Reset last update
for _ in range(5):
limiter.acquire()
# Should not have waited much (burst allows 5 instant requests)
assert mock_time[0] < 1.0
class TestRateLimiterAsync:
"""Test async rate limiter functionality."""
@pytest.mark.asyncio
async def test_rate_limiter_acquire_async(self):
"""RateLimiter.acquire_async() should work with async."""
from praisonaiagents.llm import RateLimiter
limiter = RateLimiter(requests_per_minute=60)
# Should have async acquire method
await limiter.acquire_async()
@pytest.mark.asyncio
async def test_rate_limiter_async_blocks(self):
"""Async acquire should also respect rate limits."""
from praisonaiagents.llm import RateLimiter
import asyncio
limiter = RateLimiter(requests_per_minute=6)
# Mock async sleep
sleep_times = []
original_sleep = asyncio.sleep
async def mock_sleep(s):
sleep_times.append(s)
# Don't actually sleep in tests
with patch('asyncio.sleep', mock_sleep):
await limiter.acquire_async()
await limiter.acquire_async()
# Should have tried to sleep
class TestRateLimiterAgentIntegration:
"""Test rate limiter integration with Agent."""
def test_agent_accepts_rate_limiter(self):
"""Agent should accept rate_limiter parameter."""
from praisonaiagents import Agent
from praisonaiagents.llm import RateLimiter
limiter = RateLimiter(requests_per_minute=60)
agent = Agent(
name="Test",
instructions="Test agent",
rate_limiter=limiter
)
assert agent is not None
assert getattr(agent, '_rate_limiter', None) is limiter or \
getattr(agent, 'rate_limiter', None) is limiter
def test_agent_rate_limiter_applied_to_llm_calls(self):
"""Rate limiter should be applied before LLM calls."""
from praisonaiagents import Agent
from praisonaiagents.llm import RateLimiter
limiter = RateLimiter(requests_per_minute=60)
acquire_count = [0]
original_acquire = limiter.acquire
def counting_acquire():
acquire_count[0] += 1
return original_acquire()
limiter.acquire = counting_acquire
agent = Agent(
name="Test",
instructions="Test agent",
rate_limiter=limiter
)
with patch.object(agent, '_chat_completion', return_value="mocked"):
agent.chat("Hello")
# Rate limiter should have been called
assert acquire_count[0] >= 1
class TestRateLimiterZeroOverhead:
"""Test that rate limiter has zero overhead when not used."""
def test_no_rate_limiter_no_overhead(self):
"""When no rate limiter, should not add any overhead."""
from praisonaiagents import Agent
agent = Agent(
name="Test",
instructions="Test agent"
# No rate_limiter
)
# Should not have rate limiter
assert getattr(agent, '_rate_limiter', None) is None
def test_rate_limiter_none_check_fast(self):
"""Checking for None rate limiter should be O(1)."""
from praisonaiagents.llm import RateLimiter
# This is more of a design requirement than a test
# The implementation should do: if self._rate_limiter: self._rate_limiter.acquire()
class TestRateLimiterConfiguration:
"""Test rate limiter configuration options."""
def test_rate_limiter_default_burst(self):
"""RateLimiter should have sensible default burst."""
from praisonaiagents.llm import RateLimiter
limiter = RateLimiter(requests_per_minute=60)
# Default burst should be 1 or based on rpm
assert hasattr(limiter, 'burst') or hasattr(limiter, '_burst')
def test_rate_limiter_custom_burst(self):
"""RateLimiter should accept custom burst parameter."""
from praisonaiagents.llm import RateLimiter
limiter = RateLimiter(requests_per_minute=60, burst=10)
burst = getattr(limiter, 'burst', None) or getattr(limiter, '_burst', None)
assert burst == 10
def test_rate_limiter_tpm_optional(self):
"""RateLimiter should optionally accept tokens_per_minute."""
from praisonaiagents.llm import RateLimiter
# Should not raise
limiter = RateLimiter(
requests_per_minute=60,
tokens_per_minute=100000
)
assert limiter is not None
class TestRateLimiterDeterministic:
"""Test that rate limiter is deterministic for testing."""
def test_injectable_clock(self):
"""RateLimiter should allow injecting clock for testing."""
from praisonaiagents.llm import RateLimiter
mock_time = [0.0]
limiter = RateLimiter(requests_per_minute=60)
limiter._get_time = lambda: mock_time[0]
# Should use our mock time
t1 = limiter._get_time()
mock_time[0] = 100.0
t2 = limiter._get_time()
assert t1 == 0.0
assert t2 == 100.0
def test_injectable_sleep(self):
"""RateLimiter should allow injecting sleep for testing."""
from praisonaiagents.llm import RateLimiter
sleep_calls = []
limiter = RateLimiter(requests_per_minute=60)
limiter._sleep = lambda s: sleep_calls.append(s)
# Force a wait scenario
limiter._get_time = lambda: 0.0
limiter._tokens = 0 # No tokens available
limiter._last_update = 0.0
limiter.acquire()
# Should have called our mock sleep
# (actual behavior depends on implementation)
+1 -1
View File
@@ -2422,7 +2422,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.2.2"
version = "0.3.0"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },
+4
View File
@@ -261,6 +261,10 @@ npm install praisonai
| ↳ 100+ LLM Support | [Example](examples/python/providers/openai/openai_gpt4_example.py) | [📖](https://docs.praison.ai/models) |
| ↳ Callback Agents | [Example](examples/python/general/advanced-callback-systems.py) | [📖](https://docs.praison.ai/features/callbacks) |
| ↳ Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
| ↳ Middleware System | [Example](examples/middleware/basic_middleware.py) | [📖](https://docs.praison.ai/features/middleware) |
| ↳ Configurable Model | [Example](examples/middleware/configurable_model.py) | [📖](https://docs.praison.ai/features/configurable-model) |
| ↳ Rate Limiter | [Example](examples/middleware/rate_limiter.py) | [📖](https://docs.praison.ai/features/rate-limiter) |
| ↳ Injected Tool State | [Example](examples/middleware/injected_state.py) | [📖](https://docs.praison.ai/features/injected-state) |
| ↳ Shadow Git Checkpoints | [Example](#10-shadow-git-checkpoints) | [📖](https://docs.praison.ai/features/checkpoints) |
| ↳ Background Tasks | [Example](examples/background/basic_background.py) | [📖](https://docs.praison.ai/features/background-tasks) |
| ↳ Policy Engine | [Example](examples/policy/basic_policy.py) | [📖](https://docs.praison.ai/features/policy-engine) |
+2 -2
View File
@@ -3,8 +3,8 @@ class Praisonai < Formula
desc "AI tools for various AI applications"
homepage "https://github.com/MervinPraison/PraisonAI"
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.4.4.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.4.4.tar.gz | shasum -a 256`.split.first
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.5.0.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.5.0.tar.gz | shasum -a 256`.split.first
license "MIT"
depends_on "[email protected]"
@@ -290,10 +290,9 @@ def handle_run(args: list) -> bool:
# Run agent with persistence
try:
from praisonai.db import PraisonDB
from praisonaiagents import Agent
from praisonaiagents import Agent, db
db = PraisonDB(
db_instance = db(
database_url=config["conversation_url"],
state_url=config["state_url"],
knowledge_url=config["knowledge_url"],
@@ -302,7 +301,7 @@ def handle_run(args: list) -> bool:
agent = Agent(
name=config["agent_name"],
instructions=config["agent_instructions"],
db=db,
db=db_instance,
session_id=config["session_id"],
verbose=True
)
@@ -311,7 +310,7 @@ def handle_run(args: list) -> bool:
response = agent.chat(config["prompt"])
print(f"\nAgent: {response}")
db.close()
db_instance.close()
except Exception as e:
print(f"Error: {e}")
@@ -334,12 +333,12 @@ def handle_resume(args: list) -> bool:
return True
try:
from praisonai.db import PraisonDB
from praisonaiagents import db
db = PraisonDB(database_url=config["conversation_url"])
db_instance = db(database_url=config["conversation_url"])
# Get session history
history = db.on_agent_start(
history = db_instance.on_agent_start(
agent_name="Resume",
session_id=config["session_id"],
user_id=config["user_id"]
@@ -362,7 +361,7 @@ def handle_resume(args: list) -> bool:
agent = Agent(
name=config["agent_name"],
instructions=config["agent_instructions"],
db=db,
db=db_instance,
session_id=config["session_id"],
verbose=True
)
@@ -370,7 +369,7 @@ def handle_resume(args: list) -> bool:
response = agent.chat(config["continue_prompt"])
print(f"\nAgent: {response}")
db.close()
db_instance.close()
except Exception as e:
print(f"Error: {e}")
@@ -392,17 +391,17 @@ def handle_export(args: list) -> bool:
return True
try:
from praisonai.db import PraisonDB
from praisonaiagents import db
import json
db = PraisonDB(database_url=config["conversation_url"])
db_instance = db(database_url=config["conversation_url"])
# Export session
data = db.export_session(config["session_id"])
data = db_instance.export_session(config["session_id"])
if not data:
print(f"Error: Session {config['session_id']} not found")
db.close()
db_instance.close()
return True
# Determine output file
@@ -414,7 +413,7 @@ def handle_export(args: list) -> bool:
print(f"Exported session to: {output_file}")
print(f"Messages: {len(data.get('messages', []))}")
db.close()
db_instance.close()
except Exception as e:
print(f"Error: {e}")
@@ -437,20 +436,20 @@ def handle_import(args: list) -> bool:
return True
try:
from praisonai.db import PraisonDB
from praisonaiagents import db
import json
db = PraisonDB(database_url=config["conversation_url"])
db_instance = db(database_url=config["conversation_url"])
with open(input_file, 'r') as f:
data = json.loads(f.readline())
session_id = db.import_session(data)
session_id = db_instance.import_session(data)
print(f"Imported session: {session_id}")
print(f"Messages: {len(data.get('messages', []))}")
db.close()
db_instance.close()
except Exception as e:
print(f"Error: {e}")
@@ -469,15 +468,15 @@ def handle_migrate(args: list) -> bool:
return True
try:
from praisonai.db import PraisonDB
from praisonaiagents import db
from praisonai.persistence.migrations import MigrationManager
db = PraisonDB(
db_instance = db(
database_url=config["conversation_url"],
state_url=config["state_url"]
)
manager = MigrationManager(db)
manager = MigrationManager(db_instance)
# Check for --up or --down flags
direction = "up"
@@ -506,7 +505,7 @@ def handle_migrate(args: list) -> bool:
else:
print("No migrations to rollback")
db.close()
db_instance.close()
except Exception as e:
print(f"Error: {e}")
@@ -524,15 +523,15 @@ def handle_status(args: list) -> bool:
return True
try:
from praisonai.db import PraisonDB
from praisonaiagents import db
from praisonai.persistence.migrations import MigrationManager
db = PraisonDB(
db_instance = db(
database_url=config["conversation_url"],
state_url=config["state_url"]
)
manager = MigrationManager(db)
manager = MigrationManager(db_instance)
status = manager.get_migration_status()
print("\n=== Schema Status ===")
@@ -543,7 +542,7 @@ def handle_status(args: list) -> bool:
if status['pending_versions']:
print(f"Pending versions: {', '.join(status['pending_versions'])}")
db.close()
db_instance.close()
except Exception as e:
print(f"Error: {e}")
+9
View File
@@ -813,6 +813,15 @@ class PraisonAI:
parser.add_argument('--max-cost', type=float, help='Maximum total cost budget in USD')
parser.add_argument('--daemon', action='store_true', help=argparse.SUPPRESS) # Hidden flag for daemon mode
# Rate Limiter - control API request rate
parser.add_argument("--rpm", type=int, help="Rate limit: requests per minute for LLM calls")
parser.add_argument("--tpm", type=int, help="Rate limit: tokens per minute for LLM calls (optional)")
# Configurable Model - runtime model switching
parser.add_argument("--configurable-model", action="store_true", help="Enable runtime model switching via config parameter")
parser.add_argument("--temperature", type=float, help="Override temperature for LLM calls")
parser.add_argument("--llm-provider", type=str, help="Override LLM provider (openai, anthropic, google, etc.)")
# If we're in a test environment, parse with empty args to avoid pytest interference
if in_test_env:
args, unknown_args = parser.parse_known_args([])
+1 -1
View File
@@ -57,7 +57,7 @@ class CloudDeployer:
file.write("FROM python:3.11-slim\n")
file.write("WORKDIR /app\n")
file.write("COPY . .\n")
file.write("RUN pip install flask praisonai==2.4.4 gunicorn markdown\n")
file.write("RUN pip install flask praisonai==2.5.0 gunicorn markdown\n")
file.write("EXPOSE 8080\n")
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
+1 -1
View File
@@ -1 +1 @@
__version__ = "2.4.4"
__version__ = "2.5.0"
+1 -1
View File
@@ -12,7 +12,7 @@ dependencies = [
"rich>=13.7",
"markdown>=3.5",
"pyparsing>=3.0.0",
"praisonaiagents>=0.2.2",
"praisonaiagents>=0.3.0",
"python-dotenv>=0.19.0",
"instructor>=1.3.3",
"PyYAML>=6.0",
@@ -0,0 +1,162 @@
"""
Tests for db import deprecation warnings.
These tests verify that:
1. New simplified imports work: from praisonaiagents import db
2. New simplified imports work: from praisonai import db
3. Old imports emit DeprecationWarning: from praisonai.db import PraisonDB
"""
import warnings
import pytest
class TestNewDbImports:
"""Test that new simplified db imports work correctly."""
def test_import_db_from_praisonaiagents(self):
"""Test: from praisonaiagents import db"""
from praisonaiagents import db
assert db is not None
assert hasattr(db, 'PraisonDB')
assert callable(db) # db(...) should work as shortcut
def test_import_db_from_praisonai(self):
"""Test: from praisonai import db
Note: This returns the praisonai.db module (for backwards compat).
The recommended pattern is: from praisonaiagents import db
"""
from praisonai import db
assert db is not None
# praisonai.db is a module with lazy-loaded classes
assert hasattr(db, 'PraisonDB')
def test_db_callable_shortcut(self):
"""Test that db(...) works as shortcut for db.PraisonDB(...)"""
from praisonaiagents import db
# Should not raise - just creates the adapter
# (actual DB connection is lazy)
instance = db(database_url="sqlite:///test.db")
assert instance is not None
def test_import_agent_and_db_together(self):
"""Test: from praisonaiagents import Agent, db"""
from praisonaiagents import Agent, db
assert Agent is not None
assert db is not None
def test_import_all_common_symbols(self):
"""Test: from praisonaiagents import Agent, Agents, Task, tool, db"""
from praisonaiagents import Agent, Agents, Task, tool, db
assert Agent is not None
assert Agents is not None
assert Task is not None
assert tool is not None
assert db is not None
class TestDeprecationWarnings:
"""Test that old imports emit DeprecationWarning."""
def test_praisondb_import_warns(self):
"""Test: from praisonai.db import PraisonDB emits warning"""
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
from praisonai.db import PraisonDB
# Should have at least one deprecation warning
deprecation_warnings = [
x for x in w
if issubclass(x.category, DeprecationWarning)
]
assert len(deprecation_warnings) >= 1
# Check warning message content
msg = str(deprecation_warnings[0].message)
assert "deprecated" in msg.lower()
assert "praisonaiagents" in msg or "praisonai" in msg
def test_postgresdb_import_warns(self):
"""Test: from praisonai.db import PostgresDB emits warning"""
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
from praisonai.db import PostgresDB
deprecation_warnings = [
x for x in w
if issubclass(x.category, DeprecationWarning)
]
assert len(deprecation_warnings) >= 1
def test_sqlitedb_import_warns(self):
"""Test: from praisonai.db import SQLiteDB emits warning"""
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
from praisonai.db import SQLiteDB
deprecation_warnings = [
x for x in w
if issubclass(x.category, DeprecationWarning)
]
assert len(deprecation_warnings) >= 1
def test_redisdb_import_warns(self):
"""Test: from praisonai.db import RedisDB emits warning"""
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
from praisonai.db import RedisDB
deprecation_warnings = [
x for x in w
if issubclass(x.category, DeprecationWarning)
]
assert len(deprecation_warnings) >= 1
class TestBackwardsCompatibility:
"""Test that old imports still work (just with warnings)."""
def test_praisondb_still_works(self):
"""Old PraisonDB import should still work."""
with warnings.catch_warnings():
warnings.simplefilter("ignore")
from praisonai.db import PraisonDB
# Should be usable
instance = PraisonDB(database_url="sqlite:///test.db")
assert instance is not None
def test_postgresdb_still_works(self):
"""Old PostgresDB import should still work."""
with warnings.catch_warnings():
warnings.simplefilter("ignore")
from praisonai.db import PostgresDB
# Should be a class
assert PostgresDB is not None
class TestNoHeavyImports:
"""Test that db import doesn't load heavy dependencies."""
def test_db_import_is_lightweight(self):
"""Importing db should not import heavy DB clients."""
import sys
# Clear any cached imports
modules_before = set(sys.modules.keys())
from praisonaiagents import db
modules_after = set(sys.modules.keys())
new_modules = modules_after - modules_before
# Should not have imported heavy DB clients
heavy_modules = ['psycopg2', 'pymysql', 'redis', 'qdrant_client']
for heavy in heavy_modules:
assert heavy not in new_modules, f"Heavy module {heavy} was imported"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
+4 -4
View File
@@ -4197,7 +4197,7 @@ requires-dist = [
{ name = "praisonai-tools", marker = "extra == 'autogen'", specifier = ">=0.0.22" },
{ name = "praisonai-tools", marker = "extra == 'autogen-v4'", specifier = ">=0.0.22" },
{ name = "praisonai-tools", marker = "extra == 'crewai'", specifier = ">=0.0.22" },
{ name = "praisonaiagents", specifier = ">=0.2.2" },
{ name = "praisonaiagents", specifier = ">=0.3.0" },
{ name = "pyautogen", marker = "extra == 'autogen'", specifier = "==0.2.29" },
{ name = "pydantic", marker = "extra == 'chat'", specifier = "<=2.10.1" },
{ name = "pydantic", marker = "extra == 'code'", specifier = "<=2.10.1" },
@@ -4252,7 +4252,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.2.2"
version = "0.3.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "aiohttp" },
@@ -4261,9 +4261,9 @@ dependencies = [
{ name = "pydantic" },
{ name = "rich" },
]
sdist = { url = "https://files.pythonhosted.org/packages/23/0d/860b0fa64c0074c7e82d6fdacdf567b2b69da58602aaf749c8612d20b61d/praisonaiagents-0.2.2.tar.gz", hash = "sha256:4174a92b5a38a5597bc748ce4a0aa692d66575331dcb1839acdae4d491bc7333", size = 733706, upload-time = "2025-12-26T00:59:53.844Z" }
sdist = { url = "https://files.pythonhosted.org/packages/fc/79/08c9f2d8c12f9521f9fcc3df2510bd0d09377fd68b27413c64e501c1d47b/praisonaiagents-0.3.0.tar.gz", hash = "sha256:6fac4b754ec924ba804ef09ebef9919f276861d9f6d8c66cc33721063cb0d04c", size = 741299, upload-time = "2025-12-26T02:18:18.235Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/b4/4d/e862f14fb05fbd9e470a1d7a8c0c9920688758ecc5dbffa9ec597ec690f1/praisonaiagents-0.2.2-py3-none-any.whl", hash = "sha256:0fdfd131cb0ce7859b8a2cd033da9d3cb84c9265a77cca3509cdbf70152612c8", size = 681243, upload-time = "2025-12-26T00:59:51.892Z" },
{ url = "https://files.pythonhosted.org/packages/18/4e/2165feb9897d9b17522d7653b9111e40262bcf5172fc6fb549e5f0b2413f/praisonaiagents-0.3.0-py3-none-any.whl", hash = "sha256:fdbe3cf290e274ffd518b932d0ad8fd516b1600324bde2bf1970c5efd26c4292", size = 689391, upload-time = "2025-12-26T02:18:16.247Z" },
]
[[package]]