Compare commits

...
2 Commits
Author SHA1 Message Date
Mervin Praison 57b637948f Release v3.3.1 2026-01-08 12:28:06 +00:00
Mervin Praison 757d402cae Release v3.3.0 2026-01-08 10:33:27 +00:00
84 changed files with 11495 additions and 851 deletions

No files matched your search

+2 -2
View File
@@ -1930,11 +1930,11 @@ agent = Agent(
auto_save="my-project"
)
# Load history from last 5 sessions
# Load history from past sessions via context management
agent = Agent(
name="Assistant",
memory=True,
history_in_context=5
context=True, # Enable context management for history
)
```
+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>=3.2.1" \
"praisonai>=3.3.1" \
"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>=3.2.1" \
"praisonai>=3.3.1" \
"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>=3.2.1" \
"praisonai>=3.3.1" \
"praisonai[ui]" \
"praisonai[crewai]"
+52 -49
View File
@@ -2,10 +2,10 @@
"""
Basic Background Tasks Example for PraisonAI Agents.
This example demonstrates how to use the background task system to:
1. Run tasks asynchronously
2. Track task progress
3. Cancel running tasks
This example demonstrates how to use the background task system with Agent:
1. Create an Agent with background task support
2. Run agent tasks asynchronously
3. Track task progress
4. Handle task completion
Usage:
@@ -13,74 +13,77 @@ Usage:
"""
import asyncio
import time
from praisonaiagents import Agent
from praisonaiagents.background import BackgroundRunner, BackgroundConfig, TaskStatus
async def main():
print("=" * 60)
print("Background Tasks Demo")
print("Agent-Centric Background Tasks Demo")
print("=" * 60)
# Create a background runner with config
config = BackgroundConfig(max_concurrent_tasks=3)
runner = BackgroundRunner(config=config)
runner = BackgroundRunner(config=BackgroundConfig(max_concurrent_tasks=3))
# Create an Agent with background task support
agent = Agent(
name="AsyncAssistant",
instructions="You are a helpful research assistant.",
background=runner
)
print("\n--- Agent with Background Support Created ---")
print(f"Agent: {agent.name}")
print(f"Background runner: {agent.background is not None}")
# Define some example tasks
def slow_task(duration: float, name: str) -> str:
"""A slow task that simulates work."""
print(f" 🔄 Task '{name}' starting (will take {duration}s)")
time.sleep(duration)
return f"Task '{name}' completed after {duration}s"
async def research_task(topic: str) -> str:
"""Simulate a research task."""
print(f" 🔄 Researching '{topic}'...")
await asyncio.sleep(1.5)
return f"Research on '{topic}' completed with 5 key findings"
async def async_task(duration: float, name: str) -> str:
"""An async task that simulates work."""
print(f" 🔄 Async task '{name}' starting (will take {duration}s)")
await asyncio.sleep(duration)
return f"Async task '{name}' completed after {duration}s"
async def analysis_task(data: str) -> str:
"""Simulate an analysis task."""
print(f" 🔄 Analyzing data...")
await asyncio.sleep(1.0)
return f"Analysis complete: {data}"
# Submit tasks
print("\n--- Submitting Tasks ---")
# Submit tasks via agent's background runner
print("\n--- Submitting Background Tasks ---")
task1 = await runner.submit(slow_task, args=(1.0, "quick-job"), name="task1")
print(f"✅ Submitted task 1: {task1.id[:8]}")
task1 = await agent.background.submit(
research_task,
args=("AI trends 2025",),
name="research_ai"
)
print(f"✅ Submitted research task: {task1.id[:8]}")
task2 = await runner.submit(slow_task, args=(2.0, "medium-job"), name="task2")
print(f"✅ Submitted task 2: {task2.id[:8]}")
task3 = await runner.submit(async_task, args=(1.5, "async-job"), name="task3")
print(f"✅ Submitted async task 3: {task3.id[:8]}")
task2 = await agent.background.submit(
analysis_task,
args=("market data",),
name="analyze_market"
)
print(f"✅ Submitted analysis task: {task2.id[:8]}")
# List all tasks
print("\n--- Task Status ---")
for task in runner.tasks:
for task in agent.background.tasks:
print(f" [{task.id[:8]}] {task.status.value}")
# Wait for task 1 to complete
print("\n--- Waiting for Task 1 ---")
await task1.wait(timeout=5.0)
print(f"✅ Task 1 result: {task1.result}")
# Check status of all tasks
print("\n--- Updated Task Status ---")
for task in [task1, task2, task3]:
status = "✅" if task.status == TaskStatus.COMPLETED else "🔄"
print(f" {status} [{task.id[:8]}] {task.status.value}")
# Wait for all remaining tasks
print("\n--- Waiting for All Tasks ---")
# Wait for tasks to complete
print("\n--- Waiting for Tasks ---")
await task1.wait(timeout=10.0)
await task2.wait(timeout=10.0)
await task3.wait(timeout=10.0)
# Final status
print("\n--- Final Results ---")
for task in [task1, task2, task3]:
if task.result:
print(f" ✅ {task.result}")
# Final results
print("\n--- Results ---")
for task in [task1, task2]:
status = "✅" if task.status == TaskStatus.COMPLETED else "❌"
print(f" {status} {task.result}")
# Cleanup
runner.clear_completed()
agent.background.clear_completed()
print("\n✅ Cleared completed tasks")
print("\n" + "=" * 60)
+35 -29
View File
@@ -2,10 +2,10 @@
"""
Basic Checkpoints Example for PraisonAI Agents.
This example demonstrates how to use the shadow git checkpointing system to:
1. Save checkpoints before making changes
2. List available checkpoints
3. View diffs between checkpoints
This example demonstrates how to use the shadow git checkpointing system with Agent:
1. Create an Agent with checkpoint support
2. Save checkpoints before making changes
3. List available checkpoints
4. Restore to a previous checkpoint
Usage:
@@ -16,6 +16,7 @@ import asyncio
import tempfile
from pathlib import Path
from praisonaiagents import Agent
from praisonaiagents.checkpoints import CheckpointService
@@ -23,7 +24,7 @@ async def main():
# Create a temporary workspace for demonstration
with tempfile.TemporaryDirectory() as workspace:
print("=" * 60)
print("Shadow Git Checkpointing Demo")
print("Agent-Centric Checkpointing Demo")
print("=" * 60)
print(f"\nWorkspace: {workspace}")
@@ -32,24 +33,35 @@ async def main():
test_file.write_text("# Initial content\nprint('Hello, World!')\n")
print(f"\n✅ Created initial file: {test_file.name}")
# Initialize the checkpoint service
service = CheckpointService(workspace_dir=workspace)
# Create checkpoint service
checkpoints = CheckpointService(workspace_dir=workspace)
# Create an Agent with checkpoint support
agent = Agent(
name="RefactorBot",
instructions="You are a code refactoring assistant.",
checkpoints=checkpoints
)
print("\n--- Agent with Checkpoint Support Created ---")
print(f"Agent: {agent.name}")
print(f"Checkpoints enabled: {agent.checkpoints is not None}")
try:
# Initialize the shadow git repository
print("\n--- Initializing Checkpoint Service ---")
await service.initialize()
await agent.checkpoints.initialize()
print("✅ Shadow git repository initialized")
# Save first checkpoint
print("\n--- Saving Checkpoint 1: Initial State ---")
result1 = await service.save("Initial state")
result1 = await agent.checkpoints.save("Initial state")
print(f"✅ Checkpoint saved: {result1.checkpoint.id[:8]}")
print(f" Message: {result1.checkpoint.message}")
# Make some changes
print("\n--- Making Changes ---")
test_file.write_text("# Modified content\nprint('Hello, PraisonAI!')\nprint('Checkpoints are awesome!')\n")
# Make some changes (simulating agent modifications)
print("\n--- Simulating Agent Changes ---")
test_file.write_text("# Refactored content\nprint('Hello, PraisonAI!')\nprint('Checkpoints are awesome!')\n")
print("✅ Modified example.py")
# Create a new file
@@ -58,38 +70,32 @@ async def main():
print("✅ Created utils.py")
# Save second checkpoint
print("\n--- Saving Checkpoint 2: After Changes ---")
result2 = await service.save("Added features")
print("\n--- Saving Checkpoint 2: After Refactoring ---")
result2 = await agent.checkpoints.save("Added features")
print(f"✅ Checkpoint saved: {result2.checkpoint.id[:8]}")
# List all checkpoints
print("\n--- Listing All Checkpoints ---")
checkpoints = await service.list_checkpoints()
for i, cp in enumerate(checkpoints, 1):
print(f" {i}. [{cp.id[:8]}] {cp.message} ({cp.timestamp})")
checkpoints_list = await agent.checkpoints.list_checkpoints()
for i, cp in enumerate(checkpoints_list, 1):
print(f" {i}. [{cp.id[:8]}] {cp.message}")
# Show diff
print("\n--- Current Diff (uncommitted changes) ---")
diff = await service.diff()
print("\n--- Current Diff ---")
diff = await agent.checkpoints.diff()
if diff.files:
for f in diff.files:
print(f" {f.status}: {f.path}")
else:
print(" No uncommitted changes")
# Make more changes to show diff
test_file.write_text("# Even more changes\nprint('This will be reverted')\n")
print("\n✅ Made more changes to example.py")
diff = await service.diff()
print("\n--- New Diff ---")
if diff.files:
for f in diff.files:
print(f" {f.status}: {f.path}")
# Make more changes
test_file.write_text("# Bad changes\nprint('This will be reverted')\n")
print("\n✅ Made problematic changes to example.py")
# Restore to checkpoint 2
print(f"\n--- Restoring to Checkpoint: {result2.checkpoint.id[:8]} ---")
await service.restore(result2.checkpoint.id)
await agent.checkpoints.restore(result2.checkpoint.id)
print("✅ Restored successfully!")
# Verify restoration
+74 -100
View File
@@ -2,140 +2,114 @@
"""
Basic Context Compaction Example for PraisonAI Agents.
This example demonstrates how to use context compaction to:
1. Manage conversation context length
2. Use different compaction strategies
3. Track compaction results
4. Preserve important messages
This example demonstrates how to use context compaction with Agent:
1. Create an Agent with context= parameter for compaction
2. Use different compaction strategies via ManagerConfig
3. Monitor compaction results
Usage:
python basic_compaction.py
"""
from praisonaiagents.compaction import (
ContextCompactor, CompactionConfig, CompactionStrategy, CompactionResult
)
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
def main():
print("=" * 60)
print("Context Compaction Demo")
print("Agent-Centric Context Compaction Demo")
print("=" * 60)
# ==========================================================================
# Sample Messages
# Available Strategies
# ==========================================================================
messages = [
{"role": "system", "content": "You are a helpful AI assistant."},
{"role": "user", "content": "Hello! Can you help me with Python?"},
{"role": "assistant", "content": "Of course! I'd be happy to help you with Python. What would you like to know?"},
{"role": "user", "content": "How do I read a file in Python?"},
{"role": "assistant", "content": "To read a file in Python, you can use the built-in open() function. Here's an example:\n\nwith open('filename.txt', 'r') as f:\n content = f.read()\n\nThis opens the file, reads its content, and automatically closes it."},
{"role": "user", "content": "What about writing to a file?"},
{"role": "assistant", "content": "To write to a file, use the 'w' mode:\n\nwith open('filename.txt', 'w') as f:\n f.write('Hello, World!')\n\nUse 'a' mode to append instead of overwriting."},
{"role": "user", "content": "Can you explain list comprehensions?"},
{"role": "assistant", "content": "List comprehensions are a concise way to create lists. Instead of:\n\nresult = []\nfor x in range(10):\n result.append(x * 2)\n\nYou can write:\n\nresult = [x * 2 for x in range(10)]"},
print("\n--- Available Compaction Strategies ---")
strategies = [
("truncate", "Remove oldest messages first"),
("sliding_window", "Keep most recent messages within limit"),
("summarize", "Replace old messages with summary"),
("smart", "Intelligently select messages to keep"),
]
print(f"\n--- Original Messages ---")
print(f" Message count: {len(messages)}")
for name, desc in strategies:
print(f" {name}: {desc}")
# ==========================================================================
# Basic Compaction
# Agent with Context Compaction via context= param
# ==========================================================================
print("\n--- Basic Compaction (Truncate Strategy) ---")
print("\n--- Creating Agent with Context Compaction ---")
compactor = ContextCompactor(
max_tokens=50, # Low limit to trigger compaction
strategy=CompactionStrategy.TRUNCATE,
preserve_system=True,
preserve_recent=2
# Use context= parameter with ManagerConfig for compaction
agent = Agent(
name="LongChat",
instructions="You are a helpful assistant for extended conversations.",
context=ManagerConfig(
auto_compact=True,
compact_threshold=0.8,
strategy="sliding_window",
)
)
# Check stats before
stats = compactor.get_stats(messages)
print(f" Total tokens: {stats['total_tokens']}")
print(f" Max tokens: {stats['max_tokens']}")
print(f" Needs compaction: {stats['needs_compaction']}")
# Compact
compacted, result = compactor.compact(messages)
print(f"\n After compaction:")
print(f" Messages: {len(messages)} -> {len(compacted)}")
print(f" Tokens: {result.original_tokens} -> {result.compacted_tokens}")
print(f" Saved: {result.tokens_saved} tokens")
print(f" Compression: {result.compression_ratio:.1%}")
print(f"Agent: {agent.name}")
print(f"Context manager: {agent.context_manager is not None}")
if agent.context_manager:
print(f"Auto-compact: {agent.context_manager.config.auto_compact}")
print(f"Threshold: {agent.context_manager.config.compact_threshold}")
print(f"Strategy: {agent.context_manager.config.strategy}")
# ==========================================================================
# Sliding Window Strategy
# Different Strategy Examples via context= param
# ==========================================================================
print("\n--- Sliding Window Strategy ---")
print("\n--- Strategy Examples ---")
compactor = ContextCompactor(
max_tokens=50,
strategy=CompactionStrategy.SLIDING,
preserve_system=True,
preserve_recent=3
# Sliding window strategy
sliding_agent = Agent(
name="SlidingAgent",
instructions="You are helpful.",
context=ManagerConfig(
auto_compact=True,
strategy="sliding_window",
)
)
print(f" Sliding: strategy={sliding_agent.context_manager.config.strategy}")
compacted, result = compactor.compact(messages)
print(f" Messages kept: {result.messages_kept}")
print(f" Messages removed: {result.messages_removed}")
print(f" Strategy: {result.strategy_used.value}")
# ==========================================================================
# Summarize Strategy
# ==========================================================================
print("\n--- Summarize Strategy ---")
compactor = ContextCompactor(
max_tokens=50,
strategy=CompactionStrategy.SUMMARIZE,
preserve_system=True,
preserve_recent=2
# Summarize strategy
summarize_agent = Agent(
name="SummarizeAgent",
instructions="You are helpful.",
context=ManagerConfig(
auto_compact=True,
strategy="summarize",
)
)
print(f" Summarize: strategy={summarize_agent.context_manager.config.strategy}")
compacted, result = compactor.compact(messages)
print(f" Messages kept: {result.messages_kept}")
print(f" Strategy: {result.strategy_used.value}")
# Show what messages remain
print(f"\n Remaining message roles:")
for msg in compacted:
role = msg.get("role", "unknown")
content_preview = msg.get("content", "")[:50]
print(f" [{role}] {content_preview}...")
# ==========================================================================
# No Compaction Needed
# ==========================================================================
print("\n--- No Compaction Needed ---")
compactor = ContextCompactor(
max_tokens=10000, # High limit
strategy=CompactionStrategy.TRUNCATE
# Smart strategy
smart_agent = Agent(
name="SmartAgent",
instructions="You are helpful.",
context=ManagerConfig(
auto_compact=True,
strategy="smart",
)
)
compacted, result = compactor.compact(messages)
print(f" Was compacted: {result.was_compacted}")
print(f" Messages unchanged: {len(compacted) == len(messages)}")
print(f" Smart: strategy={smart_agent.context_manager.config.strategy}")
# ==========================================================================
# Result Serialization
# Context Manager Stats Demo
# ==========================================================================
print("\n--- Result Serialization ---")
print("\n--- Context Manager Stats Demo ---")
compactor = ContextCompactor(max_tokens=50)
_, result = compactor.compact(messages)
data = result.to_dict()
print(f" Original tokens: {data['original_tokens']}")
print(f" Compacted tokens: {data['compacted_tokens']}")
print(f" Compression ratio: {data['compression_ratio']:.2f}")
# Access context manager stats
if agent.context_manager:
stats = agent.context_manager.get_stats()
print(f" Model: {stats.get('model', 'N/A')}")
print(f" Model limit: {stats.get('model_limit', 'N/A')}")
print(f" Output reserve: {stats.get('output_reserve', 'N/A')}")
print(f" Usable budget: {stats.get('usable_budget', 'N/A')}")
else:
print(" Context manager not initialized")
print("\n" + "=" * 60)
print("Demo Complete!")
+41 -56
View File
@@ -1,87 +1,72 @@
"""
Fast Context with Agent - Basic Example
Fast Context - Basic Example
This is the simplest way to use Fast Context with an Agent.
Just set fast_context=True and the agent will automatically
have access to rapid code search capabilities.
FastContext provides rapid code search capabilities.
Use it directly from the context module for code search.
Fast Context provides:
FastContext provides:
- 10-20x faster code search than traditional methods
- Parallel execution (8 concurrent searches)
- Automatic caching for repeated queries
- Seamless integration with Agent workflows
Note: FastContext is now accessed via the context module directly,
not as Agent parameters. For context management with agents,
use Agent(context=True) or Agent(context=ManagerConfig(...)).
"""
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
def main():
print("=" * 60)
print("Agent with Fast Context - Basic Example")
print("Agent with Context Management - Basic Example")
print("=" * 60)
# Create an agent with Fast Context enabled
# This is the simplest configuration
# Create an agent with context management enabled
# This is the recommended way to use context features
agent = Agent(
name="CodeAssistant",
instructions="You are a helpful code assistant.",
fast_context=True, # Enable Fast Context
fast_context_path=".", # Search in current directory
context=True, # Enable context management with defaults
verbose=False
)
print("\n✓ Agent created with Fast Context enabled")
print(f" - fast_context_enabled: {agent.fast_context_enabled}")
print(f" - fast_context_path: {agent.fast_context_path}")
print(f" - fast_context_model: {agent.fast_context_model}")
print("\n✓ Agent created with context management enabled")
print(f" - context_manager: {agent.context_manager is not None}")
if agent.context_manager:
print(f" - auto_compact: {agent.context_manager.config.auto_compact}")
print(f" - strategy: {agent.context_manager.config.strategy}")
# Use delegate_to_fast_context to search code
# For FastContext code search, use the module directly
print("\n" + "-" * 60)
print("Searching for 'class Agent' in the codebase...")
print("Using FastContext for code search...")
print("-" * 60)
context = agent.delegate_to_fast_context("class Agent")
if context:
print(f"\n✓ Found relevant code!")
print(f" Context length: {len(context)} characters")
print("\nPreview of found context:")
print("-" * 40)
# Show first 500 characters
preview = context[:500]
print(preview)
if len(context) > 500:
print("...")
else:
print("\n✗ No relevant code found")
# Another search example
print("\n" + "-" * 60)
print("Searching for 'def chat' in the codebase...")
print("-" * 60)
context = agent.delegate_to_fast_context("def chat")
if context:
print(f"\n✓ Found relevant code!")
print(f" Context length: {len(context)} characters")
else:
print("\n✗ No relevant code found")
# Access FastContext directly for more control
print("\n" + "-" * 60)
print("Direct FastContext access for advanced usage...")
print("-" * 60)
fc = agent.fast_context
if fc:
result = fc.search("import")
print(f"\n✓ Direct search for 'import':")
try:
from praisonaiagents.context.fast import FastContext
fc = FastContext(
workspace_path=".",
model="gpt-4o-mini",
max_turns=4,
max_parallel=8,
)
result = fc.search("class Agent")
print(f"\n✓ Search for 'class Agent':")
print(f" Files found: {result.total_files}")
print(f" Search time: {result.search_time_ms}ms")
print(f" From cache: {result.from_cache}")
if result.total_files > 0:
context = fc.get_context_for_agent("class Agent")
print(f" Context length: {len(context)} characters")
except ImportError:
print("\n⚠ FastContext not available (optional feature)")
except Exception as e:
print(f"\n⚠ FastContext search failed: {e}")
print("\n" + "=" * 60)
print("Fast Context makes code search instant for your agents!")
print("Context management is now unified via context= param!")
print("=" * 60)
+61 -66
View File
@@ -1,103 +1,98 @@
"""
Example 4: Agent Integration with Fast Context
Example 4: Agent Integration with Context Management
This example shows how to use Fast Context integrated with the Agent class
for automatic code search delegation.
This example shows how to use context management with the Agent class.
Context management is now unified via the context= parameter.
Benefits:
- Seamless integration with existing Agent workflows
- Automatic context injection for code-related queries
- Configurable via Agent parameters or environment variables
- Automatic token tracking and optimization
- Configurable via context= parameter with ManagerConfig
"""
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
WORKSPACE = "/Users/praison/praisonai-package/src/praisonai-agents"
def main():
print("=" * 70)
print("Agent Integration with Fast Context")
print("Agent Integration with Context Management")
print("=" * 70)
# Create an agent with Fast Context enabled
print("\n1. Creating Agent with Fast Context")
# Create an agent with context management enabled
print("\n1. Creating Agent with Context Management")
print("-" * 40)
config = ManagerConfig(
auto_compact=True,
compact_threshold=0.8,
strategy="smart",
monitor_enabled=True,
monitor_path="./context_debug.txt",
)
agent = Agent(
name="CodeAssistant",
instructions="You are a helpful code assistant that can search and analyze code.",
fast_context=True,
fast_context_path=WORKSPACE,
fast_context_model="gpt-4o-mini",
fast_context_max_turns=4,
fast_context_parallelism=8,
fast_context_timeout=30.0,
instructions="You are a helpful code assistant.",
context=config,
verbose=False
)
print(f" Agent name: {agent.name}")
print(f" Fast Context enabled: {agent.fast_context_enabled}")
print(f" Fast Context path: {agent.fast_context_path}")
print(f" Fast Context model: {agent.fast_context_model}")
print(f" Max turns: {agent.fast_context_max_turns}")
print(f" Parallelism: {agent.fast_context_parallelism}")
print(f" Context manager: {agent.context_manager is not None}")
if agent.context_manager:
print(f" Auto-compact: {agent.context_manager.config.auto_compact}")
print(f" Threshold: {agent.context_manager.config.compact_threshold}")
print(f" Strategy: {agent.context_manager.config.strategy}")
# Use delegate_to_fast_context for code search
print("\n2. Delegating Code Search to Fast Context")
# For FastContext code search, use the module directly
print("\n2. Using FastContext for Code Search")
print("-" * 40)
queries = [
"find the Agent class definition",
"search for authentication handlers",
"locate memory management code",
]
for query in queries:
print(f"\n Query: '{query}'")
context = agent.delegate_to_fast_context(query)
try:
from praisonaiagents.context.fast import FastContext
if context:
# Show preview of context
lines = context.split('\n')
preview_lines = lines[:5]
print(f" Found context ({len(context)} chars):")
for line in preview_lines:
print(f" {line[:60]}{'...' if len(line) > 60 else ''}")
if len(lines) > 5:
print(f" ... and {len(lines) - 5} more lines")
else:
print(" No relevant code found")
# Access the FastContext instance directly
print("\n3. Direct FastContext Access")
print("-" * 40)
fc = agent.fast_context
if fc:
result = fc.search("def chat")
print(f" Direct search for 'def chat':")
print(f" Files found: {result.total_files}")
print(f" Search time: {result.search_time_ms}ms")
fc = FastContext(
workspace_path=WORKSPACE,
model="gpt-4o-mini",
max_turns=4,
max_parallel=8,
timeout=30.0,
)
for f in result.files[:3]:
print(f" - {f.path}")
queries = [
"find the Agent class definition",
"search for chat method",
]
for query in queries:
print(f"\n Query: '{query}'")
result = fc.search(query)
print(f" Files found: {result.total_files}")
print(f" Search time: {result.search_time_ms}ms")
except ImportError:
print(" FastContext not available (optional feature)")
except Exception as e:
print(f" FastContext error: {e}")
# Show configuration options
print("\n4. Configuration Options")
print("\n3. Configuration Options")
print("-" * 40)
print(" Agent parameters:")
print(" - fast_context=True # Enable Fast Context")
print(" - fast_context_path='.' # Workspace path")
print(" - fast_context_model='gpt-4o-mini' # Search model")
print(" - fast_context_max_turns=4 # Max search turns")
print(" - fast_context_parallelism=8 # Parallel calls")
print(" - fast_context_timeout=30.0 # Timeout (seconds)")
print(" Agent context= parameter:")
print(" - context=True # Enable with defaults")
print(" - context=ManagerConfig(...) # Custom configuration")
print(" - context=False # Disabled (default)")
print("\n Environment variables:")
print(" - FAST_CONTEXT_MODEL")
print(" - FAST_CONTEXT_MAX_TURNS")
print("\n ManagerConfig options:")
print(" - auto_compact=True # Auto-optimize context")
print(" - compact_threshold=0.8 # Trigger at 80%")
print(" - strategy='smart' # Optimization strategy")
print(" - monitor_enabled=True # Write snapshots")
print(" - monitor_path='./context.txt' # Snapshot path")
print(" - FAST_CONTEXT_PARALLELISM")
print(" - FAST_CONTEXT_TIMEOUT")
print(" - FAST_CONTEXT_CACHE")
+110 -2
View File
@@ -2,11 +2,22 @@
"""
Context Budgeting Example
Demonstrates how to use the ContextBudgeter to allocate token budgets
across different context segments for various models.
Demonstrates how to use context budgeting with Agent via context= param,
and also shows low-level ContextBudgeter usage for advanced scenarios.
Agent-Centric Quick Start:
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
agent = Agent(
instructions="You are helpful.",
context=ManagerConfig(output_reserve=16000), # Custom output reserve
)
"""
from praisonaiagents import Agent
from praisonaiagents.context import (
ManagerConfig,
ContextBudgeter,
BudgetAllocation,
get_model_limit,
@@ -14,6 +25,31 @@ from praisonaiagents.context import (
)
def agent_centric_example():
"""Agent-centric usage - recommended approach."""
print("=" * 60)
print("Agent-Centric Context Budgeting")
print("=" * 60)
# Simple: Enable context management with defaults
agent = Agent(
instructions="You are a helpful assistant.",
context=True,
)
print(f"Agent with context=True created")
print(f"Context manager: {agent.context_manager is not None}")
# Custom: Specify output reserve
agent2 = Agent(
instructions="You are a code assistant.",
context=ManagerConfig(output_reserve=16000),
)
print(f"Agent with custom output_reserve created")
if agent2.context_manager:
stats = agent2.context_manager.get_stats()
print(f"Output reserve: {stats.get('output_reserve', 'N/A')}")
def main():
print("=" * 60)
print("Context Budgeting Example")
@@ -92,6 +128,78 @@ def main():
budget_dict = budgeter.to_dict()
print(f"Budget as dict: {budget_dict}")
# Example 6: Overflow handling strategies overview
print("\n6. Overflow Handling Strategies")
print("-" * 40)
strategies = [
("truncate", "Remove oldest messages first", "Fast, simple", "Loses early context"),
("sliding_window", "Keep N most recent messages", "Preserves recent", "Loses early context"),
("prune_tools", "Truncate old tool outputs", "Keeps messages", "May lose tool details"),
("summarize", "Replace old messages with summary", "Preserves meaning", "Slower, uses API"),
("smart", "Combine strategies intelligently", "Best balance", "More complex"),
]
print(f"{'Strategy':<15} | {'Description':<30} | {'Pros':<18} | {'Cons'}")
print("-" * 90)
for name, desc, pros, cons in strategies:
print(f"{name:<15} | {desc:<30} | {pros:<18} | {cons}")
# Example 7: Simulated overflow playbook
print("\n7. Overflow Playbook (Thresholds → Actions)")
print("-" * 40)
playbook = [
(0.70, "INFO", "Monitor usage, no action needed"),
(0.80, "NOTICE", "Consider optimization soon"),
(0.90, "WARNING", "Trigger auto-compact if enabled"),
(0.95, "CRITICAL", "Aggressive optimization required"),
(1.00, "OVERFLOW", "Immediate truncation to prevent API error"),
]
print(f"{'Usage':<8} | {'Level':<10} | {'Action'}")
print("-" * 60)
for threshold, level, action in playbook:
print(f"{threshold:>6.0%} | {level:<10} | {action}")
# Example 8: Strategy comparison table
print("\n8. Strategy Comparison")
print("-" * 40)
print("Strategy | Preserves | Loses | Best For")
print("-" * 70)
print("truncate | Recent msgs | Old msgs | Simple chats")
print("sliding_window | Last N msgs | Earlier msgs | Long conversations")
print("prune_tools | All messages | Tool details | Tool-heavy agents")
print("summarize | Context meaning| Exact wording | Important history")
print("smart | Balanced | Minimal | Production use")
# Example 9: Default strategy for interactive mode
print("\n9. Default Strategy: Interactive Mode")
print("-" * 40)
print("Interactive mode default: context=False (zero overhead)")
print("To enable context management in interactive mode:")
print(" praisonai chat --context # Enable with defaults")
print(" Or in code:")
print(" agent = Agent(instructions='...', context=True)")
print("")
print("When enabled, defaults are:")
print(" - auto_compact: True")
print(" - compact_threshold: 0.8 (80% usage)")
print(" - strategy: smart")
print(" - output_reserve: model-specific (8K-16K)")
# Example 10: Default strategy for auto-agents mode
print("\n10. Default Strategy: Auto-Agents Mode")
print("-" * 40)
print("Auto-agents mode default: context=False (zero overhead)")
print("To enable for multi-agent workflows:")
print(" from praisonaiagents import PraisonAIAgents")
print(" agents = PraisonAIAgents(agents=[...], context=True)")
print("")
print("Recommended config for long tasks:")
print(" context=ManagerConfig(")
print(" auto_compact=True,")
print(" compact_threshold=0.7, # Earlier trigger for safety")
print(" strategy='smart',")
print(" output_reserve=16384, # Larger reserve for complex outputs")
print(" )")
print("\n" + "=" * 60)
print("✓ Context budgeting examples complete!")
print("=" * 60)
+38 -2
View File
@@ -2,13 +2,27 @@
"""
Context Monitoring Example
Demonstrates how to use the ContextMonitor to write runtime context
snapshots to disk for debugging and analysis.
Demonstrates how to use context monitoring with Agent via context= param,
and also shows low-level ContextMonitor usage for advanced scenarios.
Agent-Centric Quick Start:
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
agent = Agent(
instructions="You are helpful.",
context=ManagerConfig(
monitor_enabled=True,
monitor_path="./debug/context.json",
),
)
"""
import os
import tempfile
from praisonaiagents import Agent
from praisonaiagents.context import (
ManagerConfig,
ContextMonitor,
ContextLedger,
BudgetAllocation,
@@ -19,6 +33,28 @@ from praisonaiagents.context import (
)
def agent_centric_example():
"""Agent-centric usage - recommended approach."""
print("=" * 60)
print("Agent-Centric Context Monitoring")
print("=" * 60)
with tempfile.TemporaryDirectory() as tmpdir:
monitor_path = os.path.join(tmpdir, "context.json")
agent = Agent(
instructions="You are a helpful assistant.",
context=ManagerConfig(
monitor_enabled=True,
monitor_path=monitor_path,
monitor_format="json",
),
)
print(f"Agent with monitoring enabled")
print(f"Monitor path: {monitor_path}")
print(f"Context manager: {agent.context_manager is not None}")
def main():
print("=" * 60)
print("Context Monitoring Example")
+37 -2
View File
@@ -2,11 +2,25 @@
"""
Context Optimization Example
Demonstrates how to use various optimization strategies to reduce
context size when approaching model limits.
Demonstrates how to use context optimization with Agent via context= param,
and also shows low-level optimizer usage for advanced scenarios.
Agent-Centric Quick Start:
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
agent = Agent(
instructions="You are helpful.",
context=ManagerConfig(
auto_compact=True,
strategy="smart", # truncate, sliding_window, summarize, smart
),
)
"""
from praisonaiagents import Agent
from praisonaiagents.context import (
ManagerConfig,
get_optimizer,
OptimizerStrategy,
TruncateOptimizer,
@@ -17,6 +31,27 @@ from praisonaiagents.context import (
)
def agent_centric_example():
"""Agent-centric usage - recommended approach."""
print("=" * 60)
print("Agent-Centric Context Optimization")
print("=" * 60)
# Enable auto-compaction with smart strategy
agent = Agent(
instructions="You are a helpful assistant.",
context=ManagerConfig(
auto_compact=True,
compact_threshold=0.8,
strategy="smart",
),
)
print("Agent with auto-compaction enabled")
print(f"Strategy: smart")
print(f"Threshold: 0.8 (80% of context budget)")
print(f"Context manager: {agent.context_manager is not None}")
def main():
print("=" * 60)
print("Context Optimization Example")
+38 -2
View File
@@ -2,13 +2,24 @@
"""
Multi-Agent Context Isolation Example
Demonstrates how to use MultiAgentLedger and MultiAgentMonitor
for managing context across multiple agents with isolation.
Demonstrates how to use multi-agent context with PraisonAIAgents via context= param,
and also shows low-level MultiAgentLedger usage for advanced scenarios.
Agent-Centric Quick Start:
from praisonaiagents import Agent, PraisonAIAgents
from praisonaiagents.context import ManagerConfig
agents = PraisonAIAgents(
agents=[agent1, agent2],
context=ManagerConfig(policy="isolated"), # or "shared"
)
"""
import os
import tempfile
from praisonaiagents import Agent, PraisonAIAgents
from praisonaiagents.context import (
ManagerConfig,
MultiAgentLedger,
MultiAgentMonitor,
ContextLedgerManager,
@@ -17,6 +28,31 @@ from praisonaiagents.context import (
)
def agent_centric_example():
"""Agent-centric usage - recommended approach."""
print("=" * 60)
print("Agent-Centric Multi-Agent Context")
print("=" * 60)
# Create agents with context enabled
researcher = Agent(
name="Researcher",
instructions="You are a research specialist.",
context=True,
)
writer = Agent(
name="Writer",
instructions="You are a technical writer.",
context=True,
)
# Multi-agent setup with isolated context policy
# (each agent maintains its own context)
print("Created agents with context=True")
print("Use PraisonAIAgents(agents=[...], context=ManagerConfig(policy='isolated'))")
print("for multi-agent context isolation")
def main():
print("=" * 60)
print("Multi-Agent Context Isolation Example")
@@ -13,10 +13,10 @@ from praisonaiagents import Agent
from praisonaiagents.context import (
ContextBudgeter,
ContextLedgerManager,
ContextMonitor,
get_optimizer,
OptimizerStrategy,
estimate_messages_tokens,
format_percent,
)
@@ -37,7 +37,7 @@ def main():
# Get budget info
budget = budgeter.allocate()
print(f"\nModel: gpt-4o-mini")
print("\nModel: gpt-4o-mini")
print(f"Context budget: {budget.usable:,} tokens")
print(f"Output reserve: {budget.output_reserve:,} tokens")
@@ -76,7 +76,7 @@ def main():
current_tokens = ledger.get_total()
utilization = budgeter.get_utilization(current_tokens)
print(f"[Context: {current_tokens} tokens, {utilization:.1%} used]")
print(f"[Context: {current_tokens} tokens, {format_percent(utilization)} used]")
# Check if optimization needed
if utilization > 0.8:
@@ -101,7 +101,7 @@ def main():
final_utilization = budgeter.get_utilization(final_tokens)
print(f"Total tokens used: {final_tokens:,}")
print(f"Context utilization: {final_utilization:.1%}")
print(f"Context utilization: {format_percent(final_utilization)}")
print(f"Messages in conversation: {len(conversation)}")
print(f"Remaining capacity: {budgeter.get_remaining(final_tokens):,} tokens")
+217
View File
@@ -0,0 +1,217 @@
#!/usr/bin/env python3
"""
Agent-Centric Context Management Example
Demonstrates the recommended way to use context management with the new
`context=` parameter. This is the simplest and most powerful approach.
The `context=` parameter accepts:
- False (default): Zero overhead, no context management
- True: Enable with safe defaults (auto-compact, smart strategy)
- ManagerConfig: Custom configuration object
- ContextManager: Pre-configured manager instance
Requires: OPENAI_API_KEY environment variable
"""
import os
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
def example_1_simple_enable():
"""Simplest usage - just enable context management."""
print("=" * 60)
print("Example 1: Simple Enable (context=True)")
print("=" * 60)
# Create agent with context management enabled
agent = Agent(
instructions="You are a helpful assistant. Keep responses concise.",
llm="gpt-4o-mini",
context=True, # Enable with safe defaults
)
# The agent now automatically:
# - Tracks token usage
# - Optimizes when approaching limits (80% threshold)
# - Uses smart optimization strategy
print("Agent created with context management enabled")
print(f"Context manager initialized: {agent.context_manager is not None}")
if agent.context_manager:
print(f"Model: {agent.context_manager.model}")
print(f"Auto-compact: {agent.context_manager.config.auto_compact}")
print(f"Threshold: {agent.context_manager.config.compact_threshold}")
print(f"Strategy: {agent.context_manager.config.strategy}")
# Make a chat call
if os.environ.get("OPENAI_API_KEY"):
response = agent.chat("What is 2+2?")
print(f"\nResponse: {response}")
else:
print("\nSkipping API call (no OPENAI_API_KEY)")
print()
def example_2_custom_config():
"""Custom configuration with ManagerConfig."""
print("=" * 60)
print("Example 2: Custom Configuration (context=ManagerConfig)")
print("=" * 60)
# Create custom configuration
config = ManagerConfig(
auto_compact=True,
compact_threshold=0.7, # Trigger at 70% instead of 80%
strategy="smart",
monitor_enabled=True,
monitor_path="./context_debug.txt",
monitor_format="human",
redact_sensitive=True,
output_reserve=8000,
)
# Create agent with custom config
agent = Agent(
instructions="You are a code assistant.",
llm="gpt-4o-mini",
context=config,
)
print("Agent created with custom context config")
print(f"Threshold: {agent.context_manager.config.compact_threshold}")
print(f"Monitor enabled: {agent.context_manager.config.monitor_enabled}")
print(f"Monitor path: {agent.context_manager.config.monitor_path}")
if os.environ.get("OPENAI_API_KEY"):
response = agent.chat("Write a hello world in Python")
print(f"\nResponse: {response[:200]}...")
else:
print("\nSkipping API call (no OPENAI_API_KEY)")
print()
def example_3_disabled():
"""Disabled context management - zero overhead."""
print("=" * 60)
print("Example 3: Disabled (context=False)")
print("=" * 60)
# Create agent without context management (default)
agent = Agent(
instructions="You are a simple assistant.",
llm="gpt-4o-mini",
context=False, # Explicit disable (same as default)
)
print("Agent created without context management")
print(f"Context manager: {agent.context_manager}")
print("Zero overhead: No token tracking, no optimization")
if os.environ.get("OPENAI_API_KEY"):
response = agent.chat("Hello!")
print(f"\nResponse: {response}")
else:
print("\nSkipping API call (no OPENAI_API_KEY)")
print()
def example_4_access_stats():
"""Access context statistics and history."""
print("=" * 60)
print("Example 4: Access Context Stats")
print("=" * 60)
agent = Agent(
instructions="You are helpful.",
llm="gpt-4o-mini",
context=True,
)
manager = agent.context_manager
# Access budget info via internal attribute
budget = manager._budget
print(f"Model limit: {budget.model_limit:,} tokens")
print(f"Output reserve: {budget.output_reserve:,} tokens")
print(f"Usable budget: {budget.usable:,} tokens")
# Access config
print("\nConfiguration:")
print(f" Auto-compact: {manager.config.auto_compact}")
print(f" Threshold: {manager.config.compact_threshold}")
print(f" Strategy: {manager.config.strategy}")
print()
def example_5_multi_agent():
"""Multi-agent with shared context management."""
print("=" * 60)
print("Example 5: Multi-Agent Context Management")
print("=" * 60)
from praisonaiagents import PraisonAIAgents, Task
config = ManagerConfig(
auto_compact=True,
compact_threshold=0.8,
)
# Create agents
researcher = Agent(
name="Researcher",
instructions="You research topics thoroughly.",
llm="gpt-4o-mini",
)
writer = Agent(
name="Writer",
instructions="You write clear summaries.",
llm="gpt-4o-mini",
)
# Create tasks
task1 = Task(description="Research AI trends", agent=researcher)
task2 = Task(description="Write a summary", agent=writer)
# Create multi-agent system with context management
agents = PraisonAIAgents(
agents=[researcher, writer],
tasks=[task1, task2],
context=config, # Shared context config
)
print("Multi-agent system created")
print(f"Context manager: {agents.context_manager is not None}")
if agents.context_manager:
print("Type: MultiAgentContextManager")
print("Provides per-agent isolation with shared config")
print()
def main():
print("\n" + "=" * 60)
print("Agent-Centric Context Management Examples")
print("=" * 60 + "\n")
example_1_simple_enable()
example_2_custom_config()
example_3_disabled()
example_4_access_stats()
example_5_multi_agent()
print("=" * 60)
print("✓ All examples complete!")
print("=" * 60)
if __name__ == "__main__":
main()
+210
View File
@@ -0,0 +1,210 @@
#!/usr/bin/env python3
"""
Example: Context Compression Demo
This example demonstrates:
1. Context compression for large retrieved content
2. Token budget management
3. Agent-centric knowledge retrieval with unique codes
Requirements:
- pip install praisonaiagents[knowledge]
- OPENAI_API_KEY environment variable
Usage:
python compression_demo.py
"""
import os
import tempfile
import shutil
from praisonaiagents import Agent
# Unique verification code that CANNOT be guessed
SECRET_CODE = "COMPRESS-8K4M2"
def create_verbose_document(temp_dir: str) -> str:
"""Create a verbose document with a hidden unique code."""
filepath = os.path.join(temp_dir, 'verbose_policy.txt')
# Create a document with lots of filler text but one unique code
with open(filepath, 'w') as f:
f.write("""
ACME CORPORATION COMPREHENSIVE POLICY MANUAL
=============================================
CHAPTER 1: INTRODUCTION
-----------------------
This comprehensive policy manual outlines all procedures, guidelines, and
requirements for employees of Acme Corporation. All employees are expected
to read and understand this document in its entirety. This manual supersedes
all previous versions and is effective immediately upon publication.
The purpose of this manual is to provide clear guidance on company policies,
procedures, and expectations. It covers a wide range of topics including but
not limited to workplace conduct, benefits, leave policies, and security
protocols.
CHAPTER 2: WORKPLACE CONDUCT
----------------------------
All employees are expected to maintain professional conduct at all times.
This includes treating colleagues with respect, maintaining a clean workspace,
and adhering to the company dress code. Violations of workplace conduct
policies may result in disciplinary action up to and including termination.
Professional behavior includes punctuality, meeting deadlines, and effective
communication with team members. Employees should report any concerns about
workplace conduct to their immediate supervisor or Human Resources.
CHAPTER 3: BENEFITS AND COMPENSATION
------------------------------------
Acme Corporation offers a comprehensive benefits package to all full-time
employees. This includes health insurance, dental coverage, vision care,
and a 401(k) retirement plan with company matching up to 6% of salary.
Employees are eligible for benefits after completing 90 days of employment.
Open enrollment occurs annually in November. Changes to benefits outside
of open enrollment require a qualifying life event.
CHAPTER 4: LEAVE POLICIES
-------------------------
Employees accrue paid time off (PTO) based on years of service:
- 0-2 years: 15 days per year
- 3-5 years: 20 days per year
- 6+ years: 25 days per year
Sick leave is provided separately at 10 days per year. Unused sick leave
may be carried over to the following year up to a maximum of 30 days.
CHAPTER 5: SECURITY PROTOCOLS
-----------------------------
All employees must follow strict security protocols to protect company
assets and confidential information. This includes using strong passwords,
locking workstations when away, and reporting any suspicious activity.
IMPORTANT: The master security verification code is: COMPRESS-8K4M2
This code is required for accessing secure areas and systems. Do not share
this code with unauthorized personnel. Misuse of security credentials is
grounds for immediate termination.
CHAPTER 6: EMERGENCY PROCEDURES
-------------------------------
In case of emergency, employees should follow the evacuation procedures
posted throughout the building. Assembly points are located in the parking
lot on the north side of the building.
Fire drills are conducted quarterly. All employees must participate in
fire drills and know the location of the nearest emergency exit.
CHAPTER 7: CONCLUSION
---------------------
This policy manual is a living document and may be updated periodically.
Employees will be notified of any changes via email and updated copies
will be made available on the company intranet.
Questions about any policy should be directed to Human Resources.
Document Version: 3.2.1
Last Updated: 2024-12-01
""")
return filepath
def main():
temp_dir = tempfile.mkdtemp(prefix='praison_compress_')
try:
print("=" * 60)
print("Example: Context Compression Demo")
print("=" * 60)
# Create verbose document
doc_path = create_verbose_document(temp_dir)
print(f"\nCreated verbose document: {doc_path}")
print(f"Secret code embedded: {SECRET_CODE}")
# Demonstrate compression
print("\n" + "-" * 40)
print("Compression Demo")
print("-" * 40)
try:
from praisonaiagents.rag import ContextCompressor
# Read the document
with open(doc_path, 'r') as f:
content = f.read()
print(f"\nOriginal document length: {len(content)} characters")
# Compress with different ratios
for ratio in [0.3, 0.5, 0.7]:
compressor = ContextCompressor(
max_tokens=2000,
target_ratio=ratio,
)
result = compressor.compress([content], query="security verification code")
print(f"\nCompression ratio {ratio}:")
print(f" Original tokens: {result.original_tokens}")
print(f" Compressed tokens: {result.compressed_tokens}")
print(f" Actual ratio: {result.compressed_tokens / max(result.original_tokens, 1):.2f}")
# Check if secret code is preserved
compressed_text = " ".join(result.chunks) if result.chunks else ""
if SECRET_CODE in compressed_text:
print(f" ✅ Secret code preserved in compressed output")
else:
print(f" ⚠️ Secret code may have been removed (query-focused compression)")
except ImportError as e:
print(f"Note: Compression module not available: {e}")
# Create agent with knowledge
print("\n" + "-" * 40)
print("Agent Knowledge Retrieval Test")
print("-" * 40)
agent = Agent(
name="PolicyExpert",
instructions="""You are a company policy expert.
Answer questions based ONLY on the provided knowledge context.
When asked about codes, provide the EXACT code from the documents.
Be concise but accurate.""",
knowledge=[temp_dir],
user_id="compression_demo_user",
verbose=True,
)
# Test retrieval of the unique code
print("\nQ: What is the master security verification code?")
response = agent.chat("What is the master security verification code?")
print(f"A: {response}")
if SECRET_CODE in response.upper():
print(f"\n✅ VERIFIED: Agent correctly retrieved the secret code!")
else:
print(f"\n❌ WARNING: Expected code {SECRET_CODE} not found in response")
print("\n" + "=" * 60)
print("Demo Complete")
print("=" * 60)
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
if __name__ == "__main__":
main()
@@ -0,0 +1,133 @@
#!/usr/bin/env python3
"""
Example: Knowledge Retrieval with Context-Required Questions
This example demonstrates:
1. Agent answers questions that REQUIRE retrieved context (not general knowledge)
2. Unique codes/secrets that cannot be guessed
3. Verification that context is text, not file path
Requirements:
- pip install praisonaiagents[knowledge]
- OPENAI_API_KEY environment variable
Usage:
python context_required_example.py
"""
import os
import tempfile
import shutil
from praisonaiagents import Agent
def main():
temp_dir = tempfile.mkdtemp(prefix='praison_ctx_')
try:
# Create test document with UNIQUE content that cannot be guessed
doc = os.path.join(temp_dir, 'company_policy.txt')
with open(doc, 'w') as f:
f.write("""
Acme Corporation Internal Policy Document
==========================================
Remote Work Policy:
- Employees may work from home up to 3 days per week
- Manager approval is required for remote work
- Manager approval code: ZEBRA-71
- All requests must be submitted via the HR portal
Expense Reimbursement:
- Maximum daily meal allowance: $47.50
- Mileage rate: $0.655 per mile
- Expense submission code: TIGER-42
- Receipts required for amounts over $25
IT Security:
- VPN access code prefix: FALCON-
- Password reset hotline: 555-0199
- Security incident code: HAWK-88
This document is confidential. Last updated: 2024-01-15.
""")
print("=" * 60)
print("Example: Context-Required Knowledge Retrieval")
print("=" * 60)
print(f"\nDocument location: {doc}")
print("\nThis example asks questions that CANNOT be answered")
print("without the retrieved context (unique codes/values).\n")
# Create agent with knowledge
agent = Agent(
name="PolicyExpert",
instructions="""You are a company policy expert.
Answer questions based ONLY on the provided knowledge context.
If the answer is not in the context, say 'I don't have that information.'
Be precise and quote exact values when asked about codes or numbers.""",
knowledge=[temp_dir],
user_id="demo_user",
verbose=True,
)
# Question 1: Unique code that requires context
print("-" * 40)
print("Question 1: What is the manager approval code for remote work?")
print("-" * 40)
response = agent.chat("What is the manager approval code for remote work?")
print(f"\nAnswer: {response}\n")
# Verify the answer contains the unique code
if "ZEBRA-71" in response.upper():
print("✅ VERIFIED: Agent correctly retrieved the unique code from context!\n")
else:
print("❌ WARNING: Agent did not find the code - check retrieval\n")
# Question 2: Another unique value
print("-" * 40)
print("Question 2: What is the maximum daily meal allowance?")
print("-" * 40)
response = agent.chat("What is the maximum daily meal allowance?")
print(f"\nAnswer: {response}\n")
if "47.50" in response or "$47.50" in response:
print("✅ VERIFIED: Agent correctly retrieved the specific amount!\n")
else:
print("❌ WARNING: Agent did not find the amount - check retrieval\n")
# Question 3: Security code
print("-" * 40)
print("Question 3: What is the security incident code?")
print("-" * 40)
response = agent.chat("What is the security incident code?")
print(f"\nAnswer: {response}\n")
if "HAWK-88" in response.upper():
print("✅ VERIFIED: Agent correctly retrieved the security code!\n")
else:
print("❌ WARNING: Agent did not find the code - check retrieval\n")
# Show context preview to prove it's text, not path
print("=" * 60)
print("Context Verification")
print("=" * 60)
context, _ = agent._get_knowledge_context("policy codes", use_rag=True)
print(f"\nRetrieved context preview (first 300 chars):")
print("-" * 40)
print(context[:300] if context else "NO CONTEXT")
print("-" * 40)
# Verify context is text, not path
if temp_dir in context:
print("\n❌ ERROR: Context contains file path instead of text!")
else:
print("\n✅ VERIFIED: Context contains actual text, not file path!")
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
if __name__ == "__main__":
main()
@@ -0,0 +1,199 @@
#!/usr/bin/env python3
"""
Example: Hierarchical Summaries Demo
This example demonstrates:
1. Building multi-level summaries (file → folder → project)
2. Query routing to appropriate summary level
3. Agent-centric knowledge retrieval with unique codes
Requirements:
- pip install praisonaiagents[knowledge]
- OPENAI_API_KEY environment variable
Usage:
python hierarchical_summaries_demo.py
"""
import os
import tempfile
import shutil
from praisonaiagents import Agent
# Unique verification codes per department
DEPT_CODES = {
"engineering": "ENG-7X2K9",
"marketing": "MKT-3M5N8",
"finance": "FIN-9P4Q1",
}
def create_hierarchical_corpus(temp_dir: str) -> dict:
"""Create a hierarchical corpus with departments and files."""
structure = {}
for dept, code in DEPT_CODES.items():
dept_dir = os.path.join(temp_dir, dept)
os.makedirs(dept_dir, exist_ok=True)
structure[dept] = []
# Create department overview
overview_path = os.path.join(dept_dir, "overview.txt")
with open(overview_path, 'w') as f:
f.write(f"""
{dept.upper()} DEPARTMENT OVERVIEW
{'=' * 40}
Department: {dept.title()}
Department Code: {code}
Mission Statement:
The {dept.title()} department is responsible for driving excellence
in our organization through innovative solutions and dedicated teamwork.
Key Responsibilities:
- Strategic planning and execution
- Cross-functional collaboration
- Continuous improvement initiatives
Access Code: {code}
Use this code for all department-related system access.
Last Updated: 2024-12-01
""")
structure[dept].append(overview_path)
# Create team files
for i, team in enumerate(["alpha", "beta", "gamma"]):
team_path = os.path.join(dept_dir, f"team_{team}.txt")
with open(team_path, 'w') as f:
f.write(f"""
{dept.upper()} - TEAM {team.upper()}
{'=' * 40}
Team: {team.title()}
Department: {dept.title()}
Team Size: {5 + i * 2} members
Current Projects:
- Project {team.upper()}-001: In Progress
- Project {team.upper()}-002: Planning
- Project {team.upper()}-003: Completed
Team Lead: {team.title()} Lead
Contact: {team}@acme.com
Note: For department access, use code: {code}
""")
structure[dept].append(team_path)
return structure
def main():
temp_dir = tempfile.mkdtemp(prefix='praison_hierarchy_')
try:
print("=" * 60)
print("Example: Hierarchical Summaries Demo")
print("=" * 60)
# Create hierarchical corpus
structure = create_hierarchical_corpus(temp_dir)
total_files = sum(len(files) for files in structure.values())
print(f"\nCreated hierarchical corpus:")
print(f" Departments: {len(structure)}")
print(f" Total files: {total_files}")
print(f" Location: {temp_dir}")
for dept, files in structure.items():
print(f" - {dept}/: {len(files)} files (code: {DEPT_CODES[dept]})")
# Demonstrate hierarchical summarizer
print("\n" + "-" * 40)
print("Hierarchical Summarizer Demo")
print("-" * 40)
try:
from praisonaiagents.rag import HierarchicalSummarizer
# Collect all files
all_files = []
for files in structure.values():
all_files.extend(files)
# Build hierarchy
summarizer = HierarchicalSummarizer(max_levels=3)
result = summarizer.build_hierarchy(all_files, base_path=temp_dir)
print(f"\nHierarchy built:")
print(f" Levels: {summarizer.max_levels}")
if hasattr(result, 'nodes'):
print(f" Summary nodes: {len(result.nodes)}")
# Query at different levels
print("\nQuery routing demonstration:")
queries = [
("What is the engineering department code?", "engineering"),
("Tell me about marketing team alpha", "marketing"),
("What are the finance department responsibilities?", "finance"),
]
for query, expected_dept in queries:
print(f"\n Query: {query}")
# In a full implementation, this would route to the appropriate level
print(f" Expected department: {expected_dept}")
print(f" Expected code: {DEPT_CODES[expected_dept]}")
except ImportError as e:
print(f"Note: Hierarchical summarizer not available: {e}")
except Exception as e:
print(f"Note: Hierarchy building skipped: {e}")
# Create agent with knowledge
print("\n" + "-" * 40)
print("Agent Knowledge Retrieval Test")
print("-" * 40)
agent = Agent(
name="DepartmentExpert",
instructions="""You are a department information expert.
Answer questions based ONLY on the provided knowledge context.
When asked about department codes, provide the EXACT code from the documents.
Be specific about which department you're referring to.""",
knowledge=[temp_dir],
user_id="hierarchy_demo_user",
verbose=True,
)
# Test retrieval of department codes
test_cases = [
("What is the engineering department access code?", "ENG-7X2K9"),
("What is the marketing department code?", "MKT-3M5N8"),
("What is the finance department code?", "FIN-9P4Q1"),
]
print("\nTesting retrieval of department codes:\n")
for question, expected_code in test_cases:
print(f"Q: {question}")
response = agent.chat(question)
print(f"A: {response[:200]}...")
if expected_code in response.upper():
print(f"✅ VERIFIED: Found code {expected_code}\n")
else:
print(f"❌ WARNING: Expected code {expected_code} not found\n")
print("=" * 60)
print("Demo Complete")
print("=" * 60)
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
if __name__ == "__main__":
main()
@@ -0,0 +1,177 @@
#!/usr/bin/env python3
"""
Example: Large Corpus Strategy Selection Demo
This example demonstrates:
1. Automatic strategy selection based on corpus size
2. Different retrieval strategies (DIRECT, BASIC, HYBRID, RERANKED, COMPRESSED)
3. Agent-centric knowledge retrieval with unique codes
Requirements:
- pip install praisonaiagents[knowledge]
- OPENAI_API_KEY environment variable
Usage:
python large_corpus_strategy_demo.py
"""
import os
import tempfile
import shutil
from praisonaiagents import Agent
# Unique verification codes that CANNOT be guessed
VERIFICATION_CODES = {
"project_alpha": "ALPHA-7X9K2",
"project_beta": "BETA-3M8N1",
"project_gamma": "GAMMA-5P2Q7",
"budget_code": "BUDGET-9R4T6",
"security_clearance": "CLEARANCE-2W8Y3",
}
def create_test_corpus(temp_dir: str, num_files: int = 10) -> list:
"""Create a test corpus with unique codes in each file."""
files = []
# Create main project files with unique codes
for i, (project, code) in enumerate(VERIFICATION_CODES.items()):
filepath = os.path.join(temp_dir, f"{project}_doc.txt")
with open(filepath, 'w') as f:
f.write(f"""
{project.upper().replace('_', ' ')} Documentation
{'=' * 50}
Project Overview:
This document contains confidential information about {project}.
The verification code for this project is: {code}
Key Details:
- Project started: 2024-0{i+1}-15
- Team size: {10 + i * 5} members
- Status: Active
- Priority: {'High' if i < 2 else 'Medium'}
Access Requirements:
To access this project, use verification code: {code}
All access attempts are logged for security purposes.
Last updated: 2024-12-01
""")
files.append(filepath)
# Create additional filler files to simulate larger corpus
for i in range(num_files - len(VERIFICATION_CODES)):
filepath = os.path.join(temp_dir, f"general_doc_{i}.txt")
with open(filepath, 'w') as f:
f.write(f"""
General Documentation File {i}
{'=' * 40}
This is a general documentation file containing various information
about company processes and procedures.
Section {i}.1: Overview
Lorem ipsum dolor sit amet, consectetur adipiscing elit.
Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
Section {i}.2: Procedures
Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris.
Duis aute irure dolor in reprehenderit in voluptate velit esse.
Section {i}.3: Guidelines
Excepteur sint occaecat cupidatat non proident, sunt in culpa.
Qui officia deserunt mollit anim id est laborum.
Document ID: DOC-{1000 + i}
""")
files.append(filepath)
return files
def main():
temp_dir = tempfile.mkdtemp(prefix='praison_strategy_')
try:
print("=" * 60)
print("Example: Large Corpus Strategy Selection Demo")
print("=" * 60)
# Create test corpus
files = create_test_corpus(temp_dir, num_files=15)
print(f"\nCreated {len(files)} test documents in: {temp_dir}")
print(f"Unique verification codes embedded: {len(VERIFICATION_CODES)}")
# Show strategy selection based on corpus size
print("\n" + "-" * 40)
print("Strategy Selection Demo")
print("-" * 40)
try:
from praisonaiagents.rag import select_strategy, RetrievalStrategy
from praisonaiagents.knowledge.indexing import CorpusStats
# Get corpus stats
stats = CorpusStats.from_directory(temp_dir)
print(f"\nCorpus Statistics:")
print(f" Files: {stats.file_count}")
print(f" Estimated tokens: {stats.total_tokens}")
print(f" Recommended strategy: {stats.strategy_recommendation}")
# Show strategy selection for different corpus sizes
print(f"\nStrategy selection by corpus size:")
for size in [100, 1000, 10000, 50000, 100000]:
strategy = select_strategy(corpus_tokens=size)
print(f" {size:>6} tokens -> {strategy.value}")
except ImportError as e:
print(f"Note: Strategy module not available: {e}")
# Create agent with knowledge
print("\n" + "-" * 40)
print("Agent Knowledge Retrieval Test")
print("-" * 40)
agent = Agent(
name="ProjectExpert",
instructions="""You are a project documentation expert.
Answer questions based ONLY on the provided knowledge context.
When asked about verification codes, provide the EXACT code from the documents.
If information is not in the context, say 'Information not found in documents.'""",
knowledge=[temp_dir],
user_id="strategy_demo_user",
verbose=True,
)
# Test retrieval with unique codes
test_questions = [
("What is the verification code for Project Alpha?", "ALPHA-7X9K2"),
("What is the budget code?", "BUDGET-9R4T6"),
("What is the security clearance code?", "CLEARANCE-2W8Y3"),
]
print("\nTesting retrieval of unique codes:\n")
for question, expected_code in test_questions:
print(f"Q: {question}")
response = agent.chat(question)
print(f"A: {response[:200]}...")
if expected_code in response.upper():
print(f"✅ VERIFIED: Found code {expected_code}\n")
else:
print(f"❌ WARNING: Expected code {expected_code} not found\n")
print("=" * 60)
print("Demo Complete")
print("=" * 60)
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
if __name__ == "__main__":
main()
+115
View File
@@ -0,0 +1,115 @@
#!/usr/bin/env python3
"""
Example: Using mem0 Backend with Agent Retrieval
This example demonstrates:
1. Agent with knowledge using mem0 backend (default)
2. Proper scope identifiers (user_id/agent_id)
3. Handling of metadata normalization (mem0 returns metadata=None)
4. Agent-first retrieval API
Requirements:
- pip install praisonaiagents[knowledge]
- OPENAI_API_KEY environment variable
Usage:
python mem0_backend_example.py
"""
import os
import tempfile
import shutil
from praisonaiagents import Agent
def main():
# Create temporary directory with test documents
temp_dir = tempfile.mkdtemp(prefix='praison_example_')
try:
# Create test documents with UNIQUE content that requires retrieval
doc1 = os.path.join(temp_dir, 'acme_policy.txt')
with open(doc1, 'w') as f:
f.write("""
Acme Corp Travel Policy Document
Flight booking code: FALCON-77
Hotel reservation prefix: HTL-
Maximum daily meal allowance: $52.75
Mileage reimbursement rate: $0.67 per mile
Travel approval required for trips over: $500
Emergency travel hotline: 555-ACME-911
""")
doc2 = os.path.join(temp_dir, 'project_info.txt')
with open(doc2, 'w') as f:
f.write("""
Project Phoenix Status Report
Project code: PHOENIX-2024
Lead engineer: Dr. Sarah Chen
Budget allocation: $2.4 million
Target completion: Q3 2025
Security clearance required: Level 3
Weekly standup: Tuesdays at 10am PST
""")
print("=" * 60)
print("Example: Agent with Knowledge (mem0 backend)")
print("=" * 60)
# Create agent with knowledge
# mem0 is the default backend - it requires scope identifiers
agent = Agent(
name="PolicyExpert",
instructions="You are a company policy expert. Answer questions based ONLY on the provided knowledge. Quote exact values when asked about codes or numbers.",
knowledge=[temp_dir], # Add directory with documents
user_id="example_user", # Required for mem0 backend
)
print("\n1. Testing Agent.chat() with context-required question:")
print("-" * 40)
# Question that REQUIRES the retrieved context (unique code)
response = agent.chat("What is the flight booking code?")
print("Q: What is the flight booking code?")
print(f"A: {response}")
# Verify the answer contains the unique code
if "FALCON-77" in response.upper():
print("✅ VERIFIED: Agent correctly retrieved the unique code!")
print("\n2. Testing Agent.chat() with another context-required question:")
print("-" * 40)
response = agent.chat("What is the maximum daily meal allowance?")
print("Q: What is the maximum daily meal allowance?")
print(f"A: {response}")
if "52.75" in response:
print("✅ VERIFIED: Agent correctly retrieved the specific amount!")
print("\n3. Testing Agent.retrieve() for context-only retrieval:")
print("-" * 40)
# retrieve() returns context without LLM generation
try:
context_pack = agent.retrieve("Project Phoenix")
print("Query: Project Phoenix")
print(f"Context preview: {context_pack.context[:200]}...")
print(f"Citations: {len(context_pack.citations)} found")
except Exception as e:
print(f"Note: retrieve() requires RAG module: {e}")
print("\n" + "=" * 60)
print("Example completed successfully!")
print("=" * 60)
finally:
# Cleanup
shutil.rmtree(temp_dir, ignore_errors=True)
if __name__ == "__main__":
main()
@@ -0,0 +1,111 @@
#!/usr/bin/env python3
"""
Example: Using Scope Identifiers (user_id, agent_id, run_id)
This example demonstrates:
1. Multi-user knowledge isolation with user_id
2. Multi-agent knowledge sharing with agent_id
3. Run-specific context with run_id
4. Error handling for missing scope identifiers
Requirements:
- pip install praisonaiagents[knowledge]
- OPENAI_API_KEY environment variable
Usage:
python scope_identifiers_example.py
"""
import os
import tempfile
import shutil
from praisonaiagents import Agent
def main():
temp_dir = tempfile.mkdtemp(prefix='praison_scope_')
try:
# Create test document
doc = os.path.join(temp_dir, 'shared_knowledge.txt')
with open(doc, 'w') as f:
f.write("""
Company Policy Document
1. All employees must complete security training annually.
2. Remote work is allowed with manager approval.
3. Expense reports must be submitted within 30 days.
""")
print("=" * 60)
print("Example: Scope Identifiers for Knowledge Isolation")
print("=" * 60)
# Example 1: User-scoped knowledge
print("\n1. User-scoped knowledge (user_id):")
print("-" * 40)
user1_agent = Agent(
name="HR_Assistant_User1",
instructions="You are an HR assistant.",
knowledge=[temp_dir],
user_id="user_alice", # Knowledge scoped to Alice
)
response = user1_agent.chat("What is the policy on remote work?")
print("User Alice asks: What is the policy on remote work?")
print(f"Response: {response[:200]}...")
# Example 2: Agent-scoped knowledge
print("\n2. Agent-scoped knowledge (agent_id):")
print("-" * 40)
shared_agent = Agent(
name="Policy_Expert",
instructions="You are a policy expert.",
knowledge=[temp_dir],
agent_id="policy_agent_v1", # Knowledge scoped to this agent
)
response = shared_agent.chat("What are the expense report rules?")
print("Agent policy_agent_v1 query: What are the expense report rules?")
print(f"Response: {response[:200]}...")
# Example 3: Combined scoping
print("\n3. Combined scoping (user_id + agent_id):")
print("-" * 40)
combined_agent = Agent(
name="Personal_HR_Bot",
instructions="You are a personal HR assistant.",
knowledge=[temp_dir],
user_id="user_bob",
agent_id="hr_bot_v2",
)
response = combined_agent.chat("What training is required?")
print("User Bob with hr_bot_v2: What training is required?")
print(f"Response: {response[:200]}...")
print("\n" + "=" * 60)
print("Scope Identifier Summary:")
print("=" * 60)
print("""
- user_id: Isolates knowledge per user (multi-tenant)
- agent_id: Isolates knowledge per agent type/version
- run_id: Isolates knowledge per execution run (ephemeral)
Best Practices:
1. Always provide at least one scope identifier
2. Use user_id for user-specific data
3. Use agent_id for shared agent knowledge
4. Use run_id for temporary/session-specific context
""")
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
if __name__ == "__main__":
main()
@@ -0,0 +1,133 @@
#!/usr/bin/env python3
"""
Example: Multi-Agent Scope Isolation
This example demonstrates:
1. Agent A ingests knowledge under user_id=U1
2. Agent B queries under user_id=U2 - should NOT find the knowledge
3. Agent B queries under user_id=U1 - should find the knowledge
4. Proves scope isolation works correctly
Requirements:
- pip install praisonaiagents[knowledge]
- OPENAI_API_KEY environment variable
Usage:
python scope_isolation_example.py
"""
import os
import tempfile
import shutil
from praisonaiagents import Agent, Knowledge
def main():
temp_dir = tempfile.mkdtemp(prefix='praison_scope_')
try:
# Create test document with UNIQUE secret
doc = os.path.join(temp_dir, 'secret.txt')
with open(doc, 'w') as f:
f.write("""
Project Codename: PHOENIX
Secret access code: WRENCH-992
Classification: Top Secret
Only authorized personnel with code WRENCH-992 may access this project.
""")
print("=" * 60)
print("Example: Multi-Agent Scope Isolation")
print("=" * 60)
# Step 1: Agent A ingests knowledge under user_id=U1
print("\n1. Agent A ingests knowledge under user_id='user_alice'")
print("-" * 40)
agent_a = Agent(
name="AgentA_Ingester",
instructions="You ingest and store knowledge.",
knowledge=[temp_dir],
user_id="user_alice", # Scoped to Alice
verbose=False,
)
# Ensure knowledge is processed
agent_a._ensure_knowledge_processed()
print(" Knowledge ingested by Agent A under user_alice scope")
# Step 2: Agent B queries under user_id=U2 - should NOT find
print("\n2. Agent B queries under user_id='user_bob' (different scope)")
print("-" * 40)
# Create a separate Knowledge instance for Agent B with different scope
knowledge_b = Knowledge()
search_result = knowledge_b.search("WRENCH-992", user_id="user_bob")
results = search_result.get('results', []) if isinstance(search_result, dict) else []
if len(results) == 0:
print(" ✅ CORRECT: Agent B (user_bob) found NO results")
print(" Scope isolation is working!")
else:
# Check if any result actually contains the secret
found_secret = any('wrench-992' in str(r.get('memory', '')).lower() for r in results)
if found_secret:
print(" ❌ ERROR: Agent B found the secret - scope isolation failed!")
else:
print(" ✅ CORRECT: Agent B found no matching content")
# Step 3: Agent B queries under user_id=U1 - should find
print("\n3. Agent B queries under user_id='user_alice' (same scope)")
print("-" * 40)
search_result = knowledge_b.search("WRENCH-992", user_id="user_alice")
results = search_result.get('results', []) if isinstance(search_result, dict) else []
if len(results) > 0:
memory = results[0].get('memory', '')
if 'wrench-992' in memory.lower():
print(" ✅ CORRECT: Agent B (as user_alice) found the secret!")
print(f" Retrieved: {memory[:80]}...")
else:
print(f" ⚠️ Found results but no secret: {memory[:80]}...")
else:
print(" ❌ ERROR: Agent B should have found results with user_alice scope")
# Step 4: Full Agent chat test with scope
print("\n4. Full Agent chat test with correct scope")
print("-" * 40)
agent_b = Agent(
name="AgentB_Querier",
instructions="Answer questions based on the provided knowledge. If you don't have the information, say so.",
knowledge=[temp_dir],
user_id="user_alice", # Same scope as Agent A
verbose=False,
)
response = agent_b.chat("What is the secret access code for Project PHOENIX?")
print(f" Question: What is the secret access code?")
print(f" Answer: {response}")
if "WRENCH-992" in response.upper():
print("\n ✅ VERIFIED: Agent correctly retrieved the secret code!")
else:
print("\n ⚠️ Agent may not have found the code")
print("\n" + "=" * 60)
print("Summary: Scope Isolation")
print("=" * 60)
print("""
- user_id='user_alice': Can access knowledge ingested by user_alice
- user_id='user_bob': Cannot access knowledge ingested by user_alice
- This enables multi-tenant knowledge isolation
""")
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
if __name__ == "__main__":
main()
+33 -74
View File
@@ -2,9 +2,9 @@
"""
Basic Policy Engine Example for PraisonAI Agents.
This example demonstrates how to use the policy engine to:
1. Create policies with rules
2. Check if actions are allowed
This example demonstrates how to use the policy engine with Agent:
1. Create an Agent with policy enforcement
2. Define policies with rules
3. Block dangerous operations
4. Use convenience policy creators
@@ -12,126 +12,85 @@ Usage:
python basic_policy.py
"""
from praisonaiagents import Agent
from praisonaiagents.policy import (
PolicyEngine, Policy, PolicyRule, PolicyAction, PolicyConfig,
PolicyEngine, Policy, PolicyRule, PolicyAction,
create_deny_tools_policy, create_read_only_policy
)
def main():
print("=" * 60)
print("Policy Engine Demo")
print("Agent-Centric Policy Engine Demo")
print("=" * 60)
# Create a policy engine
# Create a policy engine with rules
engine = PolicyEngine()
# ==========================================================================
# Policy 1: Block dangerous tools
# ==========================================================================
# Add policy to block dangerous tools
print("\n--- Creating Policies ---")
dangerous_tools_policy = Policy(
engine.add_policy(Policy(
name="block_dangerous",
description="Block dangerous file operations",
rules=[
PolicyRule(
action=PolicyAction.DENY,
resource="tool:delete_*",
reason="Delete operations are blocked",
name="deny_delete"
reason="Delete operations are blocked"
),
PolicyRule(
action=PolicyAction.DENY,
resource="tool:rm_*",
reason="Remove operations are blocked",
name="deny_rm"
reason="Remove operations are blocked"
),
],
priority=100
)
engine.add_policy(dangerous_tools_policy)
print(f"✅ Added policy: {dangerous_tools_policy.name}")
))
print("✅ Added policy: block_dangerous")
# ==========================================================================
# Policy 2: Allow read operations
# ==========================================================================
read_policy = Policy(
name="allow_read",
description="Allow read operations",
rules=[
PolicyRule(
action=PolicyAction.ALLOW,
resource="tool:read_*",
name="allow_read"
),
PolicyRule(
action=PolicyAction.ALLOW,
resource="tool:list_*",
name="allow_list"
),
]
)
engine.add_policy(read_policy)
print(f"✅ Added policy: {read_policy.name}")
# Add read-only policy using convenience function
engine.add_policy(create_read_only_policy())
print("✅ Added policy: read_only")
# ==========================================================================
# Test policy checks
# ==========================================================================
# Create an Agent with policy enforcement
agent = Agent(
name="SecureAgent",
instructions="You are a file management assistant.",
policy=engine
)
print("\n--- Agent with Policy Enforcement Created ---")
print(f"Agent: {agent.name}")
print(f"Policy engine: {agent.policy is not None}")
# Test policy checks via agent's policy
print("\n--- Testing Policy Checks ---")
test_cases = [
("tool:read_file", "Reading a file"),
("tool:list_directory", "Listing directory"),
("tool:delete_file", "Deleting a file"),
("tool:rm_directory", "Removing directory"),
("tool:write_file", "Writing a file"),
]
for resource, description in test_cases:
result = engine.check(resource, {})
result = agent.policy.check(resource, {})
status = "✅ ALLOWED" if result.allowed else "🚫 DENIED"
reason = f" ({result.reason})" if result.reason else ""
print(f" {status}: {description}{reason}")
# ==========================================================================
# Using convenience functions
# ==========================================================================
print("\n--- Using Convenience Functions ---")
print("\n--- Convenience Policy Creators ---")
# Create a read-only policy
read_only = create_read_only_policy()
print(f"✅ Created read-only policy with {len(read_only.rules)} rules")
# Create a deny tools policy
deny_tools = create_deny_tools_policy(
["execute_*", "shell_*", "system_*"],
reason="System commands are blocked"
)
print(f"✅ Created deny-tools policy with {len(deny_tools.rules)} rules")
print(f"✅ create_deny_tools_policy: {len(deny_tools.rules)} rules")
# ==========================================================================
# Strict mode example
# ==========================================================================
print("\n--- Strict Mode Example ---")
strict_engine = PolicyEngine(PolicyConfig(strict_mode=True))
# In strict mode, unknown resources are denied
result = strict_engine.check("tool:unknown_operation", {})
print(f" Unknown operation in strict mode: {'DENIED' if not result.allowed else 'ALLOWED'}")
# ==========================================================================
# Policy serialization
# ==========================================================================
print("\n--- Policy Serialization ---")
policy_dict = dangerous_tools_policy.to_dict()
print(f" Serialized policy: {policy_dict['name']} with {len(policy_dict['rules'])} rules")
# Recreate from dict
restored = Policy.from_dict(policy_dict)
print(f" Restored policy: {restored.name}")
read_only = create_read_only_policy()
print(f"✅ create_read_only_policy: {len(read_only.rules)} rules")
print("\n" + "=" * 60)
print("Demo Complete!")
+35 -53
View File
@@ -2,31 +2,29 @@
"""
Basic Thinking Budgets Example for PraisonAI Agents.
This example demonstrates how to use thinking budgets to:
1. Configure token budgets for extended thinking
This example demonstrates how to use thinking budgets with Agent:
1. Create an Agent with thinking budget
2. Use predefined budget levels
3. Track thinking usage
4. Adapt budgets based on complexity
Usage:
python basic_thinking.py
"""
from praisonaiagents.thinking import (
ThinkingBudget, ThinkingConfig, ThinkingUsage, ThinkingTracker
)
from praisonaiagents import Agent
from praisonaiagents.thinking import ThinkingBudget, ThinkingTracker
from praisonaiagents.thinking.budget import BudgetLevel
def main():
print("=" * 60)
print("Thinking Budgets Demo")
print("Agent-Centric Thinking Budgets Demo")
print("=" * 60)
# ==========================================================================
# Budget Levels
# ==========================================================================
print("\n--- Budget Levels ---")
print("\n--- Available Budget Levels ---")
levels = [
("Minimal", ThinkingBudget.minimal()),
@@ -39,10 +37,25 @@ def main():
for name, budget in levels:
print(f" {name}: {budget.max_tokens:,} tokens")
# ==========================================================================
# Agent with Thinking Budget
# ==========================================================================
print("\n--- Creating Agent with Thinking Budget ---")
agent = Agent(
name="DeepThinker",
instructions="You are a problem-solving assistant that thinks step by step.",
thinking_budget=ThinkingBudget.high() # 16,000 tokens for reasoning
)
print(f"Agent: {agent.name}")
print(f"Thinking budget: {agent.thinking_budget.max_tokens:,} tokens")
print(f"Adaptive: {agent.thinking_budget.adaptive}")
# ==========================================================================
# Custom Budget
# ==========================================================================
print("\n--- Custom Budget ---")
print("\n--- Custom Thinking Budget ---")
custom_budget = ThinkingBudget(
max_tokens=12000,
@@ -50,22 +63,24 @@ def main():
adaptive=True,
max_time_seconds=120.0
)
print(f" Max tokens: {custom_budget.max_tokens:,}")
print(f" Min tokens: {custom_budget.min_tokens:,}")
print(f" Adaptive: {custom_budget.adaptive}")
print(f" Time limit: {custom_budget.max_time_seconds}s")
agent_custom = Agent(
name="CustomThinker",
instructions="You solve complex problems.",
thinking_budget=custom_budget
)
print(f"Agent: {agent_custom.name}")
print(f"Max tokens: {agent_custom.thinking_budget.max_tokens:,}")
print(f"Min tokens: {agent_custom.thinking_budget.min_tokens:,}")
print(f"Time limit: {agent_custom.thinking_budget.max_time_seconds}s")
# ==========================================================================
# Complexity-Based Scaling
# ==========================================================================
print("\n--- Complexity-Based Token Allocation ---")
budget = ThinkingBudget(
max_tokens=16000,
min_tokens=2000,
adaptive=True
)
budget = agent.thinking_budget
complexities = [0.0, 0.25, 0.5, 0.75, 1.0]
for complexity in complexities:
tokens = budget.get_tokens_for_complexity(complexity)
@@ -75,34 +90,15 @@ def main():
# ==========================================================================
# Usage Tracking
# ==========================================================================
print("\n--- Usage Tracking ---")
usage = ThinkingUsage(
budget_tokens=8000,
tokens_used=5000,
time_seconds=30.0,
budget_time=60.0
)
print(f" Tokens used: {usage.tokens_used:,} / {usage.budget_tokens:,}")
print(f" Tokens remaining: {usage.tokens_remaining:,}")
print(f" Utilization: {usage.token_utilization:.1%}")
print(f" Over budget: {usage.is_over_budget}")
print(f" Time remaining: {usage.time_remaining}s")
# ==========================================================================
# Session Tracking
# ==========================================================================
print("\n--- Session Tracking ---")
tracker = ThinkingTracker()
# Simulate multiple thinking sessions
# Simulate thinking sessions
sessions = [
(8000, 4000, 20.0),
(8000, 6000, 35.0),
(16000, 12000, 60.0),
(8000, 9000, 45.0), # Over budget
]
for budget_tokens, tokens_used, time_seconds in sessions:
@@ -112,21 +108,7 @@ def main():
summary = tracker.get_summary()
print(f" Total sessions: {summary['session_count']}")
print(f" Total tokens used: {summary['total_tokens_used']:,}")
print(f" Average tokens/session: {summary['average_tokens_per_session']:,.0f}")
print(f" Average utilization: {summary['average_utilization']:.1%}")
print(f" Over-budget sessions: {summary['over_budget_count']}")
# ==========================================================================
# Serialization
# ==========================================================================
print("\n--- Serialization ---")
budget = ThinkingBudget.from_level(BudgetLevel.HIGH)
data = budget.to_dict()
print(f" Serialized: level={data['level']}, max_tokens={data['max_tokens']}")
restored = ThinkingBudget.from_dict(data)
print(f" Restored: {restored.max_tokens:,} tokens")
print("\n" + "=" * 60)
print("Demo Complete!")
@@ -78,26 +78,29 @@ def example_get_context_for_agent():
# =============================================
def example_agent_with_fast_context():
"""Use FastContext with an Agent."""
"""Use FastContext with an Agent via context= param."""
from praisonaiagents import Agent
from praisonaiagents.context.fast import FastContext
# Create agent with FastContext enabled
# Create FastContext instance
fc = FastContext(workspace_path=".")
# Search for code context
result = fc.search("find authentication code")
context = result.to_context_string() if result.files else None
# Create agent with context management enabled
agent = Agent(
name="CodeAssistant",
instructions="You are a helpful code assistant.",
fast_context=True,
fast_context_path=".", # Current directory
fast_context_model="gpt-4o-mini",
fast_context_max_turns=4,
fast_context_parallelism=8
context=True, # Enable context management
)
# Use delegate_to_fast_context for code search
context = agent.delegate_to_fast_context("find authentication code")
if context:
print("Found relevant code context:")
print(context[:300] + "...")
# Use the context in a chat
print(agent.chat(f"Based on this code:\n{context}\n\nExplain the authentication flow."))
else:
print("No relevant code found")
@@ -0,0 +1,310 @@
#!/usr/bin/env python3
"""
Large Context Knowledge Handling Example
Demonstrates the full pipeline:
1. Token Budget calculation
2. Incremental indexing with ignore patterns
3. Strategy selection based on corpus size
4. Smart retrieval with filtering
5. Contextual compression
6. Hierarchical summarization
Usage:
python examples/large_context_knowledge.py
"""
import os
import tempfile
from pathlib import Path
def create_sample_corpus(base_path: str) -> None:
"""Create a sample corpus for demonstration."""
# Create directory structure
src_dir = Path(base_path) / "src"
docs_dir = Path(base_path) / "docs"
src_dir.mkdir(parents=True, exist_ok=True)
docs_dir.mkdir(parents=True, exist_ok=True)
# Create Python files
(src_dir / "main.py").write_text("""
# Main application entry point
def main():
'''Start the application.'''
print("Hello from PraisonAI!")
if __name__ == "__main__":
main()
""")
(src_dir / "utils.py").write_text("""
# Utility functions
def helper_function(x: int) -> int:
'''Double the input value.'''
return x * 2
def format_output(text: str) -> str:
'''Format text for display.'''
return f"[OUTPUT] {text}"
""")
# Create documentation
(docs_dir / "readme.md").write_text("""
# Project Documentation
This is a sample project demonstrating large context knowledge handling.
## Features
- Token budget management
- Incremental indexing
- Smart retrieval
- Contextual compression
""")
(docs_dir / "api.md").write_text("""
# API Reference
## Functions
### main()
Entry point for the application.
### helper_function(x)
Doubles the input value.
### format_output(text)
Formats text for display.
""")
# Create .praisonignore
(Path(base_path) / ".praisonignore").write_text("""
# Ignore patterns
*.pyc
__pycache__/
.git/
""")
print(f"Created sample corpus at: {base_path}")
def demo_token_budget():
"""Demonstrate token budget calculation."""
print("\n" + "="*60)
print("1. TOKEN BUDGET CALCULATION")
print("="*60)
from praisonaiagents.rag.budget import TokenBudget, get_model_context_window
# Get context window for different models
models = ["gpt-4o", "claude-3.5-sonnet", "gemini-1.5-pro", "llama-3.1-70b"]
print("\nModel context windows:")
for model in models:
window = get_model_context_window(model)
print(f" {model}: {window:,} tokens")
# Create budget for GPT-4o
budget = TokenBudget.from_model("gpt-4o")
print(f"\nBudget for GPT-4o:")
print(f" Max tokens: {budget.model_max_tokens:,}")
print(f" Reserved for response: {budget.reserved_response_tokens:,}")
print(f" Max context tokens: {budget.max_context_tokens:,}")
# Dynamic budget calculation
available = budget.dynamic_budget(
prompt_tokens=500,
history_tokens=2000,
system_tokens=1000,
)
print(f"\nDynamic budget (with 500 prompt + 2000 history + 1000 system):")
print(f" Available for context: {available:,} tokens")
def demo_indexing(corpus_path: str):
"""Demonstrate incremental indexing."""
print("\n" + "="*60)
print("2. INCREMENTAL INDEXING")
print("="*60)
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.knowledge.indexing import CorpusStats
# Get corpus stats before indexing
stats = CorpusStats.from_directory(corpus_path)
print(f"\nCorpus statistics:")
print(f" Files: {stats.file_count}")
print(f" Estimated tokens: {stats.total_tokens:,}")
print(f" Recommended strategy: {stats.strategy_recommendation}")
# Index the corpus
knowledge = Knowledge()
result = knowledge.index(
corpus_path,
user_id="demo_user",
incremental=False,
exclude_glob=["*.pyc"],
)
print(f"\nIndexing result:")
print(f" Files indexed: {result.files_indexed}")
print(f" Files skipped: {result.files_skipped}")
print(f" Chunks created: {result.chunks_created}")
print(f" Duration: {result.duration_seconds:.2f}s")
# Incremental re-index (should skip unchanged files)
print("\nRe-indexing (incremental)...")
result2 = knowledge.index(corpus_path, user_id="demo_user", incremental=True)
print(f" Files indexed: {result2.files_indexed}")
print(f" Files skipped: {result2.files_skipped}")
def demo_strategy_selection():
"""Demonstrate strategy selection."""
print("\n" + "="*60)
print("3. STRATEGY SELECTION")
print("="*60)
from praisonaiagents.rag.strategy import select_strategy, get_strategy_description, RetrievalStrategy
from praisonaiagents.knowledge.indexing import CorpusStats
# Show strategy for different corpus sizes
sizes = [5, 50, 500, 5000, 50000, 500000]
print("\nStrategy selection by corpus size:")
for size in sizes:
stats = CorpusStats(file_count=size)
strategy = select_strategy(stats)
print(f" {size:>7} files -> {strategy.value}")
# Show strategy descriptions
print("\nStrategy descriptions:")
for strategy in RetrievalStrategy:
desc = get_strategy_description(strategy)
print(f" {strategy.value}: {desc[:60]}...")
def demo_smart_retriever():
"""Demonstrate smart retriever with filtering."""
print("\n" + "="*60)
print("4. SMART RETRIEVER")
print("="*60)
from praisonaiagents.rag.retriever import SmartRetriever, SimpleReranker
# Create retriever with reranker
reranker = SimpleReranker()
retriever = SmartRetriever(reranker=reranker)
# Demonstrate filtering
chunks = [
{"text": "Python code for data processing", "metadata": {"filename": "process.py"}},
{"text": "JavaScript frontend code", "metadata": {"filename": "app.js"}},
{"text": "Python utilities and helpers", "metadata": {"filename": "utils.py"}},
{"text": "CSS styling rules", "metadata": {"filename": "styles.css"}},
]
print("\nFiltering by include pattern (*.py):")
filtered = retriever._apply_filters(chunks, include_glob=["*.py"])
for chunk in filtered:
print(f" - {chunk['metadata']['filename']}: {chunk['text'][:40]}...")
print("\nFiltering by exclude pattern (utils*):")
filtered = retriever._apply_filters(chunks, exclude_glob=["utils*"])
for chunk in filtered:
print(f" - {chunk['metadata']['filename']}: {chunk['text'][:40]}...")
# Demonstrate reranking
print("\nReranking for query 'Python data':")
reranked = reranker.rerank("Python data", chunks)
for i, chunk in enumerate(reranked[:3], 1):
print(f" {i}. [{chunk['score']:.2f}] {chunk['text'][:40]}...")
def demo_compression():
"""Demonstrate contextual compression."""
print("\n" + "="*60)
print("5. CONTEXTUAL COMPRESSION")
print("="*60)
from praisonaiagents.rag.compressor import ContextCompressor
compressor = ContextCompressor()
# Create chunks with some redundancy
chunks = [
{"text": "Python is a high-level programming language. It is known for its simplicity and readability. Python supports multiple programming paradigms."},
{"text": "Python is a high-level programming language. It is known for its simplicity and readability. Python supports multiple programming paradigms."}, # Duplicate
{"text": "JavaScript is the language of the web. It runs in browsers and on servers with Node.js."},
{"text": "Python has extensive libraries for data science, machine learning, and web development."},
]
print(f"\nOriginal chunks: {len(chunks)}")
result = compressor.compress(chunks, "Python programming", target_tokens=200)
print(f"After compression:")
print(f" Chunks: {len(result.chunks)}")
print(f" Original tokens: {result.original_tokens}")
print(f" Compressed tokens: {result.compressed_tokens}")
print(f" Compression ratio: {result.compression_ratio:.2%}")
print(f" Method: {result.method_used}")
def demo_hierarchical_summarizer(corpus_path: str):
"""Demonstrate hierarchical summarization."""
print("\n" + "="*60)
print("6. HIERARCHICAL SUMMARIZATION")
print("="*60)
from praisonaiagents.rag.summarizer import HierarchicalSummarizer
summarizer = HierarchicalSummarizer()
# Build hierarchy
print(f"\nBuilding hierarchy for: {corpus_path}")
result = summarizer.build_hierarchy(corpus_path, levels=2)
print(f"\nHierarchy result:")
print(f" Total files: {result.total_files}")
print(f" Total nodes: {len(result.nodes)}")
print(f" Total tokens: {result.total_tokens}")
# Show summaries
print("\nFile summaries:")
for path, node in result.nodes.items():
if node.level == 1: # File level
filename = node.metadata.get("filename", os.path.basename(path))
summary = node.summary[:60] + "..." if len(node.summary) > 60 else node.summary
print(f" {filename}: {summary}")
# Query hierarchy
print("\nQuerying hierarchy for 'function':")
results = summarizer.query_hierarchy("function", corpus_path, max_results=3)
for r in results:
print(f" - {os.path.basename(r['path'])}: score={r['score']}")
def main():
"""Run all demonstrations."""
print("="*60)
print("LARGE CONTEXT KNOWLEDGE HANDLING DEMO")
print("="*60)
# Create temporary corpus
with tempfile.TemporaryDirectory() as tmpdir:
create_sample_corpus(tmpdir)
# Run demos
demo_token_budget()
demo_indexing(tmpdir)
demo_strategy_selection()
demo_smart_retriever()
demo_compression()
demo_hierarchical_summarizer(tmpdir)
print("\n" + "="*60)
print("DEMO COMPLETE")
print("="*60)
if __name__ == "__main__":
main()
@@ -267,7 +267,6 @@ class Agent:
response_template: Optional[str] = None,
allow_code_execution: Optional[bool] = False,
max_retry_limit: int = 2,
respect_context_window: bool = True,
code_execution_mode: Literal["safe", "unsafe"] = "safe",
embedder_config: Optional[Dict[str, Any]] = None,
knowledge: Optional[Union[List[str], Any]] = None,
@@ -298,13 +297,6 @@ class Agent:
planning: bool = False,
planning_tools: Optional[List[Any]] = None,
planning_reasoning: bool = False,
fast_context: bool = False,
fast_context_path: Optional[str] = None,
fast_context_model: str = "gpt-4o-mini",
fast_context_max_turns: int = 4,
fast_context_parallelism: int = 8,
fast_context_timeout: float = 30.0,
history_in_context: Optional[int] = None,
auto_save: Optional[str] = None,
skills: Optional[List[str]] = None,
skills_dirs: Optional[List[str]] = None,
@@ -313,8 +305,15 @@ class Agent:
hooks: Optional[List[Any]] = None,
llm_config: Optional[Dict[str, Any]] = None,
rate_limiter: Optional[Any] = None,
auto_summarize: bool = False,
summarize_threshold: float = 0.8
# Agent-centric feature parameters (lazy loaded)
auto_memory: Optional[Any] = None,
policy: Optional[Any] = None,
background: Optional[Any] = None,
checkpoints: Optional[Any] = None,
output_style: Optional[Any] = None,
thinking_budget: Optional[Any] = None,
# Context management (single param for all context features)
context: Optional[Union[bool, Any]] = False
):
"""Initialize an Agent instance.
@@ -363,8 +362,6 @@ class Agent:
snippets during task completion. Use with caution for security. Defaults to False.
max_retry_limit (int, optional): Maximum number of retry attempts for failed operations
before giving up. Helps handle transient errors. Defaults to 2.
respect_context_window (bool, optional): Automatically manage context window size
to prevent token limit errors with large conversations. Defaults to True.
code_execution_mode (Literal["safe", "unsafe"], optional): Safety mode for code execution.
"safe" restricts dangerous operations, "unsafe" allows full code execution. Defaults to "safe".
embedder_config (Optional[Dict[str, Any]], optional): Configuration dictionary for
@@ -466,12 +463,6 @@ class Agent:
# Store rate limiter (optional, zero overhead when None)
self._rate_limiter = rate_limiter
# Auto-summarization configuration
if auto_summarize and not (0.0 < summarize_threshold <= 1.0):
raise ValueError("summarize_threshold must be between 0.0 and 1.0")
self.auto_summarize = auto_summarize
self.summarize_threshold = summarize_threshold
# Store OpenAI client parameters for lazy initialization
self._openai_api_key = api_key
self._openai_base_url = base_url
@@ -587,7 +578,6 @@ class Agent:
self.response_template = response_template
self.allow_code_execution = allow_code_execution
self.max_retry_limit = max_retry_limit
self.respect_context_window = respect_context_window
self.code_execution_mode = code_execution_mode
self.embedder_config = embedder_config
self.knowledge = knowledge
@@ -628,7 +618,6 @@ Your Goal: {self.goal}
self.claude_memory = claude_memory
# Session management
self.history_in_context = history_in_context # Number of past sessions to include
self.auto_save = auto_save # Session name for auto-saving
# Initialize rules manager for persistent context (like Cursor/Windsurf)
@@ -709,15 +698,6 @@ Your Goal: {self.goal}
from ..rag.retrieval_config import RetrievalConfig
self._retrieval_config = RetrievalConfig()
# Fast Context configuration (lazy loaded)
self.fast_context_enabled = fast_context
self._fast_context_path = fast_context_path # Store raw, resolve lazily
self.fast_context_model = fast_context_model
self.fast_context_max_turns = fast_context_max_turns
self.fast_context_parallelism = fast_context_parallelism
self.fast_context_timeout = fast_context_timeout
self._fast_context_instance = None # Lazy loaded
# Agent Skills configuration (lazy loaded for zero performance impact)
self._skills = skills
self._skills_dirs = skills_dirs
@@ -734,6 +714,135 @@ Your Goal: {self.goal}
self._session_store = None
self._session_store_initialized = False
# Agent-centric feature instances (lazy loaded for zero performance impact)
self._auto_memory = auto_memory
self._policy = policy
self._background = background
self._checkpoints = checkpoints
self._output_style = output_style
self._thinking_budget = thinking_budget
# Context management (lazy loaded for zero overhead when disabled)
self._context_param = context # Store raw param for lazy init
self._context_manager = None # Lazy initialized on first use
self._context_manager_initialized = False
@property
def auto_memory(self):
"""AutoMemory instance for automatic memory extraction."""
return self._auto_memory
@auto_memory.setter
def auto_memory(self, value):
self._auto_memory = value
@property
def policy(self):
"""PolicyEngine instance for execution control."""
return self._policy
@policy.setter
def policy(self, value):
self._policy = value
@property
def background(self):
"""BackgroundRunner instance for async task execution."""
return self._background
@background.setter
def background(self, value):
self._background = value
@property
def checkpoints(self):
"""CheckpointService instance for file-level undo/restore."""
return self._checkpoints
@checkpoints.setter
def checkpoints(self, value):
self._checkpoints = value
@property
def output_style(self):
"""OutputStyle instance for response formatting."""
return self._output_style
@output_style.setter
def output_style(self, value):
self._output_style = value
@property
def thinking_budget(self):
"""ThinkingBudget instance for extended thinking control."""
return self._thinking_budget
@thinking_budget.setter
def thinking_budget(self, value):
self._thinking_budget = value
@property
def context_manager(self):
"""
ContextManager instance for unified context management.
Lazy initialized on first access when context=True or context=ManagerConfig.
Returns None when context=False (zero overhead).
Example:
agent = Agent(instructions="...", context=True)
# Access manager for advanced operations
if agent.context_manager:
stats = agent.context_manager.get_stats()
"""
if self._context_manager_initialized:
return self._context_manager
# Initialize based on context param type
if self._context_param is False or self._context_param is None:
# Zero overhead - no context management
self._context_manager = None
self._context_manager_initialized = True
return None
# Lazy import to avoid overhead when not used
try:
from ..context import ContextManager, ManagerConfig
except ImportError:
# Context module not available
self._context_manager = None
self._context_manager_initialized = True
return None
if self._context_param is True:
# Enable with safe defaults
self._context_manager = ContextManager(
model=self.llm if isinstance(self.llm, str) else "gpt-4o-mini",
agent_name=self.name or "Agent",
)
elif isinstance(self._context_param, ManagerConfig):
# Use provided config
self._context_manager = ContextManager(
model=self.llm if isinstance(self.llm, str) else "gpt-4o-mini",
config=self._context_param,
agent_name=self.name or "Agent",
)
elif hasattr(self._context_param, 'process'):
# Already a ContextManager instance
self._context_manager = self._context_param
else:
# Unknown type, disable
self._context_manager = None
self._context_manager_initialized = True
return self._context_manager
@context_manager.setter
def context_manager(self, value):
"""Set context manager directly."""
self._context_manager = value
self._context_manager_initialized = True
@property
def console(self):
"""Lazily initialize Rich Console only when needed."""
@@ -742,17 +851,6 @@ Your Goal: {self.goal}
self._console = Console()
return self._console
@property
def fast_context_path(self):
"""Lazily resolve fast_context_path - avoids os.getcwd() on every init."""
if self._fast_context_path is None:
self._fast_context_path = os.getcwd()
return self._fast_context_path
@fast_context_path.setter
def fast_context_path(self, value):
self._fast_context_path = value
@property
def skill_manager(self):
"""Lazily initialize SkillManager only when skills are accessed."""
@@ -939,57 +1037,6 @@ Your Goal: {self.goal}
return supports_prompt_caching(model_name)
@property
def fast_context(self):
"""Lazily initialize FastContext instance when needed.
Returns:
FastContext instance or None if not enabled
"""
if not self.fast_context_enabled:
return None
if self._fast_context_instance is None:
try:
from ..context.fast import FastContext
self._fast_context_instance = FastContext(
workspace_path=self.fast_context_path,
model=self.fast_context_model,
max_turns=self.fast_context_max_turns,
max_parallel=self.fast_context_parallelism,
timeout=self.fast_context_timeout,
verbose=self.verbose
)
except ImportError:
logging.warning("FastContext not available")
return None
return self._fast_context_instance
def delegate_to_fast_context(self, query: str) -> Optional[str]:
"""Delegate a code search query to FastContext subagent.
This method uses the FastContext subagent to rapidly search
the codebase and return relevant context.
Args:
query: Natural language search query
Returns:
Formatted context string or None if FastContext not available
"""
if not self.fast_context_enabled or self.fast_context is None:
return None
try:
result = self.fast_context.search(query)
if result.total_files > 0:
return self.fast_context.get_context_for_agent(query)
return None
except Exception as e:
logging.warning(f"FastContext search failed: {e}")
return None
@property
def rules_manager(self):
"""
@@ -1358,7 +1405,14 @@ Your Goal: {self.goal}
return "", None
if isinstance(search_results, dict) and 'results' in search_results:
knowledge_content = "\n".join([result['memory'] for result in search_results['results']])
# CRITICAL: Handle None results and get 'memory' field safely
parts = []
for result in search_results['results']:
if result is not None:
memory = result.get('memory', '') or ''
if memory:
parts.append(memory)
knowledge_content = "\n".join(parts)
else:
knowledge_content = "\n".join(search_results) if isinstance(search_results, list) else str(search_results)
@@ -3504,7 +3558,7 @@ Write the complete compiled report:"""
def start(self, prompt: str, **kwargs):
"""Start the agent with a prompt. This is a convenience method that wraps chat()."""
# Load history from past sessions if history_in_context is set
# Load history from past sessions (now handled via context= param)
self._load_history_context()
# Check if planning mode is enabled
@@ -3524,37 +3578,13 @@ Write the complete compiled report:"""
return result
def _load_history_context(self):
"""Load history from past sessions into context if history_in_context is set."""
if not self.history_in_context or not self._memory_instance:
return
"""Load history from past sessions into context.
try:
sessions = self._memory_instance.list_sessions()
if not sessions:
return
# Get the last N sessions
sessions_to_load = sessions[:self.history_in_context]
for session_info in reversed(sessions_to_load):
try:
session_data = self._memory_instance.resume_session(session_info["name"])
history = session_data.get("conversation_history", [])
# Add history to chat_history with a marker
for msg in history:
if msg not in self.chat_history:
# Mark as from history to avoid duplication
msg_copy = msg.copy()
msg_copy["_from_history"] = True
self.chat_history.append(msg_copy)
except Exception:
continue
if sessions_to_load:
logging.debug(f"Loaded history from {len(sessions_to_load)} past session(s)")
except Exception as e:
logging.debug(f"Error loading history context: {e}")
Note: This functionality is now handled via context= param with ManagerConfig.
This method is kept for backward compatibility but is a no-op.
Use context=ManagerConfig(history_sessions=N) to load past sessions.
"""
pass
def _auto_save_session(self):
"""Auto-save session if auto_save is enabled."""
@@ -127,7 +127,7 @@ def process_task_context(context_item, verbose=0, user_id=None):
return str(context_item) # Fallback for unknown types
class PraisonAIAgents:
def __init__(self, agents, tasks=None, verbose=0, completion_checker=None, max_retries=5, process="sequential", manager_llm=None, memory=False, memory_config=None, embedder=None, user_id=None, max_iter=10, stream=True, name: Optional[str] = None, planning: bool = False, planning_llm: Optional[str] = None, auto_approve_plan: bool = False, planning_tools: Optional[List] = None, planning_reasoning: bool = False, on_task_start: Optional[Callable] = None, on_task_complete: Optional[Callable] = None, variables: Optional[Dict[str, Any]] = None):
def __init__(self, agents, tasks=None, verbose=0, completion_checker=None, max_retries=5, process="sequential", manager_llm=None, memory=False, memory_config=None, embedder=None, user_id=None, max_iter=10, stream=True, name: Optional[str] = None, planning: bool = False, planning_llm: Optional[str] = None, auto_approve_plan: bool = False, planning_tools: Optional[List] = None, planning_reasoning: bool = False, on_task_start: Optional[Callable] = None, on_task_complete: Optional[Callable] = None, variables: Optional[Dict[str, Any]] = None, context: Optional[Any] = False):
# Add check at the start if memory is requested
if memory:
try:
@@ -218,6 +218,11 @@ class PraisonAIAgents:
self._state = {} # Add state storage at PraisonAIAgents level
# Context management (single param for all context features)
self._context_param = context # Store raw param for lazy init
self._context_manager = None # Lazy initialized on first use
self._context_manager_initialized = False
# Initialize planning mode
self.planning = planning
self.planning_llm = planning_llm or "gpt-4o-mini"
@@ -296,6 +301,51 @@ class PraisonAIAgents:
cleaned = cleaned[:-3].strip()
return cleaned
@property
def context_manager(self):
"""
ContextManager instance for unified context management across all agents.
Lazy initialized on first access when context=True or context=ManagerConfig.
Returns None when context=False (zero overhead).
For multi-agent scenarios, uses MultiAgentContextManager for per-agent isolation.
"""
if self._context_manager_initialized:
return self._context_manager
# Initialize based on context param type
if self._context_param is False or self._context_param is None:
# Zero overhead - no context management
self._context_manager = None
self._context_manager_initialized = True
return None
# Lazy import to avoid overhead when not used
try:
from ..context import MultiAgentContextManager, ManagerConfig
except ImportError:
# Context module not available
self._context_manager = None
self._context_manager_initialized = True
return None
if self._context_param is True:
# Enable with safe defaults for multi-agent
self._context_manager = MultiAgentContextManager()
elif isinstance(self._context_param, ManagerConfig):
# Use provided config for all agents
self._context_manager = MultiAgentContextManager(config=self._context_param)
elif hasattr(self._context_param, 'get_agent_manager'):
# Already a MultiAgentContextManager instance
self._context_manager = self._context_param
else:
# Unknown type, disable
self._context_manager = None
self._context_manager_initialized = True
return self._context_manager
def default_completion_checker(self, task, agent_output):
if task.output_json and task.result and task.result.json_dict:
return True
@@ -56,7 +56,6 @@ class AutoAgents(PraisonAIAgents):
min_reflect: int = 1,
llm: Optional[str] = None,
function_calling_llm: Optional[str] = None,
respect_context_window: bool = True,
code_execution_mode: str = "safe",
embedder_config: Optional[Dict[str, Any]] = None,
knowledge_sources: Optional[List[Any]] = None,
@@ -94,7 +93,6 @@ class AutoAgents(PraisonAIAgents):
self.min_reflect = min_reflect
self.llm = llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-5-nano')
self.function_calling_llm = function_calling_llm
self.respect_context_window = respect_context_window
self.code_execution_mode = code_execution_mode
self.embedder_config = embedder_config
self.knowledge_sources = knowledge_sources
@@ -446,7 +444,6 @@ DO NOT use strings for tasks. Each task MUST be a complete object with all four
min_reflect=self.min_reflect,
llm=self.llm,
function_calling_llm=self.function_calling_llm,
respect_context_window=self.respect_context_window,
code_execution_mode=self.code_execution_mode,
embedder_config=self.embedder_config,
knowledge=self.knowledge_sources,
@@ -41,7 +41,32 @@ Usage:
monitor.snapshot(ledger=result.ledger, budget=budget, messages=result.messages)
"""
def format_percent(value: float) -> str:
"""
Smart percentage formatting for context utilization display.
- For values < 0.1%: shows "<0.1%"
- For values < 1%: shows 2 decimal places (e.g., "0.02%")
- For values >= 1%: shows 1 decimal place (e.g., "5.3%")
Args:
value: A ratio (0.0 to 1.0+), NOT already multiplied by 100
Returns:
Formatted percentage string
"""
pct = value * 100
if pct < 0.1 and pct > 0:
return "<0.1%"
elif pct < 1.0:
return f"{pct:.2f}%"
else:
return f"{pct:.1f}%"
__all__ = [
# Utilities
"format_percent",
# Models
"ContextSegment",
"ContextLedger",
@@ -83,6 +108,23 @@ __all__ = [
"redact_sensitive",
"format_human_snapshot",
"format_json_snapshot",
"validate_monitor_path",
"should_include_content",
"load_ignore_patterns",
# Manager Facade (NEW)
"ContextManager",
"MultiAgentContextManager",
"create_context_manager",
"ManagerConfig",
"ContextPolicy",
"EstimationMode",
"ContextShareMode",
"ToolShareMode",
"OptimizationEvent",
"OptimizationEventType",
"EstimationMetrics",
"PerToolBudget",
"SnapshotHookData",
# Fast Context (legacy)
"FastContext",
"FastContextResult",
@@ -130,10 +172,19 @@ def __getattr__(name: str):
# Monitoring
if name in ("ContextMonitor", "MultiAgentMonitor", "redact_sensitive",
"format_human_snapshot", "format_json_snapshot"):
"format_human_snapshot", "format_json_snapshot",
"validate_monitor_path", "should_include_content", "load_ignore_patterns"):
from . import monitor
return getattr(monitor, name)
# Manager Facade
if name in ("ContextManager", "MultiAgentContextManager", "create_context_manager",
"ManagerConfig", "ContextPolicy", "EstimationMode", "ContextShareMode",
"ToolShareMode", "OptimizationEvent", "OptimizationEventType",
"EstimationMetrics", "PerToolBudget", "SnapshotHookData"):
from . import manager
return getattr(manager, name)
# Fast Context (legacy)
if name == "FastContext":
from praisonaiagents.context.fast.fast_context import FastContext
@@ -2,7 +2,7 @@
FastContext - High-level API for fast parallel code search.
This module provides the main FastContext class that can be used
standalone or integrated with the main Agent class.
standalone or integrated with the main Agent class via context= param.
Example usage:
# Standalone usage
@@ -12,13 +12,13 @@ Example usage:
result = fc.search("find authentication handlers")
print(result.to_context_string())
# With Agent integration
# With Agent integration via context= param
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
agent = Agent(
name="CodeAssistant",
fast_context=True,
fast_context_path="/path/to/project"
context=True, # Enable context management
)
"""
File diff suppressed because it is too large. Load diff
@@ -4,7 +4,7 @@ Context Monitor for PraisonAI Agents.
Writes runtime context to disk for inspection and debugging.
"""
from typing import Dict, List, Any, Optional, Literal
from typing import Dict, List, Any, Optional, Literal, Tuple
from pathlib import Path
from datetime import datetime
import json
@@ -15,16 +15,42 @@ from .models import (
)
def _format_pct(value: float) -> str:
"""Smart percentage formatting: <0.1% for tiny, 2 decimals for <1%, 1 decimal otherwise."""
if value < 0.1 and value > 0:
return "<0.1%"
elif value < 1.0:
return f"{value:.2f}%"
else:
return f"{value:.1f}%"
# Sensitive patterns for redaction
SENSITIVE_PATTERNS = [
r'sk-[a-zA-Z0-9]{20,}', # OpenAI API keys
r'sk-proj-[a-zA-Z0-9_-]+', # OpenAI project keys
# OpenAI API keys
r'sk-[a-zA-Z0-9]{20,}',
r'sk-proj-[a-zA-Z0-9_-]+',
# Anthropic API keys
r'sk-ant-[a-zA-Z0-9_-]+',
r'sk-ant-api\d+-[a-zA-Z0-9_-]+',
# Google API keys and tokens
r'AIza[a-zA-Z0-9_-]{30,}', # Google API keys (30+ chars after prefix)
r'ya29\.[a-zA-Z0-9_-]+', # Google OAuth tokens
r'goog-[a-zA-Z0-9_-]{20,}', # Google service tokens
# AWS credentials
r'AKIA[0-9A-Z]{16}', # AWS access keys
r'[a-zA-Z0-9]{32,}', # Generic long tokens (conservative)
r'ASIA[0-9A-Z]{16}', # AWS temporary access keys
r'aws_secret_access_key["\s:=]+["\']?[^"\'\s]+',
# Azure credentials
r'[a-zA-Z0-9]{8}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{12}', # Azure GUIDs
# Generic patterns
r'[a-zA-Z0-9]{40,}', # Generic long tokens (conservative)
r'password["\s:=]+["\']?[^"\'\s]+', # Password patterns
r'secret["\s:=]+["\']?[^"\'\s]+', # Secret patterns
r'token["\s:=]+["\']?[^"\'\s]+', # Token patterns
r'api[_-]?key["\s:=]+["\']?[^"\'\s]+', # API key patterns
r'bearer\s+[a-zA-Z0-9_-]+', # Bearer tokens
r'authorization["\s:=]+["\']?[^"\'\s]+', # Authorization headers
]
@@ -44,6 +70,153 @@ def redact_sensitive(text: str) -> str:
return result
def validate_monitor_path(
path: str,
allow_absolute: bool = False,
base_dir: Optional[str] = None,
) -> Tuple[bool, str]:
"""
Validate monitor output path for security.
Prevents path traversal attacks and enforces path restrictions.
Args:
path: Path to validate
allow_absolute: Whether to allow absolute paths
base_dir: Base directory for relative paths
Returns:
Tuple of (is_valid, error_message)
"""
from pathlib import Path as PathLib
path_obj = PathLib(path)
# Check for path traversal attempts
if '..' in path_obj.parts:
return False, "Path traversal (..) not allowed"
# Check absolute path restriction
if path_obj.is_absolute() and not allow_absolute:
return False, "Absolute paths not allowed (set allow_absolute_paths=True to override)"
# If base_dir specified, ensure path stays within it
if base_dir:
base = PathLib(base_dir).resolve()
try:
# For relative paths, resolve against base
if not path_obj.is_absolute():
full_path = (base / path_obj).resolve()
else:
full_path = path_obj.resolve()
# Check if resolved path is under base
if not str(full_path).startswith(str(base)):
return False, f"Path must be within {base_dir}"
except Exception as e:
return False, f"Path validation error: {e}"
# Check for suspicious patterns
suspicious = ['/etc/', '/var/', '/usr/', '/root/', '/home/', '~']
path_str = str(path_obj).lower()
for pattern in suspicious:
if pattern in path_str and not allow_absolute:
return False, f"Suspicious path pattern: {pattern}"
return True, ""
def should_include_content(
file_path: str,
ignore_patterns: Optional[List[str]] = None,
include_patterns: Optional[List[str]] = None,
) -> bool:
"""
Check if file content should be included based on ignore/include patterns.
Respects .praisonignore and .praisoninclude rules.
Args:
file_path: Path to check
ignore_patterns: Patterns to ignore (glob-style)
include_patterns: Patterns to include (glob-style)
Returns:
True if content should be included
"""
import fnmatch
from pathlib import Path as PathLib
path = PathLib(file_path)
name = path.name
# Default: include everything
if not ignore_patterns and not include_patterns:
return True
# Check include patterns first (whitelist)
if include_patterns:
included = False
for pattern in include_patterns:
if fnmatch.fnmatch(name, pattern) or fnmatch.fnmatch(str(path), pattern):
included = True
break
if not included:
return False
# Check ignore patterns (blacklist)
if ignore_patterns:
for pattern in ignore_patterns:
if fnmatch.fnmatch(name, pattern) or fnmatch.fnmatch(str(path), pattern):
return False
return True
def load_ignore_patterns(base_dir: str = ".") -> Tuple[List[str], List[str]]:
"""
Load ignore and include patterns from .praisonignore and .praisoninclude files.
Args:
base_dir: Directory to search for pattern files
Returns:
Tuple of (ignore_patterns, include_patterns)
"""
from pathlib import Path as PathLib
ignore_patterns = []
include_patterns = []
base = PathLib(base_dir)
# Load .praisonignore
ignore_file = base / ".praisonignore"
if ignore_file.exists():
try:
content = ignore_file.read_text()
for line in content.splitlines():
line = line.strip()
if line and not line.startswith('#'):
ignore_patterns.append(line)
except Exception:
pass
# Load .praisoninclude
include_file = base / ".praisoninclude"
if include_file.exists():
try:
content = include_file.read_text()
for line in content.splitlines():
line = line.strip()
if line and not line.startswith('#'):
include_patterns.append(line)
except Exception:
pass
return ignore_patterns, include_patterns
def format_human_snapshot(snapshot: ContextSnapshot, redact: bool = True) -> str:
"""
Format snapshot as human-readable text.
@@ -95,10 +268,10 @@ def format_human_snapshot(snapshot: ContextSnapshot, redact: bool = True) -> str
tokens = snapshot.ledger.get_segment(seg)
budget = snapshot.budget.get_segment_budget(seg)
pct = (tokens / budget * 100) if budget > 0 else 0
lines.append(f"{name:<20} | {tokens:>10,} | {budget:>10,} | {pct:>7.1f}%")
lines.append(f"{name:<20} | {tokens:>10,} | {budget:>10,} | {_format_pct(pct):>8}")
lines.append("-" * 20 + "-+-" + "-" * 10 + "-+-" + "-" * 10 + "-+-" + "-" * 8)
lines.append(f"{'TOTAL':<20} | {snapshot.ledger.total:>10,} | {snapshot.budget.usable:>10,} | {snapshot.utilization * 100:>7.1f}%")
lines.append(f"{'TOTAL':<20} | {snapshot.ledger.total:>10,} | {snapshot.budget.usable:>10,} | {_format_pct(snapshot.utilization * 100):>8}")
lines.append("")
@@ -365,17 +538,17 @@ class ContextMonitor:
utilization = ledger.total / budget.usable if budget.usable > 0 else 0
if utilization >= 1.0:
warnings.append(f"CRITICAL: Context overflow! {utilization*100:.1f}% of usable budget")
warnings.append(f"CRITICAL: Context overflow! {_format_pct(utilization*100)} of usable budget")
elif utilization >= 0.9:
warnings.append(f"WARNING: Context at {utilization*100:.1f}% of usable budget")
warnings.append(f"WARNING: Context at {_format_pct(utilization*100)} of usable budget")
elif utilization >= 0.8:
warnings.append(f"NOTICE: Context at {utilization*100:.1f}% of usable budget")
warnings.append(f"NOTICE: Context at {_format_pct(utilization*100)} of usable budget")
# Check individual segments
if budget.tool_outputs > 0:
tool_util = ledger.tool_outputs / budget.tool_outputs
if tool_util >= 0.9:
warnings.append(f"Tool outputs at {tool_util*100:.1f}% of budget")
warnings.append(f"Tool outputs at {_format_pct(tool_util*100)} of budget")
return warnings
@@ -16,6 +16,18 @@ from praisonaiagents.knowledge.chunking import Chunking
# Lazy loading for protocols and registries to avoid import overhead
_LAZY_IMPORTS = {
# Models and Protocols (new)
"SearchResultItem": ("praisonaiagents.knowledge.models", "SearchResultItem"),
"SearchResult": ("praisonaiagents.knowledge.models", "SearchResult"),
"AddResult": ("praisonaiagents.knowledge.models", "AddResult"),
"normalize_search_item": ("praisonaiagents.knowledge.models", "normalize_search_item"),
"normalize_search_result": ("praisonaiagents.knowledge.models", "normalize_search_result"),
"normalize_to_dict": ("praisonaiagents.knowledge.models", "normalize_to_dict"),
"KnowledgeStoreProtocol": ("praisonaiagents.knowledge.protocols", "KnowledgeStoreProtocol"),
"KnowledgeBackendError": ("praisonaiagents.knowledge.protocols", "KnowledgeBackendError"),
"ScopeRequiredError": ("praisonaiagents.knowledge.protocols", "ScopeRequiredError"),
"BackendNotAvailableError": ("praisonaiagents.knowledge.protocols", "BackendNotAvailableError"),
# Readers
"Document": ("praisonaiagents.knowledge.readers", "Document"),
"ReaderProtocol": ("praisonaiagents.knowledge.readers", "ReaderProtocol"),
@@ -85,6 +97,17 @@ __all__ = [
# Core
"Knowledge",
"Chunking",
# Models and Protocols
"SearchResultItem",
"SearchResult",
"AddResult",
"normalize_search_item",
"normalize_search_result",
"normalize_to_dict",
"KnowledgeStoreProtocol",
"KnowledgeBackendError",
"ScopeRequiredError",
"BackendNotAvailableError",
# Readers
"Document",
"ReaderProtocol",
@@ -0,0 +1,38 @@
"""
Knowledge Store Adapters for PraisonAI Agents.
Provides adapter implementations for various knowledge backends.
All adapters implement KnowledgeStoreProtocol.
Adapters are lazy-loaded to avoid importing heavy dependencies.
"""
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from .mem0_adapter import Mem0Adapter
# Lazy loading for adapters
_LAZY_IMPORTS = {
"Mem0Adapter": ("praisonaiagents.knowledge.adapters.mem0_adapter", "Mem0Adapter"),
}
def __getattr__(name: str):
"""Lazy load adapters."""
if name in _LAZY_IMPORTS:
module_path, attr_name = _LAZY_IMPORTS[name]
import importlib
module = importlib.import_module(module_path)
return getattr(module, attr_name)
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
def __dir__():
"""List available attributes."""
return list(_LAZY_IMPORTS.keys())
__all__ = [
"Mem0Adapter",
]
@@ -0,0 +1,459 @@
"""
Mem0 Adapter for PraisonAI Knowledge Store.
Implements KnowledgeStoreProtocol using mem0 as the backend.
Handles all normalization to ensure metadata is ALWAYS a dict (never None).
LAZY IMPORT: mem0 is only imported when this adapter is instantiated.
"""
import logging
import os
from typing import Any, Dict, List, Optional
from ..models import (
AddResult,
SearchResult,
SearchResultItem,
normalize_search_item,
)
from ..protocols import ScopeRequiredError
logger = logging.getLogger(__name__)
class Mem0Adapter:
"""
Adapter implementing KnowledgeStoreProtocol for mem0 backend.
Features:
- Lazy imports mem0 only when instantiated
- Normalizes all results (metadata None → {})
- Enforces scope requirements (user_id/agent_id/run_id)
- Maps mem0 field names to standard schema (memory → text)
- Disables telemetry by default
Usage:
adapter = Mem0Adapter(config={...})
results = adapter.search("query", user_id="user123")
"""
def __init__(
self,
config: Optional[Dict[str, Any]] = None,
disable_telemetry: bool = True,
):
"""
Initialize Mem0 adapter.
Args:
config: mem0 configuration dict
disable_telemetry: Whether to disable mem0 telemetry (default True)
"""
self._config = config or {}
self._memory = None
self._disable_telemetry = disable_telemetry
# Disable telemetry before importing mem0
if disable_telemetry:
os.environ.setdefault("ANONYMIZED_TELEMETRY", "False")
@property
def memory(self):
"""Lazy load mem0 Memory instance."""
if self._memory is None:
self._memory = self._init_mem0()
return self._memory
def _init_mem0(self):
"""Initialize mem0 Memory with config."""
try:
from mem0 import Memory
except ImportError:
raise ImportError(
"mem0 is not installed. Install with: pip install mem0ai"
)
if self._config:
return Memory.from_config(self._config)
return Memory()
def _check_scope(
self,
user_id: Optional[str],
agent_id: Optional[str],
run_id: Optional[str],
operation: str = "operation",
) -> None:
"""
Check that at least one scope identifier is provided.
mem0 requires at least one of user_id, agent_id, or run_id.
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
if not any([user_id, agent_id, run_id]):
raise ScopeRequiredError(
message=(
f"mem0 {operation} requires at least one of 'user_id', "
f"'agent_id', or 'run_id'. Please provide at least one "
f"identifier to scope the {operation}."
),
backend="mem0",
)
def _normalize_mem0_item(self, raw: Dict[str, Any]) -> SearchResultItem:
"""
Normalize a mem0 result item to SearchResultItem.
Handles mem0-specific field names and ensures metadata is dict.
Args:
raw: Raw mem0 result dict
Returns:
SearchResultItem with guaranteed schema
"""
if raw is None:
return SearchResultItem()
# Use the canonical normalization function
item = normalize_search_item(raw)
# mem0 specific: extract user_id/agent_id/run_id from top level
# and add to metadata if present
for key in ["user_id", "agent_id", "run_id"]:
if key in raw and raw[key] is not None:
item.metadata[key] = raw[key]
return item
def _normalize_mem0_results(self, raw: Any) -> SearchResult:
"""
Normalize mem0 search results to SearchResult.
Args:
raw: Raw mem0 results (dict with 'results' key or list)
Returns:
SearchResult with normalized items
"""
if raw is None:
return SearchResult()
# Handle dict with 'results' key (standard mem0 format)
if isinstance(raw, dict):
items = raw.get("results", []) or []
normalized_items = [
self._normalize_mem0_item(item)
for item in items
if item is not None
]
return SearchResult(
results=normalized_items,
metadata=raw.get("relations", {}) if "relations" in raw else {},
)
# Handle list of results
if isinstance(raw, list):
normalized_items = [
self._normalize_mem0_item(item)
for item in raw
if item is not None
]
return SearchResult(results=normalized_items)
return SearchResult()
def search(
self,
query: str,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
limit: int = 100,
filters: Optional[Dict[str, Any]] = None,
rerank: bool = True,
**kwargs: Any,
) -> SearchResult:
"""
Search for relevant content in mem0.
Args:
query: Search query string
user_id: User identifier for scoping
agent_id: Agent identifier for scoping
run_id: Run identifier for scoping
limit: Maximum number of results
filters: Optional metadata filters
rerank: Whether to use mem0 reranking (default True)
**kwargs: Additional mem0 search options
Returns:
SearchResult with normalized items (metadata always dict)
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
self._check_scope(user_id, agent_id, run_id, "search")
try:
raw_results = self.memory.search(
query=query,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
limit=limit,
filters=filters,
rerank=rerank,
**kwargs,
)
return self._normalize_mem0_results(raw_results)
except Exception as e:
logger.warning(f"mem0 search failed: {e}")
return SearchResult(metadata={"error": str(e)})
def add(
self,
content: Any,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> AddResult:
"""
Add content to mem0.
Args:
content: Content to add (string or messages list)
user_id: User identifier for scoping
agent_id: Agent identifier for scoping
run_id: Run identifier for scoping
metadata: Optional metadata to attach
**kwargs: Additional mem0 add options
Returns:
AddResult with operation status
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
self._check_scope(user_id, agent_id, run_id, "add")
try:
# Convert content to messages format for mem0
if isinstance(content, str):
messages = [{"role": "user", "content": content}]
elif isinstance(content, list):
messages = content
else:
messages = [{"role": "user", "content": str(content)}]
result = self.memory.add(
messages=messages,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
metadata=metadata,
**kwargs,
)
# Extract ID from result
result_id = ""
if isinstance(result, dict):
result_id = result.get("id", "")
elif isinstance(result, list) and result:
result_id = result[0].get("id", "") if isinstance(result[0], dict) else ""
return AddResult(
id=str(result_id),
success=True,
message="Content added successfully",
metadata={"raw_result": result} if result else {},
)
except Exception as e:
logger.error(f"mem0 add failed: {e}")
return AddResult(
success=False,
message=str(e),
metadata={"error": str(e)},
)
def get(
self,
item_id: str,
**kwargs: Any,
) -> Optional[SearchResultItem]:
"""
Get a specific item by ID from mem0.
Args:
item_id: ID of the item to retrieve
**kwargs: Additional options
Returns:
SearchResultItem if found, None otherwise
"""
try:
raw = self.memory.get(item_id)
if raw is None:
return None
return self._normalize_mem0_item(raw)
except Exception as e:
logger.warning(f"mem0 get failed: {e}")
return None
def get_all(
self,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
limit: int = 100,
**kwargs: Any,
) -> SearchResult:
"""
Get all items from mem0, optionally filtered by scope.
Args:
user_id: User identifier for scoping
agent_id: Agent identifier for scoping
run_id: Run identifier for scoping
limit: Maximum number of results
**kwargs: Additional options
Returns:
SearchResult with all matching items
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
self._check_scope(user_id, agent_id, run_id, "get_all")
try:
raw_results = self.memory.get_all(
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
**kwargs,
)
return self._normalize_mem0_results(raw_results)
except Exception as e:
logger.warning(f"mem0 get_all failed: {e}")
return SearchResult(metadata={"error": str(e)})
def update(
self,
item_id: str,
content: Any,
**kwargs: Any,
) -> AddResult:
"""
Update an existing item in mem0.
Args:
item_id: ID of the item to update
content: New content
**kwargs: Additional options
Returns:
AddResult with operation status
"""
try:
result = self.memory.update(item_id, str(content))
return AddResult(
id=item_id,
success=True,
message=result.get("message", "Updated successfully") if isinstance(result, dict) else "Updated",
)
except Exception as e:
logger.error(f"mem0 update failed: {e}")
return AddResult(
id=item_id,
success=False,
message=str(e),
metadata={"error": str(e)},
)
def delete(
self,
item_id: str,
**kwargs: Any,
) -> bool:
"""
Delete an item by ID from mem0.
Args:
item_id: ID of the item to delete
**kwargs: Additional options
Returns:
True if deleted, False otherwise
"""
try:
self.memory.delete(item_id)
return True
except Exception as e:
logger.warning(f"mem0 delete failed: {e}")
return False
def delete_all(
self,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs: Any,
) -> bool:
"""
Delete all items matching scope from mem0.
Args:
user_id: User identifier for scoping
agent_id: Agent identifier for scoping
run_id: Run identifier for scoping
**kwargs: Additional options
Returns:
True if operation succeeded
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
self._check_scope(user_id, agent_id, run_id, "delete_all")
try:
self.memory.delete_all(
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
)
return True
except Exception as e:
logger.warning(f"mem0 delete_all failed: {e}")
return False
def history(
self,
item_id: str,
**kwargs: Any,
) -> List[Dict[str, Any]]:
"""
Get history of changes for an item.
Args:
item_id: ID of the item
**kwargs: Additional options
Returns:
List of history entries
"""
try:
return self.memory.history(item_id)
except Exception as e:
logger.warning(f"mem0 history failed: {e}")
return []
@@ -0,0 +1,439 @@
"""
Incremental Indexing Infrastructure for PraisonAI Agents.
Provides:
- CorpusStats: Statistics about indexed corpus
- IndexResult: Result of indexing operation
- IgnoreMatcher: .praisonignore pattern matching
- FileTracker: File hash/mtime tracking for incremental updates
No heavy imports - only stdlib.
"""
import fnmatch
import hashlib
import json
import os
import re
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional
# Strategy thresholds based on corpus size
STRATEGY_THRESHOLDS = {
10: "direct", # < 10 files: load all
100: "basic", # < 100 files: semantic search only
1000: "hybrid", # < 1000 files: keyword + semantic
10000: "reranked", # < 10000 files: hybrid + rerank
100000: "compressed", # < 100000 files: reranked + compression
}
DEFAULT_STRATEGY = "hierarchical" # > 100000 files
def estimate_tokens_simple(text: str) -> int:
"""Simple token estimation (~4 chars per token)."""
if not text:
return 0
return len(text) // 4 + 1
@dataclass
class CorpusStats:
"""
Statistics about an indexed corpus.
Used for strategy selection and monitoring.
Attributes:
file_count: Number of files in corpus
chunk_count: Number of chunks created
total_tokens: Estimated total tokens
indexed_at: Timestamp of last indexing
path: Path to corpus root
strategy_recommendation: Recommended retrieval strategy
"""
file_count: int = 0
chunk_count: int = 0
total_tokens: int = 0
indexed_at: Optional[str] = None
path: Optional[str] = None
@property
def strategy_recommendation(self) -> str:
"""Recommend retrieval strategy based on corpus size."""
for threshold, strategy in sorted(STRATEGY_THRESHOLDS.items()):
if self.file_count < threshold:
return strategy
return DEFAULT_STRATEGY
@classmethod
def from_directory(
cls,
path: str,
extensions: Optional[List[str]] = None,
) -> "CorpusStats":
"""
Create CorpusStats from directory scan.
Args:
path: Directory path to scan
extensions: Optional list of extensions to include
Returns:
CorpusStats with file count and token estimate
"""
if extensions is None:
extensions = [
".txt", ".md", ".py", ".js", ".ts", ".json", ".yaml", ".yml",
".html", ".css", ".csv", ".xml", ".rst", ".pdf", ".docx",
]
file_count = 0
total_tokens = 0
for root, _, files in os.walk(path):
for filename in files:
ext = os.path.splitext(filename)[1].lower()
if ext in extensions:
file_count += 1
filepath = os.path.join(root, filename)
try:
size = os.path.getsize(filepath)
# Rough estimate: 4 chars per token
total_tokens += size // 4
except OSError:
pass
return cls(
file_count=file_count,
total_tokens=total_tokens,
path=path,
indexed_at=datetime.now().isoformat(),
)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"file_count": self.file_count,
"chunk_count": self.chunk_count,
"total_tokens": self.total_tokens,
"indexed_at": self.indexed_at,
"path": self.path,
"strategy_recommendation": self.strategy_recommendation,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "CorpusStats":
"""Create from dictionary."""
return cls(
file_count=data.get("file_count", 0),
chunk_count=data.get("chunk_count", 0),
total_tokens=data.get("total_tokens", 0),
indexed_at=data.get("indexed_at"),
path=data.get("path"),
)
@dataclass
class IndexResult:
"""
Result of an indexing operation.
Attributes:
success: Whether indexing succeeded
files_indexed: Number of files indexed
files_skipped: Number of files skipped (unchanged)
chunks_created: Number of chunks created
errors: List of error messages
duration_seconds: Time taken for indexing
corpus_stats: Stats about the indexed corpus
"""
success: bool = True
files_indexed: int = 0
files_skipped: int = 0
chunks_created: int = 0
errors: List[str] = field(default_factory=list)
duration_seconds: float = 0.0
corpus_stats: Optional[CorpusStats] = None
@property
def total_files(self) -> int:
"""Total files processed (indexed + skipped)."""
return self.files_indexed + self.files_skipped
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"success": self.success,
"files_indexed": self.files_indexed,
"files_skipped": self.files_skipped,
"chunks_created": self.chunks_created,
"total_files": self.total_files,
"errors": self.errors,
"duration_seconds": self.duration_seconds,
"corpus_stats": self.corpus_stats.to_dict() if self.corpus_stats else None,
}
class IgnoreMatcher:
"""
Matcher for .praisonignore patterns.
Supports gitignore-style patterns:
- * matches any sequence of characters
- ** matches any path segments
- ! negates a pattern
- # starts a comment
"""
def __init__(self, patterns: Optional[List[str]] = None):
"""
Initialize with patterns.
Args:
patterns: List of gitignore-style patterns
"""
self._include_patterns: List[str] = []
self._exclude_patterns: List[str] = []
if patterns:
for pattern in patterns:
pattern = pattern.strip()
if not pattern or pattern.startswith("#"):
continue
if pattern.startswith("!"):
self._include_patterns.append(pattern[1:])
else:
self._exclude_patterns.append(pattern)
def should_ignore(self, path: str) -> bool:
"""
Check if path should be ignored.
Args:
path: Path to check (relative or absolute)
Returns:
True if path should be ignored
"""
# Normalize path
path = path.replace("\\", "/")
basename = os.path.basename(path)
# Check exclude patterns
ignored = False
for pattern in self._exclude_patterns:
if self._matches(path, pattern) or self._matches(basename, pattern):
ignored = True
break
# Check include patterns (negation)
if ignored:
for pattern in self._include_patterns:
if self._matches(path, pattern) or self._matches(basename, pattern):
return False
return ignored
def _matches(self, path: str, pattern: str) -> bool:
"""Check if path matches pattern."""
# Handle ** for recursive matching
if "**" in pattern:
# Convert gitignore pattern to regex
# First escape special regex chars except * and /
regex_pattern = re.escape(pattern)
# Restore * and convert to regex
regex_pattern = regex_pattern.replace(r"\*\*", ".*")
regex_pattern = regex_pattern.replace(r"\*", "[^/]*")
# Match anywhere in path
try:
return bool(re.search(regex_pattern, path))
except re.error:
return False
# Handle trailing slash for directories
if pattern.endswith("/"):
pattern = pattern[:-1]
return path == pattern or path.startswith(pattern + "/") or ("/" + pattern + "/") in ("/" + path + "/")
# Standard fnmatch - check full path and basename
if fnmatch.fnmatch(path, pattern):
return True
if fnmatch.fnmatch(os.path.basename(path), pattern):
return True
# Also check if pattern matches any path segment
for segment in path.split("/"):
if fnmatch.fnmatch(segment, pattern):
return True
return False
@classmethod
def from_file(cls, filepath: str) -> "IgnoreMatcher":
"""
Load patterns from file.
Args:
filepath: Path to .praisonignore file
Returns:
IgnoreMatcher with loaded patterns
"""
patterns = []
try:
with open(filepath, "r", encoding="utf-8") as f:
patterns = f.read().splitlines()
except (OSError, IOError):
pass
return cls(patterns=patterns)
@classmethod
def from_directory(cls, directory: str) -> "IgnoreMatcher":
"""
Auto-detect and load .praisonignore from directory.
Also checks for .gitignore as fallback.
Args:
directory: Directory to search
Returns:
IgnoreMatcher with loaded patterns
"""
patterns = []
# Check for .praisonignore
praison_ignore = os.path.join(directory, ".praisonignore")
if os.path.exists(praison_ignore):
try:
with open(praison_ignore, "r", encoding="utf-8") as f:
patterns.extend(f.read().splitlines())
except (OSError, IOError):
pass
# Also check .gitignore
git_ignore = os.path.join(directory, ".gitignore")
if os.path.exists(git_ignore):
try:
with open(git_ignore, "r", encoding="utf-8") as f:
patterns.extend(f.read().splitlines())
except (OSError, IOError):
pass
return cls(patterns=patterns)
class FileTracker:
"""
Track file hashes and modification times for incremental indexing.
Persists state to a JSON file for cross-session tracking.
"""
def __init__(self, state_file: Optional[str] = None):
"""
Initialize tracker.
Args:
state_file: Path to state file for persistence
"""
self._state_file = state_file
self._tracked: Dict[str, Dict[str, Any]] = {}
def get_file_info(self, filepath: str) -> Dict[str, Any]:
"""
Get file information (hash, mtime, size).
Args:
filepath: Path to file
Returns:
Dict with path, hash, mtime, size
"""
stat = os.stat(filepath)
# Calculate hash for small files, use mtime+size for large files
if stat.st_size < 1024 * 1024: # 1MB
with open(filepath, "rb") as f:
file_hash = hashlib.md5(f.read()).hexdigest()
else:
# For large files, use mtime + size as proxy
file_hash = f"{stat.st_mtime}:{stat.st_size}"
return {
"path": filepath,
"hash": file_hash,
"mtime": stat.st_mtime,
"size": stat.st_size,
}
def mark_indexed(self, filepath: str, info: Dict[str, Any]) -> None:
"""
Mark file as indexed.
Args:
filepath: Path to file
info: File info from get_file_info()
"""
self._tracked[filepath] = info
def has_changed(self, filepath: str) -> bool:
"""
Check if file has changed since last indexing.
Args:
filepath: Path to file
Returns:
True if file is new or changed
"""
if filepath not in self._tracked:
return True
try:
current = self.get_file_info(filepath)
tracked = self._tracked[filepath]
# Compare hash
if current["hash"] != tracked["hash"]:
return True
# Compare mtime as backup
if current["mtime"] != tracked["mtime"]:
return True
return False
except OSError:
return True
def save(self) -> None:
"""Save state to file."""
if not self._state_file:
return
try:
os.makedirs(os.path.dirname(self._state_file), exist_ok=True)
with open(self._state_file, "w", encoding="utf-8") as f:
json.dump(self._tracked, f, indent=2)
except (OSError, IOError):
pass
def load(self) -> None:
"""Load state from file."""
if not self._state_file or not os.path.exists(self._state_file):
return
try:
with open(self._state_file, "r", encoding="utf-8") as f:
self._tracked = json.load(f)
except (OSError, IOError, json.JSONDecodeError):
self._tracked = {}
def clear(self) -> None:
"""Clear all tracked files."""
self._tracked = {}
if self._state_file and os.path.exists(self._state_file):
try:
os.remove(self._state_file)
except OSError:
pass
@@ -683,7 +683,7 @@ class Knowledge:
return self._process_single_input(file_path, user_id, agent_id, run_id, metadata)
def _process_single_input(self, input_path, user_id=None, agent_id=None, run_id=None, metadata=None):
"""Process a single input which can be a file path or URL."""
"""Process a single input which can be a file path, directory, or URL."""
try:
# Define supported file extensions
DOCUMENT_EXTENSIONS = {
@@ -692,11 +692,44 @@ class Knowledge:
'text': ('.txt', '.csv', '.json', '.xml', '.md', '.html', '.htm'),
'archive': '.zip'
}
# Get all supported extensions as a flat tuple
all_extensions = []
for exts in DOCUMENT_EXTENSIONS.values():
if isinstance(exts, tuple):
all_extensions.extend(exts)
else:
all_extensions.append(exts)
all_extensions = tuple(all_extensions)
# Check if input is URL
if isinstance(input_path, str) and (input_path.startswith('http://') or input_path.startswith('https://')):
self._log(f"Processing URL: {input_path}")
raise NotImplementedError("URL processing not yet implemented")
# CRITICAL FIX: Check if input is a directory - recursively process all files
if os.path.isdir(input_path):
self._log(f"Processing directory: {input_path}")
all_results = []
# Walk through directory and process all supported files
for root, dirs, files in os.walk(input_path):
for filename in files:
file_path = os.path.join(root, filename)
# Only process files with supported extensions
if filename.lower().endswith(all_extensions):
try:
result = self._process_single_input(
file_path, user_id, agent_id, run_id, metadata
)
all_results.extend(result.get('results', []))
except Exception as e:
logger.warning(f"Failed to process file {file_path}: {e}")
if not all_results:
logger.warning(f"No supported files found in directory: {input_path}")
return {'results': all_results, 'relations': []}
# Check if input ends with any supported extension
is_supported_file = any(input_path.lower().endswith(ext)
@@ -770,4 +803,158 @@ class Knowledge:
except Exception as e:
logger.error(f"Error processing input {input_path}: {str(e)}", exc_info=True)
raise
raise
def index(
self,
path: str,
incremental: bool = True,
force: bool = False,
include_glob: list = None,
exclude_glob: list = None,
user_id: str = None,
agent_id: str = None,
run_id: str = None,
):
"""
Index a directory or file for knowledge retrieval.
Supports incremental indexing - only changed files are re-indexed.
Args:
path: Directory or file path to index
incremental: If True, only index changed files (default: True)
force: If True, re-index all files regardless of changes
include_glob: List of glob patterns to include (e.g., ["*.py", "*.md"])
exclude_glob: List of glob patterns to exclude (e.g., ["*.log", "test_*"])
user_id: Optional user ID for scoping
agent_id: Optional agent ID for scoping
run_id: Optional run ID for scoping
Returns:
IndexResult with indexing statistics
"""
from .indexing import IndexResult, CorpusStats, IgnoreMatcher, FileTracker
import time as time_module
import fnmatch as fnmatch_module
start_time = time_module.time()
result = IndexResult()
# Initialize file tracker for incremental indexing
state_dir = os.path.join(os.path.dirname(path) if os.path.isfile(path) else path, ".praison")
os.makedirs(state_dir, exist_ok=True)
state_file = os.path.join(state_dir, ".index_state.json")
tracker = FileTracker(state_file=state_file)
if incremental and not force:
tracker.load()
# Initialize ignore matcher
ignore_matcher = IgnoreMatcher.from_directory(path if os.path.isdir(path) else os.path.dirname(path))
# Collect files to index
files_to_index = []
if os.path.isfile(path):
files_to_index = [path]
else:
for root, dirs, files in os.walk(path):
# Skip hidden directories
dirs[:] = [d for d in dirs if not d.startswith('.')]
for filename in files:
filepath = os.path.join(root, filename)
rel_path = os.path.relpath(filepath, path)
# Check ignore patterns
if ignore_matcher.should_ignore(rel_path):
continue
# Check include patterns
if include_glob:
matched = False
for pattern in include_glob:
if fnmatch_module.fnmatch(filename, pattern):
matched = True
break
if not matched:
continue
# Check exclude patterns
if exclude_glob:
excluded = False
for pattern in exclude_glob:
if fnmatch_module.fnmatch(filename, pattern) or fnmatch_module.fnmatch(rel_path, pattern):
excluded = True
break
if excluded:
continue
files_to_index.append(filepath)
# Index files
total_chunks = 0
for filepath in files_to_index:
try:
# Check if file has changed (for incremental indexing)
if incremental and not force and not tracker.has_changed(filepath):
result.files_skipped += 1
continue
# Index the file
add_result = self.add(
filepath,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
)
# Count chunks
if add_result and isinstance(add_result, dict):
chunks = len(add_result.get('results', []))
total_chunks += chunks
# Mark as indexed
file_info = tracker.get_file_info(filepath)
tracker.mark_indexed(filepath, file_info)
result.files_indexed += 1
except Exception as e:
result.errors.append(f"{filepath}: {str(e)}")
# Save tracker state
if incremental:
tracker.save()
# Calculate stats
result.chunks_created = total_chunks
result.duration_seconds = time_module.time() - start_time
result.corpus_stats = CorpusStats(
file_count=result.files_indexed + result.files_skipped,
chunk_count=total_chunks,
path=path,
indexed_at=datetime.now().isoformat(),
)
# Store corpus stats for later retrieval
self._corpus_stats = result.corpus_stats
return result
def get_corpus_stats(self):
"""
Get statistics about the indexed corpus.
Returns:
CorpusStats with file count, chunk count, and strategy recommendation
"""
from .indexing import CorpusStats
if hasattr(self, '_corpus_stats') and self._corpus_stats:
return self._corpus_stats
# Return empty stats if not indexed
return CorpusStats()
@@ -0,0 +1,249 @@
"""
Knowledge Models for PraisonAI Agents.
Typed dataclasses for knowledge retrieval results.
These models guarantee metadata is ALWAYS a dict (never None).
No heavy imports - only stdlib.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
@dataclass
class SearchResultItem:
"""
A single search result item with guaranteed schema.
CRITICAL: metadata is ALWAYS a dict, never None.
This is enforced at construction time.
Attributes:
id: Unique identifier for the result
text: The text content (normalized from 'memory' field if needed)
score: Relevance score (default 0.0)
metadata: Additional metadata (ALWAYS dict, never None)
source: Optional source identifier
filename: Optional filename
created_at: Optional creation timestamp
updated_at: Optional update timestamp
"""
id: str = ""
text: str = ""
score: float = 0.0
metadata: Dict[str, Any] = field(default_factory=dict)
source: Optional[str] = None
filename: Optional[str] = None
created_at: Optional[str] = None
updated_at: Optional[str] = None
def __post_init__(self):
"""Ensure metadata is never None."""
if self.metadata is None:
self.metadata = {}
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary format."""
return {
"id": self.id,
"text": self.text,
"score": self.score,
"metadata": self.metadata,
"source": self.source,
"filename": self.filename,
"created_at": self.created_at,
"updated_at": self.updated_at,
}
@dataclass
class SearchResult:
"""
Container for search results with guaranteed schema.
Attributes:
results: List of SearchResultItem objects
metadata: Additional metadata about the search (ALWAYS dict)
query: The original query string
total_count: Total number of results (may differ from len(results) if paginated)
"""
results: List[SearchResultItem] = field(default_factory=list)
metadata: Dict[str, Any] = field(default_factory=dict)
query: str = ""
total_count: Optional[int] = None
def __post_init__(self):
"""Ensure metadata is never None."""
if self.metadata is None:
self.metadata = {}
if self.results is None:
self.results = []
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary format compatible with legacy code."""
return {
"results": [r.to_dict() if hasattr(r, 'to_dict') else r for r in self.results],
"metadata": self.metadata,
"query": self.query,
"total_count": self.total_count or len(self.results),
}
def to_legacy_format(self) -> Dict[str, Any]:
"""
Convert to legacy format for backward compatibility.
Returns dict with 'results' key containing list of dicts.
"""
return {
"results": [
{
"id": r.id,
"memory": r.text, # Legacy field name
"text": r.text,
"score": r.score,
"metadata": r.metadata,
"source": r.source,
"filename": r.filename,
"created_at": r.created_at,
"updated_at": r.updated_at,
}
for r in self.results
]
}
@dataclass
class AddResult:
"""
Result of adding content to knowledge store.
Attributes:
id: ID of the added item
success: Whether the operation succeeded
message: Optional message about the operation
metadata: Additional metadata (ALWAYS dict)
"""
id: str = ""
success: bool = True
message: str = ""
metadata: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
"""Ensure metadata is never None."""
if self.metadata is None:
self.metadata = {}
def normalize_search_item(raw: Dict[str, Any]) -> SearchResultItem:
"""
Normalize a raw search result item to SearchResultItem.
This is the SINGLE canonical normalization function.
Handles mem0 format, internal format, and any other variations.
CRITICAL: Ensures metadata is ALWAYS a dict, never None.
Args:
raw: Raw result dict from any backend
Returns:
SearchResultItem with guaranteed schema
"""
if raw is None:
return SearchResultItem()
# Handle text field (mem0 uses 'memory', others use 'text')
text = raw.get("text") or raw.get("memory", "") or ""
# Handle metadata - CRITICAL: ensure never None
metadata = raw.get("metadata")
if metadata is None:
metadata = {}
elif not isinstance(metadata, dict):
metadata = {}
# Extract source/filename from metadata if not at top level
source = raw.get("source") or metadata.get("source")
filename = raw.get("filename") or metadata.get("filename")
return SearchResultItem(
id=str(raw.get("id", "")),
text=str(text),
score=float(raw.get("score", 0.0) or 0.0),
metadata=metadata,
source=source,
filename=filename,
created_at=raw.get("created_at"),
updated_at=raw.get("updated_at"),
)
def normalize_search_result(raw: Any) -> SearchResult:
"""
Normalize raw search results to SearchResult.
Handles various input formats:
- dict with 'results' key (mem0 format)
- list of results
- single result dict
- None
Args:
raw: Raw search results from any backend
Returns:
SearchResult with guaranteed schema
"""
if raw is None:
return SearchResult()
# Handle dict with 'results' key
if isinstance(raw, dict):
if "results" in raw:
items = raw.get("results", []) or []
normalized_items = [
normalize_search_item(item)
for item in items
if item is not None
]
return SearchResult(
results=normalized_items,
metadata=raw.get("metadata") or {},
query=raw.get("query", ""),
)
else:
# Single result dict
return SearchResult(
results=[normalize_search_item(raw)],
metadata={},
)
# Handle list of results
if isinstance(raw, list):
normalized_items = [
normalize_search_item(item)
for item in raw
if item is not None
]
return SearchResult(results=normalized_items)
# Fallback
return SearchResult()
def normalize_to_dict(raw: Any) -> Dict[str, Any]:
"""
Normalize raw search results to legacy dict format.
Convenience function for code that expects dict format.
Ensures all items have metadata as dict (never None).
Args:
raw: Raw search results from any backend
Returns:
Dict with 'results' key containing normalized items
"""
result = normalize_search_result(raw)
return result.to_legacy_format()
@@ -0,0 +1,211 @@
"""
Knowledge Store Protocol for PraisonAI Agents.
Defines the interface that knowledge backends must implement.
Uses typing.Protocol for structural subtyping (duck typing with type hints).
No heavy imports - only stdlib typing.
"""
from typing import Any, Dict, Optional, Protocol, runtime_checkable
from .models import AddResult, SearchResult, SearchResultItem
@runtime_checkable
class KnowledgeStoreProtocol(Protocol):
"""
Protocol defining the interface for knowledge store backends.
Implementations include:
- Mem0Adapter (mem0 backend)
- ChromaAdapter (direct chromadb)
- InternalAdapter (built-in storage)
All implementations MUST:
1. Return SearchResult/SearchResultItem with metadata as dict (never None)
2. Handle identifier scoping (user_id/agent_id/run_id) appropriately
3. Provide clear error messages for missing required parameters
"""
def search(
self,
query: str,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
limit: int = 10,
filters: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> SearchResult:
"""
Search for relevant content.
Args:
query: Search query string
user_id: Optional user identifier for scoping
agent_id: Optional agent identifier for scoping
run_id: Optional run identifier for scoping
limit: Maximum number of results
filters: Optional metadata filters
**kwargs: Backend-specific options
Returns:
SearchResult with normalized items (metadata always dict)
"""
...
def add(
self,
content: Any,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> AddResult:
"""
Add content to the knowledge store.
Args:
content: Content to add (string, file path, or structured data)
user_id: Optional user identifier for scoping
agent_id: Optional agent identifier for scoping
run_id: Optional run identifier for scoping
metadata: Optional metadata to attach
**kwargs: Backend-specific options
Returns:
AddResult with operation status
"""
...
def get(
self,
item_id: str,
**kwargs: Any,
) -> Optional[SearchResultItem]:
"""
Get a specific item by ID.
Args:
item_id: ID of the item to retrieve
**kwargs: Backend-specific options
Returns:
SearchResultItem if found, None otherwise
"""
...
def get_all(
self,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
limit: int = 100,
**kwargs: Any,
) -> SearchResult:
"""
Get all items, optionally filtered by scope.
Args:
user_id: Optional user identifier for scoping
agent_id: Optional agent identifier for scoping
run_id: Optional run identifier for scoping
limit: Maximum number of results
**kwargs: Backend-specific options
Returns:
SearchResult with all matching items
"""
...
def update(
self,
item_id: str,
content: Any,
**kwargs: Any,
) -> AddResult:
"""
Update an existing item.
Args:
item_id: ID of the item to update
content: New content
**kwargs: Backend-specific options
Returns:
AddResult with operation status
"""
...
def delete(
self,
item_id: str,
**kwargs: Any,
) -> bool:
"""
Delete an item by ID.
Args:
item_id: ID of the item to delete
**kwargs: Backend-specific options
Returns:
True if deleted, False otherwise
"""
...
def delete_all(
self,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs: Any,
) -> bool:
"""
Delete all items matching scope.
Args:
user_id: Optional user identifier for scoping
agent_id: Optional agent identifier for scoping
run_id: Optional run identifier for scoping
**kwargs: Backend-specific options
Returns:
True if operation succeeded
"""
...
class KnowledgeBackendError(Exception):
"""Base exception for knowledge backend errors."""
pass
class ScopeRequiredError(KnowledgeBackendError):
"""Raised when required scope identifiers are missing."""
def __init__(self, message: str = None, backend: str = None):
self.backend = backend
default_msg = (
f"At least one of 'user_id', 'agent_id', or 'run_id' must be provided"
f"{f' for {backend} backend' if backend else ''}."
)
super().__init__(message or default_msg)
class BackendNotAvailableError(KnowledgeBackendError):
"""Raised when a requested backend is not available."""
def __init__(self, backend: str, reason: str = None):
self.backend = backend
self.reason = reason
msg = f"Knowledge backend '{backend}' is not available"
if reason:
msg += f": {reason}"
super().__init__(msg)
@@ -26,12 +26,36 @@ _LAZY_IMPORTS = {
"ContextPack": ("praisonaiagents.rag.models", "ContextPack"),
"RAGResult": ("praisonaiagents.rag.models", "RAGResult"),
"RAGConfig": ("praisonaiagents.rag.models", "RAGConfig"),
"RetrievalStrategy": ("praisonaiagents.rag.models", "RetrievalStrategy"),
# Unified retrieval config (Agent-first)
"RetrievalConfig": ("praisonaiagents.rag.retrieval_config", "RetrievalConfig"),
"RetrievalPolicy": ("praisonaiagents.rag.retrieval_config", "RetrievalPolicy"),
"CitationsMode": ("praisonaiagents.rag.retrieval_config", "CitationsMode"),
"create_retrieval_config": ("praisonaiagents.rag.retrieval_config", "create_retrieval_config"),
# Token Budget (Phase 1)
"TokenBudget": ("praisonaiagents.rag.budget", "TokenBudget"),
"get_model_context_window": ("praisonaiagents.rag.budget", "get_model_context_window"),
"BudgetEnforcerProtocol": ("praisonaiagents.rag.budget", "BudgetEnforcerProtocol"),
"DefaultBudgetEnforcer": ("praisonaiagents.rag.budget", "DefaultBudgetEnforcer"),
"estimate_tokens": ("praisonaiagents.rag.budget", "estimate_tokens"),
# Strategy Selection (Phase 3) - RetrievalStrategy now from strategy module
"RetrievalStrategy": ("praisonaiagents.rag.strategy", "RetrievalStrategy"),
"select_strategy": ("praisonaiagents.rag.strategy", "select_strategy"),
"get_strategy_description": ("praisonaiagents.rag.strategy", "get_strategy_description"),
"STRATEGY_THRESHOLDS": ("praisonaiagents.rag.strategy", "STRATEGY_THRESHOLDS"),
# SmartRetriever (Phase 4)
"SmartRetriever": ("praisonaiagents.rag.retriever", "SmartRetriever"),
"RetrievalResult": ("praisonaiagents.rag.retriever", "RetrievalResult"),
"SimpleReranker": ("praisonaiagents.rag.retriever", "SimpleReranker"),
"RetrieverProtocol": ("praisonaiagents.rag.retriever", "RetrieverProtocol"),
"RerankerProtocol": ("praisonaiagents.rag.retriever", "RerankerProtocol"),
# Compressor (Phase 5)
"ContextCompressor": ("praisonaiagents.rag.compressor", "ContextCompressor"),
"CompressionResult": ("praisonaiagents.rag.compressor", "CompressionResult"),
"CompressorProtocol": ("praisonaiagents.rag.compressor", "CompressorProtocol"),
# Summarizer (Phase 6)
"HierarchicalSummarizer": ("praisonaiagents.rag.summarizer", "HierarchicalSummarizer"),
"SummaryNode": ("praisonaiagents.rag.summarizer", "SummaryNode"),
"HierarchyResult": ("praisonaiagents.rag.summarizer", "HierarchyResult"),
# Protocols
"ContextBuilderProtocol": ("praisonaiagents.rag.protocols", "ContextBuilderProtocol"),
"CitationFormatterProtocol": ("praisonaiagents.rag.protocols", "CitationFormatterProtocol"),
@@ -0,0 +1,350 @@
"""
Token Budget Infrastructure for PraisonAI Agents.
Provides dynamic token budget calculation for context building.
Ensures retrieved context never exceeds available token budget.
No heavy imports - only stdlib and typing.
"""
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
# Model context window mapping (tokens)
# Updated with latest model information
MODEL_CONTEXT_WINDOWS: Dict[str, int] = {
# OpenAI models
"gpt-4": 8192,
"gpt-4-32k": 32768,
"gpt-4-turbo": 128000,
"gpt-4-turbo-preview": 128000,
"gpt-4o": 128000,
"gpt-4o-mini": 128000,
"gpt-4.1": 1000000,
"gpt-4.1-mini": 1000000,
"gpt-4.1-nano": 1000000,
"gpt-3.5-turbo": 16385,
"gpt-3.5-turbo-16k": 16385,
"o1": 200000,
"o1-mini": 128000,
"o1-preview": 128000,
"o3": 200000,
"o3-mini": 200000,
"o4-mini": 200000,
# Anthropic models
"claude-3-opus": 200000,
"claude-3-sonnet": 200000,
"claude-3-haiku": 200000,
"claude-3.5-sonnet": 200000,
"claude-3.5-haiku": 200000,
"claude-3.7-sonnet": 200000,
"claude-4-sonnet": 200000,
"claude-sonnet-4": 200000,
# Google models
"gemini-pro": 32768,
"gemini-1.0-pro": 32768,
"gemini-1.5-pro": 1000000,
"gemini-1.5-flash": 1000000,
"gemini-2.0-flash": 1000000,
"gemini-2.5-pro": 1000000,
"gemini-2.5-flash": 1000000,
# Mistral models
"mistral-large": 128000,
"mistral-medium": 32768,
"mistral-small": 32768,
"mixtral-8x7b": 32768,
# Llama models
"llama-3.1-405b": 128000,
"llama-3.1-70b": 128000,
"llama-3.1-8b": 128000,
"llama-3.2-90b": 128000,
"llama-3.2-11b": 128000,
"llama-3.3-70b": 128000,
"llama-4-maverick": 1000000,
"llama-4-scout": 10000000,
# DeepSeek models
"deepseek-chat": 64000,
"deepseek-coder": 64000,
"deepseek-r1": 64000,
# Cohere models
"command-r": 128000,
"command-r-plus": 128000,
}
# Default fallback for unknown models
DEFAULT_CONTEXT_WINDOW = 8192
def get_model_context_window(model_name: Optional[str]) -> int:
"""
Get the context window size for a model.
Args:
model_name: Name of the model (e.g., "gpt-4-turbo", "claude-3-opus")
Returns:
Context window size in tokens
"""
if not model_name:
return DEFAULT_CONTEXT_WINDOW
# Normalize model name
model_lower = model_name.lower().strip()
# Direct lookup
if model_lower in MODEL_CONTEXT_WINDOWS:
return MODEL_CONTEXT_WINDOWS[model_lower]
# Prefix matching for versioned models
for prefix, tokens in MODEL_CONTEXT_WINDOWS.items():
if model_lower.startswith(prefix):
return tokens
# Pattern matching for common patterns
if "gpt-4o" in model_lower or "gpt-4-turbo" in model_lower:
return 128000
if "gpt-4" in model_lower:
return 8192
if "gpt-3.5" in model_lower:
return 16385
if "claude-3" in model_lower or "claude-4" in model_lower:
return 200000
if "gemini-1.5" in model_lower or "gemini-2" in model_lower:
return 1000000
if "gemini" in model_lower:
return 32768
if "llama-4" in model_lower:
return 1000000
if "llama-3" in model_lower:
return 128000
if "mistral" in model_lower or "mixtral" in model_lower:
return 32768
if "deepseek" in model_lower:
return 64000
if "command" in model_lower:
return 128000
return DEFAULT_CONTEXT_WINDOW
@dataclass
class TokenBudget:
"""
Dynamic token budget calculation for context building.
Calculates available tokens for retrieved context based on:
- Model's maximum context window
- Reserved tokens for response generation
- Reserved tokens for system prompt
- Reserved tokens for conversation history
Attributes:
model_max_tokens: Model's maximum context window
reserved_response_tokens: Tokens reserved for response generation
reserved_system_tokens: Tokens reserved for system prompt
reserved_history_tokens: Tokens reserved for conversation history
"""
model_max_tokens: int = 128000
reserved_response_tokens: int = 4096
reserved_system_tokens: int = 1000
reserved_history_tokens: int = 2000
@property
def max_context_tokens(self) -> int:
"""
Calculate maximum tokens available for retrieved context.
This is the static budget based on reserved allocations.
"""
available = (
self.model_max_tokens
- self.reserved_response_tokens
- self.reserved_system_tokens
- self.reserved_history_tokens
)
return max(0, available)
def dynamic_budget(
self,
prompt_tokens: int = 0,
history_tokens: int = 0,
system_tokens: int = 0,
) -> int:
"""
Calculate remaining tokens for context after accounting for actual usage.
Args:
prompt_tokens: Actual tokens used by current prompt
history_tokens: Actual tokens used by conversation history
system_tokens: Actual tokens used by system prompt
Returns:
Remaining tokens available for retrieved context
"""
used = (
prompt_tokens
+ history_tokens
+ system_tokens
+ self.reserved_response_tokens
)
remaining = self.model_max_tokens - used
return max(0, remaining)
@classmethod
def from_model(cls, model_name: Optional[str], **kwargs) -> "TokenBudget":
"""
Create TokenBudget from model name.
Args:
model_name: Name of the model
**kwargs: Override default reserved token values
Returns:
TokenBudget configured for the model
"""
context_window = get_model_context_window(model_name)
return cls(model_max_tokens=context_window, **kwargs)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"model_max_tokens": self.model_max_tokens,
"reserved_response_tokens": self.reserved_response_tokens,
"reserved_system_tokens": self.reserved_system_tokens,
"reserved_history_tokens": self.reserved_history_tokens,
"max_context_tokens": self.max_context_tokens,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "TokenBudget":
"""Create from dictionary."""
return cls(
model_max_tokens=data.get("model_max_tokens", 128000),
reserved_response_tokens=data.get("reserved_response_tokens", 4096),
reserved_system_tokens=data.get("reserved_system_tokens", 1000),
reserved_history_tokens=data.get("reserved_history_tokens", 2000),
)
@runtime_checkable
class BudgetEnforcerProtocol(Protocol):
"""
Protocol for enforcing token budgets on retrieved chunks.
Implementations must provide an enforce method that truncates
chunks to fit within the available token budget.
"""
def enforce(
self,
chunks: List[Dict[str, Any]],
budget: TokenBudget,
prompt_tokens: int = 0,
history_tokens: int = 0,
) -> List[Dict[str, Any]]:
"""
Enforce token budget on chunks.
Args:
chunks: List of chunks with 'text' and 'metadata' keys
budget: TokenBudget instance
prompt_tokens: Tokens used by current prompt
history_tokens: Tokens used by conversation history
Returns:
Truncated list of chunks that fit within budget
"""
...
def estimate_tokens(text: str) -> int:
"""
Estimate token count for text.
Uses a simple heuristic: ~4 characters per token for English text.
This is a fast approximation; for exact counts, use tiktoken.
Args:
text: Text to estimate tokens for
Returns:
Estimated token count
"""
if not text:
return 0
# Average of 4 characters per token for English
# This is a conservative estimate
return len(text) // 4 + 1
class DefaultBudgetEnforcer:
"""
Default implementation of BudgetEnforcerProtocol.
Truncates chunks to fit within available token budget.
Uses simple token estimation for speed.
"""
def __init__(self, token_estimator=None):
"""
Initialize enforcer.
Args:
token_estimator: Optional custom token estimation function
"""
self._token_estimator = token_estimator or estimate_tokens
def enforce(
self,
chunks: List[Dict[str, Any]],
budget: TokenBudget,
prompt_tokens: int = 0,
history_tokens: int = 0,
) -> List[Dict[str, Any]]:
"""
Enforce token budget on chunks.
Iteratively adds chunks until budget is exhausted.
"""
available = budget.dynamic_budget(
prompt_tokens=prompt_tokens,
history_tokens=history_tokens,
)
result = []
used_tokens = 0
for chunk in chunks:
text = chunk.get("text", "")
chunk_tokens = self._token_estimator(text)
if used_tokens + chunk_tokens <= available:
result.append(chunk)
used_tokens += chunk_tokens
else:
# Try to fit partial chunk if there's remaining space
remaining = available - used_tokens
if remaining > 100: # Only if meaningful space remains
# Truncate text to fit
chars_to_keep = remaining * 4 # Approximate
truncated_text = text[:chars_to_keep]
if truncated_text:
truncated_chunk = {
**chunk,
"text": truncated_text + "...",
"metadata": {
**chunk.get("metadata", {}),
"_truncated": True,
},
}
result.append(truncated_chunk)
break
return result
@@ -0,0 +1,277 @@
"""
Contextual Compression for PraisonAI Agents (Phase 5).
Provides context compression to fit within token budgets:
- Redundancy removal / deduplication
- Query-relevant sentence extraction
- LLM-based summarization fallback
No heavy imports at module level.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
@runtime_checkable
class CompressorProtocol(Protocol):
"""Protocol for compression implementations."""
def compress(
self,
chunks: List[Dict[str, Any]],
query: str,
target_tokens: int,
) -> List[Dict[str, Any]]:
"""Compress chunks to fit target token budget."""
...
@dataclass
class CompressionResult:
"""Result of compression operation."""
chunks: List[Dict[str, Any]] = field(default_factory=list)
original_tokens: int = 0
compressed_tokens: int = 0
compression_ratio: float = 1.0
method_used: str = "none"
def estimate_tokens(text: str) -> int:
"""Estimate token count (~4 chars per token)."""
if not text:
return 0
return len(text) // 4 + 1
class ContextCompressor:
"""
Compresses retrieved context to fit within token budget.
Compression strategies (applied in order):
1. Deduplication - remove redundant chunks
2. Sentence extraction - keep query-relevant sentences
3. Truncation - cut to fit budget
4. Summarization - LLM fallback (optional)
"""
def __init__(
self,
llm=None,
similarity_threshold: float = 0.85,
verbose: bool = False,
):
"""
Initialize compressor.
Args:
llm: Optional LLM for summarization fallback
similarity_threshold: Threshold for deduplication
verbose: Enable verbose logging
"""
self._llm = llm
self._similarity_threshold = similarity_threshold
self._verbose = verbose
def compress(
self,
chunks: List[Dict[str, Any]],
query: str,
target_tokens: int,
compression_ratio: float = 0.5,
) -> CompressionResult:
"""
Compress chunks to fit target token budget.
Args:
chunks: List of chunks with 'text' and 'metadata'
query: Original query for relevance scoring
target_tokens: Target token budget
compression_ratio: Target compression ratio (0.5 = 50% reduction)
Returns:
CompressionResult with compressed chunks
"""
result = CompressionResult()
if not chunks:
return result
# Calculate original tokens
result.original_tokens = sum(
estimate_tokens(c.get('text', '')) for c in chunks
)
# Step 1: Deduplicate
deduped = self._deduplicate(chunks)
# Step 2: Extract relevant sentences
extracted = self._extract_relevant(deduped, query)
# Step 3: Truncate to fit budget
truncated = self._truncate_to_budget(extracted, target_tokens)
# Calculate compressed tokens
result.compressed_tokens = sum(
estimate_tokens(c.get('text', '')) for c in truncated
)
if result.original_tokens > 0:
result.compression_ratio = result.compressed_tokens / result.original_tokens
result.chunks = truncated
result.method_used = "dedupe+extract+truncate"
return result
def _deduplicate(
self,
chunks: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""Remove duplicate or near-duplicate chunks."""
if not chunks:
return []
seen_texts = set()
unique = []
for chunk in chunks:
text = chunk.get('text', '').strip()
# Skip empty
if not text:
continue
# Simple dedup by exact match
text_key = text[:200].lower() # First 200 chars as key
if text_key in seen_texts:
continue
seen_texts.add(text_key)
unique.append(chunk)
return unique
def _extract_relevant(
self,
chunks: List[Dict[str, Any]],
query: str,
) -> List[Dict[str, Any]]:
"""Extract query-relevant sentences from chunks."""
if not chunks or not query:
return chunks
query_words = set(query.lower().split())
extracted = []
for chunk in chunks:
text = chunk.get('text', '')
# Split into sentences
sentences = self._split_sentences(text)
# Score sentences by relevance
relevant_sentences = []
for sentence in sentences:
sentence_words = set(sentence.lower().split())
overlap = len(query_words & sentence_words)
# Keep sentences with any overlap or important ones
if overlap > 0 or len(sentence) > 50:
relevant_sentences.append(sentence)
# Reconstruct text from relevant sentences
if relevant_sentences:
extracted.append({
**chunk,
'text': ' '.join(relevant_sentences),
'metadata': {
**chunk.get('metadata', {}),
'_compressed': True,
},
})
return extracted
def _split_sentences(self, text: str) -> List[str]:
"""Split text into sentences."""
import re
# Simple sentence splitting
sentences = re.split(r'(?<=[.!?])\s+', text)
return [s.strip() for s in sentences if s.strip()]
def _truncate_to_budget(
self,
chunks: List[Dict[str, Any]],
target_tokens: int,
) -> List[Dict[str, Any]]:
"""Truncate chunks to fit token budget."""
if not chunks:
return []
result = []
used_tokens = 0
for chunk in chunks:
text = chunk.get('text', '')
chunk_tokens = estimate_tokens(text)
if used_tokens + chunk_tokens <= target_tokens:
result.append(chunk)
used_tokens += chunk_tokens
else:
# Partial fit
remaining = target_tokens - used_tokens
if remaining > 50: # Only if meaningful space
chars_to_keep = remaining * 4
truncated_text = text[:chars_to_keep] + "..."
result.append({
**chunk,
'text': truncated_text,
'metadata': {
**chunk.get('metadata', {}),
'_truncated': True,
},
})
break
return result
def summarize(
self,
chunks: List[Dict[str, Any]],
query: str,
target_tokens: int,
) -> str:
"""
Summarize chunks using LLM (fallback method).
Args:
chunks: Chunks to summarize
query: Original query for context
target_tokens: Target token budget
Returns:
Summarized text
"""
if not self._llm or not chunks:
return ""
# Combine chunk texts
combined = "\n\n".join(c.get('text', '') for c in chunks)
# Create summarization prompt
prompt = f"""Summarize the following content to answer this query: {query}
Content:
{combined}
Provide a concise summary that captures the key information relevant to the query.
Keep the summary under {target_tokens * 4} characters."""
try:
response = self._llm.generate(prompt)
return response if isinstance(response, str) else str(response)
except Exception:
return combined[:target_tokens * 4]
@@ -57,7 +57,9 @@ def deduplicate_chunks(
continue
# Handle different result formats
text = result.get("text") or result.get("memory", "")
source = result.get("metadata", {}).get("source", "")
# CRITICAL: Handle metadata=None from mem0 - ensure always dict
metadata = result.get("metadata") or {}
source = metadata.get("source", "")
chunk_id = _chunk_hash(text, source)
@@ -146,8 +148,10 @@ def build_context(
# Build chunk text with optional source
if include_source:
source = result.get("metadata", {}).get("source", "")
filename = result.get("metadata", {}).get("filename", "")
# CRITICAL: Handle metadata=None from mem0 - ensure always dict
metadata = result.get("metadata") or {}
source = metadata.get("source", "")
filename = metadata.get("filename", "")
source_label = filename or source or f"Source {i + 1}"
chunk_text = f"[{source_label}]\n{text}"
else:
@@ -33,7 +33,8 @@ class DefaultCitationFormatter:
for i, result in enumerate(results):
# Handle different result formats
text = result.get("text") or result.get("memory", "")
metadata = result.get("metadata", {})
# CRITICAL: Handle metadata=None from mem0 - ensure always dict
metadata = result.get("metadata") or {}
score = result.get("score", 0.0)
# Extract source info
@@ -71,6 +71,23 @@ class RetrievalConfig:
citations: bool = True
citations_mode: CitationsMode = CitationsMode.APPEND
# Token budget settings (Phase 1)
model_context_window: Optional[int] = None # Auto-detect if None
reserved_response_tokens: int = 4096
dynamic_budget: bool = True # Use dynamic budget calculation
# Strategy selection (Phase 3)
strategy: str = "auto" # auto, direct, basic, hybrid, reranked, compressed, hierarchical
# Compression settings (Phase 5)
compress: bool = False
compression_ratio: float = 0.5 # Target compression ratio
# Filtering settings (Phase 4)
include_glob: Optional[list] = None # Glob patterns to include
exclude_glob: Optional[list] = None # Glob patterns to exclude
path_filter: Optional[str] = None # Regex for path filtering
# Vector store configuration
vector_store_provider: str = "chroma"
vector_store_config: Dict[str, Any] = field(default_factory=dict)
@@ -96,6 +113,54 @@ class RetrievalConfig:
</retrieved_context>"""
system_separation: bool = True # Use system prompt separation for safety
def get_token_budget(self, model_name: Optional[str] = None):
"""
Get TokenBudget instance for this config.
Args:
model_name: Optional model name for context window detection
Returns:
TokenBudget configured for this retrieval config
"""
from .budget import TokenBudget, get_model_context_window
# Use configured context window or auto-detect from model
if self.model_context_window is not None:
context_window = self.model_context_window
elif model_name:
context_window = get_model_context_window(model_name)
else:
context_window = 128000 # Default fallback
return TokenBudget(
model_max_tokens=context_window,
reserved_response_tokens=self.reserved_response_tokens,
)
def get_strategy(self, corpus_stats=None):
"""
Get retrieval strategy based on config and corpus stats.
Args:
corpus_stats: Optional CorpusStats for auto-selection
Returns:
RetrievalStrategy enum value
"""
from .strategy import select_strategy, RetrievalStrategy
if corpus_stats is None:
# No stats, use explicit strategy or default
if self.strategy and self.strategy != "auto":
try:
return RetrievalStrategy(self.strategy.lower())
except ValueError:
pass
return RetrievalStrategy.BASIC
return select_strategy(corpus_stats, override=self.strategy)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
@@ -108,6 +173,15 @@ class RetrievalConfig:
"hybrid": self.hybrid,
"citations": self.citations,
"citations_mode": self.citations_mode.value if isinstance(self.citations_mode, CitationsMode) else self.citations_mode,
"model_context_window": self.model_context_window,
"reserved_response_tokens": self.reserved_response_tokens,
"dynamic_budget": self.dynamic_budget,
"strategy": self.strategy,
"compress": self.compress,
"compression_ratio": self.compression_ratio,
"include_glob": self.include_glob,
"exclude_glob": self.exclude_glob,
"path_filter": self.path_filter,
"vector_store_provider": self.vector_store_provider,
"vector_store_config": self.vector_store_config,
"collection_name": self.collection_name,
@@ -147,6 +221,15 @@ class RetrievalConfig:
hybrid=data.get("hybrid", False),
citations=data.get("citations", True),
citations_mode=citations_mode,
model_context_window=data.get("model_context_window"),
reserved_response_tokens=data.get("reserved_response_tokens", 4096),
dynamic_budget=data.get("dynamic_budget", True),
strategy=data.get("strategy", "auto"),
compress=data.get("compress", False),
compression_ratio=data.get("compression_ratio", 0.5),
include_glob=data.get("include_glob"),
exclude_glob=data.get("exclude_glob"),
path_filter=data.get("path_filter"),
vector_store_provider=data.get("vector_store_provider", "chroma"),
vector_store_config=data.get("vector_store_config", {}),
collection_name=data.get("collection_name"),
@@ -0,0 +1,309 @@
"""
Smart Retriever for PraisonAI Agents (Phase 4).
Provides enhanced retrieval with:
- Hybrid retrieval (keyword + semantic)
- Metadata filtering
- Path/glob filtering
- Pluggable reranking
No heavy imports at module level.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
@runtime_checkable
class RetrieverProtocol(Protocol):
"""Protocol for retrieval implementations."""
def retrieve(
self,
query: str,
top_k: int = 10,
filters: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
"""Retrieve relevant chunks for query."""
...
@runtime_checkable
class RerankerProtocol(Protocol):
"""Protocol for reranking implementations."""
name: str
def rerank(
self,
query: str,
results: List[Dict[str, Any]],
top_k: Optional[int] = None,
) -> List[Dict[str, Any]]:
"""Rerank results by relevance to query."""
...
@dataclass
class RetrievalResult:
"""Result of a retrieval operation."""
chunks: List[Dict[str, Any]] = field(default_factory=list)
total_found: int = 0
strategy_used: str = "basic"
reranked: bool = False
filtered: bool = False
metadata: Dict[str, Any] = field(default_factory=dict)
class SmartRetriever:
"""
Enhanced retriever with hybrid search, filtering, and reranking.
Implements the retrieval ladder based on strategy:
- DIRECT: Return all content
- BASIC: Semantic search only
- HYBRID: Keyword + semantic search
- RERANKED: Hybrid + cross-encoder reranking
"""
def __init__(
self,
knowledge=None,
reranker: Optional[RerankerProtocol] = None,
verbose: bool = False,
):
"""
Initialize SmartRetriever.
Args:
knowledge: Knowledge instance for retrieval
reranker: Optional reranker implementation
verbose: Enable verbose logging
"""
self._knowledge = knowledge
self._reranker = reranker
self._verbose = verbose
def retrieve(
self,
query: str,
strategy: str = "basic",
top_k: int = 10,
rerank_top_k: Optional[int] = None,
include_glob: Optional[List[str]] = None,
exclude_glob: Optional[List[str]] = None,
path_filter: Optional[str] = None,
metadata_filters: Optional[Dict[str, Any]] = None,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
) -> RetrievalResult:
"""
Retrieve relevant chunks using specified strategy.
Args:
query: Search query
strategy: Retrieval strategy (direct, basic, hybrid, reranked)
top_k: Number of results to return
rerank_top_k: Number of results after reranking
include_glob: Glob patterns to include
exclude_glob: Glob patterns to exclude
path_filter: Regex for path filtering
metadata_filters: Metadata key-value filters
user_id: User ID for scoping
agent_id: Agent ID for scoping
Returns:
RetrievalResult with chunks and metadata
"""
result = RetrievalResult(strategy_used=strategy)
if not self._knowledge:
return result
# Get initial results from knowledge store
search_results = self._search(
query=query,
top_k=top_k * 3 if strategy in ("reranked", "hybrid") else top_k,
user_id=user_id,
agent_id=agent_id,
)
chunks = self._normalize_results(search_results)
result.total_found = len(chunks)
# Apply filters
if include_glob or exclude_glob or path_filter or metadata_filters:
chunks = self._apply_filters(
chunks,
include_glob=include_glob,
exclude_glob=exclude_glob,
path_filter=path_filter,
metadata_filters=metadata_filters,
)
result.filtered = True
# Apply reranking if strategy requires it
if strategy in ("reranked", "compressed", "hierarchical") and self._reranker:
chunks = self._reranker.rerank(query, chunks, top_k=rerank_top_k or top_k)
result.reranked = True
# Limit to top_k
result.chunks = chunks[:top_k]
return result
def _search(
self,
query: str,
top_k: int,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
) -> Any:
"""Perform search on knowledge store."""
try:
return self._knowledge.search(
query,
user_id=user_id,
agent_id=agent_id,
limit=top_k,
)
except Exception:
return []
def _normalize_results(self, results: Any) -> List[Dict[str, Any]]:
"""Normalize search results to list of dicts."""
if not results:
return []
# Handle dict with 'results' key
if isinstance(results, dict):
if 'results' in results:
items = results.get('results', [])
else:
items = [results]
elif isinstance(results, list):
items = results
else:
return []
normalized = []
for item in items:
if item is None:
continue
if isinstance(item, dict):
text = item.get('text') or item.get('memory', '')
metadata = item.get('metadata') or {}
normalized.append({
'text': str(text),
'metadata': metadata if isinstance(metadata, dict) else {},
'score': float(item.get('score', 0.0) or 0.0),
})
return normalized
def _apply_filters(
self,
chunks: List[Dict[str, Any]],
include_glob: Optional[List[str]] = None,
exclude_glob: Optional[List[str]] = None,
path_filter: Optional[str] = None,
metadata_filters: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
"""Apply filters to chunks."""
import fnmatch
import re
filtered = []
for chunk in chunks:
metadata = chunk.get('metadata', {})
filename = metadata.get('filename', '')
filepath = metadata.get('filepath', metadata.get('source', ''))
# Check include patterns
if include_glob:
matched = False
for pattern in include_glob:
if fnmatch.fnmatch(filename, pattern) or fnmatch.fnmatch(filepath, pattern):
matched = True
break
if not matched:
continue
# Check exclude patterns
if exclude_glob:
excluded = False
for pattern in exclude_glob:
if fnmatch.fnmatch(filename, pattern) or fnmatch.fnmatch(filepath, pattern):
excluded = True
break
if excluded:
continue
# Check path filter regex
if path_filter:
try:
if not re.search(path_filter, filepath):
continue
except re.error:
pass
# Check metadata filters
if metadata_filters:
match = True
for key, value in metadata_filters.items():
if metadata.get(key) != value:
match = False
break
if not match:
continue
filtered.append(chunk)
return filtered
class SimpleReranker:
"""
Simple reranker using keyword matching.
For production, use cross-encoder models like:
- sentence-transformers/ms-marco-MiniLM-L-6-v2
- Cohere Rerank API
"""
name = "simple"
def rerank(
self,
query: str,
results: List[Dict[str, Any]],
top_k: Optional[int] = None,
) -> List[Dict[str, Any]]:
"""Rerank results by keyword overlap with query."""
query_words = set(query.lower().split())
scored = []
for result in results:
text = result.get('text', '').lower()
text_words = set(text.split())
# Calculate overlap score
overlap = len(query_words & text_words)
original_score = result.get('score', 0.0)
# Combined score
combined_score = original_score + (overlap * 0.1)
scored.append({
**result,
'score': combined_score,
'_rerank_overlap': overlap,
})
# Sort by combined score
scored.sort(key=lambda x: x['score'], reverse=True)
if top_k:
return scored[:top_k]
return scored
@@ -0,0 +1,173 @@
"""
Retrieval Strategy Selection for PraisonAI Agents.
Provides corpus-aware strategy selection for optimal retrieval performance.
Strategies are selected based on corpus size and characteristics.
No heavy imports - only stdlib and typing.
"""
from enum import Enum
from typing import Any, Dict, Optional, TYPE_CHECKING
if TYPE_CHECKING:
from ..knowledge.indexing import CorpusStats
class RetrievalStrategy(str, Enum):
"""
Retrieval strategy levels based on corpus size and complexity.
Each level adds more sophisticated retrieval techniques:
- DIRECT: Load all content directly (tiny corpora)
- BASIC: Semantic search only (small corpora)
- HYBRID: Keyword + semantic search (medium corpora)
- RERANKED: Hybrid + cross-encoder reranking (large corpora)
- COMPRESSED: Reranked + contextual compression (very large corpora)
- HIERARCHICAL: Summaries + routing (massive corpora)
"""
DIRECT = "direct"
BASIC = "basic"
HYBRID = "hybrid"
RERANKED = "reranked"
COMPRESSED = "compressed"
HIERARCHICAL = "hierarchical"
# Strategy thresholds based on file count
STRATEGY_THRESHOLDS: Dict[int, RetrievalStrategy] = {
10: RetrievalStrategy.DIRECT, # < 10 files
100: RetrievalStrategy.BASIC, # < 100 files
1000: RetrievalStrategy.HYBRID, # < 1000 files
10000: RetrievalStrategy.RERANKED, # < 10000 files
100000: RetrievalStrategy.COMPRESSED, # < 100000 files
}
DEFAULT_STRATEGY = RetrievalStrategy.HIERARCHICAL # >= 100000 files
# Strategy descriptions for logging and UI
STRATEGY_DESCRIPTIONS: Dict[RetrievalStrategy, str] = {
RetrievalStrategy.DIRECT: "Direct loading - all content loaded into context (< 10 files)",
RetrievalStrategy.BASIC: "Basic semantic search - embedding-based retrieval (< 100 files)",
RetrievalStrategy.HYBRID: "Hybrid retrieval - keyword + semantic search (< 1000 files)",
RetrievalStrategy.RERANKED: "Reranked retrieval - hybrid + cross-encoder reranking (< 10000 files)",
RetrievalStrategy.COMPRESSED: "Compressed retrieval - reranked + contextual compression (< 100000 files)",
RetrievalStrategy.HIERARCHICAL: "Hierarchical retrieval - summaries + top-down routing (>= 100000 files)",
}
def get_strategy_description(strategy: RetrievalStrategy) -> str:
"""
Get human-readable description of a strategy.
Args:
strategy: The retrieval strategy
Returns:
Description string
"""
return STRATEGY_DESCRIPTIONS.get(strategy, f"Unknown strategy: {strategy}")
def select_strategy(
corpus_stats: "CorpusStats",
override: Optional[str] = None,
) -> RetrievalStrategy:
"""
Select optimal retrieval strategy based on corpus characteristics.
Args:
corpus_stats: Statistics about the indexed corpus
override: Optional explicit strategy override ("auto" for auto-select)
Returns:
Selected RetrievalStrategy
"""
# Handle explicit override
if override and override != "auto":
try:
return RetrievalStrategy(override.lower())
except ValueError:
pass # Fall through to auto-selection
# Auto-select based on file count
file_count = corpus_stats.file_count if corpus_stats else 0
for threshold, strategy in sorted(STRATEGY_THRESHOLDS.items()):
if file_count < threshold:
return strategy
return DEFAULT_STRATEGY
def get_strategy_config(strategy: RetrievalStrategy) -> Dict[str, Any]:
"""
Get configuration parameters for a strategy.
Args:
strategy: The retrieval strategy
Returns:
Dict of configuration parameters
"""
configs = {
RetrievalStrategy.DIRECT: {
"use_retrieval": False,
"load_all": True,
"top_k": None,
"rerank": False,
"compress": False,
"use_hierarchy": False,
},
RetrievalStrategy.BASIC: {
"use_retrieval": True,
"load_all": False,
"top_k": 5,
"rerank": False,
"compress": False,
"use_hierarchy": False,
},
RetrievalStrategy.HYBRID: {
"use_retrieval": True,
"load_all": False,
"top_k": 10,
"use_keyword_search": True,
"rerank": False,
"compress": False,
"use_hierarchy": False,
},
RetrievalStrategy.RERANKED: {
"use_retrieval": True,
"load_all": False,
"top_k": 20,
"rerank_top_k": 5,
"use_keyword_search": True,
"rerank": True,
"compress": False,
"use_hierarchy": False,
},
RetrievalStrategy.COMPRESSED: {
"use_retrieval": True,
"load_all": False,
"top_k": 30,
"rerank_top_k": 10,
"use_keyword_search": True,
"rerank": True,
"compress": True,
"compression_ratio": 0.5,
"use_hierarchy": False,
},
RetrievalStrategy.HIERARCHICAL: {
"use_retrieval": True,
"load_all": False,
"top_k": 50,
"rerank_top_k": 15,
"use_keyword_search": True,
"rerank": True,
"compress": True,
"compression_ratio": 0.3,
"use_hierarchy": True,
"hierarchy_levels": 3,
},
}
return configs.get(strategy, configs[RetrievalStrategy.BASIC])
@@ -0,0 +1,416 @@
"""
Hierarchical Summarization for PraisonAI Agents (Phase 6).
Provides hierarchical summaries for massive corpora:
- File-level summaries
- Folder-level summaries
- Project-level summary
- Top-down query routing
No heavy imports at module level.
"""
import json
import os
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
@runtime_checkable
class SummarizerProtocol(Protocol):
"""Protocol for summarization implementations."""
def summarize(self, text: str, max_tokens: int = 500) -> str:
"""Summarize text to max_tokens."""
...
@dataclass
class SummaryNode:
"""A node in the summary hierarchy."""
path: str
level: int # 0=chunk, 1=file, 2=folder, 3=project
summary: str = ""
children: List[str] = field(default_factory=list)
token_count: int = 0
created_at: Optional[str] = None
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
return {
"path": self.path,
"level": self.level,
"summary": self.summary,
"children": self.children,
"token_count": self.token_count,
"created_at": self.created_at,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "SummaryNode":
return cls(
path=data.get("path", ""),
level=data.get("level", 0),
summary=data.get("summary", ""),
children=data.get("children", []),
token_count=data.get("token_count", 0),
created_at=data.get("created_at"),
metadata=data.get("metadata", {}),
)
@dataclass
class HierarchyResult:
"""Result of hierarchy building or querying."""
nodes: Dict[str, SummaryNode] = field(default_factory=dict)
root_path: str = ""
total_files: int = 0
total_tokens: int = 0
levels: int = 3
built_at: Optional[str] = None
def estimate_tokens(text: str) -> int:
"""Estimate token count (~4 chars per token)."""
if not text:
return 0
return len(text) // 4 + 1
class HierarchicalSummarizer:
"""
Builds and queries hierarchical summaries for large corpora.
Hierarchy levels:
- Level 0: Individual chunks
- Level 1: File summaries
- Level 2: Folder summaries
- Level 3: Project summary
Query routing:
- Start at project level
- Route to relevant folders
- Drill into relevant files
- Return relevant chunks
"""
def __init__(
self,
llm=None,
persist_path: str = ".praison/summaries",
verbose: bool = False,
):
"""
Initialize summarizer.
Args:
llm: LLM for generating summaries
persist_path: Path to persist summaries
verbose: Enable verbose logging
"""
self._llm = llm
self._persist_path = persist_path
self._verbose = verbose
self._hierarchy: Dict[str, SummaryNode] = {}
def build_hierarchy(
self,
corpus_path: str,
levels: int = 3,
force: bool = False,
) -> HierarchyResult:
"""
Build summary hierarchy for corpus.
Args:
corpus_path: Path to corpus root
levels: Number of hierarchy levels (1-3)
force: Force rebuild even if cached
Returns:
HierarchyResult with built hierarchy
"""
result = HierarchyResult(
root_path=corpus_path,
levels=levels,
built_at=datetime.now().isoformat(),
)
# Check for cached hierarchy
if not force:
cached = self._load_hierarchy(corpus_path)
if cached:
return cached
# Collect all files
files = self._collect_files(corpus_path)
result.total_files = len(files)
# Build file-level summaries (Level 1)
for filepath in files:
node = self._build_file_summary(filepath, corpus_path)
if node:
self._hierarchy[filepath] = node
result.nodes[filepath] = node
# Build folder-level summaries (Level 2)
if levels >= 2:
folders = self._get_folders(files, corpus_path)
for folder in folders:
node = self._build_folder_summary(folder, corpus_path)
if node:
self._hierarchy[folder] = node
result.nodes[folder] = node
# Build project-level summary (Level 3)
if levels >= 3:
node = self._build_project_summary(corpus_path)
if node:
self._hierarchy[corpus_path] = node
result.nodes[corpus_path] = node
# Calculate total tokens
result.total_tokens = sum(n.token_count for n in result.nodes.values())
# Persist hierarchy
self._save_hierarchy(corpus_path, result)
return result
def query_hierarchy(
self,
query: str,
corpus_path: str,
max_results: int = 10,
) -> List[Dict[str, Any]]:
"""
Query hierarchy top-down to find relevant content.
Args:
query: Search query
corpus_path: Path to corpus root
max_results: Maximum results to return
Returns:
List of relevant chunks with paths
"""
# Load hierarchy if not in memory
if corpus_path not in self._hierarchy:
cached = self._load_hierarchy(corpus_path)
if cached:
self._hierarchy.update(cached.nodes)
# Start at project level
project_node = self._hierarchy.get(corpus_path)
if not project_node:
return []
# Score folders by relevance
query_words = set(query.lower().split())
relevant_paths = []
for path, node in self._hierarchy.items():
if node.level == 1: # File level
summary_words = set(node.summary.lower().split())
overlap = len(query_words & summary_words)
if overlap > 0:
relevant_paths.append((path, overlap, node))
# Sort by relevance
relevant_paths.sort(key=lambda x: x[1], reverse=True)
# Return top results
results = []
for path, score, node in relevant_paths[:max_results]:
results.append({
"path": path,
"summary": node.summary,
"score": score,
"level": node.level,
"metadata": node.metadata,
})
return results
def _collect_files(self, corpus_path: str) -> List[str]:
"""Collect all indexable files in corpus."""
files = []
extensions = {'.txt', '.md', '.py', '.js', '.ts', '.json', '.yaml', '.yml'}
for root, dirs, filenames in os.walk(corpus_path):
# Skip hidden directories
dirs[:] = [d for d in dirs if not d.startswith('.')]
for filename in filenames:
ext = os.path.splitext(filename)[1].lower()
if ext in extensions:
files.append(os.path.join(root, filename))
return files
def _get_folders(self, files: List[str], corpus_path: str) -> List[str]:
"""Get unique folder paths from files."""
folders = set()
for filepath in files:
folder = os.path.dirname(filepath)
if folder != corpus_path:
folders.add(folder)
return list(folders)
def _build_file_summary(self, filepath: str, corpus_path: str) -> Optional[SummaryNode]:
"""Build summary for a single file."""
try:
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
content = f.read()
# Simple extractive summary (first 500 chars + key lines)
summary = self._extract_summary(content, max_chars=500)
return SummaryNode(
path=filepath,
level=1,
summary=summary,
token_count=estimate_tokens(summary),
created_at=datetime.now().isoformat(),
metadata={
"filename": os.path.basename(filepath),
"relative_path": os.path.relpath(filepath, corpus_path),
},
)
except Exception:
return None
def _build_folder_summary(self, folder: str, corpus_path: str) -> Optional[SummaryNode]:
"""Build summary for a folder from child summaries."""
children = []
child_summaries = []
for path, node in self._hierarchy.items():
if node.level == 1 and os.path.dirname(path) == folder:
children.append(path)
child_summaries.append(node.summary[:200])
if not children:
return None
# Combine child summaries
combined = " | ".join(child_summaries)
summary = f"Folder contains {len(children)} files: {combined[:500]}"
return SummaryNode(
path=folder,
level=2,
summary=summary,
children=children,
token_count=estimate_tokens(summary),
created_at=datetime.now().isoformat(),
metadata={
"folder_name": os.path.basename(folder),
"file_count": len(children),
},
)
def _build_project_summary(self, corpus_path: str) -> Optional[SummaryNode]:
"""Build project-level summary from folder summaries."""
folders = []
folder_summaries = []
for path, node in self._hierarchy.items():
if node.level == 2:
folders.append(path)
folder_summaries.append(node.summary[:100])
# Also include direct file children
files = []
for path, node in self._hierarchy.items():
if node.level == 1 and os.path.dirname(path) == corpus_path:
files.append(path)
combined = " | ".join(folder_summaries[:10])
summary = f"Project with {len(folders)} folders and {len(files)} root files: {combined[:500]}"
return SummaryNode(
path=corpus_path,
level=3,
summary=summary,
children=folders + files,
token_count=estimate_tokens(summary),
created_at=datetime.now().isoformat(),
metadata={
"folder_count": len(folders),
"root_file_count": len(files),
},
)
def _extract_summary(self, content: str, max_chars: int = 500) -> str:
"""Extract summary from content."""
if not content:
return ""
# Get first paragraph or lines
lines = content.split('\n')
summary_lines = []
char_count = 0
for line in lines:
line = line.strip()
if not line:
continue
if char_count + len(line) > max_chars:
break
summary_lines.append(line)
char_count += len(line)
return ' '.join(summary_lines)
def _save_hierarchy(self, corpus_path: str, result: HierarchyResult) -> None:
"""Save hierarchy to disk."""
try:
os.makedirs(self._persist_path, exist_ok=True)
# Create safe filename from path
safe_name = corpus_path.replace('/', '_').replace('\\', '_')
filepath = os.path.join(self._persist_path, f"{safe_name}_hierarchy.json")
data = {
"root_path": result.root_path,
"total_files": result.total_files,
"total_tokens": result.total_tokens,
"levels": result.levels,
"built_at": result.built_at,
"nodes": {k: v.to_dict() for k, v in result.nodes.items()},
}
with open(filepath, 'w', encoding='utf-8') as f:
json.dump(data, f, indent=2)
except Exception:
pass
def _load_hierarchy(self, corpus_path: str) -> Optional[HierarchyResult]:
"""Load hierarchy from disk."""
try:
safe_name = corpus_path.replace('/', '_').replace('\\', '_')
filepath = os.path.join(self._persist_path, f"{safe_name}_hierarchy.json")
if not os.path.exists(filepath):
return None
with open(filepath, 'r', encoding='utf-8') as f:
data = json.load(f)
result = HierarchyResult(
root_path=data.get("root_path", ""),
total_files=data.get("total_files", 0),
total_tokens=data.get("total_tokens", 0),
levels=data.get("levels", 3),
built_at=data.get("built_at"),
)
for path, node_data in data.get("nodes", {}).items():
result.nodes[path] = SummaryNode.from_dict(node_data)
return result
except Exception:
return None
@@ -299,6 +299,11 @@ class Workflow:
status: str = "not_started" # not_started, running, completed, failed
step_statuses: Dict[str, str] = field(default_factory=dict) # {step_name: status}
# Context management (single param for all context features)
context: Optional[Any] = False # False=disabled, True=defaults, ManagerConfig, or ContextManager
_context_manager: Optional[Any] = field(default=None, repr=False)
_context_manager_initialized: bool = field(default=False, repr=False)
def to_dict(self) -> Dict[str, Any]:
return {
"name": self.name,
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "praisonaiagents"
version = "0.7.1"
version = "0.8.1"
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
readme = "README.md"
requires-python = ">=3.10"
@@ -0,0 +1,235 @@
"""
End-to-End Tests for Large Context Knowledge Handling (Phase 8).
Tests the full pipeline: indexing -> retrieval -> budget enforcement -> context building.
"""
import os
import tempfile
import pytest
class TestLargeContextE2E:
"""End-to-end tests for large context handling."""
def test_full_indexing_pipeline(self):
"""Test complete indexing pipeline."""
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.knowledge.indexing import CorpusStats, IndexResult
with tempfile.TemporaryDirectory() as tmpdir:
# Create test corpus
for i in range(5):
with open(os.path.join(tmpdir, f"doc{i}.txt"), "w") as f:
f.write(f"Document {i} content about topic {i % 3}")
# Index
knowledge = Knowledge()
result = knowledge.index(tmpdir, user_id="e2e_test", incremental=False)
assert isinstance(result, IndexResult)
assert result.success is True
assert result.files_indexed == 5
# Get stats
stats = knowledge.get_corpus_stats()
assert isinstance(stats, CorpusStats)
assert stats.file_count == 5
def test_strategy_selection_by_corpus_size(self):
"""Test strategy auto-selection based on corpus size."""
from praisonaiagents.knowledge.indexing import CorpusStats
from praisonaiagents.rag.strategy import select_strategy, RetrievalStrategy
# Small corpus -> DIRECT
stats = CorpusStats(file_count=5)
assert select_strategy(stats) == RetrievalStrategy.DIRECT
# Medium corpus -> BASIC
stats = CorpusStats(file_count=50)
assert select_strategy(stats) == RetrievalStrategy.BASIC
# Large corpus -> HYBRID
stats = CorpusStats(file_count=500)
assert select_strategy(stats) == RetrievalStrategy.HYBRID
def test_token_budget_enforcement(self):
"""Test token budget calculation and enforcement."""
from praisonaiagents.rag.budget import TokenBudget, DefaultBudgetEnforcer
budget = TokenBudget(
model_max_tokens=8000,
reserved_response_tokens=2000,
)
# Calculate available using dynamic_budget
available = budget.dynamic_budget(
system_tokens=500,
history_tokens=1000,
)
# 8000 - 2000 - 500 - 1000 = 4500
assert available == 4500
# Test enforcement
enforcer = DefaultBudgetEnforcer()
chunks = [
{"text": "A" * 1000}, # ~250 tokens
{"text": "B" * 1000}, # ~250 tokens
{"text": "C" * 1000}, # ~250 tokens
]
# Create a small budget to test enforcement
small_budget = TokenBudget(model_max_tokens=600, reserved_response_tokens=100)
enforced = enforcer.enforce(chunks, small_budget)
# Should fit within budget (500 tokens available, ~750 in chunks)
assert len(enforced) <= 3
def test_compression_pipeline(self):
"""Test contextual compression."""
from praisonaiagents.rag.compressor import ContextCompressor
compressor = ContextCompressor()
chunks = [
{"text": "Python is a programming language. It is widely used."},
{"text": "Python is a programming language. It is widely used."}, # Duplicate
{"text": "JavaScript runs in browsers."},
]
result = compressor.compress(chunks, "Python", target_tokens=500)
# Should deduplicate
assert len(result.chunks) <= 2
assert result.compression_ratio <= 1.0
def test_smart_retriever_filtering(self):
"""Test SmartRetriever with filters."""
from praisonaiagents.rag.retriever import SmartRetriever
retriever = SmartRetriever()
# Test filter application
chunks = [
{"text": "Python code", "metadata": {"filename": "main.py"}},
{"text": "JavaScript", "metadata": {"filename": "app.js"}},
{"text": "More Python", "metadata": {"filename": "utils.py"}},
]
filtered = retriever._apply_filters(chunks, include_glob=["*.py"])
assert len(filtered) == 2
filtered = retriever._apply_filters(chunks, exclude_glob=["utils*"])
assert len(filtered) == 2
def test_hierarchical_summarizer(self):
"""Test hierarchical summary building."""
from praisonaiagents.rag.summarizer import HierarchicalSummarizer
with tempfile.TemporaryDirectory() as tmpdir:
# Create nested structure
src_dir = os.path.join(tmpdir, "src")
os.makedirs(src_dir)
with open(os.path.join(src_dir, "main.py"), "w") as f:
f.write("# Main application\ndef main(): pass")
with open(os.path.join(src_dir, "utils.py"), "w") as f:
f.write("# Utilities\ndef helper(): return True")
summarizer = HierarchicalSummarizer()
result = summarizer.build_hierarchy(tmpdir, levels=2)
assert result.total_files >= 2
assert len(result.nodes) >= 2
def test_retrieval_config_integration(self):
"""Test RetrievalConfig with all new fields."""
from praisonaiagents.rag.retrieval_config import RetrievalConfig
from praisonaiagents.knowledge.indexing import CorpusStats
config = RetrievalConfig(
strategy="auto",
compress=True,
compression_ratio=0.5,
model_context_window=128000,
reserved_response_tokens=4096,
)
# Get budget
budget = config.get_token_budget()
assert budget.model_max_tokens == 128000
# Get strategy
stats = CorpusStats(file_count=500)
strategy = config.get_strategy(stats)
assert strategy.value == "hybrid"
class TestLiveAPIIntegration:
"""Live API tests (require OPENAI_API_KEY and RUN_LIVE_TESTS=1)."""
@pytest.fixture
def skip_if_no_api_key(self):
"""Skip if no API key or live tests disabled."""
if not os.environ.get("OPENAI_API_KEY"):
pytest.skip("OPENAI_API_KEY not set")
if os.environ.get("RUN_LIVE_TESTS") != "1":
pytest.skip("RUN_LIVE_TESTS not set to 1")
def test_live_indexing_and_search(self, skip_if_no_api_key):
"""Test live indexing and search with real API."""
from praisonaiagents.knowledge import Knowledge
with tempfile.TemporaryDirectory() as tmpdir:
# Create test documents
with open(os.path.join(tmpdir, "python.txt"), "w") as f:
f.write("Python is a high-level programming language known for its simplicity.")
with open(os.path.join(tmpdir, "javascript.txt"), "w") as f:
f.write("JavaScript is the language of the web, running in browsers.")
knowledge = Knowledge()
result = knowledge.index(tmpdir, user_id="live_test", incremental=False)
assert result.files_indexed == 2
# Search
results = knowledge.search("Python programming", user_id="live_test", limit=5)
# Should find results
assert results is not None
class TestPerformance:
"""Performance and import time tests."""
def test_import_time_acceptable(self):
"""Import time should be reasonable."""
import subprocess
import sys
# Measure import time
cmd = [
sys.executable, "-c",
"import time; start=time.time(); import praisonaiagents; print(time.time()-start)"
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
if result.returncode == 0:
import_time = float(result.stdout.strip())
# Should import in under 5 seconds
assert import_time < 5.0, f"Import took {import_time}s"
def test_no_heavy_deps_at_import(self):
"""Heavy dependencies should not be loaded at import."""
import subprocess
import sys
# Check if mem0 is loaded at import
cmd = [
sys.executable, "-c",
"import sys; import praisonaiagents; print('mem0' in sys.modules)"
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
if result.returncode == 0:
# mem0 should NOT be loaded at import
assert result.stdout.strip() == "False", "mem0 loaded at import time"
@@ -0,0 +1,262 @@
#!/usr/bin/env python3
"""
Live API Tests for Large Context Knowledge Handling.
These tests require a real API key and make actual API calls.
They are gated by environment variables and should only run in CI
or when explicitly enabled.
Environment Variables:
PRAISONAI_LIVE_TESTS=1 - Enable live tests
OPENAI_API_KEY - Required for LLM calls
Usage:
# Run all live tests
PRAISONAI_LIVE_TESTS=1 pytest tests/live/test_large_context_live.py -v
# Run specific test
PRAISONAI_LIVE_TESTS=1 pytest tests/live/test_large_context_live.py::test_agent_retrieves_unique_code -v
"""
import os
import tempfile
import shutil
import pytest
# Skip all tests if PRAISONAI_LIVE_TESTS is not set
pytestmark = pytest.mark.skipif(
os.environ.get("PRAISONAI_LIVE_TESTS", "").lower() not in ("1", "true", "yes"),
reason="Live tests disabled. Set PRAISONAI_LIVE_TESTS=1 to enable."
)
# Unique codes that cannot be guessed by the LLM
UNIQUE_CODES = {
"project_code": "XRAY-7K2M9",
"budget_code": "FOXTROT-3N8P1",
"security_code": "TANGO-5Q4R6",
}
@pytest.fixture
def temp_knowledge_dir():
"""Create a temporary directory with test documents."""
temp_dir = tempfile.mkdtemp(prefix='praison_live_test_')
# Create document with unique codes
doc_path = os.path.join(temp_dir, 'confidential.txt')
with open(doc_path, 'w') as f:
f.write(f"""
CONFIDENTIAL INTERNAL DOCUMENT
==============================
Project Information:
- Project Name: Phoenix Initiative
- Project Code: {UNIQUE_CODES['project_code']}
- Start Date: 2024-01-15
Budget Information:
- Total Budget: $2,500,000
- Budget Code: {UNIQUE_CODES['budget_code']}
- Fiscal Year: 2024
Security Information:
- Classification: Internal
- Security Code: {UNIQUE_CODES['security_code']}
- Access Level: Restricted
This document contains sensitive information.
Do not share outside the organization.
""")
yield temp_dir
# Cleanup
shutil.rmtree(temp_dir, ignore_errors=True)
class TestAgentKnowledgeRetrieval:
"""Test Agent retrieval of unique codes from knowledge base."""
def test_agent_retrieves_unique_code(self, temp_knowledge_dir):
"""Agent should retrieve unique project code from knowledge."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Answer based ONLY on the provided knowledge. Quote exact codes.",
knowledge=[temp_knowledge_dir],
user_id="live_test_user",
verbose=False,
)
response = agent.chat("What is the project code for Phoenix Initiative?")
assert UNIQUE_CODES['project_code'] in response.upper(), \
f"Expected code {UNIQUE_CODES['project_code']} not found in response: {response}"
def test_agent_retrieves_budget_code(self, temp_knowledge_dir):
"""Agent should retrieve unique budget code from knowledge."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Answer based ONLY on the provided knowledge. Quote exact codes.",
knowledge=[temp_knowledge_dir],
user_id="live_test_user",
verbose=False,
)
response = agent.chat("What is the budget code?")
assert UNIQUE_CODES['budget_code'] in response.upper(), \
f"Expected code {UNIQUE_CODES['budget_code']} not found in response: {response}"
def test_agent_retrieves_security_code(self, temp_knowledge_dir):
"""Agent should retrieve unique security code from knowledge."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Answer based ONLY on the provided knowledge. Quote exact codes.",
knowledge=[temp_knowledge_dir],
user_id="live_test_user",
verbose=False,
)
response = agent.chat("What is the security code?")
assert UNIQUE_CODES['security_code'] in response.upper(), \
f"Expected code {UNIQUE_CODES['security_code']} not found in response: {response}"
class TestIncrementalIndexing:
"""Test incremental indexing with real API."""
def test_incremental_indexing_skips_unchanged(self, temp_knowledge_dir):
"""Incremental indexing should skip unchanged files."""
from praisonaiagents.knowledge import Knowledge
knowledge = Knowledge()
# First index
result1 = knowledge.index(
temp_knowledge_dir,
user_id="live_test_user",
incremental=True,
)
assert result1.files_indexed >= 1, "Should index at least one file"
# Second index should skip unchanged files
result2 = knowledge.index(
temp_knowledge_dir,
user_id="live_test_user",
incremental=True,
)
assert result2.files_skipped >= 1, "Should skip unchanged files"
assert result2.files_indexed == 0, "Should not re-index unchanged files"
class TestTokenBudgeting:
"""Test token budget calculations."""
def test_token_budget_calculation(self):
"""Token budget should calculate correctly for different models."""
from praisonaiagents.rag import TokenBudget
# Test with default settings (gpt-4o-mini context window)
budget = TokenBudget(model_max_tokens=128000)
assert budget.model_max_tokens > 0, "Should have positive context window"
assert budget.reserved_response_tokens > 0, "Should have reserved response tokens"
available = budget.dynamic_budget(
system_tokens=500,
history_tokens=1000,
)
assert available > 0, "Should have available tokens"
assert available < budget.model_max_tokens, "Available should be less than total"
class TestStrategySelection:
"""Test retrieval strategy selection."""
def test_strategy_selection_by_corpus_size(self):
"""Strategy should be selected based on corpus size."""
from praisonaiagents.rag import select_strategy, RetrievalStrategy
from praisonaiagents.knowledge.indexing import CorpusStats
# Small corpus -> DIRECT
small_stats = CorpusStats(file_count=5, total_tokens=100)
small_strategy = select_strategy(small_stats)
assert small_strategy == RetrievalStrategy.DIRECT
# Medium corpus -> BASIC
medium_stats = CorpusStats(file_count=50, total_tokens=5000)
medium_strategy = select_strategy(medium_stats)
assert medium_strategy == RetrievalStrategy.BASIC
# Large corpus -> HYBRID or higher
large_stats = CorpusStats(file_count=500, total_tokens=50000)
large_strategy = select_strategy(large_stats)
assert large_strategy in (
RetrievalStrategy.HYBRID,
RetrievalStrategy.RERANKED,
RetrievalStrategy.COMPRESSED,
)
class TestContextCompression:
"""Test context compression functionality."""
def test_compression_preserves_query_relevant_content(self):
"""Compression should preserve query-relevant content."""
from praisonaiagents.rag import ContextCompressor
compressor = ContextCompressor(
similarity_threshold=0.85,
verbose=False,
)
# Chunks with metadata (as expected by the API)
chunks = [
{"text": "This is a long document with lots of filler content.", "score": 0.9},
{"text": "Lorem ipsum dolor sit amet, consectetur adipiscing elit.", "score": 0.8},
{"text": "The secret code is: COMPRESS-TEST-123.", "score": 0.95},
{"text": "More filler content follows here.", "score": 0.7},
{"text": "Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.", "score": 0.6},
]
# ContextCompressor.compress takes chunks, query, and target_tokens
result = compressor.compress(chunks, query="secret code", target_tokens=500)
# Result should be a CompressionResult
assert hasattr(result, 'chunks'), "Should have chunks attribute"
assert len(result.chunks) <= len(chunks), "Should not increase chunk count"
class TestSmartRetriever:
"""Test smart retriever functionality."""
def test_retriever_initialization(self, temp_knowledge_dir):
"""Smart retriever should initialize correctly."""
from praisonaiagents.rag import SmartRetriever
# Test basic initialization
retriever = SmartRetriever()
# Verify retriever has expected attributes
assert retriever is not None, "Retriever should initialize"
# Test with reranker
retriever_with_rerank = SmartRetriever(reranker=None)
assert retriever_with_rerank is not None, "Retriever with reranker should initialize"
if __name__ == "__main__":
# Run tests with verbose output
pytest.main([__file__, "-v", "-s"])
@@ -0,0 +1,150 @@
"""
Live integration tests for mem0 backend.
These tests require:
- OPENAI_API_KEY environment variable
- RUN_LIVE_TESTS=1 environment variable
Run with: RUN_LIVE_TESTS=1 python -m pytest tests/live/test_mem0_integration.py -v
"""
import os
import sys
import tempfile
import shutil
# Skip all tests if RUN_LIVE_TESTS is not set
import pytest
pytestmark = pytest.mark.skipif(
os.environ.get("RUN_LIVE_TESTS") != "1",
reason="Live tests disabled. Set RUN_LIVE_TESTS=1 to run."
)
class TestMem0IntegrationLive:
"""Live tests for mem0 integration."""
def test_knowledge_search_with_mem0_metadata_none(self):
"""Test that Knowledge.search handles mem0 metadata=None correctly."""
from praisonaiagents.knowledge import Knowledge
# Create temp directory with test file
temp_dir = tempfile.mkdtemp(prefix='praison_mem0_test_')
try:
test_file = os.path.join(temp_dir, 'test.txt')
with open(test_file, 'w') as f:
f.write('The capital of France is Paris. This is a test document.')
# Initialize Knowledge and add content
knowledge = Knowledge()
knowledge.add(temp_dir, user_id='test_user_live')
# Search - this should NOT raise AttributeError
results = knowledge.search('capital', user_id='test_user_live')
# Verify results structure
assert isinstance(results, dict)
assert 'results' in results
# Verify metadata handling
for result in results.get('results', []):
if result is not None:
# metadata should be accessible without error
metadata = result.get('metadata')
# Even if None, accessing it shouldn't crash
if metadata is not None:
assert isinstance(metadata, dict)
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
def test_agent_chat_with_knowledge_no_crash(self):
"""Test that Agent.chat with knowledge doesn't crash on metadata=None."""
from praisonaiagents import Agent
temp_dir = tempfile.mkdtemp(prefix='praison_agent_test_')
try:
test_file = os.path.join(temp_dir, 'test.txt')
with open(test_file, 'w') as f:
f.write('The capital of France is Paris.')
# Create agent with knowledge
agent = Agent(
name="TestAgent",
instructions="You are a helpful assistant.",
knowledge=[temp_dir],
user_id='test_user_agent',
)
# This should NOT raise "NoneType has no attribute get"
response = agent.chat("What is the capital of France?")
assert response is not None
assert isinstance(response, str)
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
def test_rag_context_builder_with_mem0_results(self):
"""Test RAG context builder handles mem0 results correctly."""
from praisonaiagents.rag.context import build_context, deduplicate_chunks
# Simulate mem0 results with metadata=None
mem0_results = [
{
"id": "1",
"memory": "Paris is the capital of France.",
"score": 0.95,
"metadata": None, # mem0 returns this!
"user_id": "test",
},
{
"id": "2",
"memory": "London is the capital of UK.",
"score": 0.85,
"metadata": {"source": "geography.txt"},
},
]
# Should not crash
deduped = deduplicate_chunks(mem0_results)
assert len(deduped) == 2
context, used = build_context(mem0_results)
assert "Paris" in context
assert "London" in context
assert len(used) == 2
class TestImportPerformance:
"""Tests for import-time performance."""
def test_core_import_does_not_load_mem0(self):
"""Test that importing praisonaiagents doesn't load mem0."""
# This test runs even without RUN_LIVE_TESTS
# Clear any existing mem0 imports
mem0_modules = [k for k in list(sys.modules.keys()) if k.startswith('mem0')]
for mod in mem0_modules:
sys.modules.pop(mod, None)
# Fresh import
import importlib
if 'praisonaiagents' in sys.modules:
importlib.reload(sys.modules['praisonaiagents'])
else:
import praisonaiagents
# mem0 should NOT be loaded
mem0_loaded = any(k.startswith('mem0') for k in sys.modules.keys())
assert not mem0_loaded, "mem0 should not be imported at package import time"
def test_knowledge_import_does_not_load_mem0(self):
"""Test that importing Knowledge doesn't load mem0."""
mem0_modules = [k for k in list(sys.modules.keys()) if k.startswith('mem0')]
for mod in mem0_modules:
sys.modules.pop(mod, None)
from praisonaiagents import Knowledge # noqa: F401
mem0_loaded = any(k.startswith('mem0') for k in sys.modules.keys())
# Note: Knowledge class definition doesn't import mem0
# mem0 is only imported when Knowledge is instantiated and used
assert not mem0_loaded, "mem0 should not be imported when importing Knowledge class"
@@ -179,7 +179,6 @@ def test_all_agent_parameters():
response_template="Test response",
allow_code_execution=True,
max_retry_limit=5,
respect_context_window=False,
code_execution_mode="unsafe",
embedder_config={"test": "config"},
knowledge=None,
@@ -22,7 +22,10 @@ class TestAgentRAGConfig:
rag_config={"include_citations": True, "top_k": 5}
)
assert agent._rag_config == {"include_citations": True, "top_k": 5}
# rag_config is merged into _retrieval_config
assert agent._retrieval_config is not None
assert agent._retrieval_config.top_k == 5
assert agent._retrieval_config.citations is True
def test_agent_without_rag_config(self):
"""Test that Agent works without rag_config."""
@@ -34,7 +37,8 @@ class TestAgentRAGConfig:
knowledge=["test.txt"]
)
assert agent._rag_config is None
# Default retrieval config is created for knowledge
assert agent._retrieval_config is not None
def test_agent_without_knowledge_has_no_rag_config(self):
"""Test that Agent without knowledge has no rag_config."""
@@ -45,7 +49,7 @@ class TestAgentRAGConfig:
instructions="Test agent"
)
assert agent._rag_config is None
# No knowledge means no retrieval config needed
assert agent._rag_instance is None
@@ -164,7 +168,7 @@ class TestAgentKnowledgeContext:
class TestAgentRAGIntegration:
"""Integration tests for Agent + RAG."""
"""Tests for Agent RAG integration."""
def test_agent_rag_config_passed_to_rag_instance(self):
"""Test that rag_config is passed to RAG instance."""
@@ -181,7 +185,7 @@ class TestAgentRAGIntegration:
}
)
# Verify config is stored
assert agent._rag_config["include_citations"] is True
assert agent._rag_config["top_k"] == 10
assert agent._rag_config["min_score"] == 0.5
# Verify config is stored in _retrieval_config
assert agent._retrieval_config.citations is True
assert agent._retrieval_config.top_k == 10
assert agent._retrieval_config.min_score == 0.5
@@ -0,0 +1,275 @@
"""
Unit tests for the context= parameter in Agent, Workflow, and PraisonAIAgents.
Tests cover:
- Parameter types (False, True, ManagerConfig, ContextManager)
- Lazy initialization (zero overhead when disabled)
- Precedence rules
- Multi-agent isolation
"""
import pytest
import time
class TestAgentContextParam:
"""Tests for Agent context= parameter."""
def test_context_false_default(self):
"""context=False (default) should have zero overhead."""
from praisonaiagents import Agent
agent = Agent(instructions="Test agent", context=False)
# Should not initialize context manager
assert agent._context_param is False
assert agent._context_manager is None
assert agent._context_manager_initialized is False
def test_context_false_no_init_on_access(self):
"""Accessing context_manager when context=False should return None."""
from praisonaiagents import Agent
agent = Agent(instructions="Test agent", context=False)
# Access should return None without error
assert agent.context_manager is None
assert agent._context_manager_initialized is True
def test_context_true_lazy_init(self):
"""context=True should lazy-init ContextManager on first access."""
from praisonaiagents import Agent
agent = Agent(instructions="Test agent", context=True)
# Not initialized until first access
assert agent._context_param is True
assert agent._context_manager_initialized is False
# Access triggers initialization
manager = agent.context_manager
assert agent._context_manager_initialized is True
assert manager is not None
def test_context_manager_config(self):
"""context=ManagerConfig should use provided config."""
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
config = ManagerConfig(
auto_compact=False,
compact_threshold=0.5,
monitor_enabled=True,
)
agent = Agent(instructions="Test agent", context=config)
manager = agent.context_manager
assert manager is not None
assert manager.config.auto_compact is False
assert manager.config.compact_threshold == 0.5
assert manager.config.monitor_enabled is True
def test_context_manager_instance(self):
"""context=ContextManager should use provided instance directly."""
from praisonaiagents import Agent
from praisonaiagents.context import ContextManager
custom_manager = ContextManager(model="gpt-4o")
agent = Agent(instructions="Test agent", context=custom_manager)
# Should use the exact same instance
assert agent.context_manager is custom_manager
def test_context_manager_setter(self):
"""context_manager setter should work correctly."""
from praisonaiagents import Agent
from praisonaiagents.context import ContextManager
agent = Agent(instructions="Test agent", context=False)
# Initially None
assert agent.context_manager is None
# Set a manager
manager = ContextManager(model="gpt-4o-mini")
agent.context_manager = manager
assert agent.context_manager is manager
assert agent._context_manager_initialized is True
class TestAgentContextPerformance:
"""Performance tests for context= parameter."""
def test_context_false_init_overhead(self):
"""context=False should add minimal overhead to init."""
from praisonaiagents import Agent
# Warm up
Agent(instructions="warmup", context=False)
# Measure init time with context=False
start = time.perf_counter()
for _ in range(10):
Agent(instructions="test", context=False)
elapsed = time.perf_counter() - start
# Should be very fast (< 50ms for 10 inits)
assert elapsed < 0.5, f"Init too slow: {elapsed}s for 10 agents"
def test_context_true_no_init_until_access(self):
"""context=True should not initialize manager until accessed."""
from praisonaiagents import Agent
agent = Agent(instructions="Test agent", context=True)
# Manager should not be initialized yet
assert agent._context_manager is None
assert agent._context_manager_initialized is False
class TestWorkflowContextParam:
"""Tests for Workflow context= parameter."""
def test_workflow_context_false_default(self):
"""Workflow context=False should be the default."""
from praisonaiagents import Workflow
workflow = Workflow(name="Test", steps=[])
assert workflow.context is False
def test_workflow_context_true(self):
"""Workflow context=True should be stored."""
from praisonaiagents import Workflow
workflow = Workflow(name="Test", steps=[], context=True)
assert workflow.context is True
class TestPraisonAIAgentsContextParam:
"""Tests for PraisonAIAgents context= parameter."""
def test_agents_context_false_default(self):
"""PraisonAIAgents context=False should have zero overhead."""
from praisonaiagents import Agent, PraisonAIAgents, Task
agent = Agent(instructions="Test agent")
task = Task(description="Test task", agent=agent)
agents = PraisonAIAgents(agents=[agent], tasks=[task], context=False)
assert agents._context_param is False
assert agents._context_manager is None
assert agents._context_manager_initialized is False
def test_agents_context_true_lazy_init(self):
"""PraisonAIAgents context=True should lazy-init MultiAgentContextManager."""
from praisonaiagents import Agent, PraisonAIAgents, Task
agent = Agent(instructions="Test agent")
task = Task(description="Test task", agent=agent)
agents = PraisonAIAgents(agents=[agent], tasks=[task], context=True)
# Not initialized until first access
assert agents._context_manager_initialized is False
# Access triggers initialization
manager = agents.context_manager
assert agents._context_manager_initialized is True
assert manager is not None
def test_agents_context_manager_config(self):
"""PraisonAIAgents context=ManagerConfig should use provided config."""
from praisonaiagents import Agent, PraisonAIAgents, Task
from praisonaiagents.context import ManagerConfig
config = ManagerConfig(
auto_compact=True,
compact_threshold=0.7,
)
agent = Agent(instructions="Test agent")
task = Task(description="Test task", agent=agent)
agents = PraisonAIAgents(agents=[agent], tasks=[task], context=config)
manager = agents.context_manager
assert manager is not None
class TestContextPrecedence:
"""Tests for context parameter precedence rules."""
def test_old_context_params_rejected(self):
"""Old context-related params should raise TypeError."""
from praisonaiagents import Agent
# Old params should be rejected (breaking change)
with pytest.raises(TypeError):
Agent(instructions="Test", auto_summarize=True)
with pytest.raises(TypeError):
Agent(instructions="Test", summarize_threshold=0.9)
with pytest.raises(TypeError):
Agent(instructions="Test", fast_context=True)
with pytest.raises(TypeError):
Agent(instructions="Test", respect_context_window=True)
with pytest.raises(TypeError):
Agent(instructions="Test", history_in_context=5)
with pytest.raises(TypeError):
Agent(instructions="Test", context_compactor="test")
def test_context_param_is_only_interface(self):
"""context= should be the only interface for context management."""
from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig
# New way: use context= param
config = ManagerConfig(auto_compact=True, compact_threshold=0.8)
agent = Agent(instructions="Test agent", context=config)
# Verify context manager is configured
assert agent.context_manager is not None
assert agent.context_manager.config.auto_compact is True
assert agent.context_manager.config.compact_threshold == 0.8
class TestContextManagerIntegration:
"""Integration tests for context manager with Agent."""
def test_context_manager_has_required_methods(self):
"""ContextManager should have required methods."""
from praisonaiagents import Agent
agent = Agent(instructions="Test agent", context=True)
manager = agent.context_manager
# Check required methods exist
assert hasattr(manager, 'process')
assert hasattr(manager, 'get_stats')
assert hasattr(manager, 'config')
def test_context_manager_model_from_agent(self):
"""ContextManager should use agent's model."""
from praisonaiagents import Agent
agent = Agent(
instructions="Test agent",
llm="gpt-4o",
context=True,
)
manager = agent.context_manager
# Manager should use agent's model
assert manager.model == "gpt-4o"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,608 @@
"""
Tests for Context Manager Facade.
Tests cover:
- ContextManager initialization and configuration
- Multi-agent orchestration with policies
- Snapshot hooks at LLM boundary
- Compression benefit checking
- Token estimation validation
- Per-tool budgets
- Privacy hardening (redaction patterns)
- Optimization history tracking
- Config precedence
"""
import json
import hashlib
from unittest.mock import patch
from praisonaiagents.context.manager import (
ContextManager,
MultiAgentContextManager,
ManagerConfig,
ContextPolicy,
EstimationMode,
ContextShareMode,
ToolShareMode,
OptimizationEvent,
OptimizationEventType,
EstimationMetrics,
PerToolBudget,
SnapshotHookData,
create_context_manager,
)
from praisonaiagents.context.models import OptimizerStrategy
class TestManagerConfig:
"""Tests for ManagerConfig."""
def test_default_config(self):
"""Test default configuration values."""
config = ManagerConfig()
assert config.auto_compact is True
assert config.compact_threshold == 0.8
assert config.strategy == OptimizerStrategy.SMART
assert config.compression_min_gain_pct == 5.0
assert config.output_reserve == 8000
assert config.estimation_mode == EstimationMode.HEURISTIC
assert config.redact_sensitive is True
def test_config_from_env(self):
"""Test loading config from environment variables."""
with patch.dict('os.environ', {
'PRAISONAI_CONTEXT_AUTO_COMPACT': 'false',
'PRAISONAI_CONTEXT_THRESHOLD': '0.7',
'PRAISONAI_CONTEXT_STRATEGY': 'truncate',
'PRAISONAI_CONTEXT_MONITOR': 'true',
}):
config = ManagerConfig.from_env()
assert config.auto_compact is False
assert config.compact_threshold == 0.7
assert config.strategy == OptimizerStrategy.TRUNCATE
assert config.monitor_enabled is True
assert config.source == "env"
def test_config_merge(self):
"""Test config merging with overrides."""
base = ManagerConfig()
merged = base.merge(
auto_compact=False,
compact_threshold=0.9,
monitor_enabled=True,
)
assert merged.auto_compact is False
assert merged.compact_threshold == 0.9
assert merged.monitor_enabled is True
# Unchanged values preserved
assert merged.strategy == OptimizerStrategy.SMART
def test_config_to_dict(self):
"""Test config serialization."""
config = ManagerConfig()
data = config.to_dict()
assert isinstance(data, dict)
assert data["auto_compact"] is True
assert data["strategy"] == "smart"
assert "tool_budgets" in data
class TestContextPolicy:
"""Tests for ContextPolicy."""
def test_default_policy(self):
"""Test default policy values."""
policy = ContextPolicy()
assert policy.share is False
assert policy.share_mode == ContextShareMode.NONE
assert policy.max_tokens == 0
assert policy.tools_share == ToolShareMode.NONE
def test_policy_to_dict(self):
"""Test policy serialization."""
policy = ContextPolicy(
share=True,
share_mode=ContextShareMode.SUMMARY,
max_tokens=5000,
)
data = policy.to_dict()
assert data["share"] is True
assert data["share_mode"] == "summary"
assert data["max_tokens"] == 5000
class TestContextManager:
"""Tests for ContextManager facade."""
def test_initialization(self):
"""Test manager initialization."""
manager = ContextManager(model="gpt-4o-mini")
assert manager.model == "gpt-4o-mini"
assert manager.config is not None
assert manager._budget is not None
def test_process_messages(self):
"""Test processing messages through pipeline."""
manager = ContextManager(model="gpt-4o-mini")
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
]
result = manager.process(messages=messages)
assert "messages" in result
assert "optimized" in result
assert "tokens_before" in result
assert "utilization" in result
def test_process_with_system_prompt(self):
"""Test processing with system prompt."""
manager = ContextManager(model="gpt-4o-mini")
messages = [{"role": "user", "content": "Hello"}]
system_prompt = "You are a helpful assistant."
result = manager.process(
messages=messages,
system_prompt=system_prompt,
)
assert result["tokens_before"] > 0
def test_capture_llm_boundary(self):
"""Test capturing exact state at LLM boundary."""
manager = ContextManager(model="gpt-4o-mini")
messages = [
{"role": "user", "content": "Hello"},
]
tools = [
{"type": "function", "function": {"name": "test_tool"}},
]
hook_data = manager.capture_llm_boundary(messages, tools)
assert isinstance(hook_data, SnapshotHookData)
assert hook_data.messages == messages
assert hook_data.tools == tools
assert len(hook_data.message_hash) == 16
assert len(hook_data.tools_hash) == 16
def test_snapshot_callback(self):
"""Test snapshot callback registration."""
manager = ContextManager(model="gpt-4o-mini")
callback_called = []
def callback(data):
callback_called.append(data)
manager.register_snapshot_callback(callback)
messages = [{"role": "user", "content": "Hello"}]
manager.capture_llm_boundary(messages, [])
assert len(callback_called) == 1
assert callback_called[0].messages == messages
def test_get_tool_budget(self):
"""Test per-tool budget retrieval."""
manager = ContextManager(model="gpt-4o-mini")
# Default budget
default = manager.get_tool_budget("unknown_tool")
assert default == manager.config.default_tool_output_max
# Set custom budget
manager.set_tool_budget("file_read", 5000, protected=True)
assert manager.get_tool_budget("file_read") == 5000
assert "file_read" in manager.config.protected_tools
def test_truncate_tool_output(self):
"""Test tool output truncation."""
manager = ContextManager(model="gpt-4o-mini")
manager.set_tool_budget("small_tool", 100)
# Large output should be truncated
large_output = "x" * 10000
truncated = manager.truncate_tool_output("small_tool", large_output)
assert len(truncated) < len(large_output)
assert "[output truncated]" in truncated
def test_optimization_history(self):
"""Test optimization history tracking."""
manager = ContextManager(model="gpt-4o-mini")
# Process some messages to generate history
messages = [{"role": "user", "content": "x" * 1000}] * 100
manager.process(messages=messages)
history = manager.get_history()
assert isinstance(history, list)
def test_get_stats(self):
"""Test getting context statistics."""
manager = ContextManager(model="gpt-4o-mini")
stats = manager.get_stats()
assert "model" in stats
assert "budget" in stats
assert "ledger" in stats
assert "utilization" in stats
def test_get_resolved_config(self):
"""Test getting resolved configuration."""
manager = ContextManager(model="gpt-4o-mini")
resolved = manager.get_resolved_config()
assert "config" in resolved
assert "precedence" in resolved
assert "effective_budget" in resolved
def test_estimate_tokens_with_cache(self):
"""Test token estimation with caching."""
manager = ContextManager(model="gpt-4o-mini")
text = "Hello world"
# First call
tokens1, _ = manager.estimate_tokens(text)
# Second call should use cache
tokens2, _ = manager.estimate_tokens(text)
assert tokens1 == tokens2
def test_reset(self):
"""Test manager reset."""
manager = ContextManager(model="gpt-4o-mini")
# Add some state
messages = [{"role": "user", "content": "Hello"}]
manager.process(messages=messages)
manager.capture_llm_boundary(messages, [])
# Reset
manager.reset()
assert manager._last_snapshot_hook is None
assert len(manager._history) == 0
class TestCompressionBenefitCheck:
"""Tests for compression benefit checking."""
def test_benefit_check_prevents_inflation(self):
"""Test that compression is reverted if not beneficial."""
config = ManagerConfig(
auto_compact=True,
compact_threshold=0.1, # Low threshold to trigger
compression_min_gain_pct=50.0, # High requirement
)
manager = ContextManager(model="gpt-4o-mini", config=config)
# Small messages that won't benefit much from compression
messages = [
{"role": "user", "content": "Hi"},
{"role": "assistant", "content": "Hello"},
]
result = manager.process(messages=messages)
# Should not have optimized due to benefit check
# (small messages don't benefit from compression)
assert result["messages"] is not None
class TestMultiAgentContextManager:
"""Tests for MultiAgentContextManager."""
def test_initialization(self):
"""Test multi-agent manager initialization."""
manager = MultiAgentContextManager()
assert manager.config is not None
assert manager.default_policy is not None
def test_get_agent_manager(self):
"""Test getting per-agent managers."""
manager = MultiAgentContextManager()
agent1 = manager.get_agent_manager("agent_1")
agent2 = manager.get_agent_manager("agent_2")
assert agent1 is not agent2
assert agent1.agent_name == "agent_1"
# Same agent returns same manager
agent1_again = manager.get_agent_manager("agent_1")
assert agent1 is agent1_again
def test_set_agent_policy(self):
"""Test setting per-agent policies."""
manager = MultiAgentContextManager()
policy = ContextPolicy(
share=True,
share_mode=ContextShareMode.SUMMARY,
)
manager.set_agent_policy("agent_1", policy)
retrieved = manager.get_agent_policy("agent_1")
assert retrieved.share is True
assert retrieved.share_mode == ContextShareMode.SUMMARY
def test_prepare_handoff_no_share(self):
"""Test handoff with no sharing."""
manager = MultiAgentContextManager()
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi"},
]
# Default policy is no share
result = manager.prepare_handoff("agent_1", "agent_2", messages)
assert result == []
def test_prepare_handoff_summary(self):
"""Test handoff with summary sharing."""
manager = MultiAgentContextManager()
policy = ContextPolicy(
share=True,
share_mode=ContextShareMode.SUMMARY,
)
manager.set_agent_policy("agent_2", policy)
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
]
result = manager.prepare_handoff("agent_1", "agent_2", messages, policy)
assert len(result) > 0
assert result[0]["role"] == "system"
assert "_handoff_summary" in result[0]
def test_prepare_handoff_full(self):
"""Test handoff with full context sharing."""
manager = MultiAgentContextManager()
policy = ContextPolicy(
share=True,
share_mode=ContextShareMode.FULL,
preserve_recent_turns=2,
)
messages = [
{"role": "system", "content": "You are helpful"},
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi"},
{"role": "user", "content": "How are you?"},
{"role": "assistant", "content": "I'm good!"},
]
result = manager.prepare_handoff("agent_1", "agent_2", messages, policy)
# Should include system + recent turns
assert len(result) > 0
def test_prepare_handoff_with_token_limit(self):
"""Test handoff respects token limits."""
manager = MultiAgentContextManager()
policy = ContextPolicy(
share=True,
share_mode=ContextShareMode.FULL,
max_tokens=100, # Very small limit
)
messages = [
{"role": "user", "content": "x" * 1000},
{"role": "assistant", "content": "y" * 1000},
]
result = manager.prepare_handoff("agent_1", "agent_2", messages, policy)
# Should be limited
from praisonaiagents.context.tokens import estimate_messages_tokens
total = estimate_messages_tokens(result)
assert total <= 100 or len(result) <= 1
def test_get_combined_stats(self):
"""Test getting combined statistics."""
manager = MultiAgentContextManager()
# Create some agents
manager.get_agent_manager("agent_1")
manager.get_agent_manager("agent_2")
stats = manager.get_combined_stats()
assert "agents" in stats
assert "agent_1" in stats["agents"]
assert "agent_2" in stats["agents"]
class TestEstimationValidation:
"""Tests for token estimation validation."""
def test_estimation_metrics(self):
"""Test estimation metrics creation."""
metrics = EstimationMetrics(
heuristic_estimate=100,
accurate_estimate=95,
error_pct=5.26,
estimator_used=EstimationMode.VALIDATED,
)
data = metrics.to_dict()
assert data["heuristic_estimate"] == 100
assert data["accurate_estimate"] == 95
assert data["error_pct"] == 5.26
def test_validated_estimation_mode(self):
"""Test validated estimation mode."""
config = ManagerConfig(
estimation_mode=EstimationMode.VALIDATED,
log_estimation_mismatch=True,
)
manager = ContextManager(model="gpt-4o-mini", config=config)
text = "Hello world"
tokens, metrics = manager.estimate_tokens(text, validate=True)
assert tokens > 0
# Metrics may or may not be available depending on tiktoken
class TestPerToolBudgets:
"""Tests for per-tool token budgets."""
def test_per_tool_budget_creation(self):
"""Test creating per-tool budgets."""
budget = PerToolBudget(
tool_name="file_read",
max_output_tokens=5000,
protected=True,
)
data = budget.to_dict()
assert data["tool_name"] == "file_read"
assert data["max_output_tokens"] == 5000
assert data["protected"] is True
def test_tool_budget_in_config(self):
"""Test tool budgets in config."""
config = ManagerConfig()
config.tool_budgets["file_read"] = PerToolBudget(
tool_name="file_read",
max_output_tokens=5000,
)
manager = ContextManager(model="gpt-4o-mini", config=config)
assert manager.get_tool_budget("file_read") == 5000
class TestOptimizationHistory:
"""Tests for optimization history tracking."""
def test_optimization_event_creation(self):
"""Test creating optimization events."""
event = OptimizationEvent(
timestamp="2024-01-01T00:00:00Z",
event_type=OptimizationEventType.AUTO_COMPACT,
strategy="smart",
tokens_before=10000,
tokens_after=5000,
tokens_saved=5000,
)
data = event.to_dict()
assert data["event_type"] == "auto_compact"
assert data["tokens_saved"] == 5000
def test_history_max_size(self):
"""Test history respects max size."""
manager = ContextManager(model="gpt-4o-mini")
manager._max_history = 5
# Add many events
for i in range(10):
manager._add_history_event(
OptimizationEventType.SNAPSHOT,
details={"index": i},
)
history = manager.get_history()
assert len(history) <= 5
class TestCreateContextManager:
"""Tests for create_context_manager factory."""
def test_create_with_defaults(self):
"""Test creating manager with defaults."""
manager = create_context_manager(model="gpt-4o-mini")
assert manager.model == "gpt-4o-mini"
assert manager.config is not None
def test_create_with_cli_overrides(self):
"""Test creating manager with CLI overrides."""
manager = create_context_manager(
model="gpt-4o-mini",
cli_overrides={
"auto_compact": False,
"monitor_enabled": True,
},
)
assert manager.config.auto_compact is False
assert manager.config.monitor_enabled is True
assert manager.config.source == "cli"
class TestSnapshotHookData:
"""Tests for SnapshotHookData."""
def test_snapshot_hook_creation(self):
"""Test creating snapshot hook data."""
messages = [{"role": "user", "content": "Hello"}]
tools = [{"type": "function", "function": {"name": "test"}}]
messages_json = json.dumps(messages, sort_keys=True)
tools_json = json.dumps(tools, sort_keys=True)
hook = SnapshotHookData(
timestamp="2024-01-01T00:00:00Z",
messages=messages,
tools=tools,
message_hash=hashlib.sha256(messages_json.encode()).hexdigest()[:16],
tools_hash=hashlib.sha256(tools_json.encode()).hexdigest()[:16],
)
data = hook.to_dict()
assert data["message_count"] == 1
assert data["tools_count"] == 1
assert len(data["message_hash"]) == 16
def test_hash_consistency(self):
"""Test that hashes are consistent for same content."""
messages = [{"role": "user", "content": "Hello"}]
manager = ContextManager(model="gpt-4o-mini")
hook1 = manager.capture_llm_boundary(messages, [])
hook2 = manager.capture_llm_boundary(messages, [])
assert hook1.message_hash == hook2.message_hash
def test_hash_changes_with_content(self):
"""Test that hashes change with different content."""
manager = ContextManager(model="gpt-4o-mini")
hook1 = manager.capture_llm_boundary(
[{"role": "user", "content": "Hello"}], []
)
hook2 = manager.capture_llm_boundary(
[{"role": "user", "content": "Goodbye"}], []
)
assert hook1.message_hash != hook2.message_hash
@@ -0,0 +1,249 @@
"""
Tests for privacy hardening in context management.
Tests cover:
- Redaction patterns for various API keys
- Path validation for monitor output
- Ignore/include pattern integration
"""
from praisonaiagents.context.monitor import (
redact_sensitive,
validate_monitor_path,
should_include_content,
load_ignore_patterns,
SENSITIVE_PATTERNS,
)
class TestRedactionPatterns:
"""Tests for sensitive data redaction."""
def test_redact_openai_key(self):
"""Test OpenAI API key redaction."""
text = "My key is sk-1234567890abcdefghijklmnop"
result = redact_sensitive(text)
assert "sk-1234567890" not in result
assert "[REDACTED]" in result
def test_redact_openai_project_key(self):
"""Test OpenAI project key redaction."""
text = "Using sk-proj-abc123_def456-xyz789"
result = redact_sensitive(text)
assert "sk-proj-" not in result
assert "[REDACTED]" in result
def test_redact_anthropic_key(self):
"""Test Anthropic API key redaction."""
text = "Anthropic key: sk-ant-abc123def456ghi789"
result = redact_sensitive(text)
assert "sk-ant-" not in result
assert "[REDACTED]" in result
def test_redact_anthropic_api_key(self):
"""Test Anthropic API key with version."""
text = "Key is sk-ant-api03-abcdefghijklmnop"
result = redact_sensitive(text)
assert "sk-ant-api" not in result
assert "[REDACTED]" in result
def test_redact_google_api_key(self):
"""Test Google API key redaction."""
text = "Google key: AIzaSyAbcdefghijklmnopqrstuvwxyz12345"
result = redact_sensitive(text)
assert "AIzaSy" not in result
assert "[REDACTED]" in result
def test_redact_google_oauth_token(self):
"""Test Google OAuth token redaction."""
text = "Token: ya29.a0AfH6SMBabcdefghijklmnop"
result = redact_sensitive(text)
assert "ya29." not in result
assert "[REDACTED]" in result
def test_redact_aws_access_key(self):
"""Test AWS access key redaction."""
text = "AWS key: AKIAIOSFODNN7EXAMPLE"
result = redact_sensitive(text)
assert "AKIA" not in result
assert "[REDACTED]" in result
def test_redact_aws_temp_key(self):
"""Test AWS temporary access key redaction."""
# AWS temp keys are exactly 20 chars: ASIA + 16 alphanumeric
text = "Temp key: ASIAIOSFODNN7EXAMPLE"
result = redact_sensitive(text)
assert "ASIAIOSFODNN7EXAMPLE" not in result
assert "[REDACTED]" in result
def test_redact_bearer_token(self):
"""Test bearer token redaction."""
text = "Authorization: Bearer abc123def456"
result = redact_sensitive(text)
assert "Bearer abc123" not in result
assert "[REDACTED]" in result
def test_redact_password_pattern(self):
"""Test password pattern redaction."""
text = 'password = "mysecretpassword123"'
result = redact_sensitive(text)
assert "mysecretpassword" not in result
assert "[REDACTED]" in result
def test_redact_api_key_pattern(self):
"""Test generic API key pattern redaction."""
text = 'api_key: "abc123def456"'
result = redact_sensitive(text)
assert "abc123def456" not in result
assert "[REDACTED]" in result
def test_preserves_normal_text(self):
"""Test that normal text is preserved."""
text = "Hello, this is a normal message without secrets."
result = redact_sensitive(text)
assert result == text
def test_multiple_redactions(self):
"""Test multiple sensitive items in one text."""
text = "Keys: sk-abc123def456ghi789jkl, sk-ant-xyz789"
result = redact_sensitive(text)
assert "sk-abc" not in result
assert "sk-ant" not in result
assert result.count("[REDACTED]") >= 2
class TestPathValidation:
"""Tests for monitor path validation."""
def test_valid_relative_path(self):
"""Test valid relative path."""
is_valid, error = validate_monitor_path("./context.txt")
assert is_valid is True
assert error == ""
def test_valid_simple_filename(self):
"""Test simple filename."""
is_valid, error = validate_monitor_path("context.txt")
assert is_valid is True
assert error == ""
def test_path_traversal_blocked(self):
"""Test path traversal is blocked."""
is_valid, error = validate_monitor_path("../../../etc/passwd")
assert is_valid is False
assert "traversal" in error.lower()
def test_absolute_path_blocked_by_default(self):
"""Test absolute paths blocked by default."""
is_valid, error = validate_monitor_path("/tmp/context.txt")
assert is_valid is False
assert "absolute" in error.lower()
def test_absolute_path_allowed_when_enabled(self):
"""Test absolute paths allowed when explicitly enabled."""
is_valid, error = validate_monitor_path(
"/tmp/context.txt",
allow_absolute=True,
)
assert is_valid is True
def test_suspicious_path_blocked(self):
"""Test suspicious paths are blocked."""
is_valid, error = validate_monitor_path("/etc/passwd")
assert is_valid is False
def test_home_path_blocked(self):
"""Test home directory paths blocked by default."""
is_valid, error = validate_monitor_path("/home/user/file.txt")
assert is_valid is False
class TestIgnoreIncludePatterns:
"""Tests for ignore/include pattern handling."""
def test_include_all_by_default(self):
"""Test all files included by default."""
result = should_include_content("test.py")
assert result is True
def test_ignore_pattern_match(self):
"""Test ignore pattern matching."""
result = should_include_content(
"secret.key",
ignore_patterns=["*.key", "*.pem"],
)
assert result is False
def test_ignore_pattern_no_match(self):
"""Test ignore pattern not matching."""
result = should_include_content(
"code.py",
ignore_patterns=["*.key", "*.pem"],
)
assert result is True
def test_include_pattern_whitelist(self):
"""Test include pattern as whitelist."""
result = should_include_content(
"code.py",
include_patterns=["*.py", "*.js"],
)
assert result is True
def test_include_pattern_excludes_others(self):
"""Test include pattern excludes non-matching."""
result = should_include_content(
"data.json",
include_patterns=["*.py", "*.js"],
)
assert result is False
def test_ignore_overrides_include(self):
"""Test ignore patterns override include."""
result = should_include_content(
"secret.py",
include_patterns=["*.py"],
ignore_patterns=["secret*"],
)
assert result is False
def test_path_pattern_match(self):
"""Test full path pattern matching."""
result = should_include_content(
"/path/to/node_modules/file.js",
ignore_patterns=["*node_modules*"],
)
assert result is False
class TestSensitivePatternsCoverage:
"""Tests to ensure all expected patterns are covered."""
def test_patterns_list_not_empty(self):
"""Test patterns list is populated."""
assert len(SENSITIVE_PATTERNS) > 0
def test_openai_patterns_present(self):
"""Test OpenAI patterns are in list."""
patterns_str = str(SENSITIVE_PATTERNS)
assert "sk-" in patterns_str
def test_anthropic_patterns_present(self):
"""Test Anthropic patterns are in list."""
patterns_str = str(SENSITIVE_PATTERNS)
assert "sk-ant" in patterns_str
def test_google_patterns_present(self):
"""Test Google patterns are in list."""
patterns_str = str(SENSITIVE_PATTERNS)
assert "AIza" in patterns_str
def test_aws_patterns_present(self):
"""Test AWS patterns are in list."""
patterns_str = str(SENSITIVE_PATTERNS)
assert "AKIA" in patterns_str
def test_bearer_pattern_present(self):
"""Test bearer token pattern is in list."""
patterns_str = str(SENSITIVE_PATTERNS)
assert "bearer" in patterns_str.lower()
@@ -182,3 +182,43 @@ class TestTokenEstimationAccuracy:
long_tokens = estimate_tokens_heuristic(long)
assert long_tokens > short_tokens * 50
class TestFormatPercent:
"""Tests for format_percent utility function."""
def test_format_percent_zero(self):
"""Test zero value."""
from praisonaiagents.context import format_percent
assert format_percent(0) == "0.00%"
def test_format_percent_tiny(self):
"""Test tiny values show <0.1%."""
from praisonaiagents.context import format_percent
# 0.0002 = 0.02% which is < 0.1%
assert format_percent(0.0002) == "<0.1%"
assert format_percent(0.0005) == "<0.1%"
assert format_percent(0.0009) == "<0.1%"
def test_format_percent_small(self):
"""Test small values show 2 decimals."""
from praisonaiagents.context import format_percent
# 0.002 = 0.2% which is < 1%
assert format_percent(0.002) == "0.20%"
assert format_percent(0.005) == "0.50%"
assert format_percent(0.009) == "0.90%"
def test_format_percent_normal(self):
"""Test normal values show 1 decimal."""
from praisonaiagents.context import format_percent
# 0.05 = 5%
assert format_percent(0.05) == "5.0%"
assert format_percent(0.5) == "50.0%"
assert format_percent(0.8) == "80.0%"
assert format_percent(1.0) == "100.0%"
def test_format_percent_overflow(self):
"""Test overflow values."""
from praisonaiagents.context import format_percent
assert format_percent(1.1) == "110.0%"
assert format_percent(1.5) == "150.0%"
@@ -0,0 +1,189 @@
"""
Unit tests for Knowledge Adapters.
Tests Mem0Adapter lazy loading and normalization.
"""
import sys
from unittest.mock import MagicMock
from praisonaiagents.knowledge.protocols import ScopeRequiredError
class TestMem0AdapterLazyLoading:
"""Tests for Mem0Adapter lazy import behavior."""
def test_mem0_not_imported_at_module_level(self):
"""Test that mem0 is not imported when adapter module is imported."""
# Remove mem0 from sys.modules if present
mem0_modules = [k for k in sys.modules.keys() if k.startswith('mem0')]
for mod in mem0_modules:
sys.modules.pop(mod, None)
# Import the adapter module
from praisonaiagents.knowledge.adapters import mem0_adapter
# mem0 should NOT be in sys.modules yet
assert 'mem0' not in sys.modules, "mem0 should not be imported at module level"
def test_adapter_class_importable(self):
"""Test that Mem0Adapter class can be imported."""
from praisonaiagents.knowledge.adapters import Mem0Adapter
assert Mem0Adapter is not None
class TestMem0AdapterNormalization:
"""Tests for Mem0Adapter result normalization."""
def test_normalize_mem0_item_with_none_metadata(self):
"""Test normalization of mem0 result with metadata=None."""
from praisonaiagents.knowledge.adapters.mem0_adapter import Mem0Adapter
adapter = Mem0Adapter.__new__(Mem0Adapter)
adapter._config = {}
adapter._memory = None
adapter._disable_telemetry = True
raw = {
"id": "test-id",
"memory": "test content",
"score": 0.95,
"metadata": None, # mem0 returns this!
"user_id": "user123",
}
item = adapter._normalize_mem0_item(raw)
assert item.id == "test-id"
assert item.text == "test content"
assert item.score == 0.95
assert item.metadata == {"user_id": "user123"} # user_id added
assert item.metadata is not None
def test_normalize_mem0_results_with_none_items(self):
"""Test normalization filters out None items."""
from praisonaiagents.knowledge.adapters.mem0_adapter import Mem0Adapter
adapter = Mem0Adapter.__new__(Mem0Adapter)
adapter._config = {}
adapter._memory = None
adapter._disable_telemetry = True
raw = {
"results": [
{"id": "1", "memory": "content 1", "metadata": None},
None,
{"id": "2", "memory": "content 2", "metadata": {"key": "val"}},
]
}
result = adapter._normalize_mem0_results(raw)
assert len(result.results) == 2
assert result.results[0].text == "content 1"
assert result.results[0].metadata == {}
assert result.results[1].metadata == {"key": "val"}
class TestMem0AdapterScopeValidation:
"""Tests for Mem0Adapter scope validation."""
def test_search_requires_scope(self):
"""Test that search raises ScopeRequiredError without scope."""
from praisonaiagents.knowledge.adapters.mem0_adapter import Mem0Adapter
adapter = Mem0Adapter.__new__(Mem0Adapter)
adapter._config = {}
adapter._memory = MagicMock()
adapter._disable_telemetry = True
try:
adapter.search("query")
assert False, "Should have raised ScopeRequiredError"
except ScopeRequiredError as e:
assert "user_id" in str(e)
assert "agent_id" in str(e)
assert "run_id" in str(e)
def test_search_with_user_id_succeeds(self):
"""Test that search succeeds with user_id."""
from praisonaiagents.knowledge.adapters.mem0_adapter import Mem0Adapter
adapter = Mem0Adapter.__new__(Mem0Adapter)
adapter._config = {}
adapter._disable_telemetry = True
# Mock the memory object
mock_memory = MagicMock()
mock_memory.search.return_value = {
"results": [
{"id": "1", "memory": "content", "metadata": None, "score": 0.9}
]
}
adapter._memory = mock_memory
result = adapter.search("query", user_id="user123")
assert len(result.results) == 1
mock_memory.search.assert_called_once()
def test_add_requires_scope(self):
"""Test that add raises ScopeRequiredError without scope."""
from praisonaiagents.knowledge.adapters.mem0_adapter import Mem0Adapter
adapter = Mem0Adapter.__new__(Mem0Adapter)
adapter._config = {}
adapter._memory = MagicMock()
adapter._disable_telemetry = True
try:
adapter.add("content")
assert False, "Should have raised ScopeRequiredError"
except ScopeRequiredError:
pass
def test_get_all_requires_scope(self):
"""Test that get_all raises ScopeRequiredError without scope."""
from praisonaiagents.knowledge.adapters.mem0_adapter import Mem0Adapter
adapter = Mem0Adapter.__new__(Mem0Adapter)
adapter._config = {}
adapter._memory = MagicMock()
adapter._disable_telemetry = True
try:
adapter.get_all()
assert False, "Should have raised ScopeRequiredError"
except ScopeRequiredError:
pass
def test_delete_all_requires_scope(self):
"""Test that delete_all raises ScopeRequiredError without scope."""
from praisonaiagents.knowledge.adapters.mem0_adapter import Mem0Adapter
adapter = Mem0Adapter.__new__(Mem0Adapter)
adapter._config = {}
adapter._memory = MagicMock()
adapter._disable_telemetry = True
try:
adapter.delete_all()
assert False, "Should have raised ScopeRequiredError"
except ScopeRequiredError:
pass
class TestProtocolCompliance:
"""Tests for KnowledgeStoreProtocol compliance."""
def test_adapter_has_required_methods(self):
"""Test that Mem0Adapter has all protocol methods."""
from praisonaiagents.knowledge.adapters.mem0_adapter import Mem0Adapter
required_methods = [
'search', 'add', 'get', 'get_all',
'update', 'delete', 'delete_all'
]
for method in required_methods:
assert hasattr(Mem0Adapter, method), f"Missing method: {method}"
@@ -0,0 +1,123 @@
"""
Tests for directory ingestion fix - ensures directories are processed correctly
and file contents (not paths) are stored.
"""
import os
import tempfile
import shutil
import pytest
class TestDirectoryIngestion:
"""Test that directories are properly processed and file contents stored."""
@pytest.fixture
def temp_knowledge_dir(self):
"""Create a temporary directory with test files."""
temp_dir = tempfile.mkdtemp(prefix='praison_test_')
# Create test file with unique content
policy_file = os.path.join(temp_dir, 'policy.txt')
with open(policy_file, 'w') as f:
f.write('Remote work policy: Manager approval code is ZEBRA-71.')
# Create another test file
guide_file = os.path.join(temp_dir, 'guide.txt')
with open(guide_file, 'w') as f:
f.write('Employee handbook: Vacation requests require code TIGER-42.')
yield temp_dir
# Cleanup
shutil.rmtree(temp_dir, ignore_errors=True)
def test_directory_ingestion_stores_text_not_path(self, temp_knowledge_dir):
"""Test that ingesting a directory stores file contents, not the directory path."""
from praisonaiagents import Knowledge
knowledge = Knowledge()
result = knowledge.add(temp_knowledge_dir, user_id='test_user')
# Should have results from files
assert 'results' in result
assert len(result['results']) > 0
# Search for unique content
search_result = knowledge.search('ZEBRA-71', user_id='test_user')
assert isinstance(search_result, dict)
assert 'results' in search_result
assert len(search_result['results']) > 0
# The memory field should contain actual text, NOT the directory path
memory = search_result['results'][0].get('memory', '')
assert 'zebra-71' in memory.lower(), f"Expected 'zebra-71' in memory, got: {memory}"
assert temp_knowledge_dir not in memory, f"Directory path should not be in memory: {memory}"
def test_directory_ingestion_processes_multiple_files(self, temp_knowledge_dir):
"""Test that all files in a directory are processed."""
from praisonaiagents import Knowledge
knowledge = Knowledge()
result = knowledge.add(temp_knowledge_dir, user_id='test_user')
# Search for content from both files
search1 = knowledge.search('ZEBRA-71', user_id='test_user')
search2 = knowledge.search('TIGER-42', user_id='test_user')
# Both should find results
assert len(search1.get('results', [])) > 0, "Should find ZEBRA-71"
assert len(search2.get('results', [])) > 0, "Should find TIGER-42"
def test_directory_ingestion_sets_metadata(self, temp_knowledge_dir):
"""Test that file metadata is properly set."""
from praisonaiagents import Knowledge
knowledge = Knowledge()
knowledge.add(temp_knowledge_dir, user_id='test_user')
search_result = knowledge.search('ZEBRA-71', user_id='test_user')
if search_result.get('results'):
metadata = search_result['results'][0].get('metadata') or {}
# Should have filename in metadata
assert 'filename' in metadata, f"Expected 'filename' in metadata: {metadata}"
assert metadata['filename'] == 'policy.txt'
class TestContextBuilderUsesText:
"""Test that context builder injects text, not paths."""
@pytest.fixture
def temp_knowledge_dir(self):
"""Create a temporary directory with test files."""
temp_dir = tempfile.mkdtemp(prefix='praison_ctx_')
policy_file = os.path.join(temp_dir, 'secret.txt')
with open(policy_file, 'w') as f:
f.write('The secret password is WRENCH-992. Do not share.')
yield temp_dir
shutil.rmtree(temp_dir, ignore_errors=True)
def test_agent_context_contains_text_not_path(self, temp_knowledge_dir):
"""Test that Agent._get_knowledge_context returns text, not path."""
from praisonaiagents import Agent
agent = Agent(
name='TestAgent',
instructions='Answer based on knowledge.',
knowledge=[temp_knowledge_dir],
user_id='test_user',
verbose=False,
)
# Ensure knowledge is processed
agent._ensure_knowledge_processed()
# Get context
context, _ = agent._get_knowledge_context('What is the password?', use_rag=True)
# Context should contain the secret, not the path
assert 'wrench-992' in context.lower(), f"Expected 'wrench-992' in context, got: {context[:200]}"
assert temp_knowledge_dir not in context, f"Path should not be in context: {context[:200]}"
@@ -0,0 +1,415 @@
"""
Unit tests for Incremental Indexing + Ignore Patterns (Phase 2).
Tests IndexResult, CorpusStats, incremental indexing, and .praisonignore support.
"""
import os
import tempfile
import time
class TestCorpusStats:
"""Tests for CorpusStats dataclass."""
def test_import_corpus_stats(self):
"""CorpusStats should be importable from knowledge module."""
from praisonaiagents.knowledge.indexing import CorpusStats
assert CorpusStats is not None
def test_default_values(self):
"""CorpusStats should have sensible defaults."""
from praisonaiagents.knowledge.indexing import CorpusStats
stats = CorpusStats()
assert stats.file_count == 0
assert stats.chunk_count == 0
assert stats.total_tokens == 0
assert stats.indexed_at is None
assert stats.strategy_recommendation == "direct"
def test_from_directory(self):
"""CorpusStats should be creatable from directory scan."""
from praisonaiagents.knowledge.indexing import CorpusStats
with tempfile.TemporaryDirectory() as tmpdir:
# Create test files
for i in range(5):
with open(os.path.join(tmpdir, f"file{i}.txt"), "w") as f:
f.write(f"Content {i} " * 100)
stats = CorpusStats.from_directory(tmpdir)
assert stats.file_count == 5
assert stats.total_tokens > 0
def test_strategy_recommendation_by_size(self):
"""CorpusStats should recommend strategy based on corpus size."""
from praisonaiagents.knowledge.indexing import CorpusStats
# Small corpus
stats = CorpusStats(file_count=5)
assert stats.strategy_recommendation == "direct"
# Medium corpus
stats = CorpusStats(file_count=50)
assert stats.strategy_recommendation == "basic"
# Large corpus
stats = CorpusStats(file_count=500)
assert stats.strategy_recommendation == "hybrid"
# Very large corpus
stats = CorpusStats(file_count=5000)
assert stats.strategy_recommendation == "reranked"
# Massive corpus
stats = CorpusStats(file_count=50000)
assert stats.strategy_recommendation == "compressed"
# Huge corpus
stats = CorpusStats(file_count=500000)
assert stats.strategy_recommendation == "hierarchical"
def test_to_dict(self):
"""CorpusStats should be serializable to dict."""
from praisonaiagents.knowledge.indexing import CorpusStats
stats = CorpusStats(file_count=10, chunk_count=50, total_tokens=5000)
d = stats.to_dict()
assert d["file_count"] == 10
assert d["chunk_count"] == 50
assert d["total_tokens"] == 5000
class TestIndexResult:
"""Tests for IndexResult dataclass."""
def test_import_index_result(self):
"""IndexResult should be importable from knowledge module."""
from praisonaiagents.knowledge.indexing import IndexResult
assert IndexResult is not None
def test_default_values(self):
"""IndexResult should have sensible defaults."""
from praisonaiagents.knowledge.indexing import IndexResult
result = IndexResult()
assert result.success is True
assert result.files_indexed == 0
assert result.files_skipped == 0
assert result.chunks_created == 0
assert result.errors == []
def test_incremental_stats(self):
"""IndexResult should track incremental indexing stats."""
from praisonaiagents.knowledge.indexing import IndexResult
result = IndexResult(
files_indexed=10,
files_skipped=5, # Unchanged files
chunks_created=50,
)
assert result.files_indexed == 10
assert result.files_skipped == 5
assert result.total_files == 15
class TestIgnoreMatcher:
"""Tests for .praisonignore pattern matching."""
def test_import_ignore_matcher(self):
"""IgnoreMatcher should be importable."""
from praisonaiagents.knowledge.indexing import IgnoreMatcher
assert IgnoreMatcher is not None
def test_basic_patterns(self):
"""IgnoreMatcher should match basic glob patterns."""
from praisonaiagents.knowledge.indexing import IgnoreMatcher
matcher = IgnoreMatcher(patterns=["*.pyc", "__pycache__", ".git"])
assert matcher.should_ignore("test.pyc") is True
assert matcher.should_ignore("__pycache__") is True
assert matcher.should_ignore(".git") is True
assert matcher.should_ignore("test.py") is False
assert matcher.should_ignore("src/main.py") is False
def test_directory_patterns(self):
"""IgnoreMatcher should handle directory patterns."""
from praisonaiagents.knowledge.indexing import IgnoreMatcher
matcher = IgnoreMatcher(patterns=["node_modules", "build/"])
assert matcher.should_ignore("node_modules/package/index.js") is True
assert matcher.should_ignore("src/node_modules/test.js") is True
assert matcher.should_ignore("build/output.js") is True
assert matcher.should_ignore("src/main.js") is False
def test_negation_patterns(self):
"""IgnoreMatcher should support negation patterns."""
from praisonaiagents.knowledge.indexing import IgnoreMatcher
matcher = IgnoreMatcher(patterns=["*.log", "!important.log"])
assert matcher.should_ignore("debug.log") is True
assert matcher.should_ignore("important.log") is False
def test_from_file(self):
"""IgnoreMatcher should load patterns from .praisonignore file."""
from praisonaiagents.knowledge.indexing import IgnoreMatcher
with tempfile.TemporaryDirectory() as tmpdir:
ignore_file = os.path.join(tmpdir, ".praisonignore")
with open(ignore_file, "w") as f:
f.write("*.pyc\n")
f.write("__pycache__\n")
f.write("# This is a comment\n")
f.write("\n") # Empty line
f.write(".git\n")
matcher = IgnoreMatcher.from_file(ignore_file)
assert matcher.should_ignore("test.pyc") is True
assert matcher.should_ignore("__pycache__") is True
assert matcher.should_ignore(".git") is True
assert matcher.should_ignore("main.py") is False
def test_from_directory(self):
"""IgnoreMatcher should auto-detect .praisonignore in directory."""
from praisonaiagents.knowledge.indexing import IgnoreMatcher
with tempfile.TemporaryDirectory() as tmpdir:
ignore_file = os.path.join(tmpdir, ".praisonignore")
with open(ignore_file, "w") as f:
f.write("*.tmp\n")
matcher = IgnoreMatcher.from_directory(tmpdir)
assert matcher.should_ignore("test.tmp") is True
assert matcher.should_ignore("test.txt") is False
class TestFileTracker:
"""Tests for file hash/mtime tracking."""
def test_import_file_tracker(self):
"""FileTracker should be importable."""
from praisonaiagents.knowledge.indexing import FileTracker
assert FileTracker is not None
def test_track_file(self):
"""FileTracker should track file hash and mtime."""
from praisonaiagents.knowledge.indexing import FileTracker
with tempfile.TemporaryDirectory() as tmpdir:
test_file = os.path.join(tmpdir, "test.txt")
with open(test_file, "w") as f:
f.write("Hello World")
tracker = FileTracker()
info = tracker.get_file_info(test_file)
assert info["path"] == test_file
assert "hash" in info
assert "mtime" in info
assert "size" in info
def test_detect_unchanged_file(self):
"""FileTracker should detect unchanged files."""
from praisonaiagents.knowledge.indexing import FileTracker
with tempfile.TemporaryDirectory() as tmpdir:
test_file = os.path.join(tmpdir, "test.txt")
with open(test_file, "w") as f:
f.write("Hello World")
tracker = FileTracker()
# First check
info1 = tracker.get_file_info(test_file)
tracker.mark_indexed(test_file, info1)
# Second check (unchanged)
assert tracker.has_changed(test_file) is False
def test_detect_changed_file(self):
"""FileTracker should detect changed files."""
from praisonaiagents.knowledge.indexing import FileTracker
with tempfile.TemporaryDirectory() as tmpdir:
test_file = os.path.join(tmpdir, "test.txt")
with open(test_file, "w") as f:
f.write("Hello World")
tracker = FileTracker()
# First check
info1 = tracker.get_file_info(test_file)
tracker.mark_indexed(test_file, info1)
# Modify file
time.sleep(0.1) # Ensure mtime changes
with open(test_file, "w") as f:
f.write("Modified Content")
# Second check (changed)
assert tracker.has_changed(test_file) is True
def test_persist_and_load(self):
"""FileTracker should persist and load state."""
from praisonaiagents.knowledge.indexing import FileTracker
with tempfile.TemporaryDirectory() as tmpdir:
state_file = os.path.join(tmpdir, ".praison_index_state.json")
test_file = os.path.join(tmpdir, "test.txt")
with open(test_file, "w") as f:
f.write("Hello World")
# Track and save
tracker1 = FileTracker(state_file=state_file)
info = tracker1.get_file_info(test_file)
tracker1.mark_indexed(test_file, info)
tracker1.save()
# Load in new tracker
tracker2 = FileTracker(state_file=state_file)
tracker2.load()
assert tracker2.has_changed(test_file) is False
class TestKnowledgeIndex:
"""Tests for Knowledge.index() method."""
def test_knowledge_has_index_method(self):
"""Knowledge should have an index() method."""
from praisonaiagents.knowledge import Knowledge
assert hasattr(Knowledge, "index")
def test_index_returns_result(self):
"""Knowledge.index() should return IndexResult."""
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.knowledge.indexing import IndexResult
with tempfile.TemporaryDirectory() as tmpdir:
# Create test file
with open(os.path.join(tmpdir, "test.txt"), "w") as f:
f.write("Test content for indexing")
knowledge = Knowledge()
result = knowledge.index(tmpdir)
assert isinstance(result, IndexResult)
assert result.files_indexed >= 1
def test_incremental_index(self):
"""Knowledge.index() should support incremental indexing."""
from praisonaiagents.knowledge import Knowledge
with tempfile.TemporaryDirectory() as tmpdir:
# Create initial file
with open(os.path.join(tmpdir, "file1.txt"), "w") as f:
f.write("Initial content")
knowledge = Knowledge()
# First index
result1 = knowledge.index(tmpdir, incremental=True)
assert result1.files_indexed == 1
# Second index (no changes)
result2 = knowledge.index(tmpdir, incremental=True)
assert result2.files_indexed == 0
assert result2.files_skipped == 1
# Add new file
with open(os.path.join(tmpdir, "file2.txt"), "w") as f:
f.write("New content")
# Third index (one new file)
result3 = knowledge.index(tmpdir, incremental=True)
assert result3.files_indexed == 1
assert result3.files_skipped == 1
def test_index_respects_ignore_patterns(self):
"""Knowledge.index() should respect exclude_glob patterns."""
from praisonaiagents.knowledge import Knowledge
with tempfile.TemporaryDirectory() as tmpdir:
# Create files
with open(os.path.join(tmpdir, "main.txt"), "w") as f:
f.write("Main content")
with open(os.path.join(tmpdir, "debug.log"), "w") as f:
f.write("Log content")
with open(os.path.join(tmpdir, "other.log"), "w") as f:
f.write("Other log content")
knowledge = Knowledge()
# Use exclude_glob to ignore .log files
result = knowledge.index(
tmpdir,
user_id="test_user",
incremental=False,
exclude_glob=["*.log"],
)
# Should only index main.txt (*.log files excluded)
assert result.files_indexed == 1
def test_index_with_include_exclude_globs(self):
"""Knowledge.index() should support include/exclude globs."""
from praisonaiagents.knowledge import Knowledge
with tempfile.TemporaryDirectory() as tmpdir:
# Create files
with open(os.path.join(tmpdir, "main.py"), "w") as f:
f.write("Python code")
with open(os.path.join(tmpdir, "test.py"), "w") as f:
f.write("Test code")
with open(os.path.join(tmpdir, "data.json"), "w") as f:
f.write('{"key": "value"}')
knowledge = Knowledge()
# Include only .py files
result = knowledge.index(
tmpdir,
include_glob=["*.py"],
user_id="test_user",
incremental=False,
)
assert result.files_indexed == 2
# Exclude test files (use fresh knowledge instance)
knowledge2 = Knowledge()
result = knowledge2.index(
tmpdir,
include_glob=["*.py"],
exclude_glob=["test*.py"],
user_id="test_user2",
incremental=False,
)
assert result.files_indexed == 1
def test_get_corpus_stats(self):
"""Knowledge should provide corpus stats."""
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.knowledge.indexing import CorpusStats
with tempfile.TemporaryDirectory() as tmpdir:
# Create test files
for i in range(3):
with open(os.path.join(tmpdir, f"file{i}.txt"), "w") as f:
f.write(f"Content {i} " * 50)
knowledge = Knowledge()
knowledge.index(tmpdir, user_id="test_user", incremental=False)
stats = knowledge.get_corpus_stats()
assert isinstance(stats, CorpusStats)
assert stats.file_count == 3
# chunk_count may be 0 if storage fails, just check stats exist
@@ -0,0 +1,234 @@
"""
Unit tests for Knowledge Models.
Tests normalization functions and dataclasses.
"""
from praisonaiagents.knowledge.models import (
SearchResultItem,
SearchResult,
AddResult,
normalize_search_item,
normalize_search_result,
normalize_to_dict,
)
class TestSearchResultItem:
"""Tests for SearchResultItem dataclass."""
def test_default_values(self):
"""Test default values are set correctly."""
item = SearchResultItem()
assert item.id == ""
assert item.text == ""
assert item.score == 0.0
assert item.metadata == {}
assert item.source is None
assert item.filename is None
def test_metadata_never_none(self):
"""Test metadata is never None - critical for mem0 compatibility."""
# Explicitly pass None
item = SearchResultItem(metadata=None)
assert item.metadata == {}
assert item.metadata is not None
def test_to_dict(self):
"""Test conversion to dict."""
item = SearchResultItem(
id="test-id",
text="test content",
score=0.95,
metadata={"key": "value"},
source="test.txt",
)
d = item.to_dict()
assert d["id"] == "test-id"
assert d["text"] == "test content"
assert d["score"] == 0.95
assert d["metadata"] == {"key": "value"}
assert d["source"] == "test.txt"
class TestSearchResult:
"""Tests for SearchResult dataclass."""
def test_default_values(self):
"""Test default values are set correctly."""
result = SearchResult()
assert result.results == []
assert result.metadata == {}
assert result.query == ""
def test_metadata_never_none(self):
"""Test metadata is never None."""
result = SearchResult(metadata=None)
assert result.metadata == {}
def test_results_never_none(self):
"""Test results is never None."""
result = SearchResult(results=None)
assert result.results == []
def test_to_legacy_format(self):
"""Test conversion to legacy format with 'memory' field."""
result = SearchResult(
results=[
SearchResultItem(id="1", text="content 1"),
SearchResultItem(id="2", text="content 2"),
]
)
legacy = result.to_legacy_format()
assert "results" in legacy
assert len(legacy["results"]) == 2
assert legacy["results"][0]["memory"] == "content 1"
assert legacy["results"][0]["text"] == "content 1"
class TestNormalizeSearchItem:
"""Tests for normalize_search_item function."""
def test_normalize_none(self):
"""Test normalizing None returns empty item."""
item = normalize_search_item(None)
assert item.id == ""
assert item.text == ""
assert item.metadata == {}
def test_normalize_mem0_format(self):
"""Test normalizing mem0 format with 'memory' field."""
raw = {
"id": "mem0-id",
"memory": "mem0 content",
"score": 0.85,
"metadata": {"source": "test.txt"},
}
item = normalize_search_item(raw)
assert item.id == "mem0-id"
assert item.text == "mem0 content"
assert item.score == 0.85
assert item.metadata == {"source": "test.txt"}
assert item.source == "test.txt"
def test_normalize_metadata_none(self):
"""Test normalizing when metadata is None - critical mem0 bug fix."""
raw = {
"id": "test-id",
"memory": "content",
"score": 0.9,
"metadata": None, # mem0 returns this!
}
item = normalize_search_item(raw)
assert item.metadata == {}
assert item.metadata is not None
def test_normalize_standard_format(self):
"""Test normalizing standard format with 'text' field."""
raw = {
"id": "std-id",
"text": "standard content",
"score": 0.75,
"metadata": {"key": "value"},
}
item = normalize_search_item(raw)
assert item.id == "std-id"
assert item.text == "standard content"
assert item.score == 0.75
def test_normalize_extracts_source_from_metadata(self):
"""Test source is extracted from metadata if not at top level."""
raw = {
"id": "1",
"text": "content",
"metadata": {"source": "doc.pdf", "filename": "doc.pdf"},
}
item = normalize_search_item(raw)
assert item.source == "doc.pdf"
assert item.filename == "doc.pdf"
class TestNormalizeSearchResult:
"""Tests for normalize_search_result function."""
def test_normalize_none(self):
"""Test normalizing None returns empty result."""
result = normalize_search_result(None)
assert result.results == []
assert result.metadata == {}
def test_normalize_dict_with_results(self):
"""Test normalizing dict with 'results' key (mem0 format)."""
raw = {
"results": [
{"id": "1", "memory": "content 1", "metadata": None},
{"id": "2", "memory": "content 2", "metadata": {"key": "val"}},
]
}
result = normalize_search_result(raw)
assert len(result.results) == 2
assert result.results[0].text == "content 1"
assert result.results[0].metadata == {} # None normalized to {}
assert result.results[1].metadata == {"key": "val"}
def test_normalize_filters_none_items(self):
"""Test None items in results list are filtered out."""
raw = {
"results": [
{"id": "1", "memory": "content"},
None,
{"id": "2", "memory": "content 2"},
]
}
result = normalize_search_result(raw)
assert len(result.results) == 2
def test_normalize_list_format(self):
"""Test normalizing list of results."""
raw = [
{"id": "1", "text": "content 1"},
{"id": "2", "text": "content 2"},
]
result = normalize_search_result(raw)
assert len(result.results) == 2
def test_normalize_single_dict(self):
"""Test normalizing single result dict."""
raw = {"id": "1", "text": "content"}
result = normalize_search_result(raw)
assert len(result.results) == 1
assert result.results[0].text == "content"
class TestNormalizeToDict:
"""Tests for normalize_to_dict function."""
def test_returns_legacy_format(self):
"""Test returns dict with 'results' key in legacy format."""
raw = {
"results": [
{"id": "1", "memory": "content", "metadata": None},
]
}
result = normalize_to_dict(raw)
assert "results" in result
assert result["results"][0]["memory"] == "content"
assert result["results"][0]["text"] == "content"
assert result["results"][0]["metadata"] == {}
class TestAddResult:
"""Tests for AddResult dataclass."""
def test_default_values(self):
"""Test default values."""
result = AddResult()
assert result.id == ""
assert result.success is True
assert result.message == ""
assert result.metadata == {}
def test_metadata_never_none(self):
"""Test metadata is never None."""
result = AddResult(metadata=None)
assert result.metadata == {}
@@ -0,0 +1,207 @@
"""
Unit tests for Token Budget Infrastructure (Phase 1).
Tests TokenBudget dataclass, model context window detection,
and dynamic budget calculation.
"""
import pytest
from dataclasses import dataclass
class TestTokenBudget:
"""Tests for TokenBudget dataclass."""
def test_import_token_budget(self):
"""TokenBudget should be importable from rag module."""
from praisonaiagents.rag.budget import TokenBudget
assert TokenBudget is not None
def test_default_values(self):
"""TokenBudget should have sensible defaults."""
from praisonaiagents.rag.budget import TokenBudget
budget = TokenBudget()
assert budget.model_max_tokens == 128000 # Default for GPT-4
assert budget.reserved_response_tokens == 4096
assert budget.reserved_system_tokens == 1000
assert budget.reserved_history_tokens == 2000
def test_max_context_tokens_calculation(self):
"""max_context_tokens should be calculated correctly."""
from praisonaiagents.rag.budget import TokenBudget
budget = TokenBudget(
model_max_tokens=100000,
reserved_response_tokens=4000,
reserved_system_tokens=1000,
reserved_history_tokens=2000,
)
# 100000 - 4000 - 1000 - 2000 = 93000
assert budget.max_context_tokens == 93000
def test_dynamic_budget_calculation(self):
"""dynamic_budget should calculate remaining tokens correctly."""
from praisonaiagents.rag.budget import TokenBudget
budget = TokenBudget(
model_max_tokens=100000,
reserved_response_tokens=4000,
)
# With 5000 prompt tokens and 3000 history tokens
remaining = budget.dynamic_budget(prompt_tokens=5000, history_tokens=3000)
# 100000 - 5000 - 3000 - 4000 = 88000
assert remaining == 88000
def test_dynamic_budget_never_negative(self):
"""dynamic_budget should never return negative values."""
from praisonaiagents.rag.budget import TokenBudget
budget = TokenBudget(model_max_tokens=10000)
# Request more than available
remaining = budget.dynamic_budget(prompt_tokens=50000, history_tokens=50000)
assert remaining == 0
def test_from_model_name(self):
"""TokenBudget should be creatable from model name."""
from praisonaiagents.rag.budget import TokenBudget
# GPT-4 Turbo
budget = TokenBudget.from_model("gpt-4-turbo")
assert budget.model_max_tokens == 128000
# GPT-3.5 Turbo
budget = TokenBudget.from_model("gpt-3.5-turbo")
assert budget.model_max_tokens == 16385
# Claude 3
budget = TokenBudget.from_model("claude-3-opus")
assert budget.model_max_tokens == 200000
# Unknown model should use fallback
budget = TokenBudget.from_model("unknown-model")
assert budget.model_max_tokens == 8192 # Safe fallback
def test_to_dict(self):
"""TokenBudget should be serializable to dict."""
from praisonaiagents.rag.budget import TokenBudget
budget = TokenBudget(model_max_tokens=50000)
d = budget.to_dict()
assert d["model_max_tokens"] == 50000
assert "reserved_response_tokens" in d
assert "max_context_tokens" in d
def test_from_dict(self):
"""TokenBudget should be deserializable from dict."""
from praisonaiagents.rag.budget import TokenBudget
d = {
"model_max_tokens": 50000,
"reserved_response_tokens": 2000,
}
budget = TokenBudget.from_dict(d)
assert budget.model_max_tokens == 50000
assert budget.reserved_response_tokens == 2000
class TestModelContextWindows:
"""Tests for model context window detection."""
def test_get_model_context_window(self):
"""get_model_context_window should return correct values."""
from praisonaiagents.rag.budget import get_model_context_window
assert get_model_context_window("gpt-4") == 8192
assert get_model_context_window("gpt-4-turbo") == 128000
assert get_model_context_window("gpt-4o") == 128000
assert get_model_context_window("gpt-4o-mini") == 128000
assert get_model_context_window("gpt-3.5-turbo") == 16385
assert get_model_context_window("claude-3-opus") == 200000
assert get_model_context_window("claude-3-sonnet") == 200000
assert get_model_context_window("claude-3-haiku") == 200000
assert get_model_context_window("gemini-pro") == 32768
assert get_model_context_window("gemini-1.5-pro") == 1000000
def test_unknown_model_fallback(self):
"""Unknown models should return safe fallback."""
from praisonaiagents.rag.budget import get_model_context_window
assert get_model_context_window("unknown-model-xyz") == 8192
assert get_model_context_window("") == 8192
assert get_model_context_window(None) == 8192
class TestBudgetEnforcement:
"""Tests for budget enforcement in context building."""
def test_budget_enforcer_protocol(self):
"""BudgetEnforcerProtocol should be importable."""
from praisonaiagents.rag.budget import BudgetEnforcerProtocol
assert BudgetEnforcerProtocol is not None
def test_default_budget_enforcer(self):
"""DefaultBudgetEnforcer should enforce token limits."""
from praisonaiagents.rag.budget import DefaultBudgetEnforcer, TokenBudget
enforcer = DefaultBudgetEnforcer()
budget = TokenBudget(model_max_tokens=1000, reserved_response_tokens=200)
# Create chunks that exceed budget
chunks = [
{"text": "word " * 100, "metadata": {}}, # ~100 tokens
{"text": "word " * 100, "metadata": {}}, # ~100 tokens
{"text": "word " * 100, "metadata": {}}, # ~100 tokens
]
result = enforcer.enforce(chunks, budget, prompt_tokens=500, history_tokens=0)
# Should have truncated to fit budget
# Available: 1000 - 200 - 500 = 300 tokens
# Each chunk is ~100 tokens, so should fit ~3 chunks
assert len(result) <= 3
class TestRetrievalConfigBudgetIntegration:
"""Tests for RetrievalConfig budget integration."""
def test_retrieval_config_has_budget_fields(self):
"""RetrievalConfig should have budget-related fields."""
from praisonaiagents.rag.retrieval_config import RetrievalConfig
config = RetrievalConfig()
# Should have max_context_tokens (already exists)
assert hasattr(config, "max_context_tokens")
# Should have new budget fields
assert hasattr(config, "model_context_window")
assert hasattr(config, "reserved_response_tokens")
assert hasattr(config, "dynamic_budget")
def test_retrieval_config_budget_defaults(self):
"""RetrievalConfig budget fields should have sensible defaults."""
from praisonaiagents.rag.retrieval_config import RetrievalConfig
config = RetrievalConfig()
# model_context_window should default to None (auto-detect)
assert config.model_context_window is None
# dynamic_budget should default to True
assert config.dynamic_budget is True
def test_retrieval_config_get_token_budget(self):
"""RetrievalConfig should provide TokenBudget instance."""
from praisonaiagents.rag.retrieval_config import RetrievalConfig
from praisonaiagents.rag.budget import TokenBudget
config = RetrievalConfig(model_context_window=100000)
budget = config.get_token_budget()
assert isinstance(budget, TokenBudget)
assert budget.model_max_tokens == 100000
@@ -0,0 +1,104 @@
"""
Unit tests for ContextCompressor (Phase 5).
"""
class TestContextCompressor:
"""Tests for ContextCompressor class."""
def test_import_context_compressor(self):
"""ContextCompressor should be importable."""
from praisonaiagents.rag.compressor import ContextCompressor
assert ContextCompressor is not None
def test_compression_result_dataclass(self):
"""CompressionResult should be a proper dataclass."""
from praisonaiagents.rag.compressor import CompressionResult
result = CompressionResult()
assert result.chunks == []
assert result.original_tokens == 0
assert result.compressed_tokens == 0
assert result.compression_ratio == 1.0
def test_compress_empty_chunks(self):
"""ContextCompressor should handle empty chunks."""
from praisonaiagents.rag.compressor import ContextCompressor
compressor = ContextCompressor()
result = compressor.compress([], "test query", 1000)
assert result.chunks == []
assert result.original_tokens == 0
def test_deduplicate_chunks(self):
"""ContextCompressor should deduplicate similar chunks."""
from praisonaiagents.rag.compressor import ContextCompressor
compressor = ContextCompressor()
chunks = [
{"text": "This is duplicate content"},
{"text": "This is duplicate content"},
{"text": "This is unique content"},
]
deduped = compressor._deduplicate(chunks)
assert len(deduped) == 2
def test_extract_relevant_sentences(self):
"""ContextCompressor should extract query-relevant sentences."""
from praisonaiagents.rag.compressor import ContextCompressor
compressor = ContextCompressor()
chunks = [
{"text": "Python is great. JavaScript is also good. Python has many libraries."},
]
extracted = compressor._extract_relevant(chunks, "Python libraries")
assert len(extracted) == 1
assert "Python" in extracted[0]["text"]
def test_truncate_to_budget(self):
"""ContextCompressor should truncate to fit budget."""
from praisonaiagents.rag.compressor import ContextCompressor
compressor = ContextCompressor()
chunks = [
{"text": "A" * 1000}, # ~250 tokens
{"text": "B" * 1000}, # ~250 tokens
]
truncated = compressor._truncate_to_budget(chunks, 300)
# Should fit within budget
total_tokens = sum(len(c["text"]) // 4 + 1 for c in truncated)
assert total_tokens <= 350 # Some buffer for truncation
def test_compress_full_pipeline(self):
"""ContextCompressor should run full compression pipeline."""
from praisonaiagents.rag.compressor import ContextCompressor
compressor = ContextCompressor()
chunks = [
{"text": "Python programming is powerful. It has many uses."},
{"text": "Python programming is powerful. It has many uses."}, # Duplicate
{"text": "JavaScript is for web development."},
]
result = compressor.compress(chunks, "Python", 500)
assert result.method_used == "dedupe+extract+truncate"
assert result.original_tokens > 0
assert len(result.chunks) <= 2 # Duplicate removed
class TestEstimateTokens:
"""Tests for token estimation."""
def test_estimate_tokens_function(self):
"""estimate_tokens should estimate correctly."""
from praisonaiagents.rag.compressor import estimate_tokens
assert estimate_tokens("") == 0
assert estimate_tokens("hello") == 2 # 5/4 + 1 = 2
assert estimate_tokens("a" * 100) == 26 # 100/4 + 1 = 26
@@ -0,0 +1,124 @@
"""
Unit tests for RAG context builder normalization.
Tests that metadata=None from mem0 is handled correctly.
"""
from praisonaiagents.rag.context import (
deduplicate_chunks,
build_context,
)
class TestDeduplicateChunksNormalization:
"""Tests for deduplicate_chunks handling of None metadata."""
def test_handles_metadata_none(self):
"""Test that metadata=None doesn't crash deduplicate_chunks."""
results = [
{"id": "1", "text": "content 1", "metadata": None},
{"id": "2", "text": "content 2", "metadata": {"source": "test.txt"}},
]
# Should not raise AttributeError
deduped = deduplicate_chunks(results)
assert len(deduped) == 2
def test_handles_none_results(self):
"""Test that None items in list are filtered."""
results = [
{"id": "1", "text": "content 1", "metadata": None},
None,
{"id": "2", "text": "content 2", "metadata": {}},
]
deduped = deduplicate_chunks(results)
assert len(deduped) == 2
def test_handles_mem0_format(self):
"""Test handling of mem0 'memory' field."""
results = [
{"id": "1", "memory": "mem0 content", "metadata": None},
]
deduped = deduplicate_chunks(results)
assert len(deduped) == 1
class TestBuildContextNormalization:
"""Tests for build_context handling of None metadata."""
def test_handles_metadata_none(self):
"""Test that metadata=None doesn't crash build_context."""
results = [
{"id": "1", "text": "content 1", "metadata": None},
{"id": "2", "text": "content 2", "metadata": {"source": "test.txt"}},
]
# Should not raise AttributeError
context, used = build_context(results)
assert len(used) == 2
assert "content 1" in context
assert "content 2" in context
def test_handles_none_results(self):
"""Test that None items in list are filtered."""
results = [
{"id": "1", "text": "content 1", "metadata": None},
None,
{"id": "2", "text": "content 2", "metadata": {}},
]
context, used = build_context(results)
assert len(used) == 2
def test_handles_mem0_format(self):
"""Test handling of mem0 'memory' field."""
results = [
{"id": "1", "memory": "mem0 content", "metadata": None, "score": 0.9},
]
context, used = build_context(results)
assert len(used) == 1
assert "mem0 content" in context
def test_source_extraction_with_none_metadata(self):
"""Test source extraction when metadata is None."""
results = [
{"id": "1", "text": "content", "metadata": None},
]
context, used = build_context(results, include_source=True)
# Should use default source label
assert "Source 1" in context or "content" in context
def test_source_extraction_with_metadata(self):
"""Test source extraction when metadata has source."""
results = [
{"id": "1", "text": "content", "metadata": {"source": "doc.pdf"}},
]
context, used = build_context(results, include_source=True)
assert "doc.pdf" in context
class TestEdgeCases:
"""Test edge cases for robustness."""
def test_empty_results(self):
"""Test empty results list."""
context, used = build_context([])
assert context == ""
assert used == []
def test_all_none_results(self):
"""Test list with all None items."""
results = [None, None, None]
deduped = deduplicate_chunks(results)
assert deduped == []
context, used = build_context(results)
assert context == ""
assert used == []
def test_empty_text_filtered(self):
"""Test that empty text items are filtered in build_context."""
results = [
{"id": "1", "text": "", "metadata": None},
{"id": "2", "text": "valid content", "metadata": None},
]
context, used = build_context(results)
assert len(used) == 1
assert "valid content" in context
@@ -0,0 +1,130 @@
"""
Unit tests for SmartRetriever (Phase 4).
"""
class TestSmartRetriever:
"""Tests for SmartRetriever class."""
def test_import_smart_retriever(self):
"""SmartRetriever should be importable."""
from praisonaiagents.rag.retriever import SmartRetriever
assert SmartRetriever is not None
def test_retrieval_result_dataclass(self):
"""RetrievalResult should be a proper dataclass."""
from praisonaiagents.rag.retriever import RetrievalResult
result = RetrievalResult()
assert result.chunks == []
assert result.total_found == 0
assert result.strategy_used == "basic"
assert result.reranked is False
def test_retriever_without_knowledge(self):
"""SmartRetriever should handle missing knowledge gracefully."""
from praisonaiagents.rag.retriever import SmartRetriever
retriever = SmartRetriever()
result = retriever.retrieve("test query")
assert result.chunks == []
assert result.total_found == 0
def test_normalize_results(self):
"""SmartRetriever should normalize various result formats."""
from praisonaiagents.rag.retriever import SmartRetriever
retriever = SmartRetriever()
# Test dict with results key
results = {"results": [{"text": "hello", "score": 0.9}]}
normalized = retriever._normalize_results(results)
assert len(normalized) == 1
assert normalized[0]["text"] == "hello"
# Test list
results = [{"text": "world", "metadata": {"source": "test"}}]
normalized = retriever._normalize_results(results)
assert len(normalized) == 1
assert normalized[0]["text"] == "world"
def test_apply_filters_include(self):
"""SmartRetriever should filter by include patterns."""
from praisonaiagents.rag.retriever import SmartRetriever
retriever = SmartRetriever()
chunks = [
{"text": "Python code", "metadata": {"filename": "main.py"}},
{"text": "JavaScript", "metadata": {"filename": "app.js"}},
]
filtered = retriever._apply_filters(chunks, include_glob=["*.py"])
assert len(filtered) == 1
assert filtered[0]["metadata"]["filename"] == "main.py"
def test_apply_filters_exclude(self):
"""SmartRetriever should filter by exclude patterns."""
from praisonaiagents.rag.retriever import SmartRetriever
retriever = SmartRetriever()
chunks = [
{"text": "Python code", "metadata": {"filename": "main.py"}},
{"text": "Test code", "metadata": {"filename": "test_main.py"}},
]
filtered = retriever._apply_filters(chunks, exclude_glob=["test_*.py"])
assert len(filtered) == 1
assert filtered[0]["metadata"]["filename"] == "main.py"
class TestSimpleReranker:
"""Tests for SimpleReranker class."""
def test_import_simple_reranker(self):
"""SimpleReranker should be importable."""
from praisonaiagents.rag.retriever import SimpleReranker
assert SimpleReranker is not None
def test_rerank_by_overlap(self):
"""SimpleReranker should rerank by keyword overlap."""
from praisonaiagents.rag.retriever import SimpleReranker
reranker = SimpleReranker()
results = [
{"text": "unrelated content", "score": 0.5},
{"text": "python programming language", "score": 0.5},
]
reranked = reranker.rerank("python programming", results)
# The one with more overlap should be first (same base score, overlap decides)
assert "python" in reranked[0]["text"].lower()
def test_rerank_top_k(self):
"""SimpleReranker should respect top_k."""
from praisonaiagents.rag.retriever import SimpleReranker
reranker = SimpleReranker()
results = [
{"text": "a", "score": 0.1},
{"text": "b", "score": 0.2},
{"text": "c", "score": 0.3},
]
reranked = reranker.rerank("test", results, top_k=2)
assert len(reranked) == 2
class TestRetrieverProtocol:
"""Tests for RetrieverProtocol."""
def test_protocol_exists(self):
"""RetrieverProtocol should be importable."""
from praisonaiagents.rag.retriever import RetrieverProtocol
assert RetrieverProtocol is not None
def test_reranker_protocol_exists(self):
"""RerankerProtocol should be importable."""
from praisonaiagents.rag.retriever import RerankerProtocol
assert RerankerProtocol is not None
@@ -0,0 +1,149 @@
"""
Unit tests for Strategy Selection (Phase 3).
Tests RetrievalStrategy enum, select_strategy(), and auto-strategy selection.
"""
class TestRetrievalStrategyEnum:
"""Tests for RetrievalStrategy enum."""
def test_import_retrieval_strategy(self):
"""RetrievalStrategy should be importable."""
from praisonaiagents.rag.strategy import RetrievalStrategy
assert RetrievalStrategy is not None
def test_all_strategy_levels(self):
"""RetrievalStrategy should have all levels."""
from praisonaiagents.rag.strategy import RetrievalStrategy
assert RetrievalStrategy.DIRECT.value == "direct"
assert RetrievalStrategy.BASIC.value == "basic"
assert RetrievalStrategy.HYBRID.value == "hybrid"
assert RetrievalStrategy.RERANKED.value == "reranked"
assert RetrievalStrategy.COMPRESSED.value == "compressed"
assert RetrievalStrategy.HIERARCHICAL.value == "hierarchical"
def test_strategy_from_string(self):
"""RetrievalStrategy should be creatable from string."""
from praisonaiagents.rag.strategy import RetrievalStrategy
assert RetrievalStrategy("direct") == RetrievalStrategy.DIRECT
assert RetrievalStrategy("basic") == RetrievalStrategy.BASIC
assert RetrievalStrategy("hybrid") == RetrievalStrategy.HYBRID
class TestSelectStrategy:
"""Tests for select_strategy() function."""
def test_import_select_strategy(self):
"""select_strategy should be importable."""
from praisonaiagents.rag.strategy import select_strategy
assert select_strategy is not None
def test_select_strategy_by_file_count(self):
"""select_strategy should select based on file count."""
from praisonaiagents.rag.strategy import select_strategy, RetrievalStrategy
from praisonaiagents.knowledge.indexing import CorpusStats
# Small corpus (< 10 files)
stats = CorpusStats(file_count=5)
assert select_strategy(stats) == RetrievalStrategy.DIRECT
# Medium corpus (< 100 files)
stats = CorpusStats(file_count=50)
assert select_strategy(stats) == RetrievalStrategy.BASIC
# Large corpus (< 1000 files)
stats = CorpusStats(file_count=500)
assert select_strategy(stats) == RetrievalStrategy.HYBRID
# Very large corpus (< 10000 files)
stats = CorpusStats(file_count=5000)
assert select_strategy(stats) == RetrievalStrategy.RERANKED
# Massive corpus (< 100000 files)
stats = CorpusStats(file_count=50000)
assert select_strategy(stats) == RetrievalStrategy.COMPRESSED
# Huge corpus (>= 100000 files)
stats = CorpusStats(file_count=500000)
assert select_strategy(stats) == RetrievalStrategy.HIERARCHICAL
def test_select_strategy_with_override(self):
"""select_strategy should respect explicit override."""
from praisonaiagents.rag.strategy import select_strategy, RetrievalStrategy
from praisonaiagents.knowledge.indexing import CorpusStats
stats = CorpusStats(file_count=5) # Would normally be DIRECT
# Override to RERANKED
result = select_strategy(stats, override="reranked")
assert result == RetrievalStrategy.RERANKED
def test_select_strategy_auto(self):
"""select_strategy with 'auto' should auto-select."""
from praisonaiagents.rag.strategy import select_strategy, RetrievalStrategy
from praisonaiagents.knowledge.indexing import CorpusStats
stats = CorpusStats(file_count=500)
result = select_strategy(stats, override="auto")
assert result == RetrievalStrategy.HYBRID
class TestStrategyMetadata:
"""Tests for strategy metadata and descriptions."""
def test_strategy_description(self):
"""Each strategy should have a description."""
from praisonaiagents.rag.strategy import get_strategy_description, RetrievalStrategy
desc = get_strategy_description(RetrievalStrategy.DIRECT)
assert "direct" in desc.lower() or "load" in desc.lower()
desc = get_strategy_description(RetrievalStrategy.HIERARCHICAL)
assert "hierarchical" in desc.lower() or "summary" in desc.lower()
def test_strategy_thresholds(self):
"""Strategy thresholds should be accessible."""
from praisonaiagents.rag.strategy import STRATEGY_THRESHOLDS
assert 10 in STRATEGY_THRESHOLDS
assert 100 in STRATEGY_THRESHOLDS
assert 1000 in STRATEGY_THRESHOLDS
class TestRetrievalConfigStrategy:
"""Tests for RetrievalConfig strategy integration."""
def test_retrieval_config_has_strategy(self):
"""RetrievalConfig should have strategy field."""
from praisonaiagents.rag.retrieval_config import RetrievalConfig
config = RetrievalConfig()
assert hasattr(config, "strategy")
assert config.strategy == "auto"
def test_retrieval_config_get_strategy(self):
"""RetrievalConfig should provide strategy selection."""
from praisonaiagents.rag.retrieval_config import RetrievalConfig
from praisonaiagents.rag.strategy import RetrievalStrategy
from praisonaiagents.knowledge.indexing import CorpusStats
config = RetrievalConfig(strategy="auto")
stats = CorpusStats(file_count=500)
strategy = config.get_strategy(stats)
assert strategy == RetrievalStrategy.HYBRID
def test_retrieval_config_explicit_strategy(self):
"""RetrievalConfig should respect explicit strategy."""
from praisonaiagents.rag.retrieval_config import RetrievalConfig
from praisonaiagents.rag.strategy import RetrievalStrategy
from praisonaiagents.knowledge.indexing import CorpusStats
config = RetrievalConfig(strategy="reranked")
stats = CorpusStats(file_count=5) # Would normally be DIRECT
strategy = config.get_strategy(stats)
assert strategy == RetrievalStrategy.RERANKED
@@ -0,0 +1,139 @@
"""
Unit tests for HierarchicalSummarizer (Phase 6).
"""
import os
import tempfile
class TestHierarchicalSummarizer:
"""Tests for HierarchicalSummarizer class."""
def test_import_hierarchical_summarizer(self):
"""HierarchicalSummarizer should be importable."""
from praisonaiagents.rag.summarizer import HierarchicalSummarizer
assert HierarchicalSummarizer is not None
def test_summary_node_dataclass(self):
"""SummaryNode should be a proper dataclass."""
from praisonaiagents.rag.summarizer import SummaryNode
node = SummaryNode(path="/test", level=1, summary="Test summary")
assert node.path == "/test"
assert node.level == 1
assert node.summary == "Test summary"
def test_summary_node_to_dict(self):
"""SummaryNode should serialize to dict."""
from praisonaiagents.rag.summarizer import SummaryNode
node = SummaryNode(path="/test", level=1, summary="Test")
data = node.to_dict()
assert data["path"] == "/test"
assert data["level"] == 1
def test_summary_node_from_dict(self):
"""SummaryNode should deserialize from dict."""
from praisonaiagents.rag.summarizer import SummaryNode
data = {"path": "/test", "level": 2, "summary": "Folder summary"}
node = SummaryNode.from_dict(data)
assert node.path == "/test"
assert node.level == 2
def test_hierarchy_result_dataclass(self):
"""HierarchyResult should be a proper dataclass."""
from praisonaiagents.rag.summarizer import HierarchyResult
result = HierarchyResult()
assert result.nodes == {}
assert result.total_files == 0
assert result.levels == 3
def test_build_hierarchy_empty_dir(self):
"""HierarchicalSummarizer should handle empty directories."""
from praisonaiagents.rag.summarizer import HierarchicalSummarizer
with tempfile.TemporaryDirectory() as tmpdir:
summarizer = HierarchicalSummarizer()
result = summarizer.build_hierarchy(tmpdir)
assert result.total_files == 0
def test_build_hierarchy_with_files(self):
"""HierarchicalSummarizer should build hierarchy for files."""
from praisonaiagents.rag.summarizer import HierarchicalSummarizer
with tempfile.TemporaryDirectory() as tmpdir:
# Create test files
with open(os.path.join(tmpdir, "main.py"), "w") as f:
f.write("def main():\n print('Hello')")
with open(os.path.join(tmpdir, "utils.py"), "w") as f:
f.write("def helper():\n return True")
summarizer = HierarchicalSummarizer()
result = summarizer.build_hierarchy(tmpdir, levels=1)
assert result.total_files == 2
assert len(result.nodes) >= 2
def test_build_hierarchy_with_folders(self):
"""HierarchicalSummarizer should build folder summaries."""
from praisonaiagents.rag.summarizer import HierarchicalSummarizer
with tempfile.TemporaryDirectory() as tmpdir:
# Create nested structure
src_dir = os.path.join(tmpdir, "src")
os.makedirs(src_dir)
with open(os.path.join(src_dir, "app.py"), "w") as f:
f.write("# Application code")
summarizer = HierarchicalSummarizer()
result = summarizer.build_hierarchy(tmpdir, levels=2)
# Should have file and folder summaries
assert result.total_files >= 1
def test_query_hierarchy(self):
"""HierarchicalSummarizer should query hierarchy."""
from praisonaiagents.rag.summarizer import HierarchicalSummarizer
with tempfile.TemporaryDirectory() as tmpdir:
# Create test files with specific content
with open(os.path.join(tmpdir, "python.py"), "w") as f:
f.write("Python programming language code")
with open(os.path.join(tmpdir, "javascript.js"), "w") as f:
f.write("JavaScript web development")
summarizer = HierarchicalSummarizer()
summarizer.build_hierarchy(tmpdir, levels=1)
results = summarizer.query_hierarchy("Python", tmpdir)
# Should find Python-related content
assert len(results) >= 0 # May be 0 if no overlap
def test_extract_summary(self):
"""HierarchicalSummarizer should extract summaries."""
from praisonaiagents.rag.summarizer import HierarchicalSummarizer
summarizer = HierarchicalSummarizer()
content = "First line.\nSecond line.\nThird line."
summary = summarizer._extract_summary(content, max_chars=50)
assert len(summary) <= 60 # Some buffer
assert "First" in summary
class TestEstimateTokensSummarizer:
"""Tests for token estimation in summarizer."""
def test_estimate_tokens(self):
"""estimate_tokens should work correctly."""
from praisonaiagents.rag.summarizer import estimate_tokens
assert estimate_tokens("") == 0
assert estimate_tokens("test") == 2
@@ -17,38 +17,40 @@ from typing import Dict, Any, List
# ============================================================================
class TestAutoSummarization:
"""Tests for automatic conversation summarization."""
"""Tests for automatic conversation summarization.
def test_summarization_config_defaults(self):
"""Test default summarization configuration."""
from praisonaiagents.agent import Agent
agent = Agent(instructions="Test agent")
# Should have auto_summarize disabled by default
assert hasattr(agent, 'auto_summarize')
assert agent.auto_summarize == False
def test_summarization_config_enabled(self):
"""Test enabling auto-summarization."""
from praisonaiagents.agent import Agent
agent = Agent(
instructions="Test agent",
auto_summarize=True,
summarize_threshold=0.8
)
assert agent.auto_summarize == True
assert agent.summarize_threshold == 0.8
def test_summarization_threshold_validation(self):
"""Test threshold must be between 0 and 1."""
Note: auto_summarize and summarize_threshold params have been removed from Agent.
Summarization is now controlled via context=ManagerConfig(auto_compact=True, compact_threshold=0.8).
These tests verify the old params are rejected.
"""
def test_old_summarization_params_rejected(self):
"""Test that old auto_summarize params are rejected."""
from praisonaiagents.agent import Agent
with pytest.raises(ValueError):
Agent(instructions="Test", auto_summarize=True, summarize_threshold=1.5)
# Old params should raise TypeError (unexpected keyword argument)
with pytest.raises(TypeError):
Agent(instructions="Test agent", auto_summarize=True)
with pytest.raises(ValueError):
Agent(instructions="Test", auto_summarize=True, summarize_threshold=-0.1)
with pytest.raises(TypeError):
Agent(instructions="Test agent", summarize_threshold=0.8)
def test_summarization_via_context_param(self):
"""Test summarization is now controlled via context= param."""
from praisonaiagents.agent import Agent
from praisonaiagents.context import ManagerConfig
# New way: use context= param with ManagerConfig
config = ManagerConfig(
auto_compact=True,
compact_threshold=0.8,
)
agent = Agent(instructions="Test agent", context=config)
# Verify context manager is configured
assert agent.context_manager is not None
assert agent.context_manager.config.auto_compact == True
assert agent.context_manager.config.compact_threshold == 0.8
def test_token_tracking_for_summarization(self):
"""Test that token tracking works for summarization decisions."""
+1 -1
View File
@@ -2422,7 +2422,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.7.1"
version = "0.8.1"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },
+2 -2
View File
@@ -1930,11 +1930,11 @@ agent = Agent(
auto_save="my-project"
)
# Load history from last 5 sessions
# Load history from past sessions via context management
agent = Agent(
name="Assistant",
memory=True,
history_in_context=5
context=True, # Enable context management for history
)
```
+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/v3.2.1.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.2.1.tar.gz | shasum -a 256`.split.first
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.3.1.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.3.1.tar.gz | shasum -a 256`.split.first
license "MIT"
depends_on "[email protected]"
@@ -122,6 +122,10 @@ def knowledge_profiler(enabled: bool, profile_out: Optional[Path], profile_top:
def knowledge_index(
sources: List[str] = typer.Argument(..., help="Source files, directories, or URLs to index"),
collection: str = typer.Option("default", "--collection", "-c", help="Collection/knowledge base name"),
user_id: Optional[str] = typer.Option(None, "--user-id", "-u", help="User ID for scoping (required for mem0 backend)"),
agent_id: Optional[str] = typer.Option(None, "--agent-id", "-a", help="Agent ID for scoping"),
run_id: Optional[str] = typer.Option(None, "--run-id", "-r", help="Run ID for scoping"),
backend: str = typer.Option("mem0", "--backend", "-b", help="Knowledge backend: mem0, chroma, internal"),
config: Optional[Path] = typer.Option(None, "--config", "-f", help="Config file path"),
verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose output"),
profile: bool = typer.Option(False, "--profile", help="Enable performance profiling"),
@@ -135,9 +139,9 @@ def knowledge_index(
Use `praisonai rag query` to answer questions using the indexed knowledge.
Examples:
praisonai knowledge index ./docs
praisonai knowledge index paper.pdf --collection research
praisonai knowledge index ./data --profile --profile-out ./profile.json
praisonai knowledge index ./docs --user-id myuser
praisonai knowledge index paper.pdf --collection research --agent-id research_agent
praisonai knowledge index ./data --backend chroma --profile
"""
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn
@@ -149,16 +153,34 @@ def knowledge_index(
# Lazy import to avoid startup cost
from praisonaiagents.knowledge import Knowledge
# Build config
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": collection,
"path": f"./.praison/knowledge/{collection}",
# Validate scope for mem0 backend
if backend == "mem0" and not any([user_id, agent_id, run_id]):
console.print("[yellow]Warning:[/yellow] mem0 backend requires at least one scope identifier.")
console.print("Use --user-id, --agent-id, or --run-id to scope your knowledge.")
console.print("Defaulting to user_id='default_user'")
user_id = "default_user"
# Build config based on backend
if backend == "chroma":
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": collection,
"path": f"./.praison/knowledge/{collection}",
}
}
}
else:
# mem0 or internal backend
knowledge_config = {
"vector_store": {
"provider": "qdrant",
"config": {
"collection_name": collection,
}
}
}
}
# Load config file if provided
if config and config.exists():
@@ -171,7 +193,7 @@ def knowledge_index(
# Initialize Knowledge
knowledge = Knowledge(config=knowledge_config, verbose=verbose)
# Index sources
# Index sources with scope identifiers
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
@@ -181,7 +203,12 @@ def knowledge_index(
for source in sources:
task = progress.add_task(f"Indexing {source}...", total=None)
try:
result = knowledge.add(source)
result = knowledge.add(
source,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
)
count = len(result.get("results", [])) if isinstance(result, dict) else 0
console.print(f"[green]✓[/green] Indexed {source}: {count} chunks")
except Exception as e:
@@ -207,6 +234,10 @@ def knowledge_index(
def knowledge_search(
query: str = typer.Argument(..., help="Search query"),
collection: str = typer.Option("default", "--collection", "-c", help="Collection to search"),
user_id: Optional[str] = typer.Option(None, "--user-id", "-u", help="User ID for scoping (required for mem0 backend)"),
agent_id: Optional[str] = typer.Option(None, "--agent-id", "-a", help="Agent ID for scoping"),
run_id: Optional[str] = typer.Option(None, "--run-id", "-r", help="Run ID for scoping"),
backend: str = typer.Option("mem0", "--backend", "-b", help="Knowledge backend: mem0, chroma, internal"),
top_k: int = typer.Option(5, "--top-k", "-k", help="Number of results to retrieve"),
hybrid: bool = typer.Option(False, "--hybrid", help="Use hybrid retrieval (dense + BM25)"),
config: Optional[Path] = typer.Option(None, "--config", "-f", help="Config file path"),
@@ -222,9 +253,9 @@ def knowledge_search(
For answers with citations, use `praisonai rag query`.
Examples:
praisonai knowledge search "capital of France"
praisonai knowledge search "capital of France" --user-id myuser
praisonai knowledge search "main findings" --collection research --top-k 10
praisonai knowledge search "Python tutorial" --hybrid
praisonai knowledge search "Python tutorial" --hybrid --backend chroma
"""
from rich.console import Console
from rich.table import Table
@@ -235,16 +266,34 @@ def knowledge_search(
try:
from praisonaiagents.knowledge import Knowledge
# Build config
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": collection,
"path": f"./.praison/knowledge/{collection}",
# Validate scope for mem0 backend
if backend == "mem0" and not any([user_id, agent_id, run_id]):
console.print("[yellow]Warning:[/yellow] mem0 backend requires at least one scope identifier.")
console.print("Use --user-id, --agent-id, or --run-id to scope your search.")
console.print("Defaulting to user_id='default_user'")
user_id = "default_user"
# Build config based on backend
if backend == "chroma":
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": collection,
"path": f"./.praison/knowledge/{collection}",
}
}
}
else:
# mem0 or internal backend
knowledge_config = {
"vector_store": {
"provider": "qdrant",
"config": {
"collection_name": collection,
}
}
}
}
# Add hybrid retrieval config if enabled
if hybrid:
@@ -263,13 +312,19 @@ def knowledge_search(
# Initialize and search
knowledge = Knowledge(config=knowledge_config, verbose=verbose)
# Use hybrid search if enabled
# Use hybrid search if enabled, pass scope identifiers
if hybrid:
if verbose:
console.print("[dim]Using hybrid retrieval (dense + BM25)...[/dim]")
results = knowledge.search(query, limit=top_k, hybrid=True)
results = knowledge.search(
query, limit=top_k, hybrid=True,
user_id=user_id, agent_id=agent_id, run_id=run_id,
)
else:
results = knowledge.search(query, limit=top_k)
results = knowledge.search(
query, limit=top_k,
user_id=user_id, agent_id=agent_id, run_id=run_id,
)
if not results:
console.print("[yellow]No results found.[/yellow]")
@@ -286,7 +341,9 @@ def knowledge_search(
for i, result in enumerate(result_list[:top_k], 1):
if isinstance(result, dict):
score = f"{result.get('score', 0):.3f}" if result.get('score') else "-"
source = result.get('metadata', {}).get('filename', result.get('metadata', {}).get('source', '-'))
# CRITICAL: Handle metadata=None from mem0
metadata = result.get('metadata') or {}
source = metadata.get('filename', metadata.get('source', '-'))
content = result.get('memory', result.get('text', ''))[:100] + "..."
else:
score = "-"
@@ -3,198 +3,116 @@ Context Management CLI Feature for PraisonAI.
Provides CLI commands and integration for context budgeting,
optimization, and monitoring.
This module uses the SDK ContextManager facade to ensure DRY
and CLI/SDK parity.
"""
import os
from typing import Dict, List, Any, Optional, Literal
from dataclasses import dataclass
from typing import Dict, List, Any, Optional
@dataclass
class ContextManagerConfig:
"""Configuration for context management."""
# Auto-compaction
auto_compact: bool = True
compact_threshold: float = 0.8
strategy: str = "smart"
# Output reserve
output_reserve: int = 8000
# Monitoring
monitor_enabled: bool = False
monitor_path: str = "./context.txt"
monitor_format: Literal["human", "json"] = "human"
monitor_frequency: Literal["turn", "tool_call", "manual", "overflow"] = "turn"
redact_sensitive: bool = True
@classmethod
def from_env(cls) -> "ContextManagerConfig":
"""Load config from environment variables."""
return cls(
auto_compact=os.getenv("PRAISONAI_CONTEXT_AUTO_COMPACT", "true").lower() == "true",
compact_threshold=float(os.getenv("PRAISONAI_CONTEXT_THRESHOLD", "0.8")),
strategy=os.getenv("PRAISONAI_CONTEXT_STRATEGY", "smart"),
output_reserve=int(os.getenv("PRAISONAI_CONTEXT_OUTPUT_RESERVE", "8000")),
monitor_enabled=os.getenv("PRAISONAI_CONTEXT_MONITOR", "false").lower() == "true",
monitor_path=os.getenv("PRAISONAI_CONTEXT_MONITOR_PATH", "./context.txt"),
monitor_format=os.getenv("PRAISONAI_CONTEXT_MONITOR_FORMAT", "human"),
monitor_frequency=os.getenv("PRAISONAI_CONTEXT_MONITOR_FREQUENCY", "turn"),
redact_sensitive=os.getenv("PRAISONAI_CONTEXT_REDACT", "true").lower() == "true",
)
def merge_cli_args(
self,
auto_compact: Optional[bool] = None,
strategy: Optional[str] = None,
threshold: Optional[float] = None,
monitor: Optional[bool] = None,
monitor_path: Optional[str] = None,
monitor_format: Optional[str] = None,
monitor_frequency: Optional[str] = None,
redact: Optional[bool] = None,
output_reserve: Optional[int] = None,
) -> "ContextManagerConfig":
"""Merge CLI arguments (highest precedence)."""
return ContextManagerConfig(
auto_compact=auto_compact if auto_compact is not None else self.auto_compact,
compact_threshold=threshold if threshold is not None else self.compact_threshold,
strategy=strategy if strategy is not None else self.strategy,
output_reserve=output_reserve if output_reserve is not None else self.output_reserve,
monitor_enabled=monitor if monitor is not None else self.monitor_enabled,
monitor_path=monitor_path if monitor_path is not None else self.monitor_path,
monitor_format=monitor_format if monitor_format is not None else self.monitor_format,
monitor_frequency=monitor_frequency if monitor_frequency is not None else self.monitor_frequency,
redact_sensitive=redact if redact is not None else self.redact_sensitive,
)
# Lazy import SDK components to avoid import overhead
def _get_sdk_manager():
"""Lazy import SDK ContextManager."""
from praisonaiagents.context import ContextManager
return ContextManager
def _get_sdk_config():
"""Lazy import SDK ManagerConfig."""
from praisonaiagents.context import ManagerConfig
return ManagerConfig
def _get_create_manager():
"""Lazy import create_context_manager factory."""
from praisonaiagents.context import create_context_manager
return create_context_manager
class ContextManagerHandler:
"""
Handles context management for interactive CLI sessions.
Integrates budgeting, composition, optimization, and monitoring.
This is a thin wrapper around the SDK ContextManager facade.
All core logic is in the SDK to ensure DRY and parity.
"""
def __init__(
self,
config: Optional[ContextManagerConfig] = None,
model: str = "gpt-4o-mini",
session_id: str = "",
agent_name: str = "Assistant",
config_file: Optional[str] = None,
cli_overrides: Optional[Dict[str, Any]] = None,
):
"""
Initialize context manager.
Initialize context manager handler.
Args:
config: Context management configuration
model: Model name for budget calculation
session_id: Session identifier
agent_name: Agent name for monitoring
config_file: Path to config.yaml
cli_overrides: CLI argument overrides
"""
self.config = config or ContextManagerConfig.from_env()
# Use SDK factory with proper precedence
create_manager = _get_create_manager()
self._manager = create_manager(
model=model,
session_id=session_id,
agent_name=agent_name,
config_file=config_file,
cli_overrides=cli_overrides,
)
# Expose model for compatibility
self.model = model
self.session_id = session_id
self.agent_name = agent_name
# Lazy-loaded components
self._budgeter = None
self._composer = None
self._monitor = None
self._ledger = None
@property
def budgeter(self):
"""Get or create budgeter."""
if self._budgeter is None:
from praisonaiagents.context import ContextBudgeter
self._budgeter = ContextBudgeter(
model=self.model,
output_reserve=self.config.output_reserve,
)
return self._budgeter
@property
def composer(self):
"""Get or create composer."""
if self._composer is None:
from praisonaiagents.context import ContextComposer
self._composer = ContextComposer(
budget=self.budgeter.allocate(),
)
return self._composer
@property
def monitor(self):
"""Get or create monitor."""
if self._monitor is None:
from praisonaiagents.context import ContextMonitor
self._monitor = ContextMonitor(
enabled=self.config.monitor_enabled,
path=self.config.monitor_path,
format=self.config.monitor_format,
frequency=self.config.monitor_frequency,
redact_sensitive=self.config.redact_sensitive,
)
return self._monitor
@property
def ledger(self):
"""Get or create ledger manager."""
if self._ledger is None:
from praisonaiagents.context import ContextLedgerManager
self._ledger = ContextLedgerManager(agent_id=self.agent_name)
self._ledger.set_budget(self.budgeter.allocate())
return self._ledger
def config(self):
"""Get SDK config for compatibility."""
return self._manager.config
def set_model(self, model: str) -> None:
"""Update model and reset budgeter."""
"""Update model and reset manager."""
self.model = model
self._budgeter = None
self._composer = None
create_manager = _get_create_manager()
self._manager = create_manager(
model=model,
session_id=self.session_id,
agent_name=self.agent_name,
)
def get_stats(self) -> Dict[str, Any]:
"""Get current context statistics."""
budget = self.budgeter.allocate()
ledger_data = self.ledger.to_dict()
return {
"model": self.model,
"model_limit": budget.model_limit,
"output_reserve": budget.output_reserve,
"usable": budget.usable,
"total_tokens": self.ledger.get_total(),
"utilization": self.ledger.get_utilization(),
"turn_count": self.ledger.get_ledger().turn_count,
"message_count": self.ledger.get_ledger().message_count,
"segments": ledger_data.get("segments", {}),
"warnings": self.ledger.get_warnings(),
"auto_compact": self.config.auto_compact,
"monitor_enabled": self.config.monitor_enabled,
}
stats = self._manager.get_stats()
# Add CLI-specific fields for compatibility
stats["auto_compact"] = self._manager.config.auto_compact
stats["monitor_enabled"] = self._manager.config.monitor_enabled
return stats
def track_history(self, messages: List[Dict[str, Any]]) -> int:
"""Track conversation history tokens."""
return self.ledger.track_history(messages)
return self._manager._ledger.track_history(messages)
def track_system_prompt(self, content: str) -> int:
"""Track system prompt tokens."""
return self.ledger.track_system_prompt(content)
return self._manager._ledger.track_system_prompt(content)
def track_tools(self, tools: List[Dict[str, Any]]) -> int:
"""Track tool schema tokens."""
return self.ledger.track_tools(tools)
return self._manager._ledger.track_tools(tools)
def check_overflow(self) -> bool:
"""Check if context is approaching limit."""
return self.budgeter.check_overflow(
self.ledger.get_total(),
threshold=self.config.compact_threshold,
)
return self._manager._ledger.get_utilization() >= self._manager.config.compact_threshold
def should_auto_compact(self) -> bool:
"""Check if auto-compaction should trigger."""
return self.config.auto_compact and self.check_overflow()
return self._manager.config.auto_compact and self.check_overflow()
def optimize(
self,
@@ -213,35 +131,43 @@ class ContextManagerHandler:
"""
from praisonaiagents.context import get_optimizer, OptimizerStrategy
strategy_name = strategy or self.config.strategy
strategy_name = strategy or self._manager.config.strategy.value
try:
strategy_enum = OptimizerStrategy(strategy_name)
except ValueError:
strategy_enum = OptimizerStrategy.SMART
optimizer = get_optimizer(strategy_enum)
target = int(self.budgeter.usable * self.config.compact_threshold)
target = int(self._manager._budget.usable * self._manager.config.compact_threshold)
return optimizer.optimize(messages, target, self.ledger.get_ledger())
return optimizer.optimize(messages, target, self._manager._ledger.get_ledger())
def process(
self,
messages: List[Dict[str, Any]],
system_prompt: str = "",
tools: Optional[List[Dict[str, Any]]] = None,
) -> Dict[str, Any]:
"""
Process messages through context pipeline.
Uses SDK ContextManager.process() for full pipeline.
"""
return self._manager.process(
messages=messages,
system_prompt=system_prompt,
tools=tools,
)
def write_snapshot(
self,
messages: List[Dict[str, Any]],
trigger: str = "manual",
) -> Optional[str]:
"""
Write context snapshot to disk.
Args:
messages: Current messages
trigger: What triggered the snapshot
Returns:
Path to written file, or None
"""
return self.monitor.snapshot(
ledger=self.ledger.get_ledger(),
budget=self.budgeter.allocate(),
"""Write context snapshot to disk."""
return self._manager._monitor.snapshot(
ledger=self._manager._ledger.get_ledger(),
budget=self._manager._budget,
messages=messages,
session_id=self.session_id,
agent_name=self.agent_name,
@@ -251,28 +177,60 @@ class ContextManagerHandler:
def enable_monitor(self) -> None:
"""Enable context monitoring."""
self.config.monitor_enabled = True
self.monitor.enable()
self._manager.config.monitor_enabled = True
self._manager._monitor.enable()
def disable_monitor(self) -> None:
"""Disable context monitoring."""
self.config.monitor_enabled = False
self.monitor.disable()
self._manager.config.monitor_enabled = False
self._manager._monitor.disable()
def set_monitor_path(self, path: str) -> None:
"""Set monitor output path."""
self.config.monitor_path = path
self.monitor.set_path(path)
self._manager.config.monitor_path = path
self._manager._monitor.set_path(path)
def set_monitor_format(self, format: str) -> None:
"""Set monitor output format."""
self.config.monitor_format = format
self.monitor.set_format(format)
self._manager.config.monitor_format = format
self._manager._monitor.set_format(format)
def set_monitor_frequency(self, frequency: str) -> None:
"""Set monitor update frequency."""
self.config.monitor_frequency = frequency
self.monitor.set_frequency(frequency)
self._manager.config.monitor_frequency = frequency
self._manager._monitor.set_frequency(frequency)
def get_history(self) -> List[Dict[str, Any]]:
"""Get optimization history."""
return self._manager.get_history()
def get_resolved_config(self) -> Dict[str, Any]:
"""Get resolved configuration with precedence info."""
return self._manager.get_resolved_config()
def capture_llm_boundary(
self,
messages: List[Dict[str, Any]],
tools: List[Dict[str, Any]],
):
"""Capture exact state at LLM call boundary."""
return self._manager.capture_llm_boundary(messages, tools)
# Compatibility properties for legacy code
@property
def budgeter(self):
"""Get budgeter for compatibility."""
return self._manager._budgeter
@property
def ledger(self):
"""Get ledger for compatibility."""
return self._manager._ledger
@property
def monitor(self):
"""Get monitor for compatibility."""
return self._manager._monitor
def handle_context_command(
@@ -296,6 +254,8 @@ def handle_context_command(
/context format <fmt> - Set format (human/json)
/context frequency <f>- Set frequency
/context compact - Trigger optimization
/context history - Show optimization history
/context config - Show resolved configuration
"""
args = args.strip().lower() if args else ""
parts = args.split(maxsplit=1)
@@ -352,8 +312,12 @@ def handle_context_command(
console.print(f"[yellow]Frequency must be one of: {', '.join(valid)}[/yellow]")
elif subcommand == "compact":
_compact_context(console, context_manager, session_state)
elif subcommand == "history":
_show_optimization_history(console, context_manager)
elif subcommand == "config":
_show_resolved_config(console, context_manager)
else:
console.print("[yellow]Unknown subcommand. Use: show, stats, budget, dump, on, off, path, format, frequency, compact[/yellow]")
console.print("[yellow]Unknown subcommand. Use: show, stats, budget, dump, on, off, path, format, frequency, compact, history, config[/yellow]")
def _show_context_summary(console, manager: ContextManagerHandler) -> None:
@@ -458,6 +422,79 @@ def _compact_context(console, manager: ContextManagerHandler, session_state: Dic
console.print("[yellow]No optimization needed[/yellow]")
def _show_optimization_history(console, manager: ContextManagerHandler) -> None:
"""Show optimization history."""
history = manager.get_history()
if not history:
console.print("[dim]No optimization events recorded yet.[/dim]")
return
console.print("\n[bold cyan]Optimization History[/bold cyan]")
console.print(f"{'Time':<24} {'Event':<20} {'Tokens':>12} {'Saved':>10}")
console.print("-" * 70)
for event in history[-20:]: # Show last 20 events
timestamp = event.get("timestamp", "")[:19] # Trim to datetime
event_type = event.get("event_type", "unknown")
tokens_before = event.get("tokens_before", 0)
tokens_saved = event.get("tokens_saved", 0)
tokens_str = f"{tokens_before:,}" if tokens_before else "-"
saved_str = f"-{tokens_saved:,}" if tokens_saved else "-"
console.print(f"{timestamp:<24} {event_type:<20} {tokens_str:>12} {saved_str:>10}")
console.print("")
console.print(f"[dim]Showing last {min(20, len(history))} of {len(history)} events[/dim]")
def _show_resolved_config(console, manager: ContextManagerHandler) -> None:
"""Show resolved configuration with precedence info."""
resolved = manager.get_resolved_config()
config = resolved.get("config", {})
console.print("\n[bold cyan]Resolved Configuration[/bold cyan]")
console.print(f"[dim]Precedence: {resolved.get('precedence', 'CLI > ENV > config.yaml > defaults')}[/dim]")
console.print(f"[dim]Source: {config.get('source', 'defaults')}[/dim]")
console.print("")
console.print("[bold]Auto-Compaction:[/bold]")
console.print(f" auto_compact: {config.get('auto_compact', True)}")
console.print(f" compact_threshold: {config.get('compact_threshold', 0.8)}")
console.print(f" strategy: {config.get('strategy', 'smart')}")
console.print(f" compression_min_gain: {config.get('compression_min_gain_pct', 5.0)}%")
console.print("")
console.print("[bold]Budget:[/bold]")
console.print(f" output_reserve: {config.get('output_reserve', 8000):,}")
console.print(f" default_tool_max: {config.get('default_tool_output_max', 10000):,}")
console.print("")
console.print("[bold]Estimation:[/bold]")
console.print(f" estimation_mode: {config.get('estimation_mode', 'heuristic')}")
console.print(f" log_mismatch: {config.get('log_estimation_mismatch', False)}")
console.print("")
console.print("[bold]Monitoring:[/bold]")
console.print(f" monitor_enabled: {config.get('monitor_enabled', False)}")
console.print(f" monitor_path: {config.get('monitor_path', './context.txt')}")
console.print(f" monitor_format: {config.get('monitor_format', 'human')}")
console.print(f" monitor_frequency: {config.get('monitor_frequency', 'turn')}")
console.print(f" monitor_write_mode: {config.get('monitor_write_mode', 'sync')}")
console.print(f" redact_sensitive: {config.get('redact_sensitive', True)}")
console.print("")
# Show effective budget
budget = resolved.get("effective_budget", {})
if budget:
console.print("[bold]Effective Budget:[/bold]")
console.print(f" model_limit: {budget.get('model_limit', 0):,}")
console.print(f" usable: {budget.get('usable', 0):,}")
console.print(f" history_budget: {budget.get('history_budget', 0):,}")
console.print("")
# CLI argument definitions for argparse integration
CONTEXT_CLI_ARGS = [
{
@@ -522,4 +559,39 @@ CONTEXT_CLI_ARGS = [
"default": None,
"help": "Output token reserve",
},
{
"flags": ["--context-estimation-mode"],
"dest": "context_estimation_mode",
"choices": ["heuristic", "accurate", "validated"],
"default": None,
"help": "Token estimation mode",
},
{
"flags": ["--context-log-mismatch"],
"dest": "context_log_mismatch",
"action": "store_true",
"default": None,
"help": "Log token estimation mismatches",
},
{
"flags": ["--context-snapshot-timing"],
"dest": "context_snapshot_timing",
"choices": ["pre_optimization", "post_optimization", "both"],
"default": None,
"help": "When to capture context snapshots",
},
{
"flags": ["--context-write-mode"],
"dest": "context_write_mode",
"choices": ["sync", "async"],
"default": None,
"help": "Monitor write mode (sync or async)",
},
{
"flags": ["--context-show-config"],
"dest": "context_show_config",
"action": "store_true",
"default": False,
"help": "Show resolved context configuration and exit",
},
]
@@ -388,8 +388,9 @@ If the query mentions a folder name that exists, include files from that folder.
Modified configuration
"""
if flag_value:
config['fast_context'] = True
config['fast_context_path'] = flag_value
# Note: fast_context params removed from Agent - FastContext is now accessed
# via the context module directly. This flag enables CLI-level fast context.
config['_fast_context_path'] = flag_value # Internal CLI flag
return config
def execute(self, query: str = None, path: str = None, **kwargs) -> str:
@@ -0,0 +1,438 @@
"""
CLI commands for Knowledge Management (Phase 7).
Provides CLI parity for knowledge indexing and retrieval operations.
Commands:
- index: Index a directory or file into the knowledge store
- stats: Show corpus statistics
- search: Search the knowledge store with strategy/rerank/compress options
- summarize: Build hierarchical summaries for large corpora
- clear: Clear the knowledge store
"""
import json
import os
from typing import Optional
from enum import Enum
try:
import typer
HAS_TYPER = True
except ImportError:
HAS_TYPER = False
class StrategyChoice(str, Enum):
"""Retrieval strategy choices."""
auto = "auto"
direct = "direct"
basic = "basic"
hybrid = "hybrid"
reranked = "reranked"
compressed = "compressed"
hierarchical = "hierarchical"
def create_knowledge_app():
"""Create the knowledge CLI app."""
if not HAS_TYPER:
return None
app = typer.Typer(
name="knowledge",
help="Knowledge management commands for indexing and retrieval.",
no_args_is_help=True,
)
@app.command("index")
def index_command(
path: str = typer.Argument(..., help="Path to directory or file to index"),
incremental: bool = typer.Option(True, "--incremental/--full", help="Use incremental indexing"),
include: Optional[str] = typer.Option(None, "--include", "-i", help="Glob patterns to include (comma-separated)"),
exclude: Optional[str] = typer.Option(None, "--exclude", "-e", help="Glob patterns to exclude (comma-separated)"),
user_id: Optional[str] = typer.Option(None, "--user-id", help="User ID for scoping"),
verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose output"),
):
"""Index a directory or file into the knowledge store."""
try:
from praisonaiagents.knowledge import Knowledge
except ImportError:
typer.echo("Error: praisonaiagents not installed", err=True)
raise typer.Exit(1)
if not os.path.exists(path):
typer.echo(f"Error: Path does not exist: {path}", err=True)
raise typer.Exit(1)
# Parse glob patterns
include_glob = include.split(",") if include else None
exclude_glob = exclude.split(",") if exclude else None
if verbose:
typer.echo(f"Indexing: {path}")
typer.echo(f" Incremental: {incremental}")
if include_glob:
typer.echo(f" Include: {include_glob}")
if exclude_glob:
typer.echo(f" Exclude: {exclude_glob}")
knowledge = Knowledge()
result = knowledge.index(
path,
incremental=incremental,
include_glob=include_glob,
exclude_glob=exclude_glob,
user_id=user_id or "cli_user",
)
typer.echo(f"✓ Indexed {result.files_indexed} files")
if result.files_skipped > 0:
typer.echo(f" Skipped {result.files_skipped} unchanged files")
if result.errors:
typer.echo(f" Errors: {len(result.errors)}", err=True)
if verbose:
for err in result.errors[:5]:
typer.echo(f" - {err}", err=True)
typer.echo(f" Duration: {result.duration_seconds:.2f}s")
@app.command("stats")
def stats_command(
path: Optional[str] = typer.Argument(None, help="Path to get stats for"),
json_output: bool = typer.Option(False, "--json", help="Output as JSON"),
):
"""Show corpus statistics."""
try:
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.knowledge.indexing import CorpusStats
except ImportError:
typer.echo("Error: praisonaiagents not installed", err=True)
raise typer.Exit(1)
if path and os.path.isdir(path):
stats = CorpusStats.from_directory(path)
else:
knowledge = Knowledge()
stats = knowledge.get_corpus_stats()
if json_output:
typer.echo(json.dumps(stats.to_dict(), indent=2))
else:
typer.echo("Corpus Statistics:")
typer.echo(f" Files: {stats.file_count}")
typer.echo(f" Chunks: {stats.chunk_count}")
typer.echo(f" Tokens: {stats.total_tokens}")
typer.echo(f" Recommended strategy: {stats.strategy_recommendation}")
@app.command("search")
def search_command(
query: str = typer.Argument(..., help="Search query"),
top_k: int = typer.Option(5, "--top-k", "-k", help="Number of results to return"),
strategy: StrategyChoice = typer.Option(StrategyChoice.auto, "--strategy", "-s", help="Retrieval strategy"),
rerank: bool = typer.Option(False, "--rerank/--no-rerank", help="Enable reranking of results"),
compress: bool = typer.Option(False, "--compress/--no-compress", help="Enable context compression"),
compression_ratio: float = typer.Option(0.5, "--compression-ratio", help="Target compression ratio (0.0-1.0)"),
max_context_tokens: Optional[int] = typer.Option(None, "--max-context-tokens", help="Maximum tokens for context"),
include: Optional[str] = typer.Option(None, "--include", "-i", help="Glob patterns to include (comma-separated)"),
exclude: Optional[str] = typer.Option(None, "--exclude", "-e", help="Glob patterns to exclude (comma-separated)"),
path_filter: Optional[str] = typer.Option(None, "--path-filter", help="Filter by path prefix"),
min_score: float = typer.Option(0.0, "--min-score", help="Minimum relevance score (0.0-1.0)"),
user_id: Optional[str] = typer.Option(None, "--user-id", help="User ID for scoping"),
json_output: bool = typer.Option(False, "--json", help="Output as JSON"),
verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose output"),
):
"""Search the knowledge store with advanced retrieval options.
Examples:
praisonai knowledge search "What is the API key?"
praisonai knowledge search "config" --strategy hybrid --rerank
praisonai knowledge search "auth" --compress --max-context-tokens 2000
praisonai knowledge search "*.py" --include "*.py" --exclude "test_*"
"""
try:
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.rag import (
SmartRetriever, ContextCompressor, RetrievalResult,
)
except ImportError as e:
typer.echo(f"Error: praisonaiagents not installed: {e}", err=True)
raise typer.Exit(1)
# Parse glob patterns
include_glob = include.split(",") if include else None
exclude_glob = exclude.split(",") if exclude else None
if verbose:
typer.echo(f"Query: {query}")
typer.echo(f"Strategy: {strategy.value}")
typer.echo(f"Rerank: {rerank}, Compress: {compress}")
if max_context_tokens:
typer.echo(f"Max context tokens: {max_context_tokens}")
knowledge = Knowledge()
# Basic search
results = knowledge.search(
query,
user_id=user_id or "cli_user",
limit=top_k * 3 if rerank or compress else top_k, # Fetch more for reranking
)
items = results.get("results", []) if isinstance(results, dict) else results
# Apply path filter
if path_filter and items:
items = [item for item in items if path_filter in str(item.get("metadata", {}).get("path", ""))]
# Apply include/exclude patterns
if include_glob and items:
import fnmatch
filtered = []
for item in items:
path = str(item.get("metadata", {}).get("path", ""))
for pattern in include_glob:
if fnmatch.fnmatch(path, f"*{pattern}*") or fnmatch.fnmatch(os.path.basename(path), pattern):
filtered.append(item)
break
items = filtered
if exclude_glob and items:
import fnmatch
filtered = []
for item in items:
path = str(item.get("metadata", {}).get("path", ""))
excluded = False
for pattern in exclude_glob:
if fnmatch.fnmatch(path, f"*{pattern}*") or fnmatch.fnmatch(os.path.basename(path), pattern):
excluded = True
break
if not excluded:
filtered.append(item)
items = filtered
# Apply min_score filter
if min_score > 0 and items:
items = [item for item in items if item.get("score", 0) >= min_score]
# Apply reranking
if rerank and items:
try:
reranker = SmartRetriever()
# Convert to RetrievalResult format
retrieval_results = [
RetrievalResult(
text=item.get("text", item.get("memory", "")),
score=item.get("score", 0),
metadata=item.get("metadata", {}),
)
for item in items
]
reranked = reranker.rerank(query, retrieval_results, top_k=top_k)
items = [
{"text": r.text, "score": r.score, "metadata": r.metadata}
for r in reranked
]
if verbose:
typer.echo(f"Reranked {len(items)} results")
except Exception as e:
if verbose:
typer.echo(f"Reranking skipped: {e}", err=True)
# Apply compression
if compress and items:
try:
compressor = ContextCompressor(
max_tokens=max_context_tokens or 4000,
target_ratio=compression_ratio,
)
texts = [item.get("text", item.get("memory", "")) for item in items]
compressed = compressor.compress(texts, query=query)
if verbose:
typer.echo(f"Compressed: {compressed.original_tokens} -> {compressed.compressed_tokens} tokens")
# Update items with compressed text
if compressed.chunks:
items = items[:len(compressed.chunks)]
for i, chunk in enumerate(compressed.chunks):
if i < len(items):
items[i]["text"] = chunk
except Exception as e:
if verbose:
typer.echo(f"Compression skipped: {e}", err=True)
# Limit to top_k
items = items[:top_k]
if json_output:
typer.echo(json.dumps(items, indent=2, default=str))
else:
if not items:
typer.echo("No results found.")
return
typer.echo(f"Found {len(items)} results:\n")
for i, item in enumerate(items, 1):
text = item.get("text", item.get("memory", ""))[:200]
score = item.get("score", 0)
metadata = item.get("metadata", {})
path = metadata.get("path", "")
typer.echo(f"{i}. [score: {score:.3f}]")
if path:
typer.echo(f" Source: {os.path.basename(path)}")
typer.echo(f" {text}...")
typer.echo()
@app.command("summarize")
def summarize_command(
path: str = typer.Argument(..., help="Path to directory to summarize"),
levels: int = typer.Option(3, "--levels", "-l", help="Number of hierarchy levels (1=file, 2=folder, 3=project)"),
incremental: bool = typer.Option(True, "--incremental/--rebuild", help="Incremental update vs full rebuild"),
include: Optional[str] = typer.Option(None, "--include", "-i", help="Glob patterns to include (comma-separated)"),
exclude: Optional[str] = typer.Option(None, "--exclude", "-e", help="Glob patterns to exclude (comma-separated)"),
output: Optional[str] = typer.Option(None, "--output", "-o", help="Output path for summary artifacts"),
user_id: Optional[str] = typer.Option(None, "--user-id", help="User ID for scoping"),
json_output: bool = typer.Option(False, "--json", help="Output as JSON"),
verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose output"),
):
"""Build hierarchical summaries for large corpora.
Creates multi-level summaries (file → folder → project) for efficient
navigation and query routing in large knowledge bases.
Examples:
praisonai knowledge summarize ./docs
praisonai knowledge summarize ./src --levels 2 --rebuild
praisonai knowledge summarize ./project --output ./summaries
"""
try:
from praisonaiagents.rag import HierarchicalSummarizer
from praisonaiagents.knowledge.indexing import CorpusStats
except ImportError as e:
typer.echo(f"Error: praisonaiagents not installed: {e}", err=True)
raise typer.Exit(1)
if not os.path.exists(path):
typer.echo(f"Error: Path does not exist: {path}", err=True)
raise typer.Exit(1)
if not os.path.isdir(path):
typer.echo(f"Error: Path must be a directory: {path}", err=True)
raise typer.Exit(1)
# Parse glob patterns
include_glob = include.split(",") if include else None
exclude_glob = exclude.split(",") if exclude else None
if verbose:
typer.echo(f"Summarizing: {path}")
typer.echo(f" Levels: {levels}")
typer.echo(f" Incremental: {incremental}")
# Get corpus stats first
stats = CorpusStats.from_directory(path)
if verbose:
typer.echo(f" Files: {stats.file_count}")
typer.echo(f" Estimated tokens: {stats.total_tokens}")
# Build hierarchical summaries
summarizer = HierarchicalSummarizer(max_levels=levels)
try:
# Collect files
import fnmatch as fnmatch_module
files_to_summarize = []
for root, dirs, files in os.walk(path):
dirs[:] = [d for d in dirs if not d.startswith('.')]
for filename in files:
filepath = os.path.join(root, filename)
rel_path = os.path.relpath(filepath, path)
# Apply include patterns
if include_glob:
matched = False
for pattern in include_glob:
if fnmatch_module.fnmatch(filename, pattern):
matched = True
break
if not matched:
continue
# Apply exclude patterns
if exclude_glob:
excluded = False
for pattern in exclude_glob:
if fnmatch_module.fnmatch(filename, pattern) or fnmatch_module.fnmatch(rel_path, pattern):
excluded = True
break
if excluded:
continue
files_to_summarize.append(filepath)
if verbose:
typer.echo(f" Processing {len(files_to_summarize)} files...")
# Build hierarchy
result = summarizer.build_hierarchy(
files_to_summarize,
base_path=path,
)
# Save if output path specified
if output:
os.makedirs(output, exist_ok=True)
output_file = os.path.join(output, "hierarchy.json")
summarizer.save(output_file)
if verbose:
typer.echo(f" Saved to: {output_file}")
if json_output:
typer.echo(json.dumps(result.to_dict() if hasattr(result, 'to_dict') else {"nodes": len(result.nodes) if hasattr(result, 'nodes') else 0}, indent=2))
else:
typer.echo("✓ Built hierarchical summaries")
typer.echo(f" Levels: {levels}")
typer.echo(f" Files processed: {len(files_to_summarize)}")
if hasattr(result, 'nodes'):
typer.echo(f" Summary nodes: {len(result.nodes)}")
if output:
typer.echo(f" Output: {output}")
except Exception as e:
typer.echo(f"Error building summaries: {e}", err=True)
if verbose:
import traceback
typer.echo(traceback.format_exc(), err=True)
raise typer.Exit(1)
@app.command("clear")
def clear_command(
confirm: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation"),
user_id: Optional[str] = typer.Option(None, "--user-id", help="User ID to clear"),
):
"""Clear the knowledge store."""
if not confirm:
confirm = typer.confirm("Are you sure you want to clear the knowledge store?")
if not confirm:
typer.echo("Aborted.")
raise typer.Exit(0)
try:
from praisonaiagents.knowledge import Knowledge
except ImportError:
typer.echo("Error: praisonaiagents not installed", err=True)
raise typer.Exit(1)
knowledge = Knowledge()
knowledge.reset()
typer.echo("✓ Knowledge store cleared.")
return app
# Create the app instance
knowledge_app = create_knowledge_app()
+1 -1
View File
@@ -3648,7 +3648,7 @@ Provide ONLY the commit message, no explanations."""
if getattr(self.args, 'history', None):
agent_config["memory"] = True # History requires memory
agent_config["history_in_context"] = self.args.history
# Note: history_in_context param removed - history loading now via context= param
print(f"[bold cyan]History enabled - loading context from last {self.args.history} session(s)[/bold cyan]")
# Claude Memory Tool (Anthropic only)
+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==3.2.1 gunicorn markdown\n")
file.write("RUN pip install flask praisonai==3.3.1 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__ = "3.2.1"
__version__ = "3.3.1"
+1 -1
View File
@@ -12,7 +12,7 @@ dependencies = [
"rich>=13.7",
"markdown>=3.5",
"pyparsing>=3.0.0",
"praisonaiagents>=0.7.1",
"praisonaiagents>=0.8.1",
"python-dotenv>=0.19.0",
"instructor>=1.3.3",
"PyYAML>=6.0",
@@ -0,0 +1,113 @@
"""
Unit tests for Knowledge CLI (Phase 7).
"""
import os
import tempfile
import pytest
class TestKnowledgeCLI:
"""Tests for knowledge CLI commands."""
def test_import_knowledge_cli(self):
"""knowledge_cli should be importable."""
from praisonai.cli.features.knowledge_cli import create_knowledge_app
assert create_knowledge_app is not None
def test_create_knowledge_app(self):
"""create_knowledge_app should return a typer app."""
from praisonai.cli.features.knowledge_cli import create_knowledge_app
app = create_knowledge_app()
# May be None if typer not installed
if app is not None:
assert hasattr(app, "command")
def test_knowledge_app_has_commands(self):
"""Knowledge app should have expected commands."""
from praisonai.cli.features.knowledge_cli import knowledge_app
if knowledge_app is None:
pytest.skip("typer not installed")
# Check registered commands exist
command_names = [cmd.name for cmd in knowledge_app.registered_commands]
assert "index" in command_names
assert "stats" in command_names
assert "search" in command_names
assert "clear" in command_names
class TestKnowledgeCLIIntegration:
"""Integration tests for knowledge CLI with typer runner."""
@pytest.fixture
def cli_runner(self):
"""Get typer CLI test runner."""
try:
from typer.testing import CliRunner
return CliRunner()
except ImportError:
pytest.skip("typer not installed")
@pytest.fixture
def knowledge_app(self):
"""Get knowledge app."""
from praisonai.cli.features.knowledge_cli import knowledge_app
if knowledge_app is None:
pytest.skip("typer not installed")
return knowledge_app
def test_index_command_help(self, cli_runner, knowledge_app):
"""Index command should show help."""
result = cli_runner.invoke(knowledge_app, ["index", "--help"])
assert result.exit_code == 0
assert "index" in result.stdout.lower() or "path" in result.stdout.lower()
def test_stats_command_help(self, cli_runner, knowledge_app):
"""Stats command should show help."""
result = cli_runner.invoke(knowledge_app, ["stats", "--help"])
assert result.exit_code == 0
assert "stats" in result.stdout.lower() or "corpus" in result.stdout.lower()
def test_search_command_help(self, cli_runner, knowledge_app):
"""Search command should show help."""
result = cli_runner.invoke(knowledge_app, ["search", "--help"])
assert result.exit_code == 0
assert "search" in result.stdout.lower() or "query" in result.stdout.lower()
def test_index_nonexistent_path(self, cli_runner, knowledge_app):
"""Index should fail for nonexistent path."""
result = cli_runner.invoke(knowledge_app, ["index", "/nonexistent/path"])
assert result.exit_code != 0
# Error may be in stdout or output
output = (result.stdout + (result.output if hasattr(result, 'output') else '')).lower()
assert "error" in output or "not exist" in output or result.exit_code == 1
def test_stats_with_directory(self, cli_runner, knowledge_app):
"""Stats should work with a directory."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create test files
with open(os.path.join(tmpdir, "test.txt"), "w") as f:
f.write("Test content")
result = cli_runner.invoke(knowledge_app, ["stats", tmpdir])
assert result.exit_code == 0
output = result.stdout.lower()
assert "files" in output or "corpus" in output or "1" in output
def test_stats_json_output(self, cli_runner, knowledge_app):
"""Stats should support JSON output."""
with tempfile.TemporaryDirectory() as tmpdir:
with open(os.path.join(tmpdir, "test.txt"), "w") as f:
f.write("Test content")
result = cli_runner.invoke(knowledge_app, ["stats", tmpdir, "--json"])
assert result.exit_code == 0
# Should be valid JSON
import json
try:
data = json.loads(result.stdout)
assert "file_count" in data
except json.JSONDecodeError:
pytest.fail("Output is not valid JSON")
+4 -4
View File
@@ -4276,7 +4276,7 @@ requires-dist = [
{ name = "praisonai-tools", marker = "extra == 'autogen'", specifier = ">=0.1.0" },
{ name = "praisonai-tools", marker = "extra == 'autogen-v4'", specifier = ">=0.1.0" },
{ name = "praisonai-tools", marker = "extra == 'crewai'", specifier = ">=0.1.0" },
{ name = "praisonaiagents", specifier = ">=0.7.1" },
{ name = "praisonaiagents", specifier = ">=0.8.1" },
{ 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" },
@@ -4324,7 +4324,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.7.1"
version = "0.8.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "aiohttp" },
@@ -4333,9 +4333,9 @@ dependencies = [
{ name = "pydantic" },
{ name = "rich" },
]
sdist = { url = "https://files.pythonhosted.org/packages/0b/47/a5630fa4e12e0d28685ed7dd846e41720ce441ea388363c69710f7c86de4/praisonaiagents-0.7.1.tar.gz", hash = "sha256:9458d9b15c404e8efb44c20c0a5088587d7f9b77c23d2fa3b48262a1e01d6110", size = 865614, upload-time = "2026-01-07T21:50:12.162Z" }
sdist = { url = "https://files.pythonhosted.org/packages/d1/a1/3577e22f24fd64c870d1b5e0cb56f7addd88c45d0c2fbb649a63cbf509d5/praisonaiagents-0.8.1.tar.gz", hash = "sha256:f17f6e8bdecdebb7c0c234ebc7103070c4546b31d41e878572208c75e6ef19ef", size = 897661, upload-time = "2026-01-08T12:24:08.607Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/58/b3/176f12c97887e2a7872104f2acf89508db4d8ea28d7555c470e8617add81/praisonaiagents-0.7.1-py3-none-any.whl", hash = "sha256:e1d670f7b2298127292af8a7ccfdc7a942a8610dc3bd55110d9497fe6114862e", size = 827576, upload-time = "2026-01-07T21:50:09.998Z" },
{ url = "https://files.pythonhosted.org/packages/71/01/d0afaf98e78da1d1293f6f1a53e6b2f7b14c8ba1908b70eed16e5b138427/praisonaiagents-0.8.1-py3-none-any.whl", hash = "sha256:b898f8f1f38ae8d010dfba69f443f146faa2d685c997aa5a4992e786dc115654", size = 866252, upload-time = "2026-01-08T12:24:05.975Z" },
]
[[package]]