mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 22:52:58 -10:00
Compare commits
37
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f834d42448 | ||
|
|
705fae31f8 | ||
|
|
b4270173d4 | ||
|
|
6256744348 | ||
|
|
15d597da37 | ||
|
|
04d5e07e1b | ||
|
|
754386022c | ||
|
|
1e27a73633 | ||
|
|
2adfe7e832 | ||
|
|
1ad58ca029 | ||
|
|
2f9677abb2 | ||
|
|
9991e00f52 | ||
|
|
86aeaa4d89 | ||
|
|
a66db1836a | ||
|
|
e90d922318 | ||
|
|
43b9552f9b | ||
|
|
d5c77c244b | ||
|
|
bcb6957dac | ||
|
|
fe78664261 | ||
|
|
d9a7c1a01e | ||
|
|
6a122c3671 | ||
|
|
1f1cfa974b | ||
|
|
34d60ae8fa | ||
|
|
095653d78a | ||
|
|
7f6a483e62 | ||
|
|
c6883def69 | ||
|
|
d3c4a2afad | ||
|
|
034138ebc9 | ||
|
|
70b2df135b | ||
|
|
3532600887 | ||
|
|
13ba1ed954 | ||
|
|
4a1fe173e8 | ||
|
|
147bcce448 | ||
|
|
b5e8052809 | ||
|
|
8a93bc5096 | ||
|
|
f65b83614c | ||
|
|
25000d9b26 |
No files matched your search
@@ -60,16 +60,16 @@ jobs:
|
||||
- name: Run PraisonAI integration tests
|
||||
run: |
|
||||
cd PraisonAI/src/praisonai
|
||||
python -m pytest tests/integration/test_aiui_* -v --timeout=300
|
||||
python -m pytest tests/integration/test_aiui_* -v
|
||||
|
||||
- name: Run PraisonAIUI integration tests
|
||||
run: |
|
||||
cd PraisonAIUI
|
||||
if [ -f tests/integration/test_agentic_roundtrip.py ]; then
|
||||
python -m pytest tests/integration/test_agentic_roundtrip.py -v --timeout=300
|
||||
python -m pytest tests/integration/test_agentic_roundtrip.py -v
|
||||
fi
|
||||
if [ -f tests/test_feature_sdk_backends.py ]; then
|
||||
python -m pytest tests/test_feature_sdk_backends.py -v --timeout=300
|
||||
python -m pytest tests/test_feature_sdk_backends.py -v
|
||||
fi
|
||||
|
||||
- name: Test Pattern C CLI
|
||||
@@ -152,7 +152,7 @@ jobs:
|
||||
run: |
|
||||
cd PraisonAI/src/praisonai
|
||||
if [ -n "$OPENAI_API_KEY" ]; then
|
||||
python -m pytest tests/integration/test_aiui_host_agentic.py -v --timeout=600
|
||||
python -m pytest tests/integration/test_aiui_host_agentic.py -v
|
||||
else
|
||||
echo "OPENAI_API_KEY not available, skipping agentic tests"
|
||||
fi
|
||||
@@ -0,0 +1,31 @@
|
||||
name: Sync Claude Secrets to PraisonAIUI
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
sync:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Sync secrets to PraisonAIUI
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.SECRETS_ADMIN_PAT }}
|
||||
TARGET: MervinPraison/PraisonAIUI
|
||||
CLAUDE_APP_ID: ${{ secrets.CLAUDE_APP_ID }}
|
||||
CLAUDE_APP_PRIVATE_KEY: ${{ secrets.CLAUDE_APP_PRIVATE_KEY }}
|
||||
CLAUDE_CODE_OAUTH_TOKEN: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
|
||||
GH_TOKEN_VALUE: ${{ secrets.GH_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
if [ -z "${GH_TOKEN:-}" ]; then
|
||||
echo "SECRETS_ADMIN_PAT is not configured in this repository"
|
||||
exit 1
|
||||
fi
|
||||
echo -n "$CLAUDE_APP_ID" | gh secret set CLAUDE_APP_ID --repo "$TARGET"
|
||||
echo -n "$CLAUDE_APP_PRIVATE_KEY" | gh secret set CLAUDE_APP_PRIVATE_KEY --repo "$TARGET"
|
||||
echo -n "$CLAUDE_CODE_OAUTH_TOKEN" | gh secret set CLAUDE_CODE_OAUTH_TOKEN --repo "$TARGET"
|
||||
echo -n "$GH_TOKEN_VALUE" | gh secret set GH_TOKEN --repo "$TARGET"
|
||||
echo "Synced CLAUDE_APP_ID, CLAUDE_APP_PRIVATE_KEY, CLAUDE_CODE_OAUTH_TOKEN, GH_TOKEN → $TARGET"
|
||||
@@ -0,0 +1,10 @@
|
||||
# Agent Instructions
|
||||
|
||||
You are working on the PraisonAI project.
|
||||
|
||||
## Project Guidelines
|
||||
- Follow the existing code style and conventions
|
||||
- Be concise and helpful in responses
|
||||
- Test implementation thoroughly
|
||||
- Ensure backward compatibility with existing APIs
|
||||
- Follow protocol-driven design: core protocols in praisonaiagents/, heavy implementations in praisonai/
|
||||
@@ -240,6 +240,7 @@ from praisonaiagents import Memory, MemoryBackend, MemoryConfig
|
||||
|
||||
Methods:
|
||||
|
||||
* <code title="class Memory">Memory.<a href="./src/praisonai-agents/praisonaiagents/memory/memory.py">build_cache_optimized_context</a>(task_descr: str, user_id: Optional[str] = None, additional: str = '', max_items: int = 3, include_cache_boundary: bool = True, include_in_output: Optional[bool] = None) -> Dict[str, str]</code>
|
||||
* <code title="class Memory">Memory.<a href="./src/praisonai-agents/praisonaiagents/memory/memory.py">build_context_for_task</a>(task_descr: str, user_id: Optional[str] = None, additional: str = '', max_items: int = 3, include_in_output: Optional[bool] = None) -> str</code>
|
||||
* <code title="class Memory">Memory.<a href="./src/praisonai-agents/praisonaiagents/memory/memory.py">calculate_quality_metrics</a>(output: str, expected_output: str, llm: Optional[str] = None, custom_prompt: Optional[str] = None) -> Dict[str, float]</code>
|
||||
* <code title="class Memory">Memory.<a href="./src/praisonai-agents/praisonaiagents/memory/memory.py">close_connections</a>()</code>
|
||||
@@ -698,7 +699,7 @@ Methods:
|
||||
|
||||
Types:
|
||||
```python
|
||||
from praisonaiagents import AgentAppConfig, AgentAppProtocol, AgentFlow, AgentManager, AgentOSConfig, AgentOSProtocol, AgentTeam, AutoApproveBackend, EmbeddingResult, aembedding, aembeddings, embedding, embeddings, get_dimensions
|
||||
from praisonaiagents import AgentAppConfig, AgentAppProtocol, AgentFlow, AgentManager, AgentOSConfig, AgentOSProtocol, AgentTeam, AutoApproveBackend, EmbeddingResult, __version__, aembedding, aembeddings, embedding, embeddings, get_dimensions
|
||||
```
|
||||
|
||||
Methods:
|
||||
@@ -914,6 +915,9 @@ Methods:
|
||||
* <code title="cli">praisonai gateway install <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway logs <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway mint-link <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway pause <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway reconnect <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway resume <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway send <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway start <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway status <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.6.55" \
|
||||
"praisonai>=4.6.59" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.6.55" \
|
||||
"praisonai>=4.6.59" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.6.55" \
|
||||
"praisonai>=4.6.59" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -32,7 +32,7 @@ workflow = AgentFlow(
|
||||
),
|
||||
],
|
||||
# Workflow-level consolidated params
|
||||
output=WorkflowOutputConfig(output="verbose", stream=True),
|
||||
output=WorkflowOutputConfig(verbose=True, stream=True),
|
||||
planning=WorkflowPlanningConfig(enabled=True, reasoning=True),
|
||||
memory=WorkflowMemoryConfig(backend="file"),
|
||||
# Agent-like params
|
||||
|
||||
@@ -5,12 +5,12 @@ Demonstrates web search with consolidated params.
|
||||
Presets: duckduckgo, tavily, google, bing, serper, search_only, fetch_only
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents import Agent, WebConfig
|
||||
|
||||
# Basic: Enable web search with preset
|
||||
# Basic: Enable web search with preset (search only — fetch needs Claude)
|
||||
agent = Agent(
|
||||
instructions="You are a research assistant with web access.",
|
||||
web="duckduckgo", # Presets: duckduckgo, tavily, google, bing, serper
|
||||
web=WebConfig(search=True, fetch=False, search_provider="duckduckgo"),
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -6,10 +6,18 @@ It includes self-reflection for improving code quality and multiple specialized
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, Task, AgentTeam
|
||||
from praisonaiagents.tools import CodeInterpreter
|
||||
|
||||
# Initialize code execution tool
|
||||
code_interpreter = CodeInterpreter()
|
||||
def run_python_code(code: str) -> str:
|
||||
"""Execute Python code safely and return stdout or errors."""
|
||||
import io
|
||||
import contextlib
|
||||
try:
|
||||
buf = io.StringIO()
|
||||
with contextlib.redirect_stdout(buf):
|
||||
exec(code, {"__builtins__": __builtins__})
|
||||
return buf.getvalue() or "Executed successfully (no output)"
|
||||
except Exception as e:
|
||||
return f"Execution error: {e}"
|
||||
|
||||
# Custom code analysis tool
|
||||
def analyze_code_complexity(code: str) -> str:
|
||||
@@ -32,9 +40,7 @@ code_writer = Agent(
|
||||
goal="Write clean, efficient, and well-documented Python code",
|
||||
backstory="You are a senior Python developer with 10+ years of experience in writing production-quality code.",
|
||||
instructions="Write Python code that follows PEP 8 standards, includes proper error handling, and has clear documentation.",
|
||||
reflection=True,
|
||||
|
||||
|
||||
output="silent",
|
||||
)
|
||||
|
||||
code_reviewer = Agent(
|
||||
@@ -52,7 +58,7 @@ code_executor = Agent(
|
||||
goal="Safely execute and test Python code",
|
||||
backstory="You are responsible for running code and reporting results or errors.",
|
||||
instructions="Execute the provided code safely and report the output or any errors encountered.",
|
||||
tools=[code_interpreter]
|
||||
tools=[run_python_code]
|
||||
)
|
||||
|
||||
# Example tasks
|
||||
@@ -87,7 +93,7 @@ def demonstrate_code_agent():
|
||||
workflow = AgentTeam(
|
||||
agents=[code_writer, code_reviewer, code_executor],
|
||||
tasks=[write_task, review_task, execute_task],
|
||||
process="sequential", output="verbose"
|
||||
process="sequential", output="silent"
|
||||
)
|
||||
|
||||
# Run the workflow
|
||||
@@ -106,10 +112,8 @@ general_code_agent = Agent(
|
||||
3. Explain code concepts
|
||||
4. Refactor code for better performance
|
||||
5. Add tests and documentation""",
|
||||
tools=[code_interpreter, analyze_code_complexity],
|
||||
reflection=True,
|
||||
|
||||
|
||||
tools=[run_python_code, analyze_code_complexity],
|
||||
output="silent",
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -41,21 +41,25 @@ def calculate_statistics(numbers: list) -> str:
|
||||
"""Calculate basic statistics for a list of numbers."""
|
||||
if not numbers:
|
||||
return "No numbers provided"
|
||||
|
||||
n = len(numbers)
|
||||
mean = sum(numbers) / n
|
||||
sorted_nums = sorted(numbers)
|
||||
try:
|
||||
nums = [float(x) for x in numbers]
|
||||
except (TypeError, ValueError):
|
||||
return "Error: all values must be numbers"
|
||||
|
||||
n = len(nums)
|
||||
mean = sum(nums) / n
|
||||
sorted_nums = sorted(nums)
|
||||
median = sorted_nums[n//2] if n % 2 else (sorted_nums[n//2-1] + sorted_nums[n//2]) / 2
|
||||
variance = sum((x - mean)**2 for x in numbers) / n
|
||||
variance = sum((x - mean)**2 for x in nums) / n
|
||||
std_dev = math.sqrt(variance)
|
||||
|
||||
|
||||
return f"""Statistics:
|
||||
- Count: {n}
|
||||
- Mean: {mean:.4f}
|
||||
- Median: {median:.4f}
|
||||
- Standard Deviation: {std_dev:.4f}
|
||||
- Min: {min(numbers)}
|
||||
- Max: {max(numbers)}"""
|
||||
- Min: {min(nums)}
|
||||
- Max: {max(nums)}"""
|
||||
|
||||
# Create specialized math agents
|
||||
calculation_agent = Agent(
|
||||
@@ -66,8 +70,7 @@ calculation_agent = Agent(
|
||||
instructions="Perform the requested calculations step by step, showing your work clearly.",
|
||||
tools=[basic_calculator, solve_quadratic, calculate_statistics],
|
||||
reflection=True,
|
||||
|
||||
|
||||
output="silent",
|
||||
)
|
||||
|
||||
problem_solver_agent = Agent(
|
||||
@@ -82,8 +85,7 @@ problem_solver_agent = Agent(
|
||||
4. Show each step of the solution
|
||||
5. Verify the answer makes sense""",
|
||||
reflection=True,
|
||||
|
||||
|
||||
output="silent",
|
||||
)
|
||||
|
||||
math_tutor_agent = Agent(
|
||||
@@ -91,7 +93,8 @@ math_tutor_agent = Agent(
|
||||
role="Mathematics tutor and explainer",
|
||||
goal="Explain mathematical concepts clearly and help users understand math",
|
||||
backstory="You are a patient math tutor who can explain concepts at any level from elementary to advanced.",
|
||||
instructions="Explain mathematical concepts using simple language, provide examples, and check for understanding."
|
||||
instructions="Explain mathematical concepts using simple language, provide examples, and check for understanding.",
|
||||
output="silent",
|
||||
)
|
||||
|
||||
# Example workflow for solving complex problems
|
||||
@@ -126,7 +129,8 @@ def solve_math_problem_workflow():
|
||||
workflow = AgentTeam(
|
||||
agents=[problem_solver_agent, calculation_agent, math_tutor_agent],
|
||||
tasks=[understand_task, solve_task, verify_task],
|
||||
process="sequential", output="verbose"
|
||||
process="sequential",
|
||||
output="silent",
|
||||
)
|
||||
|
||||
return workflow.start()
|
||||
@@ -148,8 +152,7 @@ math_assistant = Agent(
|
||||
Always show your work step by step.""",
|
||||
tools=[basic_calculator, solve_quadratic, calculate_statistics],
|
||||
reflection=True,
|
||||
|
||||
|
||||
output="silent",
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example: Bot In-Flight Run Control
|
||||
|
||||
Demonstrates how to use SessionRunControl to provide better UX for bots
|
||||
during long-running agent operations. Shows busy feedback, pending message
|
||||
handling, and /stop command support.
|
||||
|
||||
This example shows the "before" and "after" behavior described in issue #1914.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
# Mock agent for demonstration
|
||||
class MockAgent:
|
||||
"""Simple mock agent that simulates long-running tasks."""
|
||||
|
||||
def __init__(self, name="assistant"):
|
||||
self.name = name
|
||||
self.chat_history = []
|
||||
self._interrupt_controller = None
|
||||
|
||||
def chat(self, prompt: str) -> str:
|
||||
"""Simulate a long-running chat response."""
|
||||
print(f"[Agent] Starting to process: '{prompt[:50]}...'")
|
||||
|
||||
# Simulate long processing with interrupt checking
|
||||
for i in range(30): # 3 second task (30 x 0.1s)
|
||||
time.sleep(0.1)
|
||||
|
||||
# Check for interruption
|
||||
if (self._interrupt_controller and
|
||||
hasattr(self._interrupt_controller, 'is_set') and
|
||||
self._interrupt_controller.is_set()):
|
||||
reason = getattr(self._interrupt_controller, 'reason', 'unknown')
|
||||
print(f"[Agent] ⚠️ Interrupted: {reason}")
|
||||
raise InterruptedError(f"Task cancelled: {reason}")
|
||||
|
||||
response = f"Completed analysis of: {prompt}"
|
||||
print(f"[Agent] ✅ Finished: {response}")
|
||||
return response
|
||||
|
||||
|
||||
async def demo_without_run_control():
|
||||
"""Demonstrate the old behavior (silent blocking)."""
|
||||
print("\n" + "="*60)
|
||||
print("🔴 BEFORE: Without Run Control (Silent Blocking)")
|
||||
print("="*60)
|
||||
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
# Standard session manager (no run control)
|
||||
session_mgr = BotSessionManager()
|
||||
agent = MockAgent("research-agent")
|
||||
user_id = "user123"
|
||||
|
||||
print("\n1. User sends long task:")
|
||||
print(" User: 'research quantum computing trends'")
|
||||
|
||||
# Start first task (this will run for 3 seconds)
|
||||
task1 = asyncio.create_task(
|
||||
session_mgr.chat(agent, user_id, "research quantum computing trends")
|
||||
)
|
||||
|
||||
# Simulate user sending follow-up messages during execution
|
||||
await asyncio.sleep(0.5)
|
||||
print("\n2. User sends follow-up (this will block silently):")
|
||||
print(" User: 'actually focus on business impact'")
|
||||
|
||||
task2 = asyncio.create_task(
|
||||
session_mgr.chat(agent, user_id, "actually focus on business impact")
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
print("\n3. User tries to send /stop (this will also block):")
|
||||
print(" User: '/stop'")
|
||||
|
||||
task3 = asyncio.create_task(
|
||||
session_mgr.chat(agent, user_id, "/stop")
|
||||
)
|
||||
|
||||
print("\n⏳ Waiting for all messages to complete...")
|
||||
results = await asyncio.gather(task1, task2, task3)
|
||||
|
||||
print("\n📊 Results (all processed sequentially with no feedback):")
|
||||
for i, result in enumerate(results, 1):
|
||||
print(f" {i}. {result}")
|
||||
|
||||
print("\n❌ Problems with this approach:")
|
||||
print(" - No immediate feedback on follow-up messages")
|
||||
print(" - No way to cancel long-running tasks")
|
||||
print(" - User doesn't know if bot is working or broken")
|
||||
|
||||
|
||||
async def demo_with_run_control():
|
||||
"""Demonstrate the new behavior with run control."""
|
||||
print("\n" + "="*60)
|
||||
print("🟢 AFTER: With Run Control (Better UX)")
|
||||
print("="*60)
|
||||
|
||||
from praisonai.bots._run_control import SessionRunControl
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
from praisonai.bots._commands import handle_stop_command
|
||||
|
||||
# Session manager with run control
|
||||
run_control = SessionRunControl(
|
||||
busy_mode="queue",
|
||||
busy_ack_template="⏳ {action} — will process after current task finishes"
|
||||
)
|
||||
session_mgr = BotSessionManager(run_control=run_control)
|
||||
agent = MockAgent("research-agent")
|
||||
user_id = "user123"
|
||||
|
||||
print("\n1. User sends long task:")
|
||||
print(" User: 'research quantum computing trends'")
|
||||
|
||||
# Start first task
|
||||
task1 = asyncio.create_task(
|
||||
session_mgr.chat_with_run_control(agent, user_id, "research quantum computing trends")
|
||||
)
|
||||
|
||||
# Immediate follow-up messages get acknowledgments
|
||||
await asyncio.sleep(0.5)
|
||||
print("\n2. User sends follow-up (gets immediate feedback):")
|
||||
print(" User: 'actually focus on business impact'")
|
||||
|
||||
result2 = await session_mgr.chat_with_run_control(
|
||||
agent, user_id, "actually focus on business impact"
|
||||
)
|
||||
print(f" Bot: {result2['response']}")
|
||||
print(f" Metadata: {result2['metadata']}")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
print("\n3. User decides to stop current task:")
|
||||
print(" User: '/stop'")
|
||||
|
||||
stop_response = await handle_stop_command(user_id, run_control)
|
||||
print(f" Bot: {stop_response}")
|
||||
|
||||
# Wait for first task to complete (it should be cancelled)
|
||||
try:
|
||||
result1 = await task1
|
||||
print(f"\n📊 Task 1 result: {result1['response']}")
|
||||
print(f" Metadata: {result1['metadata']}")
|
||||
except asyncio.CancelledError:
|
||||
print("\n📊 Task 1 was cancelled as expected")
|
||||
|
||||
print("\n4. User sends fresh request after stopping:")
|
||||
print(" User: 'what's the weather like?'")
|
||||
|
||||
result4 = await session_mgr.chat_with_run_control(
|
||||
agent, user_id, "what's the weather like?"
|
||||
)
|
||||
print(f" Bot: {result4['response']}")
|
||||
|
||||
print("\n✅ Benefits of this approach:")
|
||||
print(" - Immediate feedback on all messages")
|
||||
print(" - Can cancel long-running tasks with /stop")
|
||||
print(" - Pending messages are queued and processed in order")
|
||||
print(" - User always knows what's happening")
|
||||
|
||||
|
||||
async def demo_interrupt_mode():
|
||||
"""Demonstrate interrupt mode."""
|
||||
print("\n" + "="*60)
|
||||
print("⚡ BONUS: Interrupt Mode (Cancel and Restart)")
|
||||
print("="*60)
|
||||
|
||||
from praisonai.bots._run_control import SessionRunControl
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
# Use interrupt mode instead of queue mode
|
||||
run_control = SessionRunControl(busy_mode="interrupt")
|
||||
session_mgr = BotSessionManager(run_control=run_control)
|
||||
agent = MockAgent("research-agent")
|
||||
user_id = "user123"
|
||||
|
||||
print("\n1. User starts long task:")
|
||||
print(" User: 'research renewable energy for 30 pages'")
|
||||
|
||||
task1 = asyncio.create_task(
|
||||
session_mgr.chat_with_run_control(agent, user_id, "research renewable energy for 30 pages")
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
print("\n2. User changes mind (interrupts and restarts):")
|
||||
print(" User: 'actually just summarize solar panel efficiency'")
|
||||
|
||||
result2 = await session_mgr.chat_with_run_control(
|
||||
agent, user_id, "actually just summarize solar panel efficiency"
|
||||
)
|
||||
print(f" Bot: {result2['response']}")
|
||||
|
||||
# Wait for original task (should be interrupted)
|
||||
try:
|
||||
result1 = await task1
|
||||
print(f"\n📊 Original task result: {result1['response']}")
|
||||
if result1['metadata'].get('interrupted'):
|
||||
print(" ✅ Original task was properly interrupted")
|
||||
except:
|
||||
print("\n📊 Original task was cancelled")
|
||||
|
||||
print("\n✅ Interrupt mode is great for:")
|
||||
print(" - When users frequently change direction")
|
||||
print(" - Interactive sessions where latest intent matters most")
|
||||
print(" - Real-time collaboration scenarios")
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run all demonstrations."""
|
||||
print("🚀 Bot In-Flight Run Control Demo")
|
||||
print("Solving the silent lock problem described in issue #1914")
|
||||
|
||||
# Show the problem
|
||||
await demo_without_run_control()
|
||||
|
||||
# Show the solution
|
||||
await demo_with_run_control()
|
||||
|
||||
# Show advanced feature
|
||||
await demo_interrupt_mode()
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("🎯 Summary")
|
||||
print("="*60)
|
||||
print("SessionRunControl solves the silent blocking problem by providing:")
|
||||
print("1. Immediate feedback for mid-run messages")
|
||||
print("2. /stop command support via InterruptController")
|
||||
print("3. Configurable policies (queue/interrupt/steer)")
|
||||
print("4. Run generation tracking to prevent race conditions")
|
||||
print("5. Pending message merging and ordering")
|
||||
print("\nThis makes bot interactions feel responsive and predictable!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -6,7 +6,6 @@ to different next steps based on output content.
|
||||
"""
|
||||
|
||||
from praisonaiagents import AgentFlow, Task
|
||||
from praisonaiagents import AgentFlowManager
|
||||
|
||||
# Create a workflow with branching
|
||||
workflow = AgentFlow(
|
||||
@@ -36,16 +35,9 @@ workflow = AgentFlow(
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create manager and register workflow
|
||||
manager = WorkflowManager()
|
||||
manager.workflows["Decision Workflow"] = workflow
|
||||
|
||||
# Execute - will branch to success_handler since 42 is positive
|
||||
result = manager.execute(
|
||||
"Decision Workflow",
|
||||
default_llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
result = workflow.run("", llm="gpt-4o-mini", verbose=True)
|
||||
print("Workflow completed!")
|
||||
for step_result in result["results"]:
|
||||
print(f" {step_result['step']}: {step_result['status']}")
|
||||
for step in result.get("steps", []):
|
||||
name = step.get("step", step.get("name", "unknown"))
|
||||
status = step.get("status", "completed")
|
||||
print(f" {name}: {status}")
|
||||
@@ -6,9 +6,8 @@ functions that return StepResult with stop_workflow=True.
|
||||
"""
|
||||
|
||||
from praisonaiagents import (
|
||||
Workflow, Task, WorkflowContext, StepResult
|
||||
AgentFlow, Task, WorkflowContext, StepResult
|
||||
)
|
||||
from praisonaiagents import AgentFlowManager
|
||||
|
||||
# Custom validator that can stop the workflow
|
||||
def validate_data(context: WorkflowContext) -> StepResult:
|
||||
@@ -50,20 +49,13 @@ workflow = AgentFlow(
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
manager = WorkflowManager()
|
||||
manager.workflows["Data Processing"] = workflow
|
||||
|
||||
# Test with invalid data (will stop early)
|
||||
print("=== Testing with invalid data ===")
|
||||
result = manager.execute("Data Processing", default_llm="gpt-4o-mini")
|
||||
|
||||
for step_result in result["results"]:
|
||||
print(f" {step_result['step']}: {step_result['output']}")
|
||||
|
||||
# Test with valid data (will complete)
|
||||
result = workflow.run("", llm="gpt-4o-mini", verbose=True)
|
||||
for step in result.get("steps", []):
|
||||
print(f" {step.get('step', 'step')}: {step.get('output', step.get('status'))}")
|
||||
|
||||
print("\n=== Testing with valid data ===")
|
||||
workflow.variables["data"] = {"value": 42}
|
||||
result = manager.execute("Data Processing", default_llm="gpt-4o-mini")
|
||||
|
||||
for step_result in result["results"]:
|
||||
print(f" {step_result['step']}: {step_result['status']}")
|
||||
result = workflow.run("", llm="gpt-4o-mini", verbose=True)
|
||||
for step in result.get("steps", []):
|
||||
print(f" {step.get('step', 'step')}: {step.get('status', 'completed')}")
|
||||
@@ -6,7 +6,6 @@ Each item is processed individually with the loop variable.
|
||||
"""
|
||||
|
||||
from praisonaiagents import AgentFlow, Task
|
||||
from praisonaiagents import AgentFlowManager
|
||||
|
||||
# Create a workflow with loop
|
||||
workflow = AgentFlow(
|
||||
@@ -30,18 +29,12 @@ workflow = AgentFlow(
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create manager and register workflow
|
||||
manager = WorkflowManager()
|
||||
manager.workflows["Process Items"] = workflow
|
||||
|
||||
# Execute - will process each fruit in the list
|
||||
result = manager.execute(
|
||||
"Process Items",
|
||||
default_llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
result = workflow.run("", llm="gpt-4o-mini", verbose=True)
|
||||
print("Workflow completed!")
|
||||
for step_result in result["results"]:
|
||||
print(f" {step_result['step']}: {step_result['status']}")
|
||||
if isinstance(step_result.get("output"), list):
|
||||
print(f" Loop processed {len(step_result['output'])} items")
|
||||
for step in result.get("steps", []):
|
||||
name = step.get("step", step.get("name", "unknown"))
|
||||
status = step.get("status", "completed")
|
||||
print(f" {name}: {status}")
|
||||
out = step.get("output")
|
||||
if isinstance(out, list):
|
||||
print(f" Loop processed {len(out)} items")
|
||||
@@ -53,7 +53,7 @@ workflow = AgentFlow(
|
||||
# Step 4: Function - formats output
|
||||
format_output
|
||||
],
|
||||
default_llm="gpt-4o-mini" # Used for action-based steps
|
||||
llm="gpt-4o-mini" # Used for action-based steps
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -5,7 +5,7 @@ Demonstrates running multiple agents concurrently and
|
||||
combining their results with an aggregator agent.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, Workflow
|
||||
from praisonaiagents import Agent, AgentFlow
|
||||
from praisonaiagents.workflows import parallel
|
||||
|
||||
# Create parallel research agents
|
||||
|
||||
@@ -5,7 +5,7 @@ Demonstrates decision-based routing where an agent classifier
|
||||
routes to specialized handler agents based on the request type.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, Workflow
|
||||
from praisonaiagents import Agent, AgentFlow
|
||||
from praisonaiagents.workflows import route
|
||||
|
||||
# Create classifier agent
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
# Agent Instantiation — Root Cause Analysis
|
||||
|
||||
**Date:** 2026-06-17
|
||||
**Version:** praisonaiagents 1.6.58 (dev)
|
||||
**Methodology:** 10 warmup + 1000 measured iterations, `time.perf_counter()`
|
||||
|
||||
## Summary
|
||||
|
||||
Instantiation regressed from **~3.2 μs** (Dec 2025) to **~31 μs** (~10×). After fixes, steady-state is **~16 μs** (~2× historical). PraisonAI is competitive with Agno (~5 μs) on warmed benchmarks using `model=`.
|
||||
|
||||
## Before vs After Fixes
|
||||
|
||||
| Config | Historical | Pre-fix | Post-fix |
|
||||
|--------|------------|---------|----------|
|
||||
| Minimal (`name`, `output="silent"`) | 3.20 μs | 31.36 μs | **16.35 μs** |
|
||||
| `model="gpt-4o-mini"` | 3.77 μs | 32.19 μs | **16.86 μs** |
|
||||
| `llm="gpt-4o-mini"` (deprecated) | 3.77 μs | 46.48 μs | **~18 μs** (warn once) |
|
||||
| LiteLLM `openai/gpt-4o-mini` cold | 7.31 μs | **3.7 s** | **53.63 μs** |
|
||||
| LiteLLM warmed instantiate | — | 60 μs | **19.83 μs** |
|
||||
| Import `Agent` | — | 20.8 ms | **29.6 ms** |
|
||||
|
||||
## Root Causes (validated)
|
||||
|
||||
| ID | Cause | Status | Fix applied |
|
||||
|----|-------|--------|-------------|
|
||||
| H1 | Config loader cold start (+9 ms first init) | CONFIRMED | Fast path when no config defaults (`_defaults_has_any_values`) |
|
||||
| H2 | LoopGuard on every init (~2 μs) | CONFIRMED | Lazy `_ensure_loop_guard()` on first tool call |
|
||||
| H3 | Approval presets | REJECTED | No change needed |
|
||||
| H4 | `llm=` deprecation every init (+16 μs) | CONFIRMED | Warn once per process in `warn_deprecated_param` |
|
||||
| H5 | Init bloat / config resolution (~25 μs) | CONFIRMED | Partially addressed via loader fast path |
|
||||
| H6 | LiteLLM eager `LLM()` import (+3.7 s) | CONFIRMED | Deferred `_llm_init_params` + lazy `llm_instance` property |
|
||||
| H7 | `ExecutionConfig` deprecation every init | CONFIRMED | Fixed stack walk; internal SDK path cached |
|
||||
|
||||
## cProfile — Post-fix (500 warmed `model=` inits)
|
||||
|
||||
| Function | Time (ms) | Per call |
|
||||
|----------|-----------|----------|
|
||||
| `Agent.__init__` | 8 | ~16 μs |
|
||||
| `apply_config_defaults` | 1 | ~2 μs |
|
||||
| `ExecutionConfig.__post_init__` | 1 | ~2 μs (first call only walks stack) |
|
||||
| `_warnings.warn` | **0** | removed from hot path |
|
||||
|
||||
## Fixes implemented
|
||||
|
||||
1. **`praisonaiagents/config/feature_configs.py`** — Fix `ExecutionConfig.__post_init__` deprecation guard; cache internal SDK detection.
|
||||
2. **`praisonaiagents/utils/deprecation.py`** — Warn once per process for deprecated params.
|
||||
3. **`praisonaiagents/config/loader.py`** — Skip `apply_config_defaults` lookups when config has no defaults.
|
||||
4. **`praisonaiagents/agent/agent.py`** — Lazy LoopGuard, lazy LLM instance, deferred LiteLLM params.
|
||||
5. **`benchmarks/simple_benchmark.py`** — Use `model=`, add 10-iteration warmup.
|
||||
6. **`benchmarks/performance_benchmark.py`** — Use `model=` instead of deprecated `llm=`.
|
||||
|
||||
## Remaining gap (~16 μs vs ~3 μs)
|
||||
|
||||
The residual ~13 μs is primarily **`Agent.__init__` size** (~1200 lines of consolidated config resolution added since v0.1.5). Further optimisation requires a dedicated slim fast-path for `Agent(name=..., output="silent")` without full config object churn.
|
||||
|
||||
## How to reproduce
|
||||
|
||||
```bash
|
||||
cd praisonai-agents
|
||||
PYTHONWARNINGS='ignore::DeprecationWarning' python benchmarks/simple_benchmark.py
|
||||
PYTHONWARNINGS='ignore::DeprecationWarning' python benchmarks/performance_benchmark.py
|
||||
```
|
||||
@@ -120,8 +120,8 @@ def run_benchmark(iterations: int = 100) -> List[ParamResult]:
|
||||
('backstory', 'I am an AI assistant.', 'backstory="..."'),
|
||||
|
||||
# LLM configuration
|
||||
('llm', 'gpt-4o-mini', 'llm="gpt-4o-mini"'),
|
||||
('llm', 'openai/gpt-4o-mini', 'llm="openai/gpt-4o-mini" (LiteLLM)'),
|
||||
('model', 'gpt-4o-mini', 'model="gpt-4o-mini"'),
|
||||
('model', 'openai/gpt-4o-mini', 'model="openai/gpt-4o-mini" (LiteLLM)'),
|
||||
|
||||
# Tools
|
||||
('tools', [sample_tool], 'tools=[sample_tool]'),
|
||||
@@ -298,7 +298,7 @@ def save_results(results: List[ParamResult], filename: str = 'PARAM_IMPACT_RESUL
|
||||
f.write('### Best Practices\n\n')
|
||||
f.write('1. Use `output="silent"` for production workloads\n')
|
||||
f.write('2. Only enable `context=True` when needed for long conversations\n')
|
||||
f.write('3. Use `llm="gpt-4o-mini"` (OpenAI SDK) for fastest instantiation\n')
|
||||
f.write('3. Use `model="gpt-4o-mini"` for fastest instantiation\n')
|
||||
f.write('4. Avoid `memory=True` unless persistent memory is required\n')
|
||||
|
||||
f.write('\n## How to Reproduce\n\n')
|
||||
|
||||
@@ -185,7 +185,7 @@ def instantiate_praisonai_agent():
|
||||
return PraisonAgent(
|
||||
name="Test Agent",
|
||||
instructions="Be concise, reply with one sentence.",
|
||||
llm="gpt-4o-mini",
|
||||
model="gpt-4o-mini",
|
||||
tools=tools,
|
||||
output="silent"
|
||||
)
|
||||
@@ -196,7 +196,7 @@ def instantiate_praisonai_litellm_agent():
|
||||
return PraisonAgent(
|
||||
name="Test Agent",
|
||||
instructions="Be concise, reply with one sentence.",
|
||||
llm="openai/gpt-4o-mini",
|
||||
model="openai/gpt-4o-mini",
|
||||
tools=tools,
|
||||
output="silent"
|
||||
)
|
||||
|
||||
@@ -14,6 +14,7 @@ import argparse
|
||||
|
||||
|
||||
ITERATIONS = 100
|
||||
WARMUP = 10
|
||||
|
||||
|
||||
def sample_tool(city: Literal['nyc', 'sf']):
|
||||
@@ -23,8 +24,10 @@ def sample_tool(city: Literal['nyc', 'sf']):
|
||||
return 'sunny'
|
||||
|
||||
|
||||
def measure_instantiation(create_fn, iterations=ITERATIONS):
|
||||
def measure_instantiation(create_fn, iterations=ITERATIONS, warmup=WARMUP):
|
||||
"""Measure average instantiation time in microseconds."""
|
||||
for _ in range(warmup):
|
||||
create_fn()
|
||||
times = []
|
||||
for _ in range(iterations):
|
||||
start = time.perf_counter()
|
||||
@@ -48,11 +51,11 @@ def run_benchmark():
|
||||
from praisonaiagents import Agent as PraisonAgent
|
||||
|
||||
results['PraisonAI'] = measure_instantiation(
|
||||
lambda: PraisonAgent(name='Test', llm='gpt-4o-mini', output="silent")
|
||||
lambda: PraisonAgent(name='Test', model='gpt-4o-mini', output="silent")
|
||||
)
|
||||
|
||||
results['PraisonAI (LiteLLM)'] = measure_instantiation(
|
||||
lambda: PraisonAgent(name='Test', llm='openai/gpt-4o-mini', output="silent")
|
||||
lambda: PraisonAgent(name='Test', model='openai/gpt-4o-mini', output="silent")
|
||||
)
|
||||
|
||||
# Other frameworks for comparison
|
||||
|
||||
@@ -114,6 +114,9 @@ def _get_lazy_cache():
|
||||
# ============================================================================
|
||||
|
||||
_LAZY_IMPORTS = {
|
||||
# Version information (lazy loaded to avoid file I/O on import)
|
||||
'__version__': ('praisonaiagents._version', '__version__'),
|
||||
|
||||
# Tools (moved from eager imports for lazy loading)
|
||||
'Tools': ('praisonaiagents.tools.tools', 'Tools'),
|
||||
'BaseTool': ('praisonaiagents.tools.base', 'BaseTool'),
|
||||
@@ -775,6 +778,9 @@ def warmup(include_litellm: bool = False, include_openai: bool = True) -> dict:
|
||||
# ============================================================================
|
||||
|
||||
__all__ = [
|
||||
# Version information
|
||||
'__version__',
|
||||
|
||||
# Core classes - the essentials
|
||||
'Agent',
|
||||
'AgentTeam', # Primary class for multi-agent coordination (v1.0+)
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
"""
|
||||
Version utilities for praisonaiagents package.
|
||||
|
||||
This module provides version information by reading from pyproject.toml
|
||||
to avoid duplication and ensure the single source of truth.
|
||||
"""
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def get_version() -> str:
|
||||
"""
|
||||
Get the version string from pyproject.toml.
|
||||
|
||||
Returns:
|
||||
str: Version string (e.g., "1.6.52")
|
||||
|
||||
Raises:
|
||||
RuntimeError: If version cannot be found or parsed
|
||||
"""
|
||||
try:
|
||||
# Get the package root (praisonaiagents directory)
|
||||
package_root = Path(__file__).parent
|
||||
# Go up to src/praisonai-agents directory
|
||||
pyproject_path = package_root.parent / "pyproject.toml"
|
||||
|
||||
if not pyproject_path.exists():
|
||||
raise RuntimeError(f"pyproject.toml not found at {pyproject_path}")
|
||||
|
||||
content = pyproject_path.read_text()
|
||||
|
||||
# Look for version = "X.Y.Z" pattern
|
||||
match = re.search(r'version\s*=\s*["\']([^"\']+)["\']', content)
|
||||
if not match:
|
||||
raise RuntimeError("Version not found in pyproject.toml")
|
||||
|
||||
return match.group(1)
|
||||
|
||||
except Exception as e:
|
||||
# Fallback to "unknown" to avoid breaking imports
|
||||
import warnings
|
||||
warnings.warn(
|
||||
f"Failed to read version from pyproject.toml: {e}. Using 'unknown'.",
|
||||
RuntimeWarning
|
||||
)
|
||||
return "unknown"
|
||||
|
||||
|
||||
# Cache the version to avoid reading the file multiple times
|
||||
__version__ = get_version()
|
||||
@@ -313,6 +313,38 @@ class Agent(SteeringMixin, SandboxMixin, UnifiedExecutionMixin, ToolExecutionMix
|
||||
cls._default_model_checked = True
|
||||
return cls._default_model
|
||||
|
||||
def _ensure_llm_instance(self):
|
||||
"""Lazy-create LLM instance from deferred init params (avoids import at Agent())."""
|
||||
if self._llm_instance is not None:
|
||||
return self._llm_instance
|
||||
if self._llm_init_params:
|
||||
try:
|
||||
from ..llm.llm import LLM
|
||||
self._llm_instance = LLM(**self._llm_init_params)
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"LLM features requested but dependencies not installed. "
|
||||
"Please install with: pip install \"praisonaiagents[llm]\""
|
||||
) from e
|
||||
return self._llm_instance
|
||||
|
||||
@property
|
||||
def llm_instance(self):
|
||||
return self._ensure_llm_instance()
|
||||
|
||||
@llm_instance.setter
|
||||
def llm_instance(self, value):
|
||||
self._llm_instance = value
|
||||
if value is not None:
|
||||
self._llm_init_params = None
|
||||
|
||||
def _ensure_loop_guard(self):
|
||||
"""Lazy-create loop guard on first tool execution."""
|
||||
if self._loop_guard is None:
|
||||
from ..escalation.loop_guard import LoopGuard, LoopGuardConfig
|
||||
self._loop_guard = LoopGuard(LoopGuardConfig(enabled=True))
|
||||
return self._loop_guard
|
||||
|
||||
@classmethod
|
||||
def _configure_logging(cls):
|
||||
"""Configure logging settings once for all agent instances."""
|
||||
@@ -718,16 +750,16 @@ class Agent(SteeringMixin, SandboxMixin, UnifiedExecutionMixin, ToolExecutionMix
|
||||
# Precedence: Explicit params > Config file > Built-in defaults
|
||||
# Only applies when param is None (not explicitly set)
|
||||
# ============================================================
|
||||
from ..config.loader import apply_config_defaults, get_default
|
||||
from ..config.loader import apply_config_defaults, get_default, _defaults_has_any_values
|
||||
|
||||
# Apply config defaults for LLM if not explicitly set
|
||||
if llm is None and model is None:
|
||||
if llm is None and model is None and _defaults_has_any_values():
|
||||
config_model = get_default("model")
|
||||
if config_model:
|
||||
llm = config_model
|
||||
|
||||
# Apply config defaults for base_url if not explicitly set
|
||||
if base_url is None:
|
||||
if base_url is None and _defaults_has_any_values():
|
||||
config_base_url = get_default("base_url")
|
||||
if config_base_url:
|
||||
base_url = config_base_url
|
||||
@@ -1478,10 +1510,8 @@ class Agent(SteeringMixin, SandboxMixin, UnifiedExecutionMixin, ToolExecutionMix
|
||||
# Initialize autonomy features (agent-centric escalation/doom-loop)
|
||||
self._init_autonomy(autonomy, verification_hooks=verification_hooks)
|
||||
|
||||
# Initialize loop guard for all execution modes (not just autonomous)
|
||||
# Provides idempotency-aware guardrails for tool execution loops
|
||||
from ..escalation.loop_guard import LoopGuard, LoopGuardConfig
|
||||
self._loop_guard = LoopGuard(LoopGuardConfig(enabled=True))
|
||||
# Initialize loop guard lazily on first tool execution (zero init overhead)
|
||||
self._loop_guard = None
|
||||
|
||||
# If instructions are provided, use them to set role, goal, and backstory
|
||||
if instructions:
|
||||
@@ -1527,6 +1557,8 @@ class Agent(SteeringMixin, SandboxMixin, UnifiedExecutionMixin, ToolExecutionMix
|
||||
self.interrupt_controller = interrupt_controller
|
||||
# Check for model name in environment variable if not provided
|
||||
self._using_custom_llm = False
|
||||
self._llm_instance = None
|
||||
self._llm_init_params = None
|
||||
# Flag to track if final result has been displayed to prevent duplicates
|
||||
self._final_display_shown = False
|
||||
|
||||
@@ -1565,99 +1597,67 @@ class Agent(SteeringMixin, SandboxMixin, UnifiedExecutionMixin, ToolExecutionMix
|
||||
self.base_url = base_url
|
||||
self.api_key = api_key
|
||||
|
||||
# If base_url is provided, always create a custom LLM instance
|
||||
# If base_url is provided, defer custom LLM instance creation
|
||||
if base_url:
|
||||
try:
|
||||
from ..llm.llm import LLM
|
||||
# Handle different llm parameter types with base_url
|
||||
if isinstance(llm, dict):
|
||||
# Merge base_url and api_key into the dict
|
||||
llm_config = llm.copy()
|
||||
llm_config['base_url'] = base_url
|
||||
if api_key:
|
||||
llm_config['api_key'] = api_key
|
||||
if auth:
|
||||
llm_config['auth'] = auth
|
||||
llm_config['metrics'] = metrics
|
||||
llm_config['max_iter'] = max_iter
|
||||
self.llm_instance = LLM(**llm_config)
|
||||
self.llm = llm.get('model', Agent._get_default_model())
|
||||
else:
|
||||
# Create LLM with model string and base_url (cached for performance)
|
||||
model_name = llm or Agent._get_default_model()
|
||||
self.llm_instance = LLM(
|
||||
model=model_name,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
auth=auth,
|
||||
metrics=metrics,
|
||||
max_iter=max_iter,
|
||||
web_search=web_search,
|
||||
web_fetch=web_fetch,
|
||||
prompt_caching=prompt_caching,
|
||||
claude_memory=claude_memory
|
||||
)
|
||||
self.llm = model_name
|
||||
self._using_custom_llm = True
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"LLM features requested but dependencies not installed. "
|
||||
"Please install with: pip install \"praisonaiagents[llm]\""
|
||||
) from e
|
||||
if isinstance(llm, dict):
|
||||
llm_config = llm.copy()
|
||||
llm_config['base_url'] = base_url
|
||||
if api_key:
|
||||
llm_config['api_key'] = api_key
|
||||
if auth:
|
||||
llm_config['auth'] = auth
|
||||
llm_config['metrics'] = metrics
|
||||
llm_config['max_iter'] = max_iter
|
||||
self._llm_init_params = llm_config
|
||||
self.llm = llm.get('model', Agent._get_default_model())
|
||||
else:
|
||||
model_name = llm or Agent._get_default_model()
|
||||
self._llm_init_params = {
|
||||
'model': model_name,
|
||||
'base_url': base_url,
|
||||
'api_key': api_key,
|
||||
'auth': auth,
|
||||
'metrics': metrics,
|
||||
'max_iter': max_iter,
|
||||
'web_search': web_search,
|
||||
'web_fetch': web_fetch,
|
||||
'prompt_caching': prompt_caching,
|
||||
'claude_memory': claude_memory,
|
||||
}
|
||||
self.llm = model_name
|
||||
self._using_custom_llm = True
|
||||
# If the user passes a dictionary (for advanced configuration)
|
||||
elif isinstance(llm, dict) and "model" in llm:
|
||||
try:
|
||||
from ..llm.llm import LLM
|
||||
# Add api_key if provided and not in dict
|
||||
if api_key and 'api_key' not in llm:
|
||||
llm = llm.copy()
|
||||
llm['api_key'] = api_key
|
||||
# Add auth if provided and not in dict
|
||||
if auth and 'auth' not in llm:
|
||||
llm = llm.copy()
|
||||
llm['auth'] = auth
|
||||
# Add metrics parameter
|
||||
llm = llm.copy()
|
||||
llm['metrics'] = metrics
|
||||
llm['max_iter'] = max_iter
|
||||
self.llm_instance = LLM(**llm) # Pass all dict items as kwargs
|
||||
self._using_custom_llm = True
|
||||
self.llm = llm.get('model', Agent._get_default_model())
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"LLM features requested but dependencies not installed. "
|
||||
"Please install with: pip install \"praisonaiagents[llm]\""
|
||||
) from e
|
||||
# If the user passes a string with a slash (provider/model)
|
||||
llm_config = llm.copy()
|
||||
if api_key and 'api_key' not in llm_config:
|
||||
llm_config['api_key'] = api_key
|
||||
if auth and 'auth' not in llm_config:
|
||||
llm_config['auth'] = auth
|
||||
llm_config['metrics'] = metrics
|
||||
llm_config['max_iter'] = max_iter
|
||||
self._llm_init_params = llm_config
|
||||
self._using_custom_llm = True
|
||||
self.llm = llm_config.get('model', Agent._get_default_model())
|
||||
# If the user passes a string with a slash (provider/model) — defer LiteLLM import
|
||||
elif isinstance(llm, str) and "/" in llm:
|
||||
try:
|
||||
from ..llm.llm import LLM
|
||||
# Pass the entire string so LiteLLM can parse provider/model
|
||||
llm_params = {'model': llm}
|
||||
if api_key:
|
||||
llm_params['api_key'] = api_key
|
||||
if auth:
|
||||
llm_params['auth'] = auth
|
||||
llm_params['metrics'] = metrics
|
||||
llm_params['web_search'] = web_search
|
||||
llm_params['web_fetch'] = web_fetch
|
||||
llm_params['prompt_caching'] = prompt_caching
|
||||
llm_params['claude_memory'] = claude_memory
|
||||
llm_params['max_iter'] = max_iter
|
||||
self.llm_instance = LLM(**llm_params)
|
||||
self._using_custom_llm = True
|
||||
self.llm = llm
|
||||
|
||||
# Ensure tools are properly accessible when using custom LLM
|
||||
if tools:
|
||||
logging.debug(f"Tools passed to Agent with custom LLM: {tools}")
|
||||
# Store the tools for later use
|
||||
self.tools = tools
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"LLM features requested but dependencies not installed. "
|
||||
"Please install with: pip install \"praisonaiagents[llm]\""
|
||||
) from e
|
||||
llm_params = {'model': llm}
|
||||
if api_key:
|
||||
llm_params['api_key'] = api_key
|
||||
if auth:
|
||||
llm_params['auth'] = auth
|
||||
llm_params['metrics'] = metrics
|
||||
llm_params['web_search'] = web_search
|
||||
llm_params['web_fetch'] = web_fetch
|
||||
llm_params['prompt_caching'] = prompt_caching
|
||||
llm_params['claude_memory'] = claude_memory
|
||||
llm_params['max_iter'] = max_iter
|
||||
self._llm_init_params = llm_params
|
||||
self._using_custom_llm = True
|
||||
self.llm = llm
|
||||
|
||||
if tools:
|
||||
logging.debug(f"Tools passed to Agent with custom LLM: {tools}")
|
||||
self.tools = tools
|
||||
# Otherwise, fall back to OpenAI environment/name (cached for performance)
|
||||
else:
|
||||
self.llm = llm or Agent._get_default_model()
|
||||
|
||||
@@ -87,12 +87,41 @@ Your Goal: {self.goal}"""
|
||||
|
||||
# Add memory context if memory is enabled
|
||||
if self._memory_instance:
|
||||
memory_context = self.get_memory_context()
|
||||
if memory_context:
|
||||
system_prompt += f"\n\n## Memory (Information you remember about the user)\n{memory_context}"
|
||||
# Display memory info to user if verbose
|
||||
if self.verbose:
|
||||
self._display_memory_info()
|
||||
# Use cache-optimized context if model supports prompt caching and cache boundary is requested
|
||||
use_cache_boundaries = hasattr(self, '_model_supports_prompt_caching') and self._model_supports_prompt_caching()
|
||||
if use_cache_boundaries and hasattr(self._memory_instance, 'build_cache_optimized_context'):
|
||||
try:
|
||||
# Get the goal/task description for context building
|
||||
task_desc = getattr(self, 'goal', '') or 'general assistance'
|
||||
cache_result = self._memory_instance.build_cache_optimized_context(
|
||||
task_descr=task_desc,
|
||||
include_cache_boundary=True
|
||||
)
|
||||
memory_context = cache_result.get('stable_prefix', '')
|
||||
cache_boundary = cache_result.get('cache_boundary', '')
|
||||
if memory_context:
|
||||
system_prompt += f"\n\n## Memory (Information you remember about the user)\n{memory_context}"
|
||||
if cache_boundary:
|
||||
system_prompt += cache_boundary
|
||||
# Display memory info to user if verbose
|
||||
if self.verbose:
|
||||
self._display_memory_info()
|
||||
except (AttributeError, KeyError, TypeError) as e:
|
||||
# Fall back to standard memory context if cache optimization fails
|
||||
memory_context = self.get_memory_context()
|
||||
if memory_context:
|
||||
system_prompt += f"\n\n## Memory (Information you remember about the user)\n{memory_context}"
|
||||
# Display memory info to user if verbose
|
||||
if self.verbose:
|
||||
self._display_memory_info()
|
||||
else:
|
||||
# Standard memory context without cache boundaries
|
||||
memory_context = self.get_memory_context()
|
||||
if memory_context:
|
||||
system_prompt += f"\n\n## Memory (Information you remember about the user)\n{memory_context}"
|
||||
# Display memory info to user if verbose
|
||||
if self.verbose:
|
||||
self._display_memory_info()
|
||||
|
||||
# Add learn context if learn is enabled (auto-inject when memory="learn")
|
||||
learn_context = self.get_learn_context()
|
||||
@@ -452,6 +481,17 @@ Your Goal: {self.goal}"""
|
||||
else:
|
||||
cleaned_tools.append(tool)
|
||||
|
||||
# Sort tools deterministically for cache stability
|
||||
def sort_formatted_tools(tools_list):
|
||||
"""Sort already-formatted tool schemas by function name."""
|
||||
def sort_key(tool):
|
||||
if isinstance(tool, dict) and tool.get('type') == 'function':
|
||||
return str(tool.get('function', {}).get('name') or '')
|
||||
return ''
|
||||
return sorted(tools_list, key=sort_key)
|
||||
|
||||
cleaned_tools = sort_formatted_tools(cleaned_tools)
|
||||
|
||||
# Cache the formatted tools with LRU eviction, including tool search metadata
|
||||
self._cache_put(
|
||||
self._formatted_tools_cache,
|
||||
|
||||
@@ -231,9 +231,10 @@ class ToolExecutionMixin:
|
||||
arguments.update(res.output.modified_data)
|
||||
|
||||
# Loop guard check - prevent tool execution loops with graduated response
|
||||
if hasattr(self, '_loop_guard') and self._loop_guard:
|
||||
if hasattr(self, '_ensure_loop_guard'):
|
||||
loop_guard = self._ensure_loop_guard()
|
||||
from ..escalation.loop_guard import GuardAction
|
||||
decision = self._loop_guard.check(function_name, arguments, is_pre_execution=True)
|
||||
decision = loop_guard.check(function_name, arguments, is_pre_execution=True)
|
||||
|
||||
if decision.action == GuardAction.WARN:
|
||||
# Inject warning into tool result so LLM sees guidance
|
||||
@@ -394,11 +395,12 @@ class ToolExecutionMixin:
|
||||
self._doom_loop_tracker.mark_progress(f"tool:{function_name}")
|
||||
|
||||
# Record tool execution in loop guard
|
||||
if hasattr(self, '_loop_guard') and self._loop_guard:
|
||||
if hasattr(self, '_ensure_loop_guard'):
|
||||
loop_guard = self._ensure_loop_guard()
|
||||
is_success = result is not None and not (isinstance(result, dict) and result.get('error'))
|
||||
self._loop_guard.record(function_name, arguments, is_success)
|
||||
loop_guard.record(function_name, arguments, is_success)
|
||||
# Handle warning injection for WARN decisions
|
||||
decision = self._loop_guard.check(function_name, arguments, is_pre_execution=False)
|
||||
decision = loop_guard.check(function_name, arguments, is_pre_execution=False)
|
||||
if decision.action.value == "warn":
|
||||
if isinstance(result, str):
|
||||
result = f"{result}\n\n[loop-guard] {decision.message}"
|
||||
|
||||
@@ -78,8 +78,8 @@ class UnifiedExecutionMixin:
|
||||
_trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
|
||||
|
||||
# Reset loop guard for new chat turn (always-on protection)
|
||||
if hasattr(self, '_loop_guard') and self._loop_guard:
|
||||
self._loop_guard.reset_turn()
|
||||
if hasattr(self, '_ensure_loop_guard'):
|
||||
self._ensure_loop_guard().reset_turn()
|
||||
|
||||
try:
|
||||
# CLI Backend routing - delegate entire turn if configured
|
||||
|
||||
@@ -40,6 +40,21 @@ class TaskStatus(Enum):
|
||||
# Set up logger
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _launch_auth_token() -> Optional[str]:
|
||||
return os.environ.get("PRAISONAI_LAUNCH_AUTH_TOKEN")
|
||||
|
||||
|
||||
def _authorise_launch_request(request) -> bool:
|
||||
token = _launch_auth_token()
|
||||
if not token:
|
||||
return True
|
||||
auth = request.headers.get("Authorization", "")
|
||||
if auth.startswith("Bearer ") and auth[7:] == token:
|
||||
return True
|
||||
return request.headers.get("X-Auth-Token") == token
|
||||
|
||||
|
||||
# Agent server registry for thread-safe server management
|
||||
|
||||
|
||||
@@ -245,7 +260,20 @@ Expected Output: {task.expected_output}."""
|
||||
# Add memory context if available
|
||||
if task.memory:
|
||||
try:
|
||||
memory_context = task.memory.build_context_for_task(task.description)
|
||||
# Use cache-optimized context if available for better prompt caching
|
||||
if hasattr(task.memory, 'build_cache_optimized_context'):
|
||||
try:
|
||||
cache_result = task.memory.build_cache_optimized_context(
|
||||
task_descr=task.description,
|
||||
include_cache_boundary=False # Don't include boundary for agent task context
|
||||
)
|
||||
memory_context = cache_result.get('stable_prefix', '')
|
||||
except Exception as e:
|
||||
# Fall back to standard context building
|
||||
logger.debug(f"Cache-optimized context failed for task '{task.description or task.id}', falling back to standard: {e}")
|
||||
memory_context = task.memory.build_context_for_task(task.description)
|
||||
else:
|
||||
memory_context = task.memory.build_context_for_task(task.description)
|
||||
if memory_context:
|
||||
# Log detailed memory context for debugging
|
||||
logger.debug(f"Memory context for task '{task.description}': {memory_context}")
|
||||
@@ -652,7 +680,7 @@ class AgentTeam(SpawnAnnounceProtocol):
|
||||
_stream = _output_config.stream
|
||||
else:
|
||||
_verbose = 0
|
||||
_stream = True
|
||||
_stream = False
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Resolve EXECUTION param using canonical resolver
|
||||
@@ -1993,6 +2021,8 @@ class AgentTeam(SpawnAnnounceProtocol):
|
||||
# Define the endpoint handler
|
||||
@app.post(path)
|
||||
async def handle_query(request: Request, query_data: Optional[AgentQuery] = None):
|
||||
if not _authorise_launch_request(request):
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
# Handle both direct JSON with query field and form data
|
||||
if query_data is None:
|
||||
try:
|
||||
@@ -2071,7 +2101,9 @@ class AgentTeam(SpawnAnnounceProtocol):
|
||||
|
||||
# Add GET endpoint to list available agents
|
||||
@app.get(f"{path}/list")
|
||||
async def list_agents():
|
||||
async def list_agents(request: Request):
|
||||
if not _authorise_launch_request(request):
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
return {
|
||||
"agents": [
|
||||
{"name": agent.display_name, "id": agent.display_name.lower().replace(' ', '_')}
|
||||
@@ -2086,6 +2118,8 @@ class AgentTeam(SpawnAnnounceProtocol):
|
||||
# Create a closure to capture the agent instance
|
||||
def create_agent_handler(agent):
|
||||
async def handle_single_agent(request: Request):
|
||||
if not _authorise_launch_request(request):
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
try:
|
||||
request_data = await request.json()
|
||||
query = request_data.get("query", "")
|
||||
|
||||
@@ -53,3 +53,4 @@ class AgentAppConfig:
|
||||
workers: int = 1
|
||||
timeout: int = 60
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
api_key: Optional[str] = None
|
||||
@@ -194,9 +194,13 @@ class ApprovalRegistry:
|
||||
|
||||
def is_yaml_approved(self, tool_name: str) -> bool:
|
||||
try:
|
||||
return tool_name in self._yaml_approved_tools.get()
|
||||
if tool_name not in self._yaml_approved_tools.get():
|
||||
return False
|
||||
except LookupError:
|
||||
return False
|
||||
if self.get_risk_level(tool_name) == "critical":
|
||||
return False
|
||||
return True
|
||||
|
||||
# ── Env-var check ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@@ -91,6 +91,16 @@ class BotConfig:
|
||||
# When set, stale sessions older than this are auto-reaped.
|
||||
session_ttl: int = 0
|
||||
|
||||
# In-flight run control settings for mid-run message handling
|
||||
# busy_mode: Policy for handling messages during active runs
|
||||
# - "queue": Queue message for next turn (default)
|
||||
# - "interrupt": Cancel current run and start new one
|
||||
# - "steer": Inject message into current run via steering
|
||||
busy_mode: str = "queue"
|
||||
|
||||
# Template for busy acknowledgment messages (use {action} placeholder)
|
||||
busy_ack: str = "⏳ {action} — will be considered next"
|
||||
|
||||
# Workspace settings for file operation containment and security
|
||||
workspace_dir: Optional[str] = None # default: ~/.praisonai/workspaces/<scope>/<session_key>
|
||||
workspace_access: str = "rw" # "rw" (read-write) | "ro" (read-only) | "none" (copy-on-write sandbox)
|
||||
@@ -136,6 +146,8 @@ class BotConfig:
|
||||
"ack_emoji": self.ack_emoji,
|
||||
"done_emoji": self.done_emoji,
|
||||
"session_ttl": self.session_ttl,
|
||||
"busy_mode": self.busy_mode,
|
||||
"busy_ack": self.busy_ack,
|
||||
"workspace_dir": self.workspace_dir,
|
||||
"workspace_access": self.workspace_access,
|
||||
"workspace_scope": self.workspace_scope,
|
||||
|
||||
@@ -33,6 +33,10 @@ __all__ = [
|
||||
"CompactionStrategy",
|
||||
# Results
|
||||
"CompactionResult",
|
||||
# Protocols
|
||||
"ToolResultPrunerProtocol",
|
||||
"MessageFormatterProtocol",
|
||||
"SummaryBuilderProtocol",
|
||||
]
|
||||
|
||||
|
||||
@@ -54,4 +58,16 @@ def __getattr__(name: str):
|
||||
from .result import CompactionResult
|
||||
return CompactionResult
|
||||
|
||||
if name == "ToolResultPrunerProtocol":
|
||||
from .protocols import ToolResultPrunerProtocol
|
||||
return ToolResultPrunerProtocol
|
||||
|
||||
if name == "MessageFormatterProtocol":
|
||||
from .protocols import MessageFormatterProtocol
|
||||
return MessageFormatterProtocol
|
||||
|
||||
if name == "SummaryBuilderProtocol":
|
||||
from .protocols import SummaryBuilderProtocol
|
||||
return SummaryBuilderProtocol
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -5,12 +5,13 @@ Manages context window by compacting messages when needed.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import List, Dict, Any, Optional, Callable, Awaitable
|
||||
from typing import List, Dict, Any, Optional, Callable, Awaitable, Tuple
|
||||
import asyncio
|
||||
|
||||
from .config import CompactionConfig, COMPACTION_PREFIX, SUMMARY_TEMPLATE
|
||||
from .strategy import CompactionStrategy
|
||||
from .result import CompactionResult
|
||||
from .protocols import ToolResultPrunerProtocol, MessageFormatterProtocol, SummaryBuilderProtocol
|
||||
|
||||
|
||||
class ContextCompactor:
|
||||
@@ -34,7 +35,10 @@ class ContextCompactor:
|
||||
preserve_system: bool = True,
|
||||
preserve_recent: int = 5,
|
||||
config: Optional[CompactionConfig] = None,
|
||||
llm_summarize_fn: Optional[Callable[[str], Awaitable[str]]] = None
|
||||
llm_summarize_fn: Optional[Callable[[str], Awaitable[str]]] = None,
|
||||
tool_pruner: Optional[ToolResultPrunerProtocol] = None,
|
||||
message_formatter: Optional[MessageFormatterProtocol] = None,
|
||||
summary_builder: Optional[SummaryBuilderProtocol] = None
|
||||
):
|
||||
"""
|
||||
Initialize the compactor.
|
||||
@@ -45,30 +49,43 @@ class ContextCompactor:
|
||||
strategy: Compaction strategy to use
|
||||
preserve_system: Keep system messages
|
||||
preserve_recent: Number of recent messages to preserve
|
||||
<<<<<<< HEAD
|
||||
config: Optional CompactionConfig to override defaults
|
||||
config: Optional CompactionConfig for advanced settings
|
||||
llm_summarize_fn: Async function to call LLM for summarization
|
||||
tool_pruner: Protocol implementation for tool result pruning
|
||||
message_formatter: Protocol implementation for message formatting
|
||||
summary_builder: Protocol implementation for summary building
|
||||
"""
|
||||
# Initialize config first
|
||||
self.config = config or CompactionConfig()
|
||||
# Use provided config or create default
|
||||
self.config = config or CompactionConfig(
|
||||
max_tokens=max_tokens,
|
||||
target_tokens=target_tokens or int(max_tokens * 0.75),
|
||||
preserve_system=preserve_system,
|
||||
preserve_recent=preserve_recent
|
||||
)
|
||||
|
||||
# Use config values if config provided, otherwise use constructor args
|
||||
if config is not None:
|
||||
self.max_tokens = config.max_tokens
|
||||
self.target_tokens = config.target_tokens
|
||||
self.preserve_system = config.preserve_system
|
||||
self.preserve_recent = config.preserve_recent
|
||||
else:
|
||||
self.max_tokens = max_tokens
|
||||
self.target_tokens = target_tokens or int(max_tokens * 0.75)
|
||||
self.preserve_system = preserve_system
|
||||
self.preserve_recent = preserve_recent
|
||||
|
||||
# Set instance attributes from config to ensure consistency
|
||||
self.max_tokens = self.config.max_tokens
|
||||
self.target_tokens = self.config.target_tokens
|
||||
self.strategy = strategy
|
||||
self.preserve_system = self.config.preserve_system
|
||||
self.preserve_recent = self.config.preserve_recent
|
||||
self.llm_summarize_fn = llm_summarize_fn
|
||||
|
||||
# Track previous summaries for iterative update
|
||||
# Protocol implementations (defaults to None - no heavy implementations in core)
|
||||
self.tool_pruner = tool_pruner
|
||||
self.message_formatter = message_formatter
|
||||
self.summary_builder = summary_builder
|
||||
|
||||
# Anti-thrashing state tracking
|
||||
self._last_savings_pct: float = 100.0 # Start high to allow first compaction
|
||||
self._low_savings_streak: int = 0
|
||||
|
||||
# Iterative summary state
|
||||
self._previous_summary: Optional[str] = None
|
||||
self._previous_summary_global_idx: int = 0
|
||||
|
||||
# Reset state tracking for each compact() call
|
||||
self._used_previous_summary: bool = False
|
||||
|
||||
def estimate_tokens(self, text: str) -> int:
|
||||
"""Estimate token count for text."""
|
||||
@@ -96,12 +113,26 @@ class ContextCompactor:
|
||||
return sum(self.count_message_tokens(m) for m in messages)
|
||||
|
||||
def needs_compaction(self, messages: List[Dict[str, Any]]) -> bool:
|
||||
"""Check if messages need compaction."""
|
||||
return self.count_total_tokens(messages) > self.max_tokens
|
||||
"""
|
||||
Check if messages need compaction with anti-thrashing protection.
|
||||
|
||||
Returns False if:
|
||||
- Token count is below threshold
|
||||
- We've had too many consecutive low-savings attempts
|
||||
"""
|
||||
if self.count_total_tokens(messages) <= self.max_tokens:
|
||||
return False
|
||||
|
||||
# Anti-thrashing: skip if we've had too many low-savings attempts
|
||||
if self._low_savings_streak >= self.config.max_consecutive_low_savings:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def compact(
|
||||
self,
|
||||
messages: List[Dict[str, Any]]
|
||||
messages: List[Dict[str, Any]],
|
||||
focus_topic: str = ""
|
||||
) -> tuple[List[Dict[str, Any]], CompactionResult]:
|
||||
"""
|
||||
Compact messages to fit within token limit (synchronous version).
|
||||
@@ -111,31 +142,49 @@ class ContextCompactor:
|
||||
|
||||
Args:
|
||||
messages: List of messages to compact
|
||||
focus_topic: Optional topic to focus on during summarization
|
||||
|
||||
Returns:
|
||||
Tuple of (compacted messages, result)
|
||||
"""
|
||||
# Reset state for this compaction call
|
||||
self._used_previous_summary = False
|
||||
|
||||
original_tokens = self.count_total_tokens(messages)
|
||||
|
||||
if original_tokens <= self.max_tokens:
|
||||
return messages, CompactionResult(
|
||||
result = CompactionResult(
|
||||
original_tokens=original_tokens,
|
||||
compacted_tokens=original_tokens,
|
||||
messages_removed=0,
|
||||
messages_kept=len(messages),
|
||||
strategy_used=self.strategy
|
||||
)
|
||||
result.calculate_savings_pct()
|
||||
return messages, result
|
||||
|
||||
# Apply tool result deduplication pre-pass if enabled
|
||||
processed_messages = messages
|
||||
tool_results_pruned = 0
|
||||
if self.config.tool_prune_before_summarise and self.tool_pruner:
|
||||
processed_messages, tool_results_pruned = self.tool_pruner.prune(
|
||||
messages,
|
||||
self.config.max_tool_result_size
|
||||
)
|
||||
|
||||
# Skip early exit - proceed with full compaction strategy
|
||||
# Anti-thrashing check will be applied after strategy runs
|
||||
|
||||
if self.strategy == CompactionStrategy.TRUNCATE:
|
||||
compacted = self._truncate(messages)
|
||||
compacted = self._truncate(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.SLIDING:
|
||||
compacted = self._sliding_window(messages)
|
||||
compacted = self._sliding_window(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.SUMMARIZE:
|
||||
compacted = self._summarize(messages)
|
||||
compacted = self._summarize(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.SMART:
|
||||
compacted = self._smart_compact(messages)
|
||||
compacted = self._smart_compact(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.PRUNE:
|
||||
compacted = self._prune(messages)
|
||||
compacted = self._prune(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.LLM_SUMMARIZE:
|
||||
if self.llm_summarize_fn:
|
||||
# For sync calls with LLM function, we need to run async
|
||||
@@ -144,17 +193,17 @@ class ContextCompactor:
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
# If in async context, fallback to naive summarization
|
||||
compacted = self._summarize(messages)
|
||||
compacted = self._summarize(processed_messages)
|
||||
except RuntimeError:
|
||||
# No running loop, safe to create one
|
||||
compacted = asyncio.run(self._llm_summarize_async(messages))
|
||||
compacted = asyncio.run(self._llm_summarize_async(processed_messages, focus_topic))
|
||||
except Exception:
|
||||
# Fallback to naive summarization if async fails
|
||||
compacted = self._summarize(messages)
|
||||
compacted = self._summarize(processed_messages)
|
||||
else:
|
||||
compacted = self._llm_summarize(messages)
|
||||
compacted = self._llm_summarize(processed_messages, focus_topic)
|
||||
else:
|
||||
compacted = self._truncate(messages)
|
||||
compacted = self._truncate(processed_messages)
|
||||
|
||||
compacted_tokens = self.count_total_tokens(compacted)
|
||||
|
||||
@@ -163,49 +212,75 @@ class ContextCompactor:
|
||||
compacted_tokens=compacted_tokens,
|
||||
messages_removed=len(messages) - len(compacted),
|
||||
messages_kept=len(compacted),
|
||||
strategy_used=self.strategy
|
||||
strategy_used=self.strategy,
|
||||
tool_results_pruned=tool_results_pruned,
|
||||
previous_summary_reused=getattr(self, '_used_previous_summary', False)
|
||||
)
|
||||
result.calculate_savings_pct()
|
||||
|
||||
# Update anti-thrashing tracking based on actual results
|
||||
self._last_savings_pct = result.savings_pct
|
||||
if result.savings_pct >= self.config.min_savings_pct:
|
||||
self._low_savings_streak = 0
|
||||
else:
|
||||
self._low_savings_streak += 1
|
||||
|
||||
return compacted, result
|
||||
|
||||
async def compact_async(
|
||||
self,
|
||||
messages: List[Dict[str, Any]]
|
||||
messages: List[Dict[str, Any]],
|
||||
focus_topic: str = ""
|
||||
) -> tuple[List[Dict[str, Any]], CompactionResult]:
|
||||
"""
|
||||
Compact messages to fit within token limit (asynchronous version).
|
||||
|
||||
Args:
|
||||
messages: List of messages to compact
|
||||
focus_topic: Optional topic to focus on during summarization
|
||||
|
||||
Returns:
|
||||
Tuple of (compacted messages, result)
|
||||
"""
|
||||
# Reset state for this compaction call
|
||||
self._used_previous_summary = False
|
||||
|
||||
original_tokens = self.count_total_tokens(messages)
|
||||
|
||||
if original_tokens <= self.max_tokens:
|
||||
return messages, CompactionResult(
|
||||
result = CompactionResult(
|
||||
original_tokens=original_tokens,
|
||||
compacted_tokens=original_tokens,
|
||||
messages_removed=0,
|
||||
messages_kept=len(messages),
|
||||
strategy_used=self.strategy
|
||||
)
|
||||
result.calculate_savings_pct()
|
||||
return messages, result
|
||||
|
||||
# Apply tool result deduplication pre-pass if enabled
|
||||
processed_messages = messages
|
||||
tool_results_pruned = 0
|
||||
if self.config.tool_prune_before_summarise and self.tool_pruner:
|
||||
processed_messages, tool_results_pruned = self.tool_pruner.prune(
|
||||
messages,
|
||||
self.config.max_tool_result_size
|
||||
)
|
||||
|
||||
if self.strategy == CompactionStrategy.TRUNCATE:
|
||||
compacted = self._truncate(messages)
|
||||
compacted = self._truncate(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.SLIDING:
|
||||
compacted = self._sliding_window(messages)
|
||||
compacted = self._sliding_window(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.SUMMARIZE:
|
||||
compacted = self._summarize(messages)
|
||||
compacted = self._summarize(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.SMART:
|
||||
compacted = self._smart_compact(messages)
|
||||
compacted = self._smart_compact(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.PRUNE:
|
||||
compacted = self._prune(messages)
|
||||
compacted = self._prune(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.LLM_SUMMARIZE:
|
||||
compacted = await self._llm_summarize_async(messages)
|
||||
compacted = await self._llm_summarize_async(processed_messages, focus_topic)
|
||||
else:
|
||||
compacted = self._truncate(messages)
|
||||
compacted = self._truncate(processed_messages)
|
||||
|
||||
compacted_tokens = self.count_total_tokens(compacted)
|
||||
|
||||
@@ -214,11 +289,35 @@ class ContextCompactor:
|
||||
compacted_tokens=compacted_tokens,
|
||||
messages_removed=len(messages) - len(compacted),
|
||||
messages_kept=len(compacted),
|
||||
strategy_used=self.strategy
|
||||
strategy_used=self.strategy,
|
||||
tool_results_pruned=tool_results_pruned,
|
||||
previous_summary_reused=getattr(self, '_used_previous_summary', False)
|
||||
)
|
||||
result.calculate_savings_pct()
|
||||
|
||||
# Update anti-thrashing tracking based on actual results
|
||||
self._last_savings_pct = result.savings_pct
|
||||
if result.savings_pct >= self.config.min_savings_pct:
|
||||
self._low_savings_streak = 0
|
||||
else:
|
||||
self._low_savings_streak += 1
|
||||
|
||||
return compacted, result
|
||||
|
||||
def _prune_tool_results(self, messages: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], int]:
|
||||
"""
|
||||
Delegate tool result pruning to injected protocol implementation.
|
||||
|
||||
If no tool pruner is provided, returns messages unchanged.
|
||||
This maintains backward compatibility while enforcing protocol-driven design.
|
||||
|
||||
Returns:
|
||||
Tuple of (processed messages, number of tool results pruned)
|
||||
"""
|
||||
if self.tool_pruner:
|
||||
return self.tool_pruner.prune(messages, self.config.max_tool_result_size)
|
||||
return messages, 0
|
||||
|
||||
def _truncate(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Truncate oldest messages."""
|
||||
result = []
|
||||
@@ -360,72 +459,19 @@ class ContextCompactor:
|
||||
result.extend(recent)
|
||||
return result
|
||||
|
||||
def _llm_summarize(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
def _llm_summarize(self, messages: List[Dict[str, Any]], focus_topic: str = "") -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Use LLM to summarize older messages with anti-injection framing.
|
||||
Use LLM to summarize older messages with iterative support.
|
||||
|
||||
Note: This implementation uses structured templates and anti-injection prefixes
|
||||
to prevent the model from treating summarized content as active instructions.
|
||||
"""
|
||||
result = []
|
||||
Supports iterative summarization - if we have a previous summary,
|
||||
we only summarize the new messages since that summary.
|
||||
Also includes anti-injection framing when configured.
|
||||
|
||||
# Keep system messages (excluding previous compacted summaries)
|
||||
system_msgs = [
|
||||
m for m in messages
|
||||
if m.get("role") == "system" and not m.get("_compacted")
|
||||
]
|
||||
other_msgs = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
result.extend(system_msgs)
|
||||
|
||||
# Keep recent messages
|
||||
recent = other_msgs[-self.preserve_recent:]
|
||||
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
|
||||
|
||||
if older:
|
||||
if self.config.structured_template:
|
||||
structured = self._build_structured_summary(older)
|
||||
else:
|
||||
# Fallback to simple summary
|
||||
summary_parts = []
|
||||
for msg in older:
|
||||
role = msg.get("role", "unknown")
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str) and content:
|
||||
summary_parts.append(f"[{role}]: {content[:150]}...")
|
||||
structured = "\n".join(summary_parts[:10])
|
||||
|
||||
if structured:
|
||||
# Apply anti-injection prefix
|
||||
prefixed = f"{self.config.compaction_prefix}\n\n{structured}"
|
||||
|
||||
# Handle iterative update if enabled
|
||||
if self.config.iterative_update and self._previous_summary:
|
||||
structured = self._merge_summaries(self._previous_summary, structured)
|
||||
prefixed = f"{self.config.compaction_prefix}\n\n{structured}"
|
||||
|
||||
# Store for next iteration
|
||||
if self.config.iterative_update:
|
||||
self._previous_summary = structured
|
||||
|
||||
result.append({
|
||||
"role": "system",
|
||||
"content": prefixed,
|
||||
"_compacted": True,
|
||||
"_original_count": len(older),
|
||||
"_anti_injection": True,
|
||||
})
|
||||
|
||||
result.extend(recent)
|
||||
return result
|
||||
|
||||
async def _llm_summarize_async(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Use LLM to intelligently summarize older messages.
|
||||
|
||||
This method invokes the agent's LLM to create a meaningful summary
|
||||
that preserves key facts, identifiers, and the user's intent.
|
||||
Args:
|
||||
messages: List of messages to compact
|
||||
focus_topic: Optional topic to focus on during summarization
|
||||
"""
|
||||
self._used_previous_summary = False
|
||||
result = []
|
||||
|
||||
# Keep system messages
|
||||
@@ -436,33 +482,169 @@ class ContextCompactor:
|
||||
|
||||
# Keep recent messages
|
||||
recent = other_msgs[-self.preserve_recent:]
|
||||
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
|
||||
|
||||
# Determine what to summarize
|
||||
total_original_messages = len(messages)
|
||||
if (self.config.enable_iterative_summary and
|
||||
self._previous_summary and
|
||||
total_original_messages > self._previous_summary_global_idx):
|
||||
# Iterative: summarize only new messages since previous summary
|
||||
# Calculate how many old messages to summarize based on global position
|
||||
messages_since_summary = total_original_messages - self._previous_summary_global_idx
|
||||
new_older_messages = max(0, messages_since_summary - self.preserve_recent)
|
||||
if new_older_messages > 0:
|
||||
to_summarize = other_msgs[-messages_since_summary:-self.preserve_recent]
|
||||
else:
|
||||
to_summarize = []
|
||||
self._used_previous_summary = True
|
||||
else:
|
||||
# Fresh summary: summarize all older messages
|
||||
to_summarize = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
|
||||
self._used_previous_summary = False
|
||||
|
||||
if to_summarize or self._previous_summary:
|
||||
# Build summary content
|
||||
summary_parts = []
|
||||
|
||||
# Include previous summary if doing iterative summarization
|
||||
if self._used_previous_summary and self._previous_summary:
|
||||
summary_parts.append(f"[Previous Summary]: {self._previous_summary}")
|
||||
summary_parts.append("[New Activity]:")
|
||||
|
||||
# Add new messages to summarize
|
||||
for msg in to_summarize:
|
||||
role = msg.get("role", "unknown")
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str) and content:
|
||||
# Extract key information with focus consideration
|
||||
if focus_topic and focus_topic.lower() in content.lower():
|
||||
# Prioritize content related to focus topic
|
||||
summary_parts.append(f"[{role}] *FOCUS*: {content[:200]}...")
|
||||
else:
|
||||
summary_parts.append(f"[{role}]: {content[:150]}...")
|
||||
|
||||
if summary_parts:
|
||||
focus_hint = f"\nFocus on: {focus_topic}" if focus_topic else ""
|
||||
|
||||
if self._used_previous_summary:
|
||||
summary = (
|
||||
f"[Iterative conversation summary - update with new activity{focus_hint}]\n"
|
||||
+ "\n".join(summary_parts[:15])
|
||||
)
|
||||
else:
|
||||
summary = (
|
||||
f"[Compacted conversation history - summarize key points{focus_hint}]\n"
|
||||
+ "\n".join(summary_parts[:10])
|
||||
)
|
||||
|
||||
# Apply anti-injection prefix if configured
|
||||
if hasattr(self.config, 'compaction_prefix') and self.config.compaction_prefix:
|
||||
summary = f"{self.config.compaction_prefix}\n\n{summary}"
|
||||
|
||||
summary_msg = {
|
||||
"role": "system",
|
||||
"content": summary,
|
||||
"_compacted": True,
|
||||
"_original_count": len(to_summarize),
|
||||
"_iterative": self._used_previous_summary,
|
||||
"_focus_topic": focus_topic
|
||||
}
|
||||
result.append(summary_msg)
|
||||
|
||||
# Update tracking for future iterative summarization
|
||||
self._previous_summary = summary
|
||||
self._previous_summary_global_idx = len(messages)
|
||||
|
||||
result.extend(recent)
|
||||
return result
|
||||
|
||||
async def _llm_summarize_async(self, messages: List[Dict[str, Any]], focus_topic: str = "") -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Use LLM to intelligently summarize older messages.
|
||||
|
||||
This method invokes the agent's LLM to create a meaningful summary
|
||||
that preserves key facts, identifiers, and the user's intent.
|
||||
|
||||
Args:
|
||||
messages: List of messages to compact
|
||||
focus_topic: Optional topic to focus on during summarization
|
||||
"""
|
||||
self._used_previous_summary = False
|
||||
result = []
|
||||
|
||||
# Keep system messages
|
||||
system_msgs = [m for m in messages if m.get("role") == "system"]
|
||||
other_msgs = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
result.extend(system_msgs)
|
||||
|
||||
# Keep recent messages
|
||||
recent = other_msgs[-self.preserve_recent:]
|
||||
|
||||
# Determine what to summarize based on iterative settings
|
||||
total_original_messages = len(messages)
|
||||
if (self.config.enable_iterative_summary and
|
||||
self._previous_summary and
|
||||
total_original_messages > self._previous_summary_global_idx):
|
||||
# Iterative: summarize only new messages since previous summary
|
||||
messages_since_summary = total_original_messages - self._previous_summary_global_idx
|
||||
new_older_messages = max(0, messages_since_summary - self.preserve_recent)
|
||||
if new_older_messages > 0:
|
||||
older = other_msgs[-messages_since_summary:-self.preserve_recent]
|
||||
else:
|
||||
older = []
|
||||
self._used_previous_summary = True
|
||||
else:
|
||||
# Fresh summary: summarize all older messages
|
||||
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
|
||||
self._used_previous_summary = False
|
||||
|
||||
if older and self.llm_summarize_fn:
|
||||
try:
|
||||
# Format messages for summarization
|
||||
history_text = self._format_messages_for_summary(older)
|
||||
# Format messages for summarization using protocol
|
||||
if self.message_formatter:
|
||||
history_text = self.message_formatter.format_for_summary(older)
|
||||
else:
|
||||
# Fallback to basic formatting if no protocol implementation
|
||||
history_text = str(older)
|
||||
|
||||
# Create summarization prompt that preserves important information
|
||||
focus_hint = f" Focus especially on: {focus_topic}." if focus_topic else ""
|
||||
iterative_hint = ""
|
||||
if self._used_previous_summary and self._previous_summary:
|
||||
iterative_hint = f"\n\n[Previous Summary]: {self._previous_summary}\n[New Activity to Add]:"
|
||||
|
||||
prompt = (
|
||||
"Summarise the following conversation history. Preserve verbatim: "
|
||||
"all file paths, IDs, hashes, URLs, task references, error messages, "
|
||||
"tool outputs, and the user's requests. Be concise but complete. "
|
||||
"Focus on facts and actions taken, not general conversation.\n\n"
|
||||
f"{history_text}"
|
||||
f"Focus on facts and actions taken, not general conversation.{focus_hint}"
|
||||
f"{iterative_hint}\n\n{history_text}"
|
||||
)
|
||||
|
||||
# Call the LLM for summarization
|
||||
summary = await self.llm_summarize_fn(prompt)
|
||||
|
||||
# Add the LLM-generated summary as a system message
|
||||
summary_content = summary
|
||||
if self._used_previous_summary:
|
||||
summary_content = f"[Iterative conversation summary]\n{summary}"
|
||||
else:
|
||||
summary_content = f"[Previous conversation summary]\n{summary}"
|
||||
|
||||
result.append({
|
||||
"role": "system",
|
||||
"content": f"[Previous conversation summary]\n{summary}",
|
||||
"content": summary_content,
|
||||
"_compacted": True,
|
||||
"_original_count": len(older),
|
||||
"_llm_generated": True,
|
||||
"_iterative": self._used_previous_summary,
|
||||
"_focus_topic": focus_topic
|
||||
})
|
||||
|
||||
# Update tracking for future iterative summarization
|
||||
self._previous_summary = summary
|
||||
self._previous_summary_global_idx = len(messages)
|
||||
except Exception as e:
|
||||
# Fallback to naive summarization if LLM call fails
|
||||
import logging
|
||||
@@ -506,133 +688,41 @@ class ContextCompactor:
|
||||
"content": summary,
|
||||
"_compacted": True,
|
||||
"_original_count": len(older),
|
||||
"_fallback": True,
|
||||
})
|
||||
|
||||
result.extend(recent)
|
||||
return result
|
||||
|
||||
def _format_messages_for_summary(self, messages: List[Dict[str, Any]]) -> str:
|
||||
"""Format messages for LLM summarization."""
|
||||
formatted = []
|
||||
for i, msg in enumerate(messages):
|
||||
role = msg.get("role", "unknown")
|
||||
content = msg.get("content", "")
|
||||
|
||||
# Handle tool calls
|
||||
if msg.get("tool_calls"):
|
||||
tool_calls = msg.get("tool_calls", [])
|
||||
tool_names = [tc.get("function", {}).get("name", "unknown") for tc in tool_calls]
|
||||
formatted.append(f"{i+1}. {role}: Called tools: {', '.join(tool_names)}")
|
||||
|
||||
# Handle content
|
||||
if isinstance(content, str) and content.strip():
|
||||
# Truncate very long content but preserve structure
|
||||
if len(content) > 1000:
|
||||
content = content[:800] + "...[truncated]..." + content[-200:]
|
||||
formatted.append(f"{i+1}. {role}: {content}")
|
||||
elif isinstance(content, list):
|
||||
# Handle multi-part content
|
||||
parts = []
|
||||
for part in content[:3]: # Limit to first 3 parts
|
||||
if isinstance(part, dict):
|
||||
text = part.get("text", "")
|
||||
if text:
|
||||
parts.append(text[:200] + "..." if len(text) > 200 else text)
|
||||
if parts:
|
||||
formatted.append(f"{i+1}. {role}: {' | '.join(parts)}")
|
||||
|
||||
# Handle tool results with IDs
|
||||
if msg.get("tool_call_id"):
|
||||
tool_id = msg.get("tool_call_id", "")
|
||||
formatted.append(f"{i+1}. {role} (tool {tool_id}): {str(content)[:500]}...")
|
||||
"""
|
||||
Delegate message formatting to protocol implementation.
|
||||
|
||||
return "\n".join(formatted)
|
||||
If no formatter is provided, returns basic string representation.
|
||||
"""
|
||||
if self.message_formatter:
|
||||
return self.message_formatter.format_for_summary(messages)
|
||||
return str(messages)
|
||||
|
||||
def _build_structured_summary(self, messages: List[Dict[str, Any]]) -> str:
|
||||
"""
|
||||
Build a structured summary using the configured template.
|
||||
Delegate structured summary building to protocol implementation.
|
||||
|
||||
Args:
|
||||
messages: Messages to summarize
|
||||
|
||||
Returns:
|
||||
Structured summary string
|
||||
If no summary builder is provided, returns basic summary.
|
||||
"""
|
||||
# Extract information from messages
|
||||
active_task = "No specific task identified"
|
||||
completed = []
|
||||
in_progress = []
|
||||
pending = []
|
||||
files = set()
|
||||
remaining = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if not isinstance(content, str):
|
||||
continue
|
||||
|
||||
content_lower = content.lower()
|
||||
|
||||
# Extract file paths
|
||||
file_matches = re.findall(r'[\w/\.\-]+\.[a-zA-Z]{1,4}', content)
|
||||
files.update(file_matches[:5]) # Limit to 5 files
|
||||
|
||||
# Categorize content based on keywords and role
|
||||
content_snippet = content[:200] + ("..." if len(content) > 200 else "")
|
||||
|
||||
if role == "user":
|
||||
if any(word in content_lower for word in ["please", "can you", "help"]) or re.search(r'\bdo\b', content_lower):
|
||||
active_task = content_snippet
|
||||
elif "?" in content:
|
||||
pending.append(content_snippet)
|
||||
|
||||
elif role == "assistant":
|
||||
if any(word in content_lower for word in ["completed", "done", "finished"]):
|
||||
completed.append(content_snippet)
|
||||
elif any(word in content_lower for word in ["working", "processing", "analyzing"]):
|
||||
in_progress.append(content_snippet)
|
||||
elif any(word in content_lower for word in ["will", "plan to", "next"]):
|
||||
remaining.append(content_snippet)
|
||||
|
||||
elif role == "tool":
|
||||
# Tool results are generally completed actions
|
||||
tool_name = msg.get("name", "unknown")
|
||||
completed.append(f"Tool {tool_name}: {content_snippet}")
|
||||
|
||||
# Format using template
|
||||
return SUMMARY_TEMPLATE.format(
|
||||
active_task=active_task,
|
||||
completed="\n".join(f"- {item}" for item in completed[:3]) or "None identified",
|
||||
in_progress="\n".join(f"- {item}" for item in in_progress[:3]) or "None identified",
|
||||
pending="\n".join(f"- {item}" for item in pending[:3]) or "None identified",
|
||||
files=", ".join(list(files)[:5]) or "None mentioned",
|
||||
remaining="\n".join(f"- {item}" for item in remaining[:3]) or "None identified"
|
||||
)
|
||||
if self.summary_builder:
|
||||
return self.summary_builder.build_structured_summary(messages)
|
||||
return f"Summary of {len(messages)} messages"
|
||||
|
||||
def _merge_summaries(self, previous: str, current: str) -> str:
|
||||
"""
|
||||
Merge previous summary with current one for iterative updates.
|
||||
Delegate summary merging to protocol implementation.
|
||||
|
||||
Args:
|
||||
previous: Previous summary content
|
||||
current: Current summary content
|
||||
|
||||
Returns:
|
||||
Merged summary
|
||||
If no summary builder is provided, returns current summary only.
|
||||
"""
|
||||
# Simple merge strategy: prioritize current but preserve unique info from previous
|
||||
if not previous:
|
||||
return current
|
||||
|
||||
# Preserve context but avoid excessive growth
|
||||
MAX_PREVIOUS_LENGTH = 500
|
||||
if len(previous) > MAX_PREVIOUS_LENGTH:
|
||||
truncated_previous = previous[:MAX_PREVIOUS_LENGTH]
|
||||
return f"{current}\n\n[Previous context]: {truncated_previous}..."
|
||||
else:
|
||||
return f"{current}\n\n[Previous context]: {previous}"
|
||||
if self.summary_builder:
|
||||
return self.summary_builder.merge_summaries(previous, current)
|
||||
return current
|
||||
|
||||
def get_stats(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""Get statistics about messages."""
|
||||
|
||||
@@ -50,3 +50,28 @@ class CompactionConfig:
|
||||
compaction_prefix: str = COMPACTION_PREFIX # Override for custom framing
|
||||
structured_template: bool = True # Emit Active Task / Remaining Work sections
|
||||
iterative_update: bool = True # Merge previous summary on re-compaction
|
||||
|
||||
# Anti-thrashing protection
|
||||
min_savings_pct: float = 10.0 # Skip if projected saving < 10% (0-100 scale)
|
||||
max_consecutive_low_savings: int = 2 # Abort after N low-savings attempts
|
||||
|
||||
# Tool result optimization
|
||||
tool_prune_before_summarise: bool = True # Deduplicate tool results before summarization
|
||||
max_tool_result_size: int = 500 # Maximum size for tool results before pruning
|
||||
|
||||
# Iterative summarization
|
||||
enable_iterative_summary: bool = True # Build on previous summaries instead of starting from scratch
|
||||
|
||||
def __post_init__(self):
|
||||
"""Validate and normalize configuration values."""
|
||||
# Normalize min_savings_pct to 0-100 scale if provided as ratio
|
||||
if self.min_savings_pct < 1.0:
|
||||
self.min_savings_pct *= 100.0
|
||||
|
||||
# Validate ranges
|
||||
if not 0.0 <= self.min_savings_pct <= 100.0:
|
||||
raise ValueError("min_savings_pct must be between 0 and 100")
|
||||
if self.max_consecutive_low_savings < 0:
|
||||
raise ValueError("max_consecutive_low_savings must be >= 0")
|
||||
if self.max_tool_result_size <= 0:
|
||||
raise ValueError("max_tool_result_size must be > 0")
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
Compaction Protocols for PraisonAI Agents.
|
||||
|
||||
Protocol-driven interfaces for context compaction functionality.
|
||||
"""
|
||||
|
||||
from typing import Protocol, List, Dict, Any, Tuple, Optional
|
||||
|
||||
|
||||
class ToolResultPrunerProtocol(Protocol):
|
||||
"""Protocol for pruning and deduplicating tool results."""
|
||||
|
||||
def prune(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
max_tool_result_size: int = 500
|
||||
) -> Tuple[List[Dict[str, Any]], int]:
|
||||
"""
|
||||
Prune and deduplicate tool results to reduce token usage.
|
||||
|
||||
Args:
|
||||
messages: List of messages to process
|
||||
max_tool_result_size: Maximum size for tool results before truncation
|
||||
|
||||
Returns:
|
||||
Tuple of (processed messages, number of tool results pruned)
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class MessageFormatterProtocol(Protocol):
|
||||
"""Protocol for formatting messages for summarization."""
|
||||
|
||||
def format_for_summary(self, messages: List[Dict[str, Any]]) -> str:
|
||||
"""
|
||||
Format messages for LLM summarization.
|
||||
|
||||
Args:
|
||||
messages: Messages to format
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for summarization
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class SummaryBuilderProtocol(Protocol):
|
||||
"""Protocol for building structured summaries."""
|
||||
|
||||
def build_structured_summary(self, messages: List[Dict[str, Any]]) -> str:
|
||||
"""
|
||||
Build a structured summary from messages.
|
||||
|
||||
Args:
|
||||
messages: Messages to summarize
|
||||
|
||||
Returns:
|
||||
Structured summary string
|
||||
"""
|
||||
...
|
||||
|
||||
def merge_summaries(self, previous: str, current: str) -> str:
|
||||
"""
|
||||
Merge previous summary with current one for iterative updates.
|
||||
|
||||
Args:
|
||||
previous: Previous summary content
|
||||
current: Current summary content
|
||||
|
||||
Returns:
|
||||
Merged summary
|
||||
"""
|
||||
...
|
||||
@@ -18,6 +18,16 @@ class CompactionResult:
|
||||
strategy_used: CompactionStrategy
|
||||
summary: str = ""
|
||||
|
||||
# Anti-thrashing and optimization info
|
||||
savings_pct: float = field(default=0.0, init=False) # Percentage of tokens saved (computed)
|
||||
tool_results_pruned: int = 0 # Number of tool results deduplicated/pruned
|
||||
previous_summary_reused: bool = False # Whether iterative summary was used
|
||||
was_skipped_due_to_low_savings: bool = False # Whether compaction was skipped due to anti-thrashing
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize computed fields after construction."""
|
||||
self.calculate_savings_pct()
|
||||
|
||||
@property
|
||||
def tokens_saved(self) -> int:
|
||||
"""Get number of tokens saved."""
|
||||
@@ -30,6 +40,13 @@ class CompactionResult:
|
||||
return 0.0
|
||||
return self.compacted_tokens / self.original_tokens
|
||||
|
||||
def calculate_savings_pct(self) -> None:
|
||||
"""Calculate and set the savings percentage."""
|
||||
if self.original_tokens > 0:
|
||||
self.savings_pct = (self.tokens_saved / self.original_tokens) * 100.0
|
||||
else:
|
||||
self.savings_pct = 0.0
|
||||
|
||||
@property
|
||||
def was_compacted(self) -> bool:
|
||||
"""Check if compaction occurred."""
|
||||
@@ -45,5 +62,9 @@ class CompactionResult:
|
||||
"messages_kept": self.messages_kept,
|
||||
"strategy_used": self.strategy_used.value,
|
||||
"compression_ratio": self.compression_ratio,
|
||||
"savings_pct": self.savings_pct,
|
||||
"tool_results_pruned": self.tool_results_pruned,
|
||||
"previous_summary_reused": self.previous_summary_reused,
|
||||
"was_skipped_due_to_low_savings": self.was_skipped_due_to_low_savings,
|
||||
"summary": self.summary
|
||||
}
|
||||
@@ -774,29 +774,39 @@ class ExecutionConfig:
|
||||
self.context_compaction
|
||||
)
|
||||
|
||||
# Emit deprecation warning for default behavior change
|
||||
# Only warn if context_compaction was not explicitly set (i.e., using the default False)
|
||||
import warnings
|
||||
import inspect
|
||||
frame = inspect.currentframe()
|
||||
try:
|
||||
# Check if this is being called from user code (not internal)
|
||||
caller_frame = frame.f_back
|
||||
if (caller_frame and
|
||||
not any(path_part in caller_frame.f_code.co_filename
|
||||
for path_part in ['praisonaiagents', 'test_', '__pycache__'])):
|
||||
# This is being called from user code
|
||||
if self.context_compaction is False:
|
||||
warnings.warn(
|
||||
"ExecutionConfig.context_compaction will default to True in the next "
|
||||
"release for proactive context overflow protection. To disable, explicitly "
|
||||
"set context_compaction=False. To use the new default early, set "
|
||||
"context_compaction=True.",
|
||||
DeprecationWarning,
|
||||
stacklevel=3
|
||||
)
|
||||
finally:
|
||||
del frame
|
||||
# Emit deprecation warning once per process for default behavior change.
|
||||
# Walk the stack once — dataclass __init__ uses caller filename "<string>".
|
||||
if self.context_compaction is False:
|
||||
if getattr(ExecutionConfig, '_context_compaction_internal', False):
|
||||
return
|
||||
if getattr(ExecutionConfig, '_context_compaction_warned', False):
|
||||
return
|
||||
import warnings
|
||||
import inspect
|
||||
_INTERNAL = ('praisonaiagents', 'test_', '__pycache__')
|
||||
frame = inspect.currentframe()
|
||||
try:
|
||||
caller = frame.f_back if frame else None
|
||||
is_internal = False
|
||||
while caller is not None:
|
||||
if any(p in caller.f_code.co_filename for p in _INTERNAL):
|
||||
is_internal = True
|
||||
break
|
||||
caller = caller.f_back
|
||||
if is_internal:
|
||||
ExecutionConfig._context_compaction_internal = True
|
||||
return
|
||||
warnings.warn(
|
||||
"ExecutionConfig.context_compaction will default to True in the next "
|
||||
"release for proactive context overflow protection. To disable, explicitly "
|
||||
"set context_compaction=False. To use the new default early, set "
|
||||
"context_compaction=True.",
|
||||
DeprecationWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
ExecutionConfig._context_compaction_warned = True
|
||||
finally:
|
||||
del frame
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
@@ -1089,8 +1099,8 @@ class MultiAgentOutputConfig:
|
||||
# Verbosity level (0=silent, 1=minimal, 2+=verbose)
|
||||
verbose: int = 0
|
||||
|
||||
# Streaming output
|
||||
stream: bool = True
|
||||
# Streaming output - False by default for sync multi-agent compatibility
|
||||
stream: bool = False
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
|
||||
@@ -83,6 +83,7 @@ class ConfigValidationError(Exception):
|
||||
# Global config cache (lazy loaded)
|
||||
_config_cache: Optional[Dict[str, Any]] = None
|
||||
_config_loaded: bool = False
|
||||
_defaults_has_values: Optional[bool] = None
|
||||
|
||||
|
||||
def _default_project_plugins_dir():
|
||||
@@ -418,11 +419,27 @@ def get_enabled_plugins() -> Optional[List[str]]:
|
||||
return None # All plugins enabled
|
||||
|
||||
|
||||
def _defaults_has_any_values() -> bool:
|
||||
"""Return True if config file defines any defaults (cached after first check)."""
|
||||
global _defaults_has_values
|
||||
if _defaults_has_values is None:
|
||||
defaults = get_defaults_config()
|
||||
_defaults_has_values = any(
|
||||
getattr(defaults, key, None) is not None
|
||||
for key in (
|
||||
"model", "base_url", "memory", "knowledge", "planning",
|
||||
"reflection", "web", "output", "execution", "caching", "autonomy",
|
||||
)
|
||||
)
|
||||
return _defaults_has_values
|
||||
|
||||
|
||||
def clear_config_cache() -> None:
|
||||
"""Clear the config cache (for testing)."""
|
||||
global _config_cache, _config_loaded
|
||||
global _config_cache, _config_loaded, _defaults_has_values
|
||||
_config_cache = None
|
||||
_config_loaded = False
|
||||
_defaults_has_values = None
|
||||
|
||||
|
||||
def validate_config(config: Dict[str, Any], raise_on_error: bool = False) -> List[str]:
|
||||
@@ -603,6 +620,8 @@ def apply_config_defaults(
|
||||
if explicit_value is not None:
|
||||
# If True and we have config defaults, merge them
|
||||
if explicit_value is True:
|
||||
if not _defaults_has_any_values():
|
||||
return explicit_value
|
||||
config_defaults = get_default(param_name)
|
||||
if config_defaults and isinstance(config_defaults, dict):
|
||||
if config_class:
|
||||
@@ -612,6 +631,10 @@ def apply_config_defaults(
|
||||
return config_defaults
|
||||
return explicit_value
|
||||
|
||||
# Fast path: no config defaults file values — skip lookups
|
||||
if not _defaults_has_any_values():
|
||||
return None
|
||||
|
||||
# Check if config has defaults for this param
|
||||
config_defaults = get_default(param_name)
|
||||
if config_defaults is None:
|
||||
|
||||
@@ -322,8 +322,8 @@ CACHING_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
# =============================================================================
|
||||
|
||||
MULTI_AGENT_OUTPUT_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
"verbose": {"verbose": 2, "stream": True},
|
||||
"minimal": {"verbose": 1, "stream": True},
|
||||
"verbose": {"verbose": 2, "stream": False},
|
||||
"minimal": {"verbose": 1, "stream": False},
|
||||
"silent": {"verbose": 0, "stream": False},
|
||||
}
|
||||
|
||||
|
||||
@@ -24,6 +24,7 @@ import xml.etree.ElementTree as ET
|
||||
from ..errors import AgentErrorKind, FailoverDecision, IdleTimeoutBreaker
|
||||
# Gap 2: Tool call execution imports
|
||||
from ..tools.call_executor import ToolCall, create_tool_call_executor
|
||||
from ..tools.schema import annotation_to_json_schema, get_parameter_requirements
|
||||
# Display functions - lazy loaded to avoid importing rich at startup
|
||||
# These are only needed when output=verbose
|
||||
_display_module = None
|
||||
@@ -5886,34 +5887,74 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
function_name = function_or_name
|
||||
logging.debug(f"Attempting to generate tool definition for: {function_name}")
|
||||
|
||||
# First try to get the tool definition if it exists
|
||||
tool_def_name = f"{function_name}_definition"
|
||||
tool_def = globals().get(tool_def_name)
|
||||
logging.debug(f"Looking for {tool_def_name} in globals: {tool_def is not None}")
|
||||
|
||||
if not tool_def:
|
||||
import __main__
|
||||
tool_def = getattr(__main__, tool_def_name, None)
|
||||
logging.debug(f"Looking for {tool_def_name} in __main__: {tool_def is not None}")
|
||||
|
||||
if tool_def:
|
||||
logging.debug(f"Found tool definition: {tool_def}")
|
||||
return tool_def
|
||||
|
||||
# Try to find the function
|
||||
func = globals().get(function_name)
|
||||
logging.debug(f"Looking for {function_name} in globals: {func is not None}")
|
||||
|
||||
if not func:
|
||||
import __main__
|
||||
func = getattr(__main__, function_name, None)
|
||||
logging.debug(f"Looking for {function_name} in __main__: {func is not None}")
|
||||
|
||||
if not func or not callable(func):
|
||||
logging.debug(f"Function {function_name} not found or not callable")
|
||||
return None
|
||||
# First try to get from the tool registry (preferred method)
|
||||
try:
|
||||
from ..tools.registry import get_registry
|
||||
registry = get_registry()
|
||||
tool = registry.get(function_name)
|
||||
if tool:
|
||||
logging.debug(f"Found tool in registry: {function_name}")
|
||||
# If it's a BaseTool, get its schema
|
||||
if hasattr(tool, 'get_schema'):
|
||||
tool_def = tool.get_schema()
|
||||
# Apply same normalization as the callable path
|
||||
if (
|
||||
isinstance(tool_def, dict)
|
||||
and isinstance(tool_def.get("function"), dict)
|
||||
and isinstance(tool_def["function"].get("parameters"), dict)
|
||||
):
|
||||
tool_def = tool_def.copy()
|
||||
tool_def["function"] = tool_def["function"].copy()
|
||||
tool_def["function"]["parameters"] = self._fix_array_schemas(
|
||||
tool_def["function"]["parameters"]
|
||||
)
|
||||
return tool_def
|
||||
# If it's a callable, use it as func
|
||||
if callable(tool):
|
||||
func = tool
|
||||
else:
|
||||
logging.debug(f"Tool {function_name} in registry is not callable")
|
||||
return None
|
||||
else:
|
||||
logging.debug(f"Tool {function_name} not found in registry, falling back to globals/__main__")
|
||||
# Fall back to globals and __main__ for backward compatibility
|
||||
tool_def_name = f"{function_name}_definition"
|
||||
tool_def = globals().get(tool_def_name)
|
||||
if not tool_def:
|
||||
import __main__
|
||||
tool_def = getattr(__main__, tool_def_name, None)
|
||||
if tool_def:
|
||||
return tool_def
|
||||
|
||||
func = globals().get(function_name)
|
||||
if not func:
|
||||
import __main__
|
||||
func = getattr(__main__, function_name, None)
|
||||
if not func or not callable(func):
|
||||
logging.debug(f"Function {function_name} not found or not callable")
|
||||
return None
|
||||
except ImportError:
|
||||
logging.debug("Tool registry not available, falling back to globals/__main__")
|
||||
# Fall back to globals and __main__ when registry unavailable
|
||||
tool_def_name = f"{function_name}_definition"
|
||||
tool_def = globals().get(tool_def_name)
|
||||
if not tool_def:
|
||||
import __main__
|
||||
tool_def = getattr(__main__, tool_def_name, None)
|
||||
if tool_def:
|
||||
return tool_def
|
||||
|
||||
func = globals().get(function_name)
|
||||
if not func:
|
||||
import __main__
|
||||
func = getattr(__main__, function_name, None)
|
||||
if not func or not callable(func):
|
||||
logging.debug(f"Function {function_name} not found or not callable")
|
||||
return None
|
||||
|
||||
import inspect
|
||||
from typing import get_type_hints
|
||||
|
||||
# Handle Langchain and CrewAI tools
|
||||
if inspect.isclass(func) and hasattr(func, 'run') and not hasattr(func, '_run'):
|
||||
original_func = func
|
||||
@@ -5961,26 +6002,26 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
logging.debug(f"Parameter descriptions: {param_descriptions}")
|
||||
|
||||
# Get type hints for proper schema generation
|
||||
try:
|
||||
hints = get_type_hints(func) if getattr(func, "__annotations__", None) else {}
|
||||
except (NameError, TypeError, AttributeError):
|
||||
hints = getattr(func, "__annotations__", {}) or {}
|
||||
|
||||
for name, param in parameters_list:
|
||||
param_type = "string" # Default type
|
||||
if param.annotation != inspect.Parameter.empty:
|
||||
if param.annotation == int:
|
||||
param_type = "integer"
|
||||
elif param.annotation == float:
|
||||
param_type = "number"
|
||||
elif param.annotation == bool:
|
||||
param_type = "boolean"
|
||||
elif param.annotation == list:
|
||||
param_type = "array"
|
||||
elif param.annotation == dict:
|
||||
param_type = "object"
|
||||
# Get type annotation
|
||||
param_type = hints.get(name, param.annotation if param.annotation != inspect.Parameter.empty else str)
|
||||
|
||||
parameters["properties"][name] = {
|
||||
"type": param_type,
|
||||
"description": param_descriptions.get(name, "Parameter description not available")
|
||||
}
|
||||
# Use new schema utility for proper type handling
|
||||
prop_schema = annotation_to_json_schema(param_type)
|
||||
|
||||
if param.default == inspect.Parameter.empty:
|
||||
# Add description from docstring
|
||||
prop_schema["description"] = param_descriptions.get(name, "Parameter description not available")
|
||||
|
||||
parameters["properties"][name] = prop_schema
|
||||
|
||||
# Check if required using improved logic
|
||||
if get_parameter_requirements(sig, name):
|
||||
parameters["required"].append(name)
|
||||
|
||||
logging.debug(f"Generated parameters: {parameters}")
|
||||
|
||||
@@ -246,3 +246,74 @@ class SecurityConfig:
|
||||
def get_bind_address(self) -> str:
|
||||
"""Get recommended bind address based on config."""
|
||||
return "127.0.0.1" if self.bind_localhost_only else "0.0.0.0"
|
||||
|
||||
|
||||
def load_sse_security_from_env():
|
||||
"""Load optional SSE security settings from environment."""
|
||||
import os
|
||||
|
||||
token = os.environ.get("MCP_SSE_AUTH_TOKEN") or os.environ.get("MCP_AUTH_TOKEN")
|
||||
origins_raw = os.environ.get("MCP_SSE_ALLOWED_ORIGINS")
|
||||
if not token and not origins_raw:
|
||||
return None, None
|
||||
|
||||
origins = (
|
||||
[origin.strip() for origin in origins_raw.split(",") if origin.strip()]
|
||||
if origins_raw
|
||||
else None
|
||||
)
|
||||
config = SecurityConfig(
|
||||
require_auth=bool(token),
|
||||
validate_origin=origins is not None,
|
||||
allowed_origins=origins or SecurityConfig().allowed_origins,
|
||||
)
|
||||
return config, token
|
||||
|
||||
|
||||
def build_sse_security_app(app, security: Optional["SecurityConfig"] = None):
|
||||
"""Wrap a Starlette app with SSE security middleware when configured."""
|
||||
config = security
|
||||
token = None
|
||||
if config is None:
|
||||
config, token = load_sse_security_from_env()
|
||||
elif config.require_auth:
|
||||
import os
|
||||
|
||||
token = os.environ.get("MCP_SSE_AUTH_TOKEN") or os.environ.get("MCP_AUTH_TOKEN")
|
||||
if config is None or (not config.require_auth and not config.validate_origin):
|
||||
return app
|
||||
|
||||
class _SSESecurityMiddleware:
|
||||
def __init__(self, inner, sec: SecurityConfig, expected_token: Optional[str]):
|
||||
self._inner = inner
|
||||
self._sec = sec
|
||||
self._token = expected_token
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope.get("type") != "http":
|
||||
await self._inner(scope, receive, send)
|
||||
return
|
||||
|
||||
headers = {k.decode().lower(): v.decode() for k, v in scope.get("headers", [])}
|
||||
|
||||
if self._sec.require_auth and self._token:
|
||||
auth = headers.get("authorization", "")
|
||||
if not auth.startswith("Bearer ") or auth[7:] != self._token:
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
response = JSONResponse({"error": "Unauthorized"}, status_code=401)
|
||||
await response(scope, receive, send)
|
||||
return
|
||||
|
||||
if self._sec.validate_origin:
|
||||
origin = headers.get("origin") or None
|
||||
if not self._sec.is_origin_allowed(origin):
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
response = JSONResponse({"error": "Origin not allowed"}, status_code=403)
|
||||
await response(scope, receive, send)
|
||||
return
|
||||
|
||||
await self._inner(scope, receive, send)
|
||||
|
||||
return _SSESecurityMiddleware(app, config, token)
|
||||
@@ -242,12 +242,18 @@ class ToolsMCPServer:
|
||||
|
||||
mcp.run()
|
||||
|
||||
def run_sse(self, host: str = "0.0.0.0", port: int = 8080) -> None:
|
||||
def run_sse(
|
||||
self,
|
||||
host: str = "0.0.0.0",
|
||||
port: int = 8080,
|
||||
security: "Optional[SecurityConfig]" = None,
|
||||
) -> None:
|
||||
"""Run the MCP server using SSE transport.
|
||||
|
||||
Args:
|
||||
host: Host to bind to
|
||||
port: Port to listen on
|
||||
security: Optional security config; when unset, reads MCP_SSE_* env vars
|
||||
"""
|
||||
try:
|
||||
import uvicorn
|
||||
@@ -260,6 +266,8 @@ class ToolsMCPServer:
|
||||
f"SSE dependencies not available: {e}. "
|
||||
"Install with: pip install uvicorn starlette"
|
||||
)
|
||||
|
||||
from .mcp_security import SecurityConfig, build_sse_security_app
|
||||
|
||||
mcp = self.get_fastmcp()
|
||||
|
||||
@@ -286,7 +294,8 @@ class ToolsMCPServer:
|
||||
Mount(messages_path, app=sse_transport.handle_post_message),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
app = build_sse_security_app(app, security)
|
||||
print(f"🚀 Starting MCP server '{self.name}' with SSE transport")
|
||||
print(f"📡 SSE endpoint: http://{host}:{port}{sse_path}")
|
||||
print(f"🛠️ Available tools: {', '.join(self.get_tool_names())}")
|
||||
|
||||
@@ -115,6 +115,16 @@ class FileMemory:
|
||||
"importance_threshold": 0.7, # Min importance for long-term
|
||||
"auto_promote": True, # Auto-promote important short-term to long-term
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _sanitise_user_id(user_id: str) -> str:
|
||||
"""Reject path traversal in user_id before using it as a directory name."""
|
||||
if not user_id or not isinstance(user_id, str):
|
||||
return "default"
|
||||
if ".." in user_id or "/" in user_id or "\\" in user_id:
|
||||
raise ValueError("user_id must not contain path separators or parent references")
|
||||
safe = user_id.strip()
|
||||
return safe or "default"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -132,7 +142,7 @@ class FileMemory:
|
||||
config: Configuration overrides
|
||||
verbose: Verbosity level (0=quiet, 1=info, 2+=debug)
|
||||
"""
|
||||
self.user_id = user_id
|
||||
self.user_id = self._sanitise_user_id(user_id)
|
||||
self.verbose = verbose
|
||||
|
||||
# Set up paths
|
||||
|
||||
@@ -29,6 +29,11 @@ logger = get_logger(__name__, extra_data={"subsystem": "memory"})
|
||||
TRACE_LEVEL = 5
|
||||
logging.addLevelName(TRACE_LEVEL, 'TRACE')
|
||||
|
||||
# Cache boundary constants for prompt prefix caching optimization
|
||||
# NOTE: This is a placeholder - current implementation does not provide
|
||||
# functional prompt caching as it requires structured API metadata
|
||||
CACHE_BOUNDARY = ""
|
||||
|
||||
|
||||
|
||||
class Memory(StorageMixin, SearchMixin, MemoryCoreMixin):
|
||||
@@ -1704,6 +1709,41 @@ class Memory(StorageMixin, SearchMixin, MemoryCoreMixin):
|
||||
entities = self.search_entity(q, limit=max_items)
|
||||
user_mem = self.search_user_memory(user_id, q, limit=max_items) if user_id else []
|
||||
|
||||
# Apply stable sorting to ensure deterministic order for prompt caching
|
||||
# Sort by timestamp descending, then by content hash for stable ordering
|
||||
def _sort_memory_results(results: List[Any]) -> List[Any]:
|
||||
"""Stable sort so identical query sets produce identical context strings."""
|
||||
import hashlib
|
||||
if not results:
|
||||
return results
|
||||
|
||||
def sort_key(r):
|
||||
if not isinstance(r, dict):
|
||||
return (0.0, hashlib.sha256(str(r).encode()).hexdigest())
|
||||
|
||||
# Prefer created_at over timestamp for better determinism
|
||||
timestamp = r.get("created_at") or r.get("timestamp") or 0
|
||||
if isinstance(timestamp, str):
|
||||
try:
|
||||
from datetime import datetime
|
||||
timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00")).timestamp()
|
||||
except ValueError:
|
||||
timestamp = 0.0
|
||||
timestamp = -(float(timestamp)) # Negative for descending order
|
||||
|
||||
# Use full content hash for true deterministic ordering
|
||||
content = r.get("text", "") if isinstance(r, dict) else str(r)
|
||||
content_hash = hashlib.sha256(content.encode()).hexdigest()
|
||||
return (timestamp, content_hash)
|
||||
|
||||
return sorted(results, key=sort_key)
|
||||
|
||||
# Apply stable sorting to all memory results
|
||||
short_term = _sort_memory_results(short_term)
|
||||
long_term = _sort_memory_results(long_term)
|
||||
entities = _sort_memory_results(entities)
|
||||
user_mem = _sort_memory_results(user_mem)
|
||||
|
||||
# Add sections in order of priority
|
||||
add_section("Short-term Memory Context", short_term)
|
||||
add_section("Long-term Memory Context", long_term)
|
||||
@@ -1713,6 +1753,51 @@ class Memory(StorageMixin, SearchMixin, MemoryCoreMixin):
|
||||
|
||||
return "\n".join(lines) if lines else ""
|
||||
|
||||
def build_cache_optimized_context(
|
||||
self,
|
||||
task_descr: str,
|
||||
user_id: Optional[str] = None,
|
||||
additional: str = "",
|
||||
max_items: int = 3,
|
||||
include_cache_boundary: bool = True,
|
||||
include_in_output: Optional[bool] = None
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
Build context with cache boundary markers for prompt prefix caching optimization.
|
||||
|
||||
Returns a dictionary with 'stable_prefix' and 'cache_boundary' keys.
|
||||
The stable prefix contains deterministically ordered content that should
|
||||
remain constant between turns for effective prompt caching.
|
||||
|
||||
Args:
|
||||
task_descr: Task description for memory search
|
||||
user_id: Optional user ID for personalized memory
|
||||
additional: Additional context to include in search
|
||||
max_items: Maximum items per memory category
|
||||
include_cache_boundary: Whether to include cache boundary marker
|
||||
include_in_output: Whether to include memory content in output
|
||||
|
||||
Returns:
|
||||
Dict with 'stable_prefix' and 'cache_boundary' keys
|
||||
"""
|
||||
# Preserve original behavior: include_in_output=None means debug-level-only
|
||||
# Don't change the default behavior for backward compatibility
|
||||
|
||||
# Build the stable context using existing method
|
||||
stable_context = self.build_context_for_task(
|
||||
task_descr=task_descr,
|
||||
user_id=user_id,
|
||||
additional=additional,
|
||||
max_items=max_items,
|
||||
include_in_output=include_in_output
|
||||
)
|
||||
|
||||
result = {"stable_prefix": stable_context}
|
||||
|
||||
result["cache_boundary"] = CACHE_BOUNDARY if include_cache_boundary else ""
|
||||
|
||||
return result
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Master Reset (Everything)
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
@@ -117,6 +117,16 @@ class AgentServer:
|
||||
self._server_thread: Optional[threading.Thread] = None
|
||||
self._app = None
|
||||
self._server = None
|
||||
|
||||
def _authorise_request(self, request) -> bool:
|
||||
"""Verify bearer token when auth_token is configured."""
|
||||
token = self.config.auth_token
|
||||
if not token:
|
||||
return True
|
||||
auth = request.headers.get("Authorization", "")
|
||||
if auth.startswith("Bearer ") and auth[7:] == token:
|
||||
return True
|
||||
return request.headers.get("X-Auth-Token") == token
|
||||
|
||||
def _create_app(self):
|
||||
"""Create the ASGI application."""
|
||||
@@ -137,6 +147,8 @@ class AgentServer:
|
||||
})
|
||||
|
||||
async def events(request):
|
||||
if not self._authorise_request(request):
|
||||
return JSONResponse({"error": "Unauthorized"}, status_code=401)
|
||||
client_id = str(uuid.uuid4())
|
||||
client = SSEClient(client_id)
|
||||
self._clients[client_id] = client
|
||||
@@ -162,6 +174,8 @@ class AgentServer:
|
||||
)
|
||||
|
||||
async def publish(request):
|
||||
if not self._authorise_request(request):
|
||||
return JSONResponse({"error": "Unauthorized"}, status_code=401)
|
||||
try:
|
||||
data = await request.json()
|
||||
event_type = data.get("type", "message")
|
||||
|
||||
@@ -14,6 +14,10 @@ import time
|
||||
import threading
|
||||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||||
|
||||
from .utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory.memory import Memory
|
||||
from .knowledge.knowledge import Knowledge
|
||||
@@ -504,6 +508,20 @@ class Session:
|
||||
Returns:
|
||||
Formatted context string
|
||||
"""
|
||||
# Use cache-optimized context if available for better prompt caching
|
||||
if hasattr(self.memory, 'build_cache_optimized_context'):
|
||||
try:
|
||||
cache_result = self.memory.build_cache_optimized_context(
|
||||
task_descr=query,
|
||||
user_id=self.user_id,
|
||||
max_items=max_items,
|
||||
include_cache_boundary=False # Don't include boundary for session context
|
||||
)
|
||||
return cache_result.get('stable_prefix', '')
|
||||
except Exception as e:
|
||||
# Fall back to standard context building
|
||||
logger.debug(f"Cache-optimized context failed for session '{self.id or self.user_id}', falling back to standard: {e}")
|
||||
|
||||
return self.memory.build_context_for_task(
|
||||
task_descr=query,
|
||||
user_id=self.user_id,
|
||||
|
||||
@@ -14,6 +14,8 @@ import json
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..memory.memory import Memory
|
||||
from ..knowledge.knowledge import Knowledge
|
||||
@@ -477,6 +479,20 @@ class Session:
|
||||
Returns:
|
||||
Formatted context string
|
||||
"""
|
||||
# Use cache-optimized context if available for better prompt caching
|
||||
if hasattr(self.memory, 'build_cache_optimized_context'):
|
||||
try:
|
||||
cache_result = self.memory.build_cache_optimized_context(
|
||||
task_descr=query,
|
||||
user_id=self.user_id,
|
||||
max_items=max_items,
|
||||
include_cache_boundary=False # Don't include boundary for session context
|
||||
)
|
||||
return cache_result.get('stable_prefix', '')
|
||||
except Exception as e:
|
||||
# Fall back to standard context building
|
||||
logger.debug(f"Cache-optimized context failed for session '{self.session_id}', falling back to standard: {e}")
|
||||
|
||||
return self.memory.build_context_for_task(
|
||||
task_descr=query,
|
||||
user_id=self.user_id,
|
||||
|
||||
@@ -740,10 +740,27 @@ class Task:
|
||||
# Build context for next tasks
|
||||
if self.next_tasks:
|
||||
logger.info(f"Task {self.id}: Building context for next tasks...")
|
||||
context = self.memory.build_context_for_task(
|
||||
task_descr=task_output.raw,
|
||||
max_items=5
|
||||
)
|
||||
# Use cache-optimized context if available for better prompt caching
|
||||
if hasattr(self.memory, 'build_cache_optimized_context'):
|
||||
try:
|
||||
cache_result = self.memory.build_cache_optimized_context(
|
||||
task_descr=task_output.raw,
|
||||
max_items=5,
|
||||
include_cache_boundary=False # Don't include boundary for task context
|
||||
)
|
||||
context = cache_result.get('stable_prefix', '')
|
||||
except Exception as e:
|
||||
# Fall back to standard context building
|
||||
logger.debug(f"Cache-optimized context failed for task '{self.id}', falling back to standard: {e}")
|
||||
context = self.memory.build_context_for_task(
|
||||
task_descr=task_output.raw,
|
||||
max_items=5
|
||||
)
|
||||
else:
|
||||
context = self.memory.build_context_for_task(
|
||||
task_descr=task_output.raw,
|
||||
max_items=5
|
||||
)
|
||||
logger.info(f"Task {self.id}: Built context for next tasks: {len(context)} items")
|
||||
|
||||
logger.info(f"Task {self.id}: Memory operations complete")
|
||||
|
||||
@@ -31,6 +31,8 @@ import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
from typing import Optional, List
|
||||
|
||||
from ..approval import require_approval
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Availability cache for performance (checked once per process)
|
||||
@@ -146,6 +148,7 @@ def ast_grep_search(
|
||||
logger.exception("Unexpected error in ast_grep_search")
|
||||
return f"Error: {e}"
|
||||
|
||||
@require_approval(risk_level="high")
|
||||
def ast_grep_rewrite(
|
||||
pattern: str,
|
||||
replacement: str,
|
||||
|
||||
@@ -17,9 +17,12 @@ Usage:
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Callable, Dict, List, Optional, Type, get_type_hints
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import copy
|
||||
|
||||
from .schema import annotation_to_json_schema, get_parameter_requirements
|
||||
|
||||
|
||||
class ToolValidationError(Exception):
|
||||
"""Raised when a tool fails validation."""
|
||||
@@ -129,18 +132,17 @@ class BaseTool(ABC):
|
||||
|
||||
# Get type hint
|
||||
param_type = hints.get(param_name, Any)
|
||||
json_type = self._python_type_to_json(param_type)
|
||||
|
||||
# Build property schema
|
||||
prop_schema = {"type": json_type}
|
||||
# Use new schema utility for proper type handling
|
||||
prop_schema = annotation_to_json_schema(param_type)
|
||||
|
||||
# Add description from docstring if available
|
||||
# (Could parse docstring for param descriptions)
|
||||
|
||||
schema["properties"][param_name] = prop_schema
|
||||
|
||||
# Check if required (no default value)
|
||||
if param.default is inspect.Parameter.empty:
|
||||
# Check if required using improved logic
|
||||
if get_parameter_requirements(sig, param_name):
|
||||
schema["required"].append(param_name)
|
||||
except Exception as e:
|
||||
logging.debug(f"Could not generate schema for {self.name}: {e}")
|
||||
@@ -149,28 +151,15 @@ class BaseTool(ABC):
|
||||
|
||||
@staticmethod
|
||||
def _python_type_to_json(python_type: Type) -> str:
|
||||
"""Convert Python type to JSON Schema type."""
|
||||
type_map = {
|
||||
str: "string",
|
||||
int: "integer",
|
||||
float: "number",
|
||||
bool: "boolean",
|
||||
list: "array",
|
||||
dict: "object",
|
||||
type(None): "null"
|
||||
}
|
||||
"""Convert Python type to JSON Schema type.
|
||||
|
||||
# Handle Optional, Union, etc.
|
||||
origin = getattr(python_type, '__origin__', None)
|
||||
if origin is not None:
|
||||
# For List[X], return "array"
|
||||
if origin is list:
|
||||
return "array"
|
||||
# For Dict[X, Y], return "object"
|
||||
if origin is dict:
|
||||
return "object"
|
||||
|
||||
return type_map.get(python_type, "string")
|
||||
DEPRECATED: Use annotation_to_json_schema() from schema.py instead.
|
||||
This method is kept for backward compatibility but will be removed.
|
||||
"""
|
||||
# Legacy fallback - delegate to new schema utility and extract type
|
||||
from .schema import annotation_to_json_schema
|
||||
schema = annotation_to_json_schema(python_type)
|
||||
return schema.get("type", "string")
|
||||
|
||||
@abstractmethod
|
||||
def run(self, **kwargs) -> Any:
|
||||
@@ -442,5 +431,49 @@ def validate_tool_schema_consistency(tools: List[Any]) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def get_sorted_tool_schemas(tools: List[Any]) -> List[Dict[str, Any]]:
|
||||
"""Get tool schemas sorted by function name for deterministic ordering.
|
||||
|
||||
This ensures tool schemas are always ordered consistently for prompt caching optimization.
|
||||
|
||||
Args:
|
||||
tools: List of tool objects (BaseTool instances, callables, etc.)
|
||||
|
||||
Returns:
|
||||
List of tool schemas sorted alphabetically by function name
|
||||
|
||||
Raises:
|
||||
ToolValidationError: If validation fails
|
||||
"""
|
||||
if not tools:
|
||||
return []
|
||||
|
||||
# First validate all tools (reuse existing validation logic)
|
||||
validate_tool_schema_consistency(tools)
|
||||
|
||||
from .decorator import get_tool_schema
|
||||
schemas = []
|
||||
|
||||
for tool in tools:
|
||||
# Get schema from different tool types
|
||||
if isinstance(tool, BaseTool):
|
||||
schema = tool.get_schema()
|
||||
elif hasattr(tool, 'get_schema') and callable(getattr(tool, 'get_schema')):
|
||||
schema = tool.get_schema()
|
||||
elif callable(tool):
|
||||
schema = get_tool_schema(tool)
|
||||
else:
|
||||
continue # Skip invalid tools (validation already happened)
|
||||
|
||||
if schema:
|
||||
schemas.append(schema)
|
||||
|
||||
# Sort schemas by function name for deterministic ordering
|
||||
def sort_key(schema):
|
||||
return schema.get("function", {}).get("name", "")
|
||||
|
||||
return sorted(schemas, key=sort_key)
|
||||
|
||||
|
||||
# For backward compatibility - tools can also just be functions
|
||||
# The @tool decorator (in decorator.py) wraps functions into BaseTool instances
|
||||
@@ -31,6 +31,7 @@ import copy
|
||||
from typing import Any, Callable, Dict, Optional, Union, get_type_hints
|
||||
|
||||
from .base import BaseTool
|
||||
from .schema import annotation_to_json_schema, get_parameter_requirements
|
||||
|
||||
# Lazy load injected module functions to reduce import time
|
||||
_injected_module = None
|
||||
@@ -132,12 +133,13 @@ class FunctionTool(BaseTool):
|
||||
if is_injected_type(param_type):
|
||||
continue
|
||||
|
||||
json_type = BaseTool._python_type_to_json(param_type)
|
||||
# Use new schema utility for proper type handling
|
||||
prop_schema = annotation_to_json_schema(param_type)
|
||||
|
||||
schema["properties"][param_name] = {"type": json_type}
|
||||
schema["properties"][param_name] = prop_schema
|
||||
|
||||
# Check if required (no default value)
|
||||
if param.default is inspect.Parameter.empty:
|
||||
# Check if required using improved logic
|
||||
if get_parameter_requirements(sig, param_name):
|
||||
schema["required"].append(param_name)
|
||||
except Exception as e:
|
||||
logging.debug(f"Could not generate schema for {func.__name__}: {e}")
|
||||
@@ -342,11 +344,14 @@ def _schema_from_function(func: Callable) -> Dict[str, Any]:
|
||||
continue
|
||||
|
||||
param_type = hints.get(param_name, Any)
|
||||
json_type = BaseTool._python_type_to_json(param_type)
|
||||
|
||||
properties[param_name] = {"type": json_type}
|
||||
# Use new schema utility for proper type handling
|
||||
prop_schema = annotation_to_json_schema(param_type)
|
||||
|
||||
if param.default is inspect.Parameter.empty:
|
||||
properties[param_name] = prop_schema
|
||||
|
||||
# Check if required using improved logic
|
||||
if get_parameter_requirements(sig, param_name):
|
||||
required.append(param_name)
|
||||
except Exception as e:
|
||||
logging.debug(f"Could not generate schema for {name}: {e}")
|
||||
|
||||
@@ -532,6 +532,14 @@ def _smtp_reply_email(message_id: str, body: str) -> str:
|
||||
logger.error(f"Failed to reply to {message_id}: {e}")
|
||||
return f"Failed to reply: {e}"
|
||||
|
||||
def _imap_sanitize(value: str) -> str:
|
||||
"""Strip characters that break IMAP quoted-string search criteria."""
|
||||
cleaned = value.replace('"', "").replace("\\", "")
|
||||
if any(c in cleaned for c in "\r\n\x00"):
|
||||
raise ValueError("Invalid IMAP search value")
|
||||
return cleaned
|
||||
|
||||
|
||||
def _smtp_search_emails(
|
||||
query: Optional[str] = None,
|
||||
from_addr: Optional[str] = None,
|
||||
@@ -544,11 +552,11 @@ def _smtp_search_emails(
|
||||
# Build IMAP search criteria
|
||||
criteria = []
|
||||
if from_addr:
|
||||
criteria.append(f'FROM "{from_addr}"')
|
||||
criteria.append(f'FROM "{_imap_sanitize(from_addr)}"')
|
||||
if subject:
|
||||
criteria.append(f'SUBJECT "{subject}"')
|
||||
criteria.append(f'SUBJECT "{_imap_sanitize(subject)}"')
|
||||
if query:
|
||||
criteria.append(f'TEXT "{query}"')
|
||||
criteria.append(f'TEXT "{_imap_sanitize(query)}"')
|
||||
if not criteria:
|
||||
criteria.append("ALL")
|
||||
search_str = " ".join(criteria)
|
||||
|
||||
@@ -165,12 +165,17 @@ class MentionsParser:
|
||||
def _process_file_mention(self, file_path: str) -> Optional[str]:
|
||||
"""Process @file:path mention."""
|
||||
try:
|
||||
# Resolve path relative to workspace
|
||||
full_path = self.workspace_path / file_path
|
||||
if not full_path.exists():
|
||||
# Try as absolute path
|
||||
full_path = Path(file_path)
|
||||
|
||||
if ".." in file_path:
|
||||
self._log(f"Rejected path traversal in mention: {file_path}", logging.WARNING)
|
||||
return f"# File: {file_path}\n[Path not allowed]"
|
||||
|
||||
# Resolve only within workspace (no arbitrary absolute paths)
|
||||
full_path = (self.workspace_path / file_path).resolve()
|
||||
workspace_root = self.workspace_path.resolve()
|
||||
if not str(full_path).startswith(str(workspace_root) + os.sep) and full_path != workspace_root:
|
||||
self._log(f"File outside workspace: {file_path}", logging.WARNING)
|
||||
return f"# File: {file_path}\n[Path outside workspace]"
|
||||
|
||||
if not full_path.exists():
|
||||
self._log(f"File not found: {file_path}", logging.WARNING)
|
||||
return f"# File: {file_path}\n[File not found]"
|
||||
|
||||
@@ -57,6 +57,7 @@ _SANDBOX_BLOCKED_CALLS = frozenset({
|
||||
'exec', 'eval', 'compile', '__import__',
|
||||
'open', 'input', 'breakpoint',
|
||||
'setattr', 'delattr', 'dir', 'vars',
|
||||
'format', 'format_map',
|
||||
})
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
"""JSON Schema generation utilities for tool parameters.
|
||||
|
||||
This module provides a shared utility for converting Python type annotations
|
||||
to JSON Schema format. It handles complex types like Optional, Union, Literal,
|
||||
and Enum that were previously collapsed to simple "string" types.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.tools.schema import annotation_to_json_schema
|
||||
|
||||
schema = annotation_to_json_schema(Optional[int])
|
||||
# {"anyOf": [{"type": "integer"}, {"type": "null"}]}
|
||||
|
||||
schema = annotation_to_json_schema(Literal["fast", "deep"])
|
||||
# {"type": "string", "enum": ["fast", "deep"]}
|
||||
"""
|
||||
|
||||
import enum
|
||||
import inspect
|
||||
from typing import Any, Union, Literal, get_origin, get_args
|
||||
|
||||
|
||||
def annotation_to_json_schema(annotation: Any) -> dict:
|
||||
"""Convert Python type annotation to JSON Schema format.
|
||||
|
||||
Handles:
|
||||
- Basic types (str, int, float, bool)
|
||||
- Optional[T] -> {"anyOf": [T schema, {"type": "null"}]}
|
||||
- Union[T1, T2] -> {"anyOf": [T1 schema, T2 schema]}
|
||||
- Literal["a", "b"] -> {"type": "string", "enum": ["a", "b"]}
|
||||
- Enum subclasses -> {"type": "string", "enum": [e.value for e in MyEnum]}
|
||||
- List[T] -> {"type": "array", "items": T schema}
|
||||
- Dict[K, V] -> {"type": "object", "additionalProperties": V schema}
|
||||
|
||||
Args:
|
||||
annotation: Python type annotation to convert
|
||||
|
||||
Returns:
|
||||
dict: JSON Schema representation
|
||||
"""
|
||||
# Handle None type
|
||||
if annotation is type(None):
|
||||
return {"type": "null"}
|
||||
|
||||
# Get typing origin and args
|
||||
origin = get_origin(annotation)
|
||||
args = get_args(annotation)
|
||||
|
||||
# Handle Union types (including Optional[T] which is Union[T, None])
|
||||
if origin is Union:
|
||||
schemas = [annotation_to_json_schema(arg) for arg in args]
|
||||
|
||||
# Check if this is Optional[T] (Union[T, None])
|
||||
null_schemas = [s for s in schemas if s == {"type": "null"}]
|
||||
non_null_schemas = [s for s in schemas if s != {"type": "null"}]
|
||||
|
||||
if null_schemas and len(non_null_schemas) == 1:
|
||||
# This is Optional[T] - simplified format
|
||||
return {"anyOf": [non_null_schemas[0], {"type": "null"}]}
|
||||
else:
|
||||
# General Union type
|
||||
return {"anyOf": schemas}
|
||||
|
||||
# Handle Literal types
|
||||
if origin is Literal or (hasattr(annotation, '__origin__') and
|
||||
str(annotation.__origin__).endswith('Literal')):
|
||||
# Infer type from literal values if all are the same type
|
||||
values = list(args)
|
||||
if values:
|
||||
types_seen = set(type(v) for v in values)
|
||||
if len(types_seen) == 1:
|
||||
single_type = next(iter(types_seen))
|
||||
type_map = {str: "string", int: "integer", float: "number", bool: "boolean"}
|
||||
if single_type in type_map:
|
||||
return {"type": type_map[single_type], "enum": values}
|
||||
# Mixed types or non-primitive - use enum without type constraint
|
||||
return {"enum": values}
|
||||
|
||||
# Handle Enum subclasses
|
||||
if isinstance(annotation, type) and issubclass(annotation, enum.Enum):
|
||||
# Infer type from enum values if all are the same type
|
||||
values = [e.value for e in annotation]
|
||||
if values:
|
||||
types_seen = set(type(v) for v in values)
|
||||
if len(types_seen) == 1:
|
||||
single_type = next(iter(types_seen))
|
||||
type_map = {str: "string", int: "integer", float: "number", bool: "boolean"}
|
||||
if single_type in type_map:
|
||||
return {"type": type_map[single_type], "enum": values}
|
||||
# Mixed types or non-primitive - use enum without type constraint
|
||||
return {"enum": values}
|
||||
|
||||
# Handle List types
|
||||
if origin is list:
|
||||
if args:
|
||||
items_schema = annotation_to_json_schema(args[0])
|
||||
return {"type": "array", "items": items_schema}
|
||||
else:
|
||||
return {"type": "array"}
|
||||
|
||||
# Handle Dict types
|
||||
if origin is dict:
|
||||
if len(args) >= 2:
|
||||
value_schema = annotation_to_json_schema(args[1])
|
||||
return {"type": "object", "additionalProperties": value_schema}
|
||||
else:
|
||||
return {"type": "object"}
|
||||
|
||||
# Handle basic types
|
||||
type_map = {
|
||||
str: "string",
|
||||
int: "integer",
|
||||
float: "number",
|
||||
bool: "boolean",
|
||||
list: "array",
|
||||
dict: "object",
|
||||
type(None): "null"
|
||||
}
|
||||
|
||||
return {"type": type_map.get(annotation, "string")}
|
||||
|
||||
|
||||
def get_parameter_requirements(sig: inspect.Signature, param_name: str) -> bool:
|
||||
"""Determine if a parameter is required based on its signature.
|
||||
|
||||
A parameter is required if it has no default value, unless it's Optional.
|
||||
|
||||
Args:
|
||||
sig: Function signature
|
||||
param_name: Parameter name to check
|
||||
|
||||
Returns:
|
||||
bool: True if parameter is required
|
||||
"""
|
||||
param = sig.parameters.get(param_name)
|
||||
if not param:
|
||||
return False
|
||||
|
||||
# Parameter with default is not required
|
||||
if param.default is not inspect.Parameter.empty:
|
||||
return False
|
||||
|
||||
# Check if the type annotation indicates optional
|
||||
if param.annotation is not inspect.Parameter.empty:
|
||||
origin = get_origin(param.annotation)
|
||||
if origin is Union:
|
||||
args = get_args(param.annotation)
|
||||
# Optional[T] is Union[T, None]
|
||||
if type(None) in args:
|
||||
return False
|
||||
|
||||
# No default and not optional = required
|
||||
return True
|
||||
@@ -58,9 +58,11 @@ def _host_is_blocked(hostname: str) -> bool:
|
||||
pass
|
||||
|
||||
try:
|
||||
return _ip_blocked(ipaddress.ip_address(socket.inet_aton(host)))
|
||||
except OSError:
|
||||
pass
|
||||
for info in socket.getaddrinfo(host, None):
|
||||
if _ip_blocked(ipaddress.ip_address(info[4][0])):
|
||||
return True
|
||||
except socket.gaierror:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@@ -169,12 +171,31 @@ class SpiderTools:
|
||||
return {"error": error_msg}
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
# Make request
|
||||
response = session.get(
|
||||
url,
|
||||
timeout=timeout,
|
||||
verify=verify_ssl
|
||||
)
|
||||
# Make request with manual redirect validation (SSRF-safe)
|
||||
current_url = url
|
||||
redirect_statuses = {301, 302, 303, 307, 308}
|
||||
max_redirects = 5
|
||||
response = None
|
||||
for _ in range(max_redirects + 1):
|
||||
if not self._validate_url(current_url):
|
||||
return {"error": f"Invalid or blocked redirect target: {current_url}"}
|
||||
response = session.get(
|
||||
current_url,
|
||||
timeout=timeout,
|
||||
verify=verify_ssl,
|
||||
allow_redirects=False,
|
||||
)
|
||||
if response.status_code not in redirect_statuses:
|
||||
break
|
||||
location = response.headers.get("Location")
|
||||
if not location:
|
||||
break
|
||||
current_url = urljoin(current_url, location)
|
||||
else:
|
||||
return {"error": "Too many redirects"}
|
||||
|
||||
if response is None:
|
||||
return {"error": "Request failed"}
|
||||
response.raise_for_status()
|
||||
|
||||
# Parse HTML
|
||||
@@ -186,7 +207,7 @@ class SpiderTools:
|
||||
|
||||
# Initialize result
|
||||
result = {
|
||||
'url': url,
|
||||
'url': current_url,
|
||||
'status_code': response.status_code,
|
||||
'encoding': response.encoding,
|
||||
'headers': dict(response.headers),
|
||||
|
||||
@@ -130,17 +130,64 @@ def _crawl_with_crawl4ai(urls: List[str]) -> List[Dict[str, Any]]:
|
||||
except RuntimeError:
|
||||
return asyncio.run(_crawl_async())
|
||||
|
||||
def _is_safe_crawl_url(url: str) -> bool:
|
||||
"""Return True when a URL passes SSRF checks for crawling."""
|
||||
import urllib.parse
|
||||
import socket
|
||||
import ipaddress
|
||||
|
||||
try:
|
||||
parsed = urllib.parse.urlparse(url)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
return False
|
||||
hostname = parsed.hostname
|
||||
if not hostname:
|
||||
return False
|
||||
if os.environ.get("ALLOW_LOCAL_CRAWL") == "true":
|
||||
return True
|
||||
for info in socket.getaddrinfo(hostname, None):
|
||||
ip = ipaddress.ip_address(info[4][0])
|
||||
if (
|
||||
ip.is_loopback
|
||||
or ip.is_private
|
||||
or ip.is_link_local
|
||||
or ip.is_multicast
|
||||
or ip.is_unspecified
|
||||
):
|
||||
return False
|
||||
return True
|
||||
except (socket.gaierror, ValueError, OSError):
|
||||
return False
|
||||
|
||||
|
||||
def _crawl_with_httpx(urls: List[str]) -> List[Dict[str, Any]]:
|
||||
"""Crawl URLs using basic HTTP fetch with httpx or urllib."""
|
||||
import urllib.parse
|
||||
results = []
|
||||
|
||||
redirect_statuses = {301, 302, 303, 307, 308}
|
||||
max_redirects = 5
|
||||
|
||||
for url in urls:
|
||||
try:
|
||||
# Try httpx first
|
||||
try:
|
||||
import httpx
|
||||
with httpx.Client(follow_redirects=True, timeout=30.0) as client:
|
||||
response = client.get(url)
|
||||
current = url
|
||||
response = None
|
||||
with httpx.Client(follow_redirects=False, timeout=30.0) as client:
|
||||
for _ in range(max_redirects + 1):
|
||||
if not _is_safe_crawl_url(current):
|
||||
raise ValueError("Redirect target failed SSRF validation")
|
||||
response = client.get(current)
|
||||
if response.status_code not in redirect_statuses:
|
||||
break
|
||||
location = response.headers.get("Location")
|
||||
if not location:
|
||||
break
|
||||
current = urllib.parse.urljoin(current, location)
|
||||
else:
|
||||
raise ValueError("Too many redirects")
|
||||
if response is not None:
|
||||
response.raise_for_status()
|
||||
content = response.text
|
||||
except ImportError:
|
||||
@@ -210,37 +257,12 @@ def web_crawl(
|
||||
raw_url_list = [urls] if single_url else urls
|
||||
|
||||
# Validate URLs to prevent SSRF and Local File Read
|
||||
import urllib.parse
|
||||
import socket
|
||||
import ipaddress
|
||||
|
||||
url_list = []
|
||||
for u in raw_url_list:
|
||||
try:
|
||||
parsed = urllib.parse.urlparse(u)
|
||||
if parsed.scheme not in ('http', 'https'):
|
||||
logger.warning(f"Rejected non-http/https URL: {u}")
|
||||
continue
|
||||
hostname = parsed.hostname
|
||||
if not hostname:
|
||||
continue
|
||||
|
||||
# Allow opting out of SSRF protection for advanced use cases
|
||||
if os.environ.get("ALLOW_LOCAL_CRAWL") != "true":
|
||||
try:
|
||||
ip_str = socket.gethostbyname(hostname)
|
||||
ip = ipaddress.ip_address(ip_str)
|
||||
if ip.is_loopback or ip.is_private or ip.is_link_local or ip.is_multicast or ip.is_unspecified:
|
||||
logger.warning(f"Rejected SSRF or private IP attempt: {u}")
|
||||
continue
|
||||
except socket.gaierror:
|
||||
logger.warning(f"Could not resolve hostname for: {u}")
|
||||
continue
|
||||
|
||||
if _is_safe_crawl_url(u):
|
||||
url_list.append(u)
|
||||
except Exception as e:
|
||||
logger.warning(f"URL validation failed for {u}: {e}")
|
||||
continue
|
||||
else:
|
||||
logger.warning(f"Rejected unsafe crawl URL: {u}")
|
||||
|
||||
if not url_list:
|
||||
return {"error": "No valid or safe URLs provided. Local and non-http(s) URLs are blocked for security."}
|
||||
|
||||
@@ -12,6 +12,9 @@ import inspect
|
||||
|
||||
__all__ = ['deprecated', 'DeprecationConfig', 'check_expired_deprecations']
|
||||
|
||||
# Warn once per process for param deprecations (avoids hot-path overhead in Agent.__init__)
|
||||
_warned_params: set[str] = set()
|
||||
|
||||
class DeprecationConfig:
|
||||
"""Configuration for deprecation behavior and version management."""
|
||||
|
||||
@@ -182,6 +185,10 @@ def warn_deprecated_param(
|
||||
|
||||
message = "".join(msg_parts)
|
||||
|
||||
warn_key = f"{param_name}:{since}"
|
||||
if warn_key in _warned_params:
|
||||
return
|
||||
_warned_params.add(warn_key)
|
||||
warnings.warn(message, category=DeprecationWarning, stacklevel=stacklevel)
|
||||
|
||||
|
||||
|
||||
@@ -2952,22 +2952,36 @@ CONCISE SUMMARY:"""
|
||||
"variables": all_variables
|
||||
}
|
||||
|
||||
# Load tools from the recipe's tools.py if present
|
||||
# Load tools from the recipe's tools.py if present (opt-in only)
|
||||
tool_registry = {}
|
||||
tools_py = recipe_path / "tools.py"
|
||||
if tools_py.exists():
|
||||
try:
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location("recipe_tools", tools_py)
|
||||
recipe_module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(recipe_module)
|
||||
|
||||
# Extract callable tools
|
||||
for name in dir(recipe_module):
|
||||
if not name.startswith("_"):
|
||||
obj = getattr(recipe_module, name)
|
||||
if callable(obj):
|
||||
tool_registry[name] = obj
|
||||
import os
|
||||
_local_ok = os.environ.get("PRAISONAI_ALLOW_LOCAL_TOOLS", "").strip().lower() in (
|
||||
"1", "true", "yes", "on"
|
||||
)
|
||||
_template_ok = os.environ.get("PRAISONAI_ALLOW_TEMPLATE_TOOLS", "").strip().lower() in (
|
||||
"1", "true", "yes", "on"
|
||||
)
|
||||
if _local_ok or _template_ok:
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location("recipe_tools", tools_py)
|
||||
recipe_module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(recipe_module)
|
||||
else:
|
||||
logger.warning(
|
||||
"Skipping included recipe tools.py: set PRAISONAI_ALLOW_LOCAL_TOOLS "
|
||||
"or PRAISONAI_ALLOW_TEMPLATE_TOOLS to enable."
|
||||
)
|
||||
recipe_module = None
|
||||
if recipe_module is not None:
|
||||
# Extract callable tools
|
||||
for name in dir(recipe_module):
|
||||
if not name.startswith("_"):
|
||||
obj = getattr(recipe_module, name)
|
||||
if callable(obj):
|
||||
tool_registry[name] = obj
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tools from {tools_py}: {e}")
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "1.6.55"
|
||||
version = "1.6.59"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -5,8 +5,7 @@ testpaths = tests
|
||||
python_files = test_*.py
|
||||
python_classes = Test*
|
||||
python_functions = test_*
|
||||
addopts = -v --tb=short --timeout=60
|
||||
timeout_method = thread
|
||||
addopts = -v --tb=short
|
||||
filterwarnings =
|
||||
ignore::DeprecationWarning
|
||||
ignore::PendingDeprecationWarning
|
||||
|
||||
@@ -18,6 +18,7 @@ class TestExecutionConfigMaxIter:
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.test_llm_config = {
|
||||
"model": "gpt-4o-mini",
|
||||
"api_base": "http://test",
|
||||
"api_key": "test-key",
|
||||
"provider": "custom"
|
||||
@@ -33,21 +34,13 @@ class TestExecutionConfigMaxIter:
|
||||
config = ExecutionConfig(max_iter=15)
|
||||
assert config.max_iter == 15
|
||||
|
||||
@patch('praisonaiagents.llm.llm.LLM._chat_completion')
|
||||
def test_agent_with_execution_config_max_iter(self, mock_chat):
|
||||
def test_agent_with_execution_config_max_iter(self):
|
||||
"""Test that Agent passes ExecutionConfig.max_iter to LLM."""
|
||||
# Mock the LLM response to avoid actual API calls
|
||||
mock_response = Mock()
|
||||
mock_response.choices = [Mock()]
|
||||
mock_response.choices[0].message.content = "Hello, I'm a test response!"
|
||||
mock_response.choices[0].message.tool_calls = None
|
||||
mock_chat.return_value = mock_response
|
||||
|
||||
# Create agent with custom max_iter
|
||||
execution_config = ExecutionConfig(max_iter=15)
|
||||
|
||||
with patch('praisonaiagents.agent.agent.LLM') as MockLLM:
|
||||
with patch('praisonaiagents.llm.llm.LLM') as MockLLM:
|
||||
mock_llm_instance = Mock()
|
||||
mock_llm_instance.get_response.return_value = "Hello, I'm a test response!"
|
||||
MockLLM.return_value = mock_llm_instance
|
||||
|
||||
agent = Agent(
|
||||
@@ -57,6 +50,9 @@ class TestExecutionConfigMaxIter:
|
||||
execution=execution_config
|
||||
)
|
||||
|
||||
# LLM is deferred until first access
|
||||
_ = agent.llm_instance
|
||||
|
||||
# Verify LLM was instantiated with max_iter from ExecutionConfig
|
||||
MockLLM.assert_called()
|
||||
call_args = MockLLM.call_args
|
||||
@@ -72,7 +68,7 @@ class TestExecutionConfigMaxIter:
|
||||
def test_llm_constructor_accepts_max_iter(self):
|
||||
"""Test that LLM constructor accepts max_iter parameter."""
|
||||
llm = LLM(
|
||||
api_base="http://test",
|
||||
model="gpt-4o-mini",
|
||||
api_key="test-key",
|
||||
max_iter=25
|
||||
)
|
||||
@@ -82,7 +78,7 @@ class TestExecutionConfigMaxIter:
|
||||
def test_llm_constructor_default_max_iter(self):
|
||||
"""Test LLM constructor default max_iter value."""
|
||||
llm = LLM(
|
||||
api_base="http://test",
|
||||
model="gpt-4o-mini",
|
||||
api_key="test-key"
|
||||
)
|
||||
|
||||
@@ -90,7 +86,7 @@ class TestExecutionConfigMaxIter:
|
||||
assert hasattr(llm, 'max_iter')
|
||||
assert llm.max_iter == 20 # Default aligned with ExecutionConfig
|
||||
|
||||
@patch('praisonaiagents.agent.agent.LLM')
|
||||
@patch('praisonaiagents.llm.llm.LLM')
|
||||
def test_agent_without_execution_config_uses_llm_default(self, MockLLM):
|
||||
"""Test that Agent without ExecutionConfig doesn't override LLM max_iter."""
|
||||
mock_llm_instance = Mock()
|
||||
@@ -102,6 +98,8 @@ class TestExecutionConfigMaxIter:
|
||||
llm=self.test_llm_config
|
||||
)
|
||||
|
||||
_ = agent.llm_instance
|
||||
|
||||
# Verify LLM was called without max_iter override
|
||||
MockLLM.assert_called()
|
||||
call_args = MockLLM.call_args
|
||||
@@ -121,7 +119,7 @@ class TestExecutionConfigMaxIter:
|
||||
# Skip if presets not implemented yet
|
||||
pytest.skip("ExecutionConfig presets not implemented yet")
|
||||
|
||||
@patch('praisonaiagents.agent.agent.LLM')
|
||||
@patch('praisonaiagents.llm.llm.LLM')
|
||||
def test_different_llm_init_paths_respect_max_iter(self, MockLLM):
|
||||
"""Test that all LLM initialization paths in Agent respect max_iter."""
|
||||
mock_llm_instance = Mock()
|
||||
@@ -129,15 +127,15 @@ class TestExecutionConfigMaxIter:
|
||||
execution_config = ExecutionConfig(max_iter=30)
|
||||
|
||||
test_cases = [
|
||||
# Dict config
|
||||
{"provider": "openai", "model": "gpt-4o", "api_key": "test"},
|
||||
# String config
|
||||
"gpt-4o",
|
||||
# With base_url
|
||||
{"base_url": "http://test", "model": "test-model", "api_key": "test"}
|
||||
# Dict config — deferred LLM
|
||||
({"provider": "openai", "model": "gpt-4o", "api_key": "test"}, True),
|
||||
# Plain model string — OpenAI path, no LLM() at init
|
||||
("gpt-4o", False),
|
||||
# Dict with base_url — deferred LLM
|
||||
({"base_url": "http://test", "model": "test-model", "api_key": "test"}, True),
|
||||
]
|
||||
|
||||
for llm_config in test_cases:
|
||||
for llm_config, expects_custom_llm in test_cases:
|
||||
MockLLM.reset_mock()
|
||||
|
||||
agent = Agent(
|
||||
@@ -147,11 +145,17 @@ class TestExecutionConfigMaxIter:
|
||||
execution=execution_config
|
||||
)
|
||||
|
||||
# Verify LLM was called with max_iter
|
||||
MockLLM.assert_called()
|
||||
call_args = MockLLM.call_args
|
||||
assert 'max_iter' in call_args.kwargs
|
||||
assert call_args.kwargs['max_iter'] == 30
|
||||
if expects_custom_llm:
|
||||
assert agent._llm_init_params is not None
|
||||
assert agent._llm_init_params.get('max_iter') == 30
|
||||
_ = agent.llm_instance
|
||||
MockLLM.assert_called()
|
||||
call_args = MockLLM.call_args
|
||||
assert call_args.kwargs.get('max_iter') == 30
|
||||
else:
|
||||
assert agent._llm_init_params is None
|
||||
assert not agent._using_custom_llm
|
||||
assert agent.max_iter == 30
|
||||
|
||||
def test_max_iter_bounds_validation(self):
|
||||
"""Test validation of max_iter bounds."""
|
||||
@@ -184,7 +188,7 @@ class TestExecutionConfigMaxIter:
|
||||
autonomy_config = AutonomyConfig()
|
||||
|
||||
# This should not conflict
|
||||
with patch('praisonaiagents.agent.agent.LLM') as MockLLM:
|
||||
with patch('praisonaiagents.llm.llm.LLM') as MockLLM:
|
||||
mock_llm_instance = Mock()
|
||||
MockLLM.return_value = mock_llm_instance
|
||||
|
||||
@@ -196,6 +200,8 @@ class TestExecutionConfigMaxIter:
|
||||
autonomy=autonomy_config
|
||||
)
|
||||
|
||||
_ = agent.llm_instance
|
||||
|
||||
# Verify both configs can coexist and max_iter is propagated
|
||||
MockLLM.assert_called()
|
||||
call_args = MockLLM.call_args
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
Test script to reproduce and verify fix for multi-agent streaming issue #1882
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
# Set minimal OpenAI API key for testing (should fail gracefully)
|
||||
os.environ.setdefault("OPENAI_API_KEY", "sk-test-key-for-reproduction")
|
||||
|
||||
|
||||
def test_config_defaults():
|
||||
"""Test that the config default is now False"""
|
||||
from praisonaiagents.config.feature_configs import MultiAgentOutputConfig
|
||||
|
||||
config = MultiAgentOutputConfig()
|
||||
assert config.stream is False
|
||||
|
||||
|
||||
@pytest.mark.live
|
||||
def test_single_agent():
|
||||
"""Test single agent (should work with fallback)"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(instructions="Reply with exactly the requested text")
|
||||
result = agent.start("Reply with exactly: OK")
|
||||
assert result is not None
|
||||
|
||||
|
||||
@pytest.mark.live
|
||||
def test_multi_agent():
|
||||
"""Test multi-agent (should work after fix)"""
|
||||
from praisonaiagents import Agent, Agents
|
||||
|
||||
research_agent = Agent(instructions="Research about AI")
|
||||
summarise_agent = Agent(instructions="Summarise research agent's findings")
|
||||
agents = Agents(agents=[research_agent, summarise_agent])
|
||||
|
||||
try:
|
||||
result = agents.start("What is Python?")
|
||||
assert result is not None
|
||||
except Exception as exc:
|
||||
if "Streaming is not supported in sync OpenAIAdapter" in str(exc):
|
||||
pytest.fail("Streaming error still exists — fix did not work")
|
||||
raise
|
||||
@@ -0,0 +1,182 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to validate cache optimization integration.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add the source path to Python path
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src', 'praisonai-agents'))
|
||||
|
||||
def test_imports():
|
||||
"""Test that all imports work correctly."""
|
||||
try:
|
||||
from praisonaiagents.tools.base import get_sorted_tool_schemas
|
||||
from praisonaiagents.memory.memory import Memory, CACHE_BOUNDARY
|
||||
print("✓ Imports successful")
|
||||
return True
|
||||
except ImportError as e:
|
||||
print(f"✗ Import failed: {e}")
|
||||
return False
|
||||
|
||||
def test_tool_sorting():
|
||||
"""Test tool schema sorting functionality (internal sorting)."""
|
||||
try:
|
||||
# Test the sorting function directly
|
||||
def sort_formatted_tools(tools_list):
|
||||
"""Sort already-formatted tool schemas by function name."""
|
||||
def sort_key(tool):
|
||||
if isinstance(tool, dict) and tool.get('type') == 'function':
|
||||
return str(tool.get('function', {}).get('name') or '')
|
||||
return ''
|
||||
return sorted(tools_list, key=sort_key)
|
||||
|
||||
# Test with formatted tool schemas
|
||||
test_tools = [
|
||||
{"type": "function", "function": {"name": "zebra_tool", "description": "Z tool"}},
|
||||
{"type": "function", "function": {"name": "alpha_tool", "description": "A tool"}},
|
||||
{"type": "function", "function": {"name": "beta_tool", "description": "B tool"}},
|
||||
]
|
||||
|
||||
sorted_tools = sort_formatted_tools(test_tools)
|
||||
tool_names = [tool["function"]["name"] for tool in sorted_tools]
|
||||
|
||||
expected_order = ["alpha_tool", "beta_tool", "zebra_tool"]
|
||||
if tool_names == expected_order:
|
||||
print("✓ Tool sorting works correctly")
|
||||
return True
|
||||
else:
|
||||
print(f"✗ Tool sorting failed. Got: {tool_names}, Expected: {expected_order}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✗ Tool sorting test failed: {e}")
|
||||
return False
|
||||
|
||||
def test_cache_optimized_context():
|
||||
"""Test cache-optimized context building API structure."""
|
||||
try:
|
||||
from praisonaiagents.memory.memory import Memory
|
||||
|
||||
# Create a basic memory instance for API testing
|
||||
config = {}
|
||||
memory = Memory(config)
|
||||
|
||||
# Mock the underlying methods to avoid database dependencies in unit test
|
||||
memory.build_context_for_task = lambda task_descr, **kwargs: "test context"
|
||||
|
||||
# Test cache-optimized context building
|
||||
result = memory.build_cache_optimized_context(
|
||||
task_descr="test task",
|
||||
include_cache_boundary=True
|
||||
)
|
||||
|
||||
if isinstance(result, dict) and 'stable_prefix' in result and 'cache_boundary' in result:
|
||||
# Verify the cache boundary is now empty (our fix)
|
||||
if result['cache_boundary'] == "":
|
||||
print("✓ Cache-optimized context building works (cache boundary properly empty)")
|
||||
return True
|
||||
else:
|
||||
print(f"✗ Cache boundary should be empty but got: '{result['cache_boundary']}'")
|
||||
return False
|
||||
else:
|
||||
print(f"✗ Cache-optimized context building failed. Got: {result}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✗ Cache-optimized context test failed: {e}")
|
||||
return False
|
||||
|
||||
def test_agent_integration():
|
||||
"""Test agent integration with cache optimization."""
|
||||
try:
|
||||
from praisonaiagents.agent.chat_mixin import ChatMixin
|
||||
|
||||
# Mock the necessary methods for testing
|
||||
class TestAgent(ChatMixin):
|
||||
def __init__(self):
|
||||
self.tools = []
|
||||
self._formatted_tools_cache = {}
|
||||
self._system_prompt_cache = {}
|
||||
self._memory_instance = None
|
||||
self.use_system_prompt = True
|
||||
self.role = "test"
|
||||
self.goal = "test"
|
||||
self.backstory = "test"
|
||||
self._rules_manager_initialized = False
|
||||
self._rules_manager = None
|
||||
self._skills = None
|
||||
self._skills_dirs = None
|
||||
self.verbose = False
|
||||
|
||||
def _cache_get(self, cache, key):
|
||||
return cache.get(key)
|
||||
|
||||
def _cache_put(self, cache, key, value):
|
||||
cache[key] = value
|
||||
|
||||
def _get_tools_cache_key(self, tools):
|
||||
return str(tools)
|
||||
|
||||
def get_memory_context(self):
|
||||
return ""
|
||||
|
||||
def get_learn_context(self):
|
||||
return ""
|
||||
|
||||
def get_skills_prompt(self):
|
||||
return ""
|
||||
|
||||
agent = TestAgent()
|
||||
|
||||
# Test tool formatting with sorting - pass already formatted tools
|
||||
test_tools = [
|
||||
{"type": "function", "function": {"name": "zebra_tool", "description": "Z tool"}},
|
||||
{"type": "function", "function": {"name": "alpha_tool", "description": "A tool"}},
|
||||
]
|
||||
|
||||
formatted_tools = agent._format_tools_for_completion(test_tools)
|
||||
if len(formatted_tools) == 2:
|
||||
tool_names = [tool["function"]["name"] for tool in formatted_tools]
|
||||
if tool_names == ["alpha_tool", "zebra_tool"]:
|
||||
print("✓ Agent tool integration works correctly")
|
||||
return True
|
||||
else:
|
||||
print(f"✗ Agent tools not sorted correctly. Got: {tool_names}")
|
||||
return False
|
||||
else:
|
||||
print(f"✗ Agent tool formatting failed. Got {len(formatted_tools)} tools")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Agent integration test failed: {e}")
|
||||
return False
|
||||
|
||||
def main():
|
||||
"""Run all tests."""
|
||||
print("Testing cache optimization integration...")
|
||||
print()
|
||||
|
||||
tests = [
|
||||
test_imports,
|
||||
test_tool_sorting,
|
||||
test_cache_optimized_context,
|
||||
test_agent_integration
|
||||
]
|
||||
|
||||
results = []
|
||||
for test in tests:
|
||||
results.append(test())
|
||||
|
||||
print()
|
||||
passed = sum(results)
|
||||
total = len(results)
|
||||
|
||||
if passed == total:
|
||||
print(f"🎉 All {total} tests passed!")
|
||||
return 0
|
||||
else:
|
||||
print(f"💥 {total - passed} out of {total} tests failed!")
|
||||
return 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,147 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script for context assembly cache optimization (Issue #1808).
|
||||
|
||||
This test verifies that:
|
||||
1. Memory search results are sorted deterministically
|
||||
2. Tool schemas are sorted by function name
|
||||
3. Cache boundary markers are properly included
|
||||
4. Multiple runs produce identical context strings for caching
|
||||
"""
|
||||
|
||||
def test_memory_stable_sorting():
|
||||
"""Test that memory context is assembled deterministically."""
|
||||
print("Testing memory stable sorting...")
|
||||
|
||||
from praisonaiagents.memory.memory import Memory
|
||||
|
||||
# Create test data that simulates memory search results
|
||||
test_results = [
|
||||
{"text": "First memory", "timestamp": 1000, "id": "1"},
|
||||
{"text": "Second memory", "timestamp": 2000, "id": "2"},
|
||||
{"text": "Third memory", "timestamp": 1500, "id": "3"}
|
||||
]
|
||||
|
||||
# Create a memory instance and stub the search methods to return controlled data
|
||||
memory = Memory(config={"provider": "none"})
|
||||
|
||||
# Mock the search methods to return our test data
|
||||
original_search_short = memory.search_short_term
|
||||
original_search_long = memory.search_long_term
|
||||
original_search_entity = memory.search_entity
|
||||
original_search_user = memory.search_user_memory
|
||||
|
||||
try:
|
||||
memory.search_short_term = lambda q, limit=3: test_results.copy()
|
||||
memory.search_long_term = lambda q, limit=3: []
|
||||
memory.search_entity = lambda q, limit=3: []
|
||||
memory.search_user_memory = lambda user_id, q, limit=3: []
|
||||
|
||||
# Test multiple runs produce identical context strings
|
||||
context1 = memory.build_context_for_task("test task", include_in_output=True)
|
||||
context2 = memory.build_context_for_task("test task", include_in_output=True)
|
||||
|
||||
assert context1 == context2, "Memory context should be deterministic"
|
||||
|
||||
# Verify the order is correct in the actual context (newest first)
|
||||
lines = context1.split('\n')
|
||||
second_idx = next((i for i, line in enumerate(lines) if "Second memory" in line), -1)
|
||||
third_idx = next((i for i, line in enumerate(lines) if "Third memory" in line), -1)
|
||||
first_idx = next((i for i, line in enumerate(lines) if "First memory" in line), -1)
|
||||
|
||||
assert second_idx < third_idx < first_idx, "Memory should be ordered newest first"
|
||||
print("✓ Memory context is deterministic and correctly ordered")
|
||||
|
||||
finally:
|
||||
# Restore original methods
|
||||
memory.search_short_term = original_search_short
|
||||
memory.search_long_term = original_search_long
|
||||
memory.search_entity = original_search_entity
|
||||
memory.search_user_memory = original_search_user
|
||||
|
||||
# Test cache boundary constants are available
|
||||
from praisonaiagents.memory.memory import CACHE_BOUNDARY
|
||||
assert CACHE_BOUNDARY is not None, "CACHE_BOUNDARY should be defined"
|
||||
print("✓ Cache boundary constants are available")
|
||||
|
||||
def test_tool_schema_sorting():
|
||||
"""Test that tool schemas are sorted deterministically."""
|
||||
print("Testing tool schema sorting...")
|
||||
|
||||
from praisonaiagents.tools.base import get_sorted_tool_schemas
|
||||
from praisonaiagents.tools.decorator import tool
|
||||
|
||||
# Create test tools with different names
|
||||
@tool
|
||||
def zebra_tool():
|
||||
"""A tool starting with Z."""
|
||||
return "zebra"
|
||||
|
||||
@tool
|
||||
def apple_tool():
|
||||
"""A tool starting with A."""
|
||||
return "apple"
|
||||
|
||||
@tool
|
||||
def middle_tool():
|
||||
"""A tool in the middle."""
|
||||
return "middle"
|
||||
|
||||
tools = [zebra_tool, apple_tool, middle_tool]
|
||||
|
||||
# Get sorted schemas multiple times
|
||||
schemas1 = get_sorted_tool_schemas(tools)
|
||||
schemas2 = get_sorted_tool_schemas(tools)
|
||||
|
||||
# Results should be identical
|
||||
assert schemas1 == schemas2, "Tool schemas should be deterministic"
|
||||
|
||||
# Check they are sorted by name
|
||||
names = [schema["function"]["name"] for schema in schemas1]
|
||||
expected_names = ["apple_tool", "middle_tool", "zebra_tool"]
|
||||
assert names == expected_names, f"Tools should be sorted alphabetically, got {names}"
|
||||
print("✓ Tool schemas are sorted deterministically")
|
||||
|
||||
def test_cache_boundary_markers():
|
||||
"""Test that cache boundary markers work correctly."""
|
||||
print("Testing cache boundary markers...")
|
||||
|
||||
from praisonaiagents.memory.memory import Memory, CACHE_BOUNDARY
|
||||
|
||||
# Test that the boundary constant is properly defined
|
||||
assert "CACHE_BOUNDARY" in CACHE_BOUNDARY
|
||||
assert len(CACHE_BOUNDARY) > 0
|
||||
print("✓ Cache boundary constant is defined")
|
||||
|
||||
# Test the cache-optimized context method actually works
|
||||
memory = Memory(config={"provider": "none"})
|
||||
|
||||
# Mock search methods to return empty results for simpler test
|
||||
memory.search_short_term = lambda q, limit=3: []
|
||||
memory.search_long_term = lambda q, limit=3: []
|
||||
memory.search_entity = lambda q, limit=3: []
|
||||
memory.search_user_memory = lambda user_id, q, limit=3: []
|
||||
|
||||
# Test with cache boundary enabled
|
||||
result = memory.build_cache_optimized_context("test task", include_cache_boundary=True, include_in_output=True)
|
||||
assert "stable_prefix" in result, "Result should have stable_prefix"
|
||||
assert "cache_boundary" in result, "Result should have cache_boundary"
|
||||
assert result["cache_boundary"] == CACHE_BOUNDARY, "Cache boundary should match constant"
|
||||
|
||||
# Test with cache boundary disabled
|
||||
result2 = memory.build_cache_optimized_context("test task", include_cache_boundary=False, include_in_output=True)
|
||||
assert result2["cache_boundary"] == "", "Cache boundary should be empty when disabled"
|
||||
|
||||
print("✓ Cache optimization method works correctly")
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
try:
|
||||
test_memory_stable_sorting()
|
||||
test_tool_schema_sorting()
|
||||
test_cache_boundary_markers()
|
||||
print("\n🎉 All tests passed! Context assembly cache optimization is working correctly.")
|
||||
except Exception as e:
|
||||
print(f"\n❌ Test failed: {e}")
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Unit tests for Agent instantiation performance optimisations (lazy LLM, LoopGuard)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig
|
||||
from praisonaiagents.config.loader import clear_config_cache
|
||||
|
||||
|
||||
class TestLazyInstantiation:
|
||||
def test_agent_init_no_deprecation_warning(self):
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
Agent(name="T", output="silent")
|
||||
Agent(name="T", model="gpt-4o-mini", output="silent")
|
||||
deprecations = [w for w in caught if issubclass(w.category, DeprecationWarning)]
|
||||
assert not deprecations
|
||||
|
||||
def test_litellm_model_deferred_at_init(self):
|
||||
agent = Agent(name="T", model="openai/gpt-4o-mini", output="silent")
|
||||
assert agent._llm_instance is None
|
||||
assert agent._llm_init_params["model"] == "openai/gpt-4o-mini"
|
||||
assert agent._using_custom_llm
|
||||
|
||||
def test_loop_guard_deferred_at_init(self):
|
||||
agent = Agent(name="T", output="silent")
|
||||
assert agent._loop_guard is None
|
||||
guard = agent._ensure_loop_guard()
|
||||
assert guard is not None
|
||||
assert agent._ensure_loop_guard() is guard
|
||||
|
||||
def test_llm_instance_materializes_once(self):
|
||||
with patch("praisonaiagents.llm.llm.LLM") as MockLLM:
|
||||
MockLLM.return_value = Mock()
|
||||
agent = Agent(name="T", model="openai/gpt-4o-mini", output="silent")
|
||||
first = agent.llm_instance
|
||||
second = agent.llm_instance
|
||||
assert first is second
|
||||
MockLLM.assert_called_once()
|
||||
|
||||
def test_dict_llm_stores_max_iter_in_deferred_params(self):
|
||||
agent = Agent(
|
||||
name="T",
|
||||
llm={"model": "gpt-4o-mini", "api_key": "x"},
|
||||
execution=ExecutionConfig(max_iter=12),
|
||||
output="silent",
|
||||
)
|
||||
assert agent._llm_init_params["max_iter"] == 12
|
||||
|
||||
def test_plain_model_string_no_custom_llm(self):
|
||||
agent = Agent(name="T", model="gpt-4o-mini", output="silent")
|
||||
assert not agent._using_custom_llm
|
||||
assert agent._llm_init_params is None
|
||||
|
||||
def test_config_defaults_fast_path_when_empty(self, tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
clear_config_cache()
|
||||
agent = Agent(name="T", output="silent")
|
||||
assert agent.name == "T"
|
||||
|
||||
def test_execution_config_internal_from_tests_is_suppressed(self):
|
||||
"""Test/SDK paths suppress context_compaction spam during Agent init."""
|
||||
ExecutionConfig._context_compaction_warned = False # noqa: SLF001
|
||||
ExecutionConfig._context_compaction_internal = False # noqa: SLF001
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
ExecutionConfig()
|
||||
compaction = [w for w in caught if "context_compaction" in str(w.message)]
|
||||
assert not compaction
|
||||
assert ExecutionConfig._context_compaction_internal is True # noqa: SLF001
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Smoke tests for security audit hardening."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_workflow_include_skips_tools_without_opt_in(tmp_path, monkeypatch):
|
||||
monkeypatch.delenv("PRAISONAI_ALLOW_LOCAL_TOOLS", raising=False)
|
||||
monkeypatch.delenv("PRAISONAI_ALLOW_TEMPLATE_TOOLS", raising=False)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
child = tmp_path / "child_recipe"
|
||||
child.mkdir()
|
||||
marker = tmp_path / "marker.txt"
|
||||
(child / "tools.py").write_text(f'from pathlib import Path\nPath({str(marker)!r}).write_text("x")\n')
|
||||
(child / "workflow.yaml").write_text("name: child\nsteps: []\n")
|
||||
|
||||
from praisonaiagents.workflows.workflows import Workflow, include
|
||||
|
||||
Workflow(steps=[include("child_recipe")]).run(input="", llm="dummy/local", stream=False)
|
||||
assert not marker.exists()
|
||||
|
||||
|
||||
def test_web_crawl_rejects_redirect_to_loopback():
|
||||
pytest.importorskip("httpx")
|
||||
from praisonaiagents.tools.web_crawl_tools import _is_safe_crawl_url
|
||||
|
||||
assert _is_safe_crawl_url("http://127.0.0.1/secret") is False
|
||||
assert _is_safe_crawl_url("http://example.com/") is True
|
||||
|
||||
|
||||
def test_spider_blocks_rebound_hostname():
|
||||
from praisonaiagents.tools.spider_tools import _host_is_blocked
|
||||
|
||||
with patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("127.0.0.1", 0))]):
|
||||
assert _host_is_blocked("rebound.example") is True
|
||||
|
||||
|
||||
def test_file_memory_rejects_traversal_user_id():
|
||||
from praisonaiagents.memory.file_memory import FileMemory
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
FileMemory(user_id="../escape")
|
||||
|
||||
|
||||
def test_ast_grep_rewrite_requires_approval():
|
||||
from praisonaiagents.tools import ast_grep_tool # noqa: F401
|
||||
from praisonaiagents.approval import is_approval_required
|
||||
|
||||
assert is_approval_required("ast_grep_rewrite") is True
|
||||
|
||||
|
||||
def test_agent_server_enforces_auth_token():
|
||||
from praisonaiagents.server.server import AgentServer, ServerConfig
|
||||
|
||||
server = AgentServer(config=ServerConfig(auth_token="secret-token"))
|
||||
request = type("Req", (), {"headers": {"Authorization": "Bearer wrong"}})()
|
||||
assert server._authorise_request(request) is False
|
||||
request_ok = type("Req", (), {"headers": {"Authorization": "Bearer secret-token"}})()
|
||||
assert server._authorise_request(request_ok) is True
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Unit tests for security advisory hardening."""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_webhook_fail_closed_without_secret(monkeypatch):
|
||||
from praisonai.bots.webhook_security import webhooks_require_verification
|
||||
|
||||
monkeypatch.delenv("PRAISONAI_INSECURE_WEBHOOKS", raising=False)
|
||||
assert webhooks_require_verification() is True
|
||||
|
||||
|
||||
def test_webhook_insecure_override(monkeypatch):
|
||||
from praisonai.bots.webhook_security import webhooks_require_verification
|
||||
|
||||
monkeypatch.setenv("PRAISONAI_INSECURE_WEBHOOKS", "true")
|
||||
assert webhooks_require_verification() is False
|
||||
|
||||
|
||||
def test_jobs_agent_root_rejects_escape(tmp_path, monkeypatch):
|
||||
from praisonai.jobs.path_validation import validate_agent_file_path
|
||||
|
||||
root = tmp_path / "agents"
|
||||
root.mkdir()
|
||||
outside = tmp_path / "outside.yaml"
|
||||
outside.write_text("agents: []")
|
||||
monkeypatch.setenv("PRAISONAI_JOBS_AGENT_ROOT", str(root))
|
||||
|
||||
with pytest.raises(ValueError, match="PRAISONAI_JOBS_AGENT_ROOT"):
|
||||
validate_agent_file_path(str(outside))
|
||||
|
||||
|
||||
def test_jobs_agent_root_allows_inside(tmp_path, monkeypatch):
|
||||
from praisonai.jobs.path_validation import validate_agent_file_path
|
||||
|
||||
root = tmp_path / "agents"
|
||||
root.mkdir()
|
||||
inside = root / "team.yaml"
|
||||
inside.write_text("agents: []")
|
||||
monkeypatch.setenv("PRAISONAI_JOBS_AGENT_ROOT", str(root))
|
||||
|
||||
assert validate_agent_file_path(str(inside)) == str(inside)
|
||||
|
||||
|
||||
def test_yaml_approve_blocks_critical_tools():
|
||||
from praisonaiagents.approval.registry import ApprovalRegistry
|
||||
|
||||
registry = ApprovalRegistry()
|
||||
token = registry.set_yaml_approved_tools(["execute_code", "search_web"])
|
||||
try:
|
||||
assert registry.is_yaml_approved("search_web") is True
|
||||
assert registry.is_yaml_approved("execute_code") is False
|
||||
finally:
|
||||
registry.reset_yaml_approved_tools(token)
|
||||
|
||||
|
||||
def test_sandbox_blocks_format():
|
||||
from praisonaiagents.tools.python_tools import _validate_code_ast
|
||||
|
||||
assert _validate_code_ast('"{0.__class__}".format(42)') is not None
|
||||
assert _validate_code_ast('"".format_map({})') is not None
|
||||
|
||||
|
||||
def test_sse_security_middleware_rejects_bad_origin(monkeypatch):
|
||||
import asyncio
|
||||
|
||||
from starlette.applications import Starlette
|
||||
from starlette.responses import PlainTextResponse
|
||||
from starlette.routing import Route
|
||||
|
||||
from praisonaiagents.mcp.mcp_security import SecurityConfig, build_sse_security_app
|
||||
|
||||
monkeypatch.delenv("MCP_SSE_ALLOWED_ORIGINS", raising=False)
|
||||
monkeypatch.delenv("MCP_SSE_AUTH_TOKEN", raising=False)
|
||||
monkeypatch.delenv("MCP_AUTH_TOKEN", raising=False)
|
||||
|
||||
app = Starlette(routes=[Route("/", lambda r: PlainTextResponse("ok"))])
|
||||
config = SecurityConfig(validate_origin=True, allowed_origins=["localhost"])
|
||||
app = build_sse_security_app(app, config)
|
||||
|
||||
async def _run(origin: str):
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "GET",
|
||||
"path": "/",
|
||||
"headers": [(b"origin", origin.encode())],
|
||||
}
|
||||
status_code = 500
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
async def send(message):
|
||||
nonlocal status_code
|
||||
if message["type"] == "http.response.start":
|
||||
status_code = message["status"]
|
||||
|
||||
await app(scope, receive, send)
|
||||
return status_code
|
||||
|
||||
assert asyncio.run(_run("https://evil.com")) == 403
|
||||
assert asyncio.run(_run("http://localhost:8080")) == 200
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Tests for remaining security hardening (round 2)."""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_launch_auth_optional_by_default():
|
||||
from praisonaiagents.agents.agents import _authorise_launch_request
|
||||
|
||||
class Req:
|
||||
headers = {}
|
||||
|
||||
assert _authorise_launch_request(Req()) is True
|
||||
|
||||
|
||||
def test_launch_auth_enforced_when_token_set(monkeypatch):
|
||||
from praisonaiagents.agents.agents import _authorise_launch_request
|
||||
|
||||
monkeypatch.setenv("PRAISONAI_LAUNCH_AUTH_TOKEN", "secret")
|
||||
|
||||
class Bad:
|
||||
headers = {}
|
||||
|
||||
class Good:
|
||||
headers = {"Authorization": "Bearer secret"}
|
||||
|
||||
assert _authorise_launch_request(Bad()) is False
|
||||
assert _authorise_launch_request(Good()) is True
|
||||
|
||||
monkeypatch.delenv("PRAISONAI_LAUNCH_AUTH_TOKEN", raising=False)
|
||||
|
||||
|
||||
def test_mentions_reject_outside_workspace(tmp_path):
|
||||
from praisonaiagents.tools.mentions import MentionsParser
|
||||
|
||||
outside = tmp_path / "secret.txt"
|
||||
outside.write_text("secret", encoding="utf-8")
|
||||
ws = tmp_path / "ws"
|
||||
ws.mkdir()
|
||||
parser = MentionsParser(workspace_path=str(ws))
|
||||
context, _ = parser.process(f"@file:{outside}")
|
||||
assert "outside workspace" in context.lower() or "not allowed" in context.lower()
|
||||
|
||||
|
||||
def test_imap_sanitize_rejects_injection():
|
||||
from praisonaiagents.tools.email_tools import _imap_sanitize
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_imap_sanitize('test"\r\nBAD')
|
||||
Generated
+1
-1
@@ -3125,7 +3125,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.6.55"
|
||||
version = "1.6.59"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -95,6 +95,25 @@ async def require_workspace_owner(
|
||||
)
|
||||
|
||||
|
||||
async def require_delete_permission(
|
||||
workspace_id: str,
|
||||
user: AuthIdentity,
|
||||
session: AsyncSession,
|
||||
*,
|
||||
resource_owner_id: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Allow delete when user is admin/owner or owns the resource."""
|
||||
member_svc = MemberService(session)
|
||||
if await member_svc.has_role(workspace_id, user.id, "admin"):
|
||||
return
|
||||
if resource_owner_id and resource_owner_id == user.id:
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Insufficient permission to delete this resource",
|
||||
)
|
||||
|
||||
|
||||
def ensure_resource_in_workspace(
|
||||
resource_workspace_id: str | None,
|
||||
workspace_id: str,
|
||||
|
||||
@@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from praisonaiagents.auth import AuthIdentity
|
||||
|
||||
from ..deps import get_db, require_workspace_member
|
||||
from ..deps import get_db, require_delete_permission, require_workspace_member
|
||||
from ..schemas import AgentCreate, AgentResponse, AgentUpdate
|
||||
from ...services.agent_service import AgentService
|
||||
|
||||
@@ -96,6 +96,12 @@ async def delete_agent(
|
||||
session: AsyncSession = Depends(get_db),
|
||||
):
|
||||
svc = AgentService(session)
|
||||
agent = await svc.get(agent_id, workspace_id=workspace_id)
|
||||
if agent is None:
|
||||
raise HTTPException(status_code=404, detail="Agent not found")
|
||||
await require_delete_permission(
|
||||
workspace_id, user, session, resource_owner_id=agent.owner_id
|
||||
)
|
||||
deleted = await svc.delete(agent_id, workspace_id=workspace_id)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail="Agent not found")
|
||||
@@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from praisonaiagents.auth import AuthIdentity
|
||||
|
||||
from ..deps import get_db, require_issue_in_workspace, require_workspace_member
|
||||
from ..deps import get_db, require_delete_permission, require_issue_in_workspace, require_workspace_member
|
||||
from ...services.dependency_service import DependencyService
|
||||
from ..schemas import DependencyCreate, DependencyResponse
|
||||
|
||||
@@ -58,12 +58,21 @@ async def delete_dependency(
|
||||
session: AsyncSession = Depends(get_db),
|
||||
):
|
||||
await require_issue_in_workspace(workspace_id, issue_id, session)
|
||||
from ...services.issue_service import IssueService
|
||||
|
||||
issue_svc = IssueService(session)
|
||||
issue = await issue_svc.get(issue_id, workspace_id=workspace_id)
|
||||
if issue is None:
|
||||
raise HTTPException(status_code=404, detail="Issue not found")
|
||||
svc = DependencyService(session)
|
||||
dep = await svc.get(dep_id)
|
||||
if dep is None:
|
||||
raise HTTPException(status_code=404, detail="Dependency not found")
|
||||
if dep.issue_id != issue_id and dep.depends_on_issue_id != issue_id:
|
||||
raise HTTPException(status_code=404, detail="Dependency not found")
|
||||
await require_delete_permission(
|
||||
workspace_id, user, session, resource_owner_id=issue.creator_id
|
||||
)
|
||||
deleted = await svc.delete(dep_id)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail="Dependency not found")
|
||||
@@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from praisonaiagents.auth import AuthIdentity
|
||||
|
||||
from ..deps import ensure_resource_in_workspace, get_db, require_workspace_member
|
||||
from ..deps import ensure_resource_in_workspace, get_db, require_delete_permission, require_workspace_member
|
||||
from ..schemas import (
|
||||
CommentCreate,
|
||||
CommentResponse,
|
||||
@@ -133,6 +133,12 @@ async def delete_issue(
|
||||
session: AsyncSession = Depends(get_db),
|
||||
):
|
||||
svc = IssueService(session)
|
||||
issue = await svc.get(issue_id, workspace_id=workspace_id)
|
||||
if issue is None:
|
||||
raise HTTPException(status_code=404, detail="Issue not found")
|
||||
await require_delete_permission(
|
||||
workspace_id, user, session, resource_owner_id=issue.creator_id
|
||||
)
|
||||
deleted = await svc.delete(issue_id, workspace_id=workspace_id)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail="Issue not found")
|
||||
|
||||
@@ -12,6 +12,7 @@ from praisonaiagents.auth import AuthIdentity
|
||||
from ..deps import (
|
||||
ensure_resource_in_workspace,
|
||||
get_db,
|
||||
require_delete_permission,
|
||||
require_issue_in_workspace,
|
||||
require_workspace_member,
|
||||
)
|
||||
@@ -75,6 +76,7 @@ async def delete_label(
|
||||
if label is None:
|
||||
raise HTTPException(status_code=404, detail="Label not found")
|
||||
ensure_resource_in_workspace(label.workspace_id, workspace_id, label="Label")
|
||||
await require_delete_permission(workspace_id, user, session)
|
||||
deleted = await svc.delete(label_id)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail="Label not found")
|
||||
|
||||
@@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from praisonaiagents.auth import AuthIdentity
|
||||
|
||||
from ..deps import get_db, require_workspace_member
|
||||
from ..deps import get_db, require_delete_permission, require_workspace_member
|
||||
from ..schemas import ProjectCreate, ProjectResponse, ProjectUpdate
|
||||
from ...services.project_service import ProjectService
|
||||
|
||||
@@ -93,6 +93,12 @@ async def delete_project(
|
||||
session: AsyncSession = Depends(get_db),
|
||||
):
|
||||
svc = ProjectService(session)
|
||||
project = await svc.get(project_id, workspace_id=workspace_id)
|
||||
if project is None:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
await require_delete_permission(
|
||||
workspace_id, user, session, resource_owner_id=project.lead_id
|
||||
)
|
||||
deleted = await svc.delete(project_id, workspace_id=workspace_id)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
"""DELETE ownership checks for platform resources."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonai_platform.services.auth_service import AuthService
|
||||
from praisonai_platform.services.issue_service import IssueService
|
||||
from praisonai_platform.services.member_service import MemberService
|
||||
from praisonai_platform.services.workspace_service import WorkspaceService
|
||||
from praisonaiagents.auth import AuthIdentity
|
||||
|
||||
from praisonai_platform.api.deps import require_delete_permission
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_member_cannot_delete_others_issue(session):
|
||||
auth = AuthService(session)
|
||||
owner, _ = await auth.register("[email protected]", "pass")
|
||||
other_user, _ = await auth.register("[email protected]", "pass")
|
||||
ws = await WorkspaceService(session).create("Team", owner.id)
|
||||
await MemberService(session).add(ws.id, other_user.id, "member")
|
||||
|
||||
issue = await IssueService(session).create(
|
||||
workspace_id=ws.id,
|
||||
title="private",
|
||||
creator_id=owner.id,
|
||||
)
|
||||
|
||||
other = AuthIdentity(id=other_user.id, email=other_user.email, name=other_user.name)
|
||||
with pytest.raises(Exception) as exc:
|
||||
await require_delete_permission(
|
||||
ws.id, other, session, resource_owner_id=issue.creator_id
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_can_delete_any_issue(session):
|
||||
auth = AuthService(session)
|
||||
owner, _ = await auth.register("[email protected]", "pass")
|
||||
admin_user, _ = await auth.register("[email protected]", "pass")
|
||||
ws = await WorkspaceService(session).create("AdminTeam", owner.id)
|
||||
await MemberService(session).add(ws.id, admin_user.id, "admin")
|
||||
|
||||
issue = await IssueService(session).create(
|
||||
workspace_id=ws.id,
|
||||
title="shared",
|
||||
creator_id=owner.id,
|
||||
)
|
||||
|
||||
admin = AuthIdentity(id=admin_user.id, email=admin_user.email, name=admin_user.name)
|
||||
await require_delete_permission(
|
||||
ws.id, admin, session, resource_owner_id=issue.creator_id
|
||||
)
|
||||
@@ -1,14 +1,14 @@
|
||||
{
|
||||
"version": "1.5.87",
|
||||
"lastUpdated": "2026-06-12",
|
||||
"lastUpdated": "2026-06-17",
|
||||
"generatedBy": "praisonai._dev.parity.generator",
|
||||
"sourceOfTruth": "Python SDK (praisonaiagents)",
|
||||
"status": "PARITY_ACHIEVED",
|
||||
"summary": {
|
||||
"pythonCoreFeatures": 359,
|
||||
"pythonCoreFeatures": 360,
|
||||
"rustFeatures": 667,
|
||||
"gapCount": 80,
|
||||
"parityPercentage": 185.8
|
||||
"gapCount": 81,
|
||||
"parityPercentage": 185.3
|
||||
},
|
||||
"pythonCoreSDK": {
|
||||
"exports": [
|
||||
@@ -275,6 +275,7 @@
|
||||
"WebSearchProvider",
|
||||
"Workflow",
|
||||
"WorkflowContext",
|
||||
"__version__",
|
||||
"aembed",
|
||||
"aembedding",
|
||||
"aembeddings",
|
||||
@@ -4374,6 +4375,15 @@
|
||||
"status": "DONE",
|
||||
"category": "config"
|
||||
},
|
||||
{
|
||||
"feature": "__version__",
|
||||
"python": true,
|
||||
"rust": false,
|
||||
"priority": "P3",
|
||||
"effort": "low",
|
||||
"status": "TODO",
|
||||
"category": "other"
|
||||
},
|
||||
{
|
||||
"feature": "aembed",
|
||||
"python": true,
|
||||
|
||||
@@ -1,16 +1,16 @@
|
||||
# Rust Feature Parity Tracker
|
||||
|
||||
> **Python Features:** 359 | **Rust Features:** 667 | **Parity:** 79.4%
|
||||
> **Python Features:** 360 | **Rust Features:** 667 | **Parity:** 79.2%
|
||||
|
||||
## Summary
|
||||
|
||||
| Metric | Count |
|
||||
|--------|-------|
|
||||
| Python Core Features | 359 |
|
||||
| Python Core Features | 360 |
|
||||
| Rust Features | 667 |
|
||||
| **Actual Gap Count** | **74** |
|
||||
| **Actual Gap Count** | **75** |
|
||||
| Language Limitations (N/A) | 5 |
|
||||
| **Parity** | **79.4%** |
|
||||
| **Parity** | **79.2%** |
|
||||
|
||||
## Implemented Features
|
||||
|
||||
@@ -751,6 +751,7 @@ These Python features cannot be directly implemented in Rust due to reserved key
|
||||
- ❌ `ToolValidationError`
|
||||
- ❌ `ToolsetRegistry`
|
||||
- ❌ `ToolsetSpec`
|
||||
- ❌ `__version__`
|
||||
- ❌ `configure_structured_logging`
|
||||
- ❌ `get_default_policy`
|
||||
- ❌ `get_logger`
|
||||
|
||||
@@ -1,17 +1,17 @@
|
||||
{
|
||||
"version": "1.5.87",
|
||||
"lastUpdated": "2026-06-12",
|
||||
"lastUpdated": "2026-06-17",
|
||||
"generatedBy": "praisonai._dev.parity.generator",
|
||||
"sourceOfTruth": "Python SDK (praisonaiagents)",
|
||||
"summary": {
|
||||
"pythonCoreFeatures": 359,
|
||||
"pythonCoreFeatures": 360,
|
||||
"pythonWrapperFeatures": 99,
|
||||
"typescriptFeatures": 1192,
|
||||
"gapCount": 67,
|
||||
"gapCount": 68,
|
||||
"priorityP0": 2,
|
||||
"priorityP1": 0,
|
||||
"priorityP2": 0,
|
||||
"priorityP3": 65
|
||||
"priorityP3": 66
|
||||
},
|
||||
"pythonCoreSDK": {
|
||||
"path": "/home/runner/work/PraisonAI/PraisonAI/src/praisonai-agents/praisonaiagents",
|
||||
@@ -279,6 +279,7 @@
|
||||
"WebSearchProvider",
|
||||
"Workflow",
|
||||
"WorkflowContext",
|
||||
"__version__",
|
||||
"aembed",
|
||||
"aembedding",
|
||||
"aembeddings",
|
||||
@@ -655,6 +656,7 @@
|
||||
"ToolsetRegistry",
|
||||
"ToolsetSpec",
|
||||
"ValidationError",
|
||||
"__version__",
|
||||
"config",
|
||||
"configure_structured_logging",
|
||||
"embedding",
|
||||
@@ -5491,6 +5493,15 @@
|
||||
"status": "DONE",
|
||||
"category": "config"
|
||||
},
|
||||
{
|
||||
"feature": "__version__",
|
||||
"python": true,
|
||||
"typescript": false,
|
||||
"priority": "P3",
|
||||
"effort": "low",
|
||||
"status": "TODO",
|
||||
"category": "other"
|
||||
},
|
||||
{
|
||||
"feature": "aembed",
|
||||
"python": true,
|
||||
|
||||
@@ -1,20 +1,20 @@
|
||||
# Feature Parity Tracker
|
||||
|
||||
> **Version:** 1.5.87 | **Last Updated:** 2026-06-12
|
||||
> **Version:** 1.5.87 | **Last Updated:** 2026-06-17
|
||||
> **Source of Truth:** Python SDK (praisonaiagents)
|
||||
|
||||
## Summary
|
||||
|
||||
| Metric | Count |
|
||||
|--------|-------|
|
||||
| Python Core Features | 359 |
|
||||
| Python Core Features | 360 |
|
||||
| Python Wrapper Features | 99 |
|
||||
| TypeScript Features | 1192 |
|
||||
| **Gap Count** | **67** |
|
||||
| **Gap Count** | **68** |
|
||||
| P0 (Critical) | 2 |
|
||||
| P1 (High) | 0 |
|
||||
| P2 (Medium) | 0 |
|
||||
| P3 (Low) | 65 |
|
||||
| P3 (Low) | 66 |
|
||||
|
||||
## Gap Matrix
|
||||
|
||||
@@ -148,7 +148,7 @@
|
||||
| `enable\_telemetry` | ✅ | ✅ | low | ✅ DONE |
|
||||
| `get\_telemetry` | ✅ | ✅ | low | ✅ DONE |
|
||||
|
||||
### P3_Advanced (179 done, 65 todo)
|
||||
### P3_Advanced (179 done, 66 todo)
|
||||
|
||||
| Feature | Python | TypeScript | Effort | Status |
|
||||
|---------|--------|------------|--------|--------|
|
||||
@@ -205,6 +205,7 @@
|
||||
| `ToolsetRegistry` | ✅ | ❌ | high | ⏳ TODO |
|
||||
| `ToolsetSpec` | ✅ | ❌ | high | ⏳ TODO |
|
||||
| `ValidationError` | ✅ | ❌ | low | ⏳ TODO |
|
||||
| `\_\_version\_\_` | ✅ | ❌ | low | ⏳ TODO |
|
||||
| `configure\_structured\_logging` | ✅ | ❌ | low | ⏳ TODO |
|
||||
| `get\_default\_policy` | ✅ | ❌ | low | ⏳ TODO |
|
||||
| `get\_logger` | ✅ | ❌ | low | ⏳ TODO |
|
||||
@@ -555,7 +556,7 @@ from praisonaiagents import Memory
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>other</strong> (70 exports)</summary>
|
||||
<summary><strong>other</strong> (71 exports)</summary>
|
||||
|
||||
```python
|
||||
from praisonaiagents import AGGRESSIVE_POLICY, AgentMessageEvent, AgentRunOutcome, Agents, AsyncLearnProtocol, AutonomyConfig, BALANCED_POLICY, BudgetExceededError, CONSERVATIVE_POLICY, CliBackendConfig...
|
||||
|
||||
@@ -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/v4.6.55.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.6.55.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.6.59.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.6.59.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -632,6 +632,37 @@ class AgentsGenerator:
|
||||
return tools_dict
|
||||
|
||||
|
||||
def _load_config(self):
|
||||
"""Load configuration from agent file or agent_yaml."""
|
||||
if self.agent_yaml:
|
||||
config = yaml.safe_load(self.agent_yaml)
|
||||
else:
|
||||
if self.agent_file in ('/app/api:app', 'api:app'):
|
||||
self.agent_file = 'agents.yaml'
|
||||
try:
|
||||
with open(self.agent_file, 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
except FileNotFoundError:
|
||||
print(f"File not found: {self.agent_file}")
|
||||
return None
|
||||
|
||||
# Apply CLI config overrides to both paths (agent_yaml and agent_file)
|
||||
if self.cli_config:
|
||||
self._merge_cli_config(config, self.cli_config)
|
||||
return config
|
||||
|
||||
def _is_workflow_yaml(self, config):
|
||||
"""Check if configuration is workflow mode YAML."""
|
||||
process_type = config.get('process', 'sequential')
|
||||
has_steps = 'steps' in config
|
||||
has_workflow_config = 'workflow' in config
|
||||
workflow_type = config.get('type')
|
||||
return (
|
||||
process_type == 'workflow'
|
||||
or (has_steps and has_workflow_config)
|
||||
or workflow_type in {'job', 'hybrid'}
|
||||
)
|
||||
|
||||
def generate_crew_and_kickoff(self):
|
||||
"""
|
||||
Generates a crew of agents and initiates tasks based on the provided configuration.
|
||||
@@ -649,30 +680,10 @@ class AgentsGenerator:
|
||||
|
||||
This function first loads the agent configuration from the specified file. It then initializes the tools required for the agents based on the specified framework. If the specified framework is "autogen", it loads the LLM configuration dynamically and creates an AssistantAgent for each role in the configuration. It then adds tools to the agents if specified in the configuration. Finally, it prepares tasks for the agents based on the configuration and initiates the tasks using the crew of agents. If the specified framework is not "autogen", it creates a crew of agents and initiates tasks based on the configuration.
|
||||
"""
|
||||
if self.agent_yaml:
|
||||
config = yaml.safe_load(self.agent_yaml)
|
||||
else:
|
||||
if self.agent_file == '/app/api:app' or self.agent_file == 'api:app':
|
||||
self.agent_file = 'agents.yaml'
|
||||
try:
|
||||
with open(self.agent_file, 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
except FileNotFoundError:
|
||||
print(f"File not found: {self.agent_file}")
|
||||
return
|
||||
|
||||
# Apply CLI configuration overrides to YAML config
|
||||
if self.cli_config:
|
||||
# Merge CLI configuration with YAML config
|
||||
self._merge_cli_config(config, self.cli_config)
|
||||
|
||||
# Check if this is a workflow-mode YAML (process: workflow or has steps section)
|
||||
process_type = config.get('process', 'sequential')
|
||||
has_steps = 'steps' in config
|
||||
has_workflow_config = 'workflow' in config
|
||||
|
||||
if process_type == 'workflow' or (has_steps and has_workflow_config):
|
||||
# Route to YAMLWorkflowParser for advanced workflow patterns
|
||||
config = self._load_config()
|
||||
if config is None:
|
||||
return
|
||||
if self._is_workflow_yaml(config):
|
||||
return self._run_yaml_workflow(config)
|
||||
|
||||
# Use shared preparation logic
|
||||
@@ -694,35 +705,12 @@ class AgentsGenerator:
|
||||
Async version of generate_crew_and_kickoff.
|
||||
Generates a crew of agents and initiates tasks based on the provided configuration.
|
||||
"""
|
||||
if self.agent_yaml:
|
||||
config = yaml.safe_load(self.agent_yaml)
|
||||
else:
|
||||
if self.agent_file == '/app/api:app' or self.agent_file == 'api:app':
|
||||
self.agent_file = 'agents.yaml'
|
||||
try:
|
||||
with open(self.agent_file, 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
except FileNotFoundError:
|
||||
print(f"File not found: {self.agent_file}")
|
||||
return
|
||||
|
||||
# Apply CLI configuration overrides to YAML config
|
||||
if self.cli_config:
|
||||
# Merge CLI configuration with YAML config
|
||||
self._merge_cli_config(config, self.cli_config)
|
||||
|
||||
# Check if this is a workflow-mode YAML (process: workflow or has steps section)
|
||||
process_type = config.get('process', 'sequential')
|
||||
has_steps = 'steps' in config
|
||||
has_workflow_config = 'workflow' in config
|
||||
|
||||
if process_type == 'workflow' or (has_steps and has_workflow_config):
|
||||
config = self._load_config()
|
||||
if config is None:
|
||||
return
|
||||
if self._is_workflow_yaml(config):
|
||||
return await self._arun_yaml_workflow(config)
|
||||
else:
|
||||
return await self._arun_framework(config)
|
||||
|
||||
async def _arun_framework(self, config):
|
||||
"""Async version of _run_framework with shared preparation logic."""
|
||||
# Use shared preparation logic
|
||||
prep = self._prepare_for_run(config)
|
||||
|
||||
@@ -737,6 +725,7 @@ class AgentsGenerator:
|
||||
cli_config=getattr(self, 'cli_config', None),
|
||||
)
|
||||
|
||||
|
||||
async def _arun_yaml_workflow(self, config):
|
||||
"""
|
||||
Async version of _run_yaml_workflow using YAMLWorkflowParser.
|
||||
|
||||
@@ -7,6 +7,8 @@ It creates a FastAPI-based web service for deploying agents.
|
||||
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import os
|
||||
|
||||
from praisonaiagents import AgentOSConfig, AgentOSProtocol
|
||||
|
||||
|
||||
@@ -112,6 +114,25 @@ class AgentOS:
|
||||
|
||||
# Register routes
|
||||
self._register_routes(app)
|
||||
|
||||
launch_token = self.config.api_key or os.environ.get("PRAISONAI_AGENTOS_API_KEY")
|
||||
if launch_token:
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
expected = launch_token
|
||||
|
||||
class _AgentOSAuthMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request, call_next):
|
||||
if request.url.path in ("/health", "/"):
|
||||
return await call_next(request)
|
||||
auth = request.headers.get("Authorization", "")
|
||||
token = auth[7:] if auth.startswith("Bearer ") else request.headers.get("X-API-Key", "")
|
||||
if token != expected:
|
||||
return JSONResponse({"error": "Unauthorized"}, status_code=401)
|
||||
return await call_next(request)
|
||||
|
||||
app.add_middleware(_AgentOSAuthMiddleware)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
@@ -19,11 +19,13 @@ if TYPE_CHECKING:
|
||||
from .botos import BotOS
|
||||
from ._session import BotSessionManager
|
||||
from ._dlq import InboundDLQ, DLQEntry
|
||||
from ._ingress import InboundJournal, JournalEntry
|
||||
from ._slack_approval import SlackApproval
|
||||
from ._telegram_approval import TelegramApproval
|
||||
from ._discord_approval import DiscordApproval
|
||||
from ._webhook_approval import WebhookApproval
|
||||
from ._http_approval import HTTPApproval
|
||||
from ._streaming import StreamingConfig, StreamingMode, DraftStreamer
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy loading of bot components."""
|
||||
@@ -83,6 +85,23 @@ def __getattr__(name: str):
|
||||
if name == "DLQEntry":
|
||||
from ._dlq import DLQEntry
|
||||
return DLQEntry
|
||||
# Inbound message journal for durable processing
|
||||
if name == "InboundJournal":
|
||||
from ._ingress import InboundJournal
|
||||
return InboundJournal
|
||||
if name == "JournalEntry":
|
||||
from ._ingress import JournalEntry
|
||||
return JournalEntry
|
||||
# Streaming components
|
||||
if name == "StreamingConfig":
|
||||
from ._streaming import StreamingConfig
|
||||
return StreamingConfig
|
||||
if name == "StreamingMode":
|
||||
from ._streaming import StreamingMode
|
||||
return StreamingMode
|
||||
if name == "DraftStreamer":
|
||||
from ._streaming import DraftStreamer
|
||||
return DraftStreamer
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
__all__ = [
|
||||
@@ -90,8 +109,11 @@ __all__ = [
|
||||
"Bot", "BotOS",
|
||||
"BotSessionManager",
|
||||
"InboundDLQ", "DLQEntry",
|
||||
"InboundJournal", "JournalEntry",
|
||||
"SlackApproval", "TelegramApproval", "DiscordApproval",
|
||||
"WebhookApproval", "HTTPApproval",
|
||||
# Streaming
|
||||
"StreamingConfig", "StreamingMode", "DraftStreamer",
|
||||
# W1
|
||||
"mirror_to_session", "BotSessionManager",
|
||||
]
|
||||
@@ -12,6 +12,7 @@ from typing import Dict, Optional, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent
|
||||
from ._run_control import SessionRunControl
|
||||
|
||||
|
||||
def format_status(
|
||||
@@ -64,6 +65,7 @@ def format_help(
|
||||
"Available Commands",
|
||||
"/status - Show bot status and info",
|
||||
"/new - Reset conversation session",
|
||||
"/stop - Cancel current agent task",
|
||||
"/help - Show this help message",
|
||||
]
|
||||
if extra_commands:
|
||||
@@ -72,3 +74,65 @@ def format_help(
|
||||
lines.append(f"\nAgent: {agent_name}")
|
||||
lines.append(f"Model: {model}")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
async def handle_stop_command(
|
||||
user_id: str,
|
||||
run_control: Optional["SessionRunControl"] = None,
|
||||
) -> str:
|
||||
"""Handle /stop command to cancel current agent task.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
run_control: SessionRunControl instance for managing runs
|
||||
|
||||
Returns:
|
||||
Response message to send to user
|
||||
"""
|
||||
if run_control is None:
|
||||
return "❌ Stop command not available (run control not enabled)"
|
||||
|
||||
try:
|
||||
stopped = await run_control.stop(user_id)
|
||||
if stopped:
|
||||
return "✅ Current task cancelled. Send a new message to start fresh."
|
||||
else:
|
||||
return "ℹ️ No active task to cancel."
|
||||
except Exception as e:
|
||||
return f"❌ Error stopping task: {e}"
|
||||
|
||||
|
||||
def handle_run_status_command(
|
||||
user_id: str,
|
||||
run_control: Optional["SessionRunControl"] = None,
|
||||
) -> str:
|
||||
"""Handle request for current run status.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
run_control: SessionRunControl instance
|
||||
|
||||
Returns:
|
||||
Status message
|
||||
"""
|
||||
if run_control is None:
|
||||
return "Run control not enabled"
|
||||
|
||||
try:
|
||||
status = run_control.get_run_status(user_id)
|
||||
|
||||
if not status["is_running"]:
|
||||
pending_info = " (has queued message)" if status["has_pending"] else ""
|
||||
return f"💤 No active task{pending_info}"
|
||||
|
||||
elapsed = status.get("elapsed_seconds", 0)
|
||||
elapsed_str = f"{elapsed//60}m {elapsed%60}s" if elapsed > 60 else f"{elapsed}s"
|
||||
|
||||
pending_info = ""
|
||||
if status["has_pending"]:
|
||||
preview = status.get("pending_preview", "")
|
||||
pending_info = f"\n📝 Queued: {preview}"
|
||||
|
||||
return f"⚡ Task running for {elapsed_str}{pending_info}"
|
||||
except Exception as e:
|
||||
return f"Error getting status: {e}"
|
||||
@@ -21,6 +21,39 @@ class AgentConfigSchema(BaseModel):
|
||||
tools: List[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class StreamingConfigSchema(BaseModel):
|
||||
"""Schema for streaming reply configuration."""
|
||||
mode: str = "off" # off | draft | progress
|
||||
min_interval: float = 1.5 # Minimum seconds between edits
|
||||
min_delta: int = 120 # Minimum character delta before edit
|
||||
placeholder_text: str = "🤔 Thinking..."
|
||||
progress_prefix: str = "🤔 "
|
||||
|
||||
@field_validator("mode")
|
||||
@classmethod
|
||||
def validate_streaming_mode(cls, v: str) -> str:
|
||||
allowed = {"off", "draft", "progress"}
|
||||
if v not in allowed:
|
||||
raise ValueError(
|
||||
f"Invalid streaming mode '{v}'. Must be one of: {', '.join(sorted(allowed))}"
|
||||
)
|
||||
return v
|
||||
|
||||
@field_validator("min_interval")
|
||||
@classmethod
|
||||
def validate_min_interval(cls, v: float) -> float:
|
||||
if v < 0.1:
|
||||
raise ValueError("min_interval must be at least 0.1 seconds")
|
||||
return v
|
||||
|
||||
@field_validator("min_delta")
|
||||
@classmethod
|
||||
def validate_min_delta(cls, v: int) -> int:
|
||||
if v < 1:
|
||||
raise ValueError("min_delta must be at least 1 character")
|
||||
return v
|
||||
|
||||
|
||||
class ChannelConfigSchema(BaseModel):
|
||||
"""Schema for a single channel configuration."""
|
||||
platform: Optional[str] = None
|
||||
@@ -33,6 +66,7 @@ class ChannelConfigSchema(BaseModel):
|
||||
routes: Dict[str, str] = Field(default_factory=dict)
|
||||
webhook_url: Optional[str] = None
|
||||
webhook_port: int = 8080
|
||||
streaming: Optional[StreamingConfigSchema] = None
|
||||
|
||||
@field_validator("mode")
|
||||
@classmethod
|
||||
|
||||
@@ -253,6 +253,11 @@ class DiscordApproval:
|
||||
if author.get("bot"):
|
||||
continue
|
||||
|
||||
# Only accept replies to the approval message
|
||||
ref = msg.get("message_reference") or {}
|
||||
if ref.get("message_id") != message_id:
|
||||
continue
|
||||
|
||||
text = msg.get("content", "").strip()
|
||||
user_id = author.get("id", "unknown")
|
||||
username = author.get("username", user_id)
|
||||
|
||||
@@ -131,17 +131,20 @@ class HTTPApproval:
|
||||
|
||||
def _build_html(self, request_id: str, info: Dict[str, Any]) -> str:
|
||||
"""Build a minimal approval page."""
|
||||
tool_name = info.get("tool_name", "unknown")
|
||||
risk_level = info.get("risk_level", "unknown")
|
||||
agent_name = info.get("agent_name", "")
|
||||
import html as html_module
|
||||
|
||||
tool_name = html_module.escape(str(info.get("tool_name", "unknown")))
|
||||
risk_level = html_module.escape(str(info.get("risk_level", "unknown")))
|
||||
agent_name = html_module.escape(str(info.get("agent_name", "")))
|
||||
arguments = info.get("arguments", {})
|
||||
|
||||
args_html = ""
|
||||
for k, v in arguments.items():
|
||||
val_str = str(v)
|
||||
val_str = html_module.escape(str(v))
|
||||
if len(val_str) > 200:
|
||||
val_str = val_str[:197] + "..."
|
||||
args_html += f"<tr><td><code>{k}</code></td><td><code>{val_str}</code></td></tr>"
|
||||
key_str = html_module.escape(str(k))
|
||||
args_html += f"<tr><td><code>{key_str}</code></td><td><code>{val_str}</code></td></tr>"
|
||||
|
||||
risk_colors = {
|
||||
"critical": "#FF0000", "high": "#FF8C00",
|
||||
|
||||
@@ -0,0 +1,382 @@
|
||||
"""
|
||||
Inbound Message Journal for PraisonAI bot adapters.
|
||||
|
||||
Provides durable inbound message processing with deduplication, persistent
|
||||
ingress journal, and idempotent replay to prevent duplicate agent runs and
|
||||
silent message loss during crashes or webhook redeliveries.
|
||||
|
||||
Design constraints (per PraisonAI principles):
|
||||
- Wrapper-only — heavy implementation stays out of core SDK
|
||||
- Lazy: sqlite3 is stdlib so no extra dependency
|
||||
- Optional: BotSessionManager works exactly as before unless journal is provided
|
||||
- Bounded: TTL + max_size prevent unbounded disk growth
|
||||
- Thread-safe: per-instance threading.Lock guards SQLite writes
|
||||
- Atomic: claim/complete mechanism ensures at-least-once with effective exactly-once
|
||||
|
||||
Storage schema::
|
||||
|
||||
ingress_journal(
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
ts REAL,
|
||||
platform TEXT,
|
||||
account TEXT,
|
||||
channel_id TEXT,
|
||||
message_id TEXT,
|
||||
payload TEXT,
|
||||
status TEXT, -- 'pending', 'claimed', 'completed'
|
||||
claimed_at REAL,
|
||||
completed_at REAL,
|
||||
attempts INTEGER DEFAULT 0,
|
||||
UNIQUE(platform, account, channel_id, message_id)
|
||||
)
|
||||
|
||||
Public API:
|
||||
- InboundJournal(path, *, max_size=50_000, ttl_seconds=30*86400, claim_timeout=300)
|
||||
- receive(platform, account, channel_id, message_id, payload) -> Optional[str]
|
||||
- claim(key) -> ClaimContext
|
||||
- complete(key)
|
||||
- replay() -> int
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 30 days default — long enough for ops to investigate, short enough not to blow up disk
|
||||
_DEFAULT_TTL_SECONDS = 30 * 86400
|
||||
_DEFAULT_MAX_SIZE = 50_000
|
||||
_DEFAULT_CLAIM_TIMEOUT = 300 # 5 minutes
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class JournalEntry:
|
||||
"""A journaled inbound message."""
|
||||
|
||||
id: int
|
||||
ts: float
|
||||
platform: str
|
||||
account: str
|
||||
channel_id: str
|
||||
message_id: str
|
||||
payload: str
|
||||
status: str
|
||||
claimed_at: Optional[float]
|
||||
completed_at: Optional[float]
|
||||
attempts: int
|
||||
|
||||
|
||||
class ClaimContext:
|
||||
"""Context manager for claiming/releasing journal entries."""
|
||||
|
||||
def __init__(self, journal: "InboundJournal", key: str):
|
||||
self._journal = journal
|
||||
self._key = key
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
# No longer needed - release is handled by the outer context managers
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
# No longer needed - release is handled by the outer context managers
|
||||
pass
|
||||
|
||||
|
||||
class InboundJournal:
|
||||
"""SQLite-backed inbound message journal for durable bot processing.
|
||||
|
||||
Args:
|
||||
path: SQLite file path. Created if missing; parent dirs created.
|
||||
max_size: Maximum entries kept; oldest completed entries evicted when exceeded.
|
||||
ttl_seconds: Completed entries older than this are evicted.
|
||||
claim_timeout: Claimed entries older than this are considered stale for replay.
|
||||
|
||||
Example::
|
||||
|
||||
from praisonai.bots import InboundJournal, BotSessionManager
|
||||
|
||||
journal = InboundJournal(path="~/.praisonai/state/ingress.sqlite")
|
||||
session_mgr = BotSessionManager(platform="telegram", ingress_journal=journal)
|
||||
|
||||
# In message handler:
|
||||
key = journal.receive("telegram", "bot123", chat_id, update.message.message_id, payload)
|
||||
if key is None:
|
||||
return # duplicate redelivery
|
||||
async with journal.claim(key):
|
||||
response = await session_mgr.chat(agent, user_id, text)
|
||||
journal.complete(key)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path: Union[str, Path],
|
||||
*,
|
||||
max_size: int = _DEFAULT_MAX_SIZE,
|
||||
ttl_seconds: int = _DEFAULT_TTL_SECONDS,
|
||||
claim_timeout: int = _DEFAULT_CLAIM_TIMEOUT,
|
||||
) -> None:
|
||||
self.path = Path(path).expanduser()
|
||||
self.max_size = int(max_size)
|
||||
self.ttl_seconds = int(ttl_seconds)
|
||||
self.claim_timeout = int(claim_timeout)
|
||||
self._lock = threading.Lock()
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._init_schema()
|
||||
|
||||
# ── Schema ──────────────────────────────────────────────────────
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
conn = sqlite3.connect(str(self.path))
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA synchronous=NORMAL")
|
||||
return conn
|
||||
|
||||
def _init_schema(self) -> None:
|
||||
with self._lock, self._connect() as conn:
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS ingress_journal (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
ts REAL NOT NULL,
|
||||
platform TEXT NOT NULL,
|
||||
account TEXT NOT NULL,
|
||||
channel_id TEXT NOT NULL,
|
||||
message_id TEXT NOT NULL,
|
||||
payload TEXT NOT NULL DEFAULT '',
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
claimed_at REAL,
|
||||
completed_at REAL,
|
||||
attempts INTEGER DEFAULT 0,
|
||||
UNIQUE(platform, account, channel_id, message_id)
|
||||
)
|
||||
""")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_status ON ingress_journal(status)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_ts ON ingress_journal(ts)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_claimed_at ON ingress_journal(claimed_at)")
|
||||
|
||||
# ── Core API ────────────────────────────────────────────────────
|
||||
def receive(
|
||||
self,
|
||||
platform: str,
|
||||
account: str,
|
||||
channel_id: str,
|
||||
message_id: str,
|
||||
payload: Dict[str, Any] = None
|
||||
) -> Optional[str]:
|
||||
"""Journal an inbound message for durable processing.
|
||||
|
||||
Args:
|
||||
platform: Bot platform ("telegram", "discord", "slack", etc.)
|
||||
account: Bot account/instance identifier
|
||||
channel_id: Platform channel/chat identifier
|
||||
message_id: Platform message identifier
|
||||
payload: Message data to store (will be JSON-serialized)
|
||||
|
||||
Returns:
|
||||
str: Unique key for claim/complete, or None if duplicate
|
||||
"""
|
||||
if payload is None:
|
||||
payload = {}
|
||||
|
||||
with self._lock, self._connect() as conn:
|
||||
try:
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
self._evict_expired_locked(conn)
|
||||
|
||||
# Try to insert - will fail on duplicate (platform, account, channel_id, message_id)
|
||||
cur = conn.execute("""
|
||||
INSERT INTO ingress_journal(ts, platform, account, channel_id,
|
||||
message_id, payload, status)
|
||||
VALUES (?, ?, ?, ?, ?, ?, 'pending')
|
||||
""", (
|
||||
time.time(), platform, account, channel_id,
|
||||
message_id, json.dumps(payload)
|
||||
))
|
||||
|
||||
entry_id = cur.lastrowid
|
||||
self._evict_overflow_locked(conn)
|
||||
conn.commit()
|
||||
|
||||
# Return a key for claim/complete operations
|
||||
return f"{platform}:{account}:{channel_id}:{message_id}:{entry_id}"
|
||||
|
||||
except sqlite3.IntegrityError:
|
||||
# Duplicate message - this is expected for webhook redeliveries
|
||||
conn.rollback()
|
||||
return None
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
|
||||
@contextmanager
|
||||
def claim(self, key: str):
|
||||
"""Claim a journal entry for processing."""
|
||||
self._claim_entry(key)
|
||||
try:
|
||||
yield ClaimContext(self, key)
|
||||
finally:
|
||||
# Release claim if not completed
|
||||
self._release_claim(key)
|
||||
|
||||
@asynccontextmanager
|
||||
async def aclaim(self, key: str):
|
||||
"""Async claim a journal entry for processing."""
|
||||
self._claim_entry(key)
|
||||
try:
|
||||
yield ClaimContext(self, key)
|
||||
finally:
|
||||
# Release claim if not completed
|
||||
self._release_claim(key)
|
||||
|
||||
def complete(self, key: str) -> None:
|
||||
"""Mark a journal entry as completed."""
|
||||
entry_id = self._extract_id_from_key(key)
|
||||
|
||||
with self._lock, self._connect() as conn:
|
||||
conn.execute("""
|
||||
UPDATE ingress_journal
|
||||
SET status = 'completed', completed_at = ?
|
||||
WHERE id = ?
|
||||
""", (time.time(), entry_id))
|
||||
|
||||
# ── Internal Implementation ────────────────────────────────────
|
||||
def _claim_entry(self, key: str) -> None:
|
||||
"""Mark entry as claimed for processing."""
|
||||
entry_id = self._extract_id_from_key(key)
|
||||
|
||||
with self._lock, self._connect() as conn:
|
||||
conn.execute("""
|
||||
UPDATE ingress_journal
|
||||
SET status = 'claimed', claimed_at = ?, attempts = attempts + 1
|
||||
WHERE id = ? AND status IN ('pending', 'claimed')
|
||||
""", (time.time(), entry_id))
|
||||
|
||||
def _release_claim(self, key: str) -> None:
|
||||
"""Release claim without completing (for crash recovery)."""
|
||||
entry_id = self._extract_id_from_key(key)
|
||||
|
||||
with self._lock, self._connect() as conn:
|
||||
conn.execute("""
|
||||
UPDATE ingress_journal
|
||||
SET status = 'pending', claimed_at = NULL
|
||||
WHERE id = ? AND status = 'claimed'
|
||||
""", (entry_id,))
|
||||
|
||||
def _extract_id_from_key(self, key: str) -> int:
|
||||
"""Extract entry ID from journal key."""
|
||||
try:
|
||||
return int(key.split(":")[-1])
|
||||
except (ValueError, IndexError):
|
||||
raise ValueError(f"Invalid journal key format: {key}")
|
||||
|
||||
# ── Maintenance ─────────────────────────────────────────────────
|
||||
def _evict_expired_locked(self, conn: sqlite3.Connection) -> int:
|
||||
"""Remove completed entries older than ttl_seconds."""
|
||||
cutoff = time.time() - self.ttl_seconds
|
||||
cur = conn.execute("""
|
||||
DELETE FROM ingress_journal
|
||||
WHERE status = 'completed' AND completed_at <= ?
|
||||
""", (cutoff,))
|
||||
return int(cur.rowcount or 0)
|
||||
|
||||
def _evict_overflow_locked(self, conn: sqlite3.Connection) -> int:
|
||||
"""Remove oldest entries if over max_size. Prefers completed, then oldest pending."""
|
||||
n = conn.execute("SELECT COUNT(*) FROM ingress_journal").fetchone()[0]
|
||||
if n <= self.max_size:
|
||||
return 0
|
||||
|
||||
excess = n - self.max_size
|
||||
|
||||
# First try to delete completed entries
|
||||
completed_cur = conn.execute("""
|
||||
DELETE FROM ingress_journal WHERE id IN (
|
||||
SELECT id FROM ingress_journal
|
||||
WHERE status = 'completed'
|
||||
ORDER BY completed_at ASC
|
||||
LIMIT ?
|
||||
)
|
||||
""", (excess,))
|
||||
deleted = int(completed_cur.rowcount or 0)
|
||||
|
||||
# If we still have excess after deleting completed entries, delete oldest pending
|
||||
remaining_excess = excess - deleted
|
||||
if remaining_excess > 0:
|
||||
pending_cur = conn.execute("""
|
||||
DELETE FROM ingress_journal WHERE id IN (
|
||||
SELECT id FROM ingress_journal
|
||||
WHERE status = 'pending'
|
||||
ORDER BY ts ASC
|
||||
LIMIT ?
|
||||
)
|
||||
""", (remaining_excess,))
|
||||
deleted += int(pending_cur.rowcount or 0)
|
||||
|
||||
return deleted
|
||||
|
||||
def replay(self) -> int:
|
||||
"""Find and replay stale claimed entries. Returns count replayed."""
|
||||
stale_cutoff = time.time() - self.claim_timeout
|
||||
|
||||
with self._lock, self._connect() as conn:
|
||||
# Find stale claimed entries
|
||||
cur = conn.execute("""
|
||||
SELECT id FROM ingress_journal
|
||||
WHERE status = 'claimed' AND claimed_at <= ?
|
||||
""", (stale_cutoff,))
|
||||
|
||||
stale_ids = [row[0] for row in cur.fetchall()]
|
||||
|
||||
if stale_ids:
|
||||
# Reset them to pending for replay
|
||||
placeholders = ",".join("?" * len(stale_ids))
|
||||
conn.execute(f"""
|
||||
UPDATE ingress_journal
|
||||
SET status = 'pending', claimed_at = NULL
|
||||
WHERE id IN ({placeholders})
|
||||
""", stale_ids)
|
||||
|
||||
logger.info(f"InboundJournal: reset {len(stale_ids)} stale claimed entries for replay")
|
||||
|
||||
return len(stale_ids)
|
||||
|
||||
def size(self) -> int:
|
||||
"""Return total number of entries."""
|
||||
with self._lock, self._connect() as conn:
|
||||
return int(conn.execute("SELECT COUNT(*) FROM ingress_journal").fetchone()[0])
|
||||
|
||||
def pending_count(self) -> int:
|
||||
"""Return number of pending entries."""
|
||||
with self._lock, self._connect() as conn:
|
||||
return int(conn.execute(
|
||||
"SELECT COUNT(*) FROM ingress_journal WHERE status = 'pending'"
|
||||
).fetchone()[0])
|
||||
|
||||
def purge(self) -> int:
|
||||
"""Delete all entries. Returns count removed."""
|
||||
with self._lock, self._connect() as conn:
|
||||
n = conn.execute("SELECT COUNT(*) FROM ingress_journal").fetchone()[0]
|
||||
conn.execute("DELETE FROM ingress_journal")
|
||||
return int(n)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"InboundJournal(path={str(self.path)!r}, "
|
||||
f"size={self.size()}, pending={self.pending_count()}, "
|
||||
f"max_size={self.max_size}, ttl={self.ttl_seconds}s)"
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["InboundJournal", "JournalEntry", "ClaimContext"]
|
||||
@@ -0,0 +1,346 @@
|
||||
"""
|
||||
Per-session run control for PraisonAI bots.
|
||||
|
||||
Provides busy feedback, pending message slots, interrupt support, and
|
||||
optional steering for in-flight agent runs. Solves the silent lock
|
||||
problem where mid-run messages queue invisibly with no user feedback.
|
||||
|
||||
Usage inside a bot message handler::
|
||||
|
||||
# Initialize with busy mode policy
|
||||
run_ctrl = SessionRunControl(busy_mode="queue") # or "interrupt", "steer"
|
||||
|
||||
# Handle message with run control
|
||||
decision = await run_ctrl.submit(user_id, message_text)
|
||||
|
||||
if decision == RunDecision.RUN_NOW:
|
||||
# Start agent.chat() with interrupt controller
|
||||
response = await session_mgr.chat(agent, user_id, text)
|
||||
elif decision == RunDecision.QUEUED:
|
||||
# Return acknowledgment, message will be processed after current run
|
||||
return "⏳ Noted - will consider this after the current task finishes"
|
||||
elif decision == RunDecision.MERGED:
|
||||
# Message was merged into pending slot
|
||||
return "⏳ Added to pending message"
|
||||
|
||||
Key features:
|
||||
- Busy acknowledgment for mid-run messages instead of silent queuing
|
||||
- Per-user pending slot that merges follow-ups received mid-run
|
||||
- /stop command via InterruptController integration
|
||||
- Optional steering mode to inject messages into running turns
|
||||
- Run generation tracking to prevent canceled runs overwriting fresh state
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import Dict, Optional, Any, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents.agent.interrupt import InterruptController
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RunDecision(Enum):
|
||||
"""Decision for how to handle a submitted message."""
|
||||
RUN_NOW = "run_now" # No current run, start immediately
|
||||
QUEUED = "queued" # Current run active, message queued for next turn
|
||||
MERGED = "merged" # Message merged into existing pending slot
|
||||
INTERRUPTED = "interrupted" # Current run was interrupted, start new run
|
||||
|
||||
|
||||
class BusyMode(Enum):
|
||||
"""Policy for handling messages during active runs."""
|
||||
QUEUE = "queue" # Queue message for next turn (default)
|
||||
INTERRUPT = "interrupt" # Cancel current run and start new one
|
||||
STEER = "steer" # Inject message into current run via steering
|
||||
|
||||
|
||||
class SessionRunState:
|
||||
"""Per-user run state tracking."""
|
||||
|
||||
def __init__(self):
|
||||
self.is_running = False
|
||||
self.run_generation = 0 # Incremented on each run start
|
||||
self.pending_message: Optional[str] = None
|
||||
self.interrupt_controller: Optional["InterruptController"] = None
|
||||
self.start_time: Optional[float] = None
|
||||
self.last_activity: float = time.time() # Always track activity for cleanup
|
||||
|
||||
|
||||
class SessionRunControl:
|
||||
"""Per-session pending slot + interrupt handle around agent runs.
|
||||
|
||||
Manages run states, pending messages, and interrupt controllers
|
||||
for each user session to provide better UX around long-running
|
||||
agent operations.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
busy_mode: str = "queue",
|
||||
busy_ack_template: str = "⏳ {action} — will be considered next"
|
||||
):
|
||||
"""Initialize session run control.
|
||||
|
||||
Args:
|
||||
busy_mode: Policy for mid-run messages ("queue", "interrupt", "steer")
|
||||
busy_ack_template: Template for busy acknowledgment messages
|
||||
"""
|
||||
try:
|
||||
self._busy_mode = BusyMode(busy_mode)
|
||||
except ValueError:
|
||||
logger.warning(f"Invalid busy_mode '{busy_mode}', using 'queue'")
|
||||
self._busy_mode = BusyMode.QUEUE
|
||||
|
||||
self._busy_ack_template = busy_ack_template
|
||||
self._sessions: Dict[str, SessionRunState] = {}
|
||||
self._locks: Dict[str, asyncio.Lock] = {}
|
||||
|
||||
def _get_lock(self, user_id: str) -> asyncio.Lock:
|
||||
"""Get or create lock for user."""
|
||||
if user_id not in self._locks:
|
||||
self._locks[user_id] = asyncio.Lock()
|
||||
return self._locks[user_id]
|
||||
|
||||
def _get_session(self, user_id: str) -> SessionRunState:
|
||||
"""Get or create session state for user."""
|
||||
if user_id not in self._sessions:
|
||||
self._sessions[user_id] = SessionRunState()
|
||||
return self._sessions[user_id]
|
||||
|
||||
async def submit(self, user_id: str, text: str) -> RunDecision:
|
||||
"""Submit a message and get decision on how to handle it.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
text: Message text
|
||||
|
||||
Returns:
|
||||
RunDecision indicating how to proceed
|
||||
"""
|
||||
lock = self._get_lock(user_id)
|
||||
async with lock:
|
||||
session = self._get_session(user_id)
|
||||
session.last_activity = time.time() # Update activity timestamp
|
||||
|
||||
if not session.is_running:
|
||||
# No active run, start immediately
|
||||
session.is_running = True
|
||||
session.run_generation += 1
|
||||
session.start_time = time.time()
|
||||
|
||||
# Create interrupt controller for this run
|
||||
try:
|
||||
from praisonaiagents.agent.interrupt import InterruptController
|
||||
session.interrupt_controller = InterruptController()
|
||||
except ImportError:
|
||||
logger.warning("InterruptController not available")
|
||||
session.interrupt_controller = None
|
||||
|
||||
return RunDecision.RUN_NOW
|
||||
|
||||
# There's an active run, handle based on busy mode
|
||||
if self._busy_mode == BusyMode.INTERRUPT:
|
||||
# Cancel current run and start new one
|
||||
if session.interrupt_controller:
|
||||
session.interrupt_controller.request("user_new_message")
|
||||
|
||||
session.pending_message = None # Clear any pending
|
||||
session.run_generation += 1
|
||||
session.start_time = time.time()
|
||||
|
||||
# Create new interrupt controller
|
||||
try:
|
||||
from praisonaiagents.agent.interrupt import InterruptController
|
||||
session.interrupt_controller = InterruptController()
|
||||
except ImportError:
|
||||
session.interrupt_controller = None
|
||||
|
||||
return RunDecision.INTERRUPTED
|
||||
|
||||
elif self._busy_mode == BusyMode.STEER:
|
||||
# Try to inject into running agent via steering
|
||||
# Note: This requires the agent to have steering enabled
|
||||
# For now, we'll fall back to queueing
|
||||
logger.warning("Steer mode is not yet implemented, falling back to queue mode")
|
||||
# TODO: Implement steering integration when agent reference is available
|
||||
|
||||
# Default to queue mode (or steer fallback)
|
||||
if session.pending_message is None:
|
||||
session.pending_message = text
|
||||
return RunDecision.QUEUED
|
||||
else:
|
||||
# Merge with existing pending message
|
||||
session.pending_message = f"{session.pending_message}\n{text}"
|
||||
return RunDecision.MERGED
|
||||
|
||||
async def get_busy_ack_message(self, user_id: str, decision: RunDecision) -> str:
|
||||
"""Get appropriate busy acknowledgment message.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
decision: The decision returned by submit()
|
||||
|
||||
Returns:
|
||||
Acknowledgment message to send to user
|
||||
"""
|
||||
session = self._get_session(user_id)
|
||||
|
||||
if decision == RunDecision.QUEUED:
|
||||
action = "noted"
|
||||
elif decision == RunDecision.MERGED:
|
||||
action = "added to pending request"
|
||||
elif decision == RunDecision.INTERRUPTED:
|
||||
return "⚠️ Previous task cancelled, starting your new request"
|
||||
else:
|
||||
return "" # No ack needed for RUN_NOW
|
||||
|
||||
# Add timing info if available
|
||||
timing = ""
|
||||
if session.start_time:
|
||||
elapsed = int(time.time() - session.start_time)
|
||||
if elapsed > 60:
|
||||
timing = f" (running for {elapsed//60}m {elapsed%60}s)"
|
||||
elif elapsed > 10:
|
||||
timing = f" (running for {elapsed}s)"
|
||||
|
||||
return self._busy_ack_template.format(action=action) + timing
|
||||
|
||||
async def stop(self, user_id: str) -> bool:
|
||||
"""Stop the current run for a user.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
|
||||
Returns:
|
||||
True if a run was stopped, False if no active run
|
||||
"""
|
||||
lock = self._get_lock(user_id)
|
||||
async with lock:
|
||||
session = self._get_session(user_id)
|
||||
|
||||
if not session.is_running:
|
||||
return False
|
||||
|
||||
# Request interrupt
|
||||
if session.interrupt_controller:
|
||||
session.interrupt_controller.request("user_stop_command")
|
||||
|
||||
# Clear state
|
||||
session.is_running = False
|
||||
session.pending_message = None
|
||||
session.interrupt_controller = None
|
||||
session.start_time = None
|
||||
|
||||
return True
|
||||
|
||||
def next_pending(self, user_id: str) -> Optional[str]:
|
||||
"""Get and clear the next pending message for a user.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
|
||||
Returns:
|
||||
Pending message if any, None otherwise
|
||||
"""
|
||||
session = self._get_session(user_id)
|
||||
pending = session.pending_message
|
||||
session.pending_message = None
|
||||
return pending
|
||||
|
||||
async def finish_run(self, user_id: str, run_generation: Optional[int] = None):
|
||||
"""Mark a run as finished.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
run_generation: Generation of the run that finished (for race protection)
|
||||
"""
|
||||
lock = self._get_lock(user_id)
|
||||
async with lock:
|
||||
session = self._get_session(user_id)
|
||||
session.last_activity = time.time() # Update activity timestamp
|
||||
|
||||
# Race protection: only clear if this is the current generation
|
||||
if run_generation is None or run_generation == session.run_generation:
|
||||
session.is_running = False
|
||||
session.interrupt_controller = None
|
||||
session.start_time = None
|
||||
|
||||
def get_interrupt_controller(self, user_id: str) -> Optional["InterruptController"]:
|
||||
"""Get the interrupt controller for a user's current run.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
|
||||
Returns:
|
||||
InterruptController if run is active, None otherwise
|
||||
"""
|
||||
session = self._get_session(user_id)
|
||||
if session.is_running:
|
||||
return session.interrupt_controller
|
||||
return None
|
||||
|
||||
def get_run_status(self, user_id: str) -> Dict[str, Any]:
|
||||
"""Get run status for a user.
|
||||
|
||||
Args:
|
||||
user_id: User identifier
|
||||
|
||||
Returns:
|
||||
Dict with run status information
|
||||
"""
|
||||
session = self._get_session(user_id)
|
||||
|
||||
status = {
|
||||
"is_running": session.is_running,
|
||||
"run_generation": session.run_generation,
|
||||
"has_pending": session.pending_message is not None,
|
||||
"busy_mode": self._busy_mode.value
|
||||
}
|
||||
|
||||
if session.start_time:
|
||||
status["elapsed_seconds"] = int(time.time() - session.start_time)
|
||||
|
||||
if session.pending_message:
|
||||
status["pending_preview"] = session.pending_message[:100] + "..." if len(session.pending_message) > 100 else session.pending_message
|
||||
|
||||
return status
|
||||
|
||||
@property
|
||||
def active_sessions_count(self) -> int:
|
||||
"""Number of sessions with active runs."""
|
||||
return sum(1 for session in self._sessions.values() if session.is_running)
|
||||
|
||||
def cleanup_stale_sessions(self, max_age_seconds: int = 3600) -> int:
|
||||
"""Clean up stale sessions older than max_age_seconds.
|
||||
|
||||
Args:
|
||||
max_age_seconds: Maximum age in seconds
|
||||
|
||||
Returns:
|
||||
Number of sessions cleaned up
|
||||
"""
|
||||
if max_age_seconds <= 0:
|
||||
return 0
|
||||
|
||||
now = time.time()
|
||||
stale_users = []
|
||||
|
||||
for user_id, session in self._sessions.items():
|
||||
# Clean up based on last activity instead of only active runs
|
||||
if (now - session.last_activity) > max_age_seconds:
|
||||
stale_users.append(user_id)
|
||||
|
||||
for user_id in stale_users:
|
||||
del self._sessions[user_id]
|
||||
self._locks.pop(user_id, None)
|
||||
|
||||
if stale_users:
|
||||
logger.debug(f"SessionRunControl: cleaned up {len(stale_users)} stale sessions")
|
||||
|
||||
return len(stale_users)
|
||||
@@ -64,6 +64,8 @@ class BotSessionManager:
|
||||
platform: str = "",
|
||||
dlq: Optional[Any] = None,
|
||||
identity_resolver: Optional[Any] = None,
|
||||
ingress_journal: Optional[Any] = None,
|
||||
run_control: Optional[Any] = None,
|
||||
) -> None:
|
||||
self._histories: Dict[str, List[Dict[str, Any]]] = {}
|
||||
self._locks: Dict[str, asyncio.Lock] = {}
|
||||
@@ -80,6 +82,13 @@ class BotSessionManager:
|
||||
# session key is the resolver-returned unified user id, so the
|
||||
# same human pinging from multiple platforms shares one history.
|
||||
self._identity_resolver = identity_resolver
|
||||
# Ingress journal: optional durable message processing with dedup.
|
||||
# When set, messages are journaled before agent processing for
|
||||
# crash recovery and webhook redelivery protection.
|
||||
self._ingress_journal = ingress_journal
|
||||
self._last_journal_key = None # Store key for delayed completion
|
||||
# Run control for in-flight message handling
|
||||
self._run_control = run_control
|
||||
|
||||
def _storage_key(self, user_id: str) -> str:
|
||||
"""Resolve a raw platform user id to the in-memory/store key.
|
||||
@@ -171,6 +180,9 @@ class BotSessionManager:
|
||||
chat_id: str = "",
|
||||
thread_id: str = "",
|
||||
user_name: str = "",
|
||||
message_id: str = "",
|
||||
account: str = "",
|
||||
stream_callback: Optional[Any] = None,
|
||||
) -> str:
|
||||
"""Run ``agent.chat(prompt)`` with *user_id*-scoped history.
|
||||
|
||||
@@ -180,13 +192,42 @@ class BotSessionManager:
|
||||
Uses both a per-user lock (serialise same user) and a per-agent
|
||||
lock (prevent concurrent history swaps on a shared Agent).
|
||||
|
||||
Args:
|
||||
stream_callback: Optional async callback for streaming events.
|
||||
If provided, will be passed to agent.astart() for streaming.
|
||||
|
||||
N4 — Inbound DLQ: if a ``dlq`` was passed to ``__init__`` and
|
||||
``agent.chat()`` raises, the failing message is persisted to
|
||||
the dead-letter queue **before** the exception is re-raised.
|
||||
This makes the error visible to the caller (so the bot adapter
|
||||
can log / show the user a friendly message) while preserving
|
||||
the message for later replay.
|
||||
|
||||
Ingress Journal: if an ``ingress_journal`` was passed to ``__init__``
|
||||
and ``message_id`` is provided, the message is journaled with deduplication
|
||||
and claim/complete semantics for crash-safe, exactly-once processing.
|
||||
"""
|
||||
# Handle ingress journaling for durable message processing
|
||||
journal_key = None
|
||||
if self._ingress_journal is not None and message_id:
|
||||
payload = {
|
||||
"user_id": user_id,
|
||||
"prompt": prompt,
|
||||
"chat_id": chat_id,
|
||||
"thread_id": thread_id,
|
||||
"user_name": user_name,
|
||||
}
|
||||
journal_key = self._ingress_journal.receive(
|
||||
platform=self._platform or "unknown",
|
||||
account=account or "default",
|
||||
channel_id=chat_id or user_id,
|
||||
message_id=message_id,
|
||||
payload=payload
|
||||
)
|
||||
if journal_key is None:
|
||||
# Duplicate message - return empty response
|
||||
return ""
|
||||
|
||||
self._last_active[self._storage_key(user_id)] = time.monotonic()
|
||||
user_lock = self._get_lock(user_id)
|
||||
agent_lock = self._get_agent_lock(agent)
|
||||
@@ -213,62 +254,92 @@ class BotSessionManager:
|
||||
_clear_ctx = None # type: ignore[assignment]
|
||||
|
||||
try:
|
||||
async with user_lock:
|
||||
# Load history (may hit disk via run_in_executor for async safety)
|
||||
loop = asyncio.get_running_loop()
|
||||
user_history = await loop.run_in_executor(
|
||||
None, self._load_history, user_id
|
||||
)
|
||||
# Claim journal entry if we have one
|
||||
if journal_key is not None:
|
||||
claim_ctx = self._ingress_journal.aclaim(journal_key)
|
||||
await claim_ctx.__aenter__()
|
||||
else:
|
||||
claim_ctx = None
|
||||
|
||||
try:
|
||||
async with user_lock:
|
||||
# Load history (may hit disk via run_in_executor for async safety)
|
||||
loop = asyncio.get_running_loop()
|
||||
user_history = await loop.run_in_executor(
|
||||
None, self._load_history, user_id
|
||||
)
|
||||
|
||||
# W1 robustness: hold ``agent_lock`` across the FULL LLM call
|
||||
# (not only the history swap) so concurrent users on a shared
|
||||
# Agent instance never observe each other's chat_history.
|
||||
async with agent_lock:
|
||||
saved_history = agent.chat_history
|
||||
agent.chat_history = user_history
|
||||
try:
|
||||
# Copy current task's contextvars (incl. SessionContext)
|
||||
# into the worker thread so tools the agent invokes can
|
||||
# read platform/user metadata.
|
||||
import contextvars
|
||||
_ctx = contextvars.copy_context()
|
||||
response = await loop.run_in_executor(
|
||||
None, _ctx.run, agent.chat, prompt
|
||||
)
|
||||
# Capture updated history before restoring caller's.
|
||||
updated_history = agent.chat_history
|
||||
except Exception as exc:
|
||||
# N4: persist the failed inbound message before bubbling.
|
||||
if self._dlq is not None:
|
||||
try:
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._dlq.enqueue(
|
||||
platform=self._platform,
|
||||
user_id=user_id,
|
||||
prompt=prompt,
|
||||
error=f"{type(exc).__name__}: {exc}",
|
||||
chat_id=chat_id,
|
||||
thread_id=thread_id,
|
||||
user_name=user_name,
|
||||
# W1 robustness: hold ``agent_lock`` across the FULL LLM call
|
||||
# (not only the history swap) so concurrent users on a shared
|
||||
# Agent instance never observe each other's chat_history.
|
||||
async with agent_lock:
|
||||
saved_history = agent.chat_history
|
||||
agent.chat_history = user_history
|
||||
try:
|
||||
# Choose streaming vs non-streaming path based on callback
|
||||
if stream_callback:
|
||||
# Streaming path: use agent.astart() with stream callback
|
||||
response = await agent.astart(prompt, stream_callback=stream_callback)
|
||||
# Handle AutonomyResult when autonomy is enabled in caller mode
|
||||
if hasattr(response, 'output'):
|
||||
response = response.output
|
||||
else:
|
||||
# Legacy non-streaming path: use agent.chat() in executor
|
||||
# Copy current task's contextvars (incl. SessionContext)
|
||||
# into the worker thread so tools the agent invokes can
|
||||
# read platform/user metadata.
|
||||
import contextvars
|
||||
_ctx = contextvars.copy_context()
|
||||
response = await loop.run_in_executor(
|
||||
None, _ctx.run, agent.chat, prompt
|
||||
)
|
||||
# Capture updated history before restoring caller's.
|
||||
updated_history = agent.chat_history
|
||||
except Exception as exc:
|
||||
# N4: persist the failed inbound message before bubbling.
|
||||
if self._dlq is not None:
|
||||
try:
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._dlq.enqueue(
|
||||
platform=self._platform,
|
||||
user_id=user_id,
|
||||
prompt=prompt,
|
||||
error=f"{type(exc).__name__}: {exc}",
|
||||
chat_id=chat_id,
|
||||
thread_id=thread_id,
|
||||
user_name=user_name,
|
||||
)
|
||||
)
|
||||
)
|
||||
except Exception as dlq_exc: # pragma: no cover — defensive
|
||||
logger.error(
|
||||
"Failed to enqueue inbound DLQ entry: %s", dlq_exc
|
||||
)
|
||||
agent.chat_history = saved_history
|
||||
raise
|
||||
finally:
|
||||
agent.chat_history = saved_history
|
||||
except Exception as dlq_exc: # pragma: no cover — defensive
|
||||
logger.error(
|
||||
"Failed to enqueue inbound DLQ entry: %s", dlq_exc
|
||||
)
|
||||
agent.chat_history = saved_history
|
||||
raise
|
||||
finally:
|
||||
agent.chat_history = saved_history
|
||||
|
||||
# Persist outside the agent_lock — it's per-user and the agent
|
||||
# is no longer touched.
|
||||
await loop.run_in_executor(
|
||||
None, self._save_history, user_id, updated_history
|
||||
)
|
||||
# Persist outside the agent_lock — it's per-user and the agent
|
||||
# is no longer touched.
|
||||
await loop.run_in_executor(
|
||||
None, self._save_history, user_id, updated_history
|
||||
)
|
||||
|
||||
return response
|
||||
# Store journal key in instance for later completion after message delivery
|
||||
self._last_journal_key = journal_key
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
# Handle any remaining exceptions and ensure claim is released
|
||||
if claim_ctx is not None:
|
||||
await claim_ctx.__aexit__(type(e), e, e.__traceback__)
|
||||
raise
|
||||
else:
|
||||
# Clean exit - no exception
|
||||
if claim_ctx is not None:
|
||||
await claim_ctx.__aexit__(None, None, None)
|
||||
finally:
|
||||
# Always clear task-local session context, even if an exception occurred.
|
||||
if ctx_token is not None and _clear_ctx is not None:
|
||||
@@ -277,6 +348,129 @@ class BotSessionManager:
|
||||
except Exception as e:
|
||||
logger.debug("Failed to clear task-local session context for ctx_token=%r: %s", ctx_token, e)
|
||||
|
||||
async def chat_with_run_control(
|
||||
self,
|
||||
agent: "Agent",
|
||||
user_id: str,
|
||||
prompt: str,
|
||||
chat_id: str = "",
|
||||
thread_id: str = "",
|
||||
user_name: str = "",
|
||||
) -> Dict[str, Any]:
|
||||
"""Run agent.chat() with run control for better UX during long operations.
|
||||
|
||||
This method integrates with SessionRunControl to provide:
|
||||
- Busy acknowledgment for mid-run messages
|
||||
- Pending message slot for follow-ups
|
||||
- Interrupt support via /stop command
|
||||
- Optional steering for real-time guidance
|
||||
|
||||
Returns:
|
||||
Dict with 'response' (str) and 'metadata' (dict) keys.
|
||||
Metadata includes run control information.
|
||||
"""
|
||||
if self._run_control is None:
|
||||
# Fall back to regular chat if no run control
|
||||
response = await self.chat(agent, user_id, prompt, chat_id, thread_id, user_name)
|
||||
return {"response": response, "metadata": {"run_control": False}}
|
||||
|
||||
try:
|
||||
# Import here to avoid circular dependency
|
||||
from ._run_control import RunDecision
|
||||
except ImportError:
|
||||
# Fall back if run control not available
|
||||
response = await self.chat(agent, user_id, prompt, chat_id, thread_id, user_name)
|
||||
return {"response": response, "metadata": {"run_control": False, "error": "run_control_unavailable"}}
|
||||
|
||||
# Submit message to run control
|
||||
decision = await self._run_control.submit(user_id, prompt)
|
||||
|
||||
if decision in (RunDecision.QUEUED, RunDecision.MERGED):
|
||||
# Message was queued or merged, return acknowledgment
|
||||
ack_msg = await self._run_control.get_busy_ack_message(user_id, decision)
|
||||
return {
|
||||
"response": ack_msg,
|
||||
"metadata": {
|
||||
"run_control": True,
|
||||
"decision": decision.value,
|
||||
"queued": True
|
||||
}
|
||||
}
|
||||
|
||||
# We're running now (RUN_NOW or INTERRUPTED)
|
||||
run_generation = None
|
||||
interrupt_controller = None
|
||||
|
||||
if decision == RunDecision.RUN_NOW:
|
||||
# Get the interrupt controller for this run
|
||||
interrupt_controller = self._run_control.get_interrupt_controller(user_id)
|
||||
elif decision == RunDecision.INTERRUPTED:
|
||||
# Previous run was cancelled, get new controller
|
||||
interrupt_controller = self._run_control.get_interrupt_controller(user_id)
|
||||
|
||||
# Get current run generation for race protection
|
||||
status = self._run_control.get_run_status(user_id)
|
||||
run_generation = status.get("run_generation")
|
||||
|
||||
try:
|
||||
# Attach interrupt controller to agent if available
|
||||
original_interrupt = None
|
||||
if interrupt_controller and hasattr(agent, '_interrupt_controller'):
|
||||
original_interrupt = getattr(agent, '_interrupt_controller', None)
|
||||
agent._interrupt_controller = interrupt_controller
|
||||
|
||||
# Run the chat with the existing method
|
||||
response = await self.chat(agent, user_id, prompt, chat_id, thread_id, user_name)
|
||||
|
||||
# Check for pending messages to process next
|
||||
pending = self._run_control.next_pending(user_id)
|
||||
pending_info = {}
|
||||
if pending:
|
||||
pending_info = {
|
||||
"next_pending": pending[:100] + "..." if len(pending) > 100 else pending
|
||||
}
|
||||
|
||||
return {
|
||||
"response": response,
|
||||
"metadata": {
|
||||
"run_control": True,
|
||||
"decision": decision.value,
|
||||
"completed": True,
|
||||
"run_generation": run_generation,
|
||||
**pending_info
|
||||
}
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
# Handle interruption specifically
|
||||
if interrupt_controller and interrupt_controller.is_set():
|
||||
reason = interrupt_controller.reason or "unknown"
|
||||
return {
|
||||
"response": f"⚠️ Task cancelled: {reason}",
|
||||
"metadata": {
|
||||
"run_control": True,
|
||||
"decision": decision.value,
|
||||
"interrupted": True,
|
||||
"reason": reason,
|
||||
"run_generation": run_generation
|
||||
}
|
||||
}
|
||||
else:
|
||||
# Re-raise other exceptions
|
||||
raise
|
||||
|
||||
finally:
|
||||
# Restore original interrupt controller
|
||||
if interrupt_controller and hasattr(agent, '_interrupt_controller'):
|
||||
if original_interrupt is not None:
|
||||
agent._interrupt_controller = original_interrupt
|
||||
else:
|
||||
agent._interrupt_controller = None
|
||||
|
||||
# Mark run as finished
|
||||
if run_generation is not None:
|
||||
await self._run_control.finish_run(user_id, run_generation)
|
||||
|
||||
def reap_stale(self, max_age_seconds: int) -> int:
|
||||
"""Remove sessions older than *max_age_seconds*. Returns count reaped.
|
||||
|
||||
@@ -304,6 +498,23 @@ class BotSessionManager:
|
||||
logger.debug("BotSessionManager: reaped %d stale sessions", len(stale))
|
||||
return len(stale)
|
||||
|
||||
def complete_last_journal_entry(self) -> bool:
|
||||
"""Complete the last journal entry if one exists.
|
||||
|
||||
Call this after successfully delivering a message to the platform
|
||||
to ensure the journal entry is marked as completed.
|
||||
|
||||
Returns True if an entry was completed, False if no entry was pending.
|
||||
"""
|
||||
if self._last_journal_key is not None and self._ingress_journal is not None:
|
||||
try:
|
||||
self._ingress_journal.complete(self._last_journal_key)
|
||||
self._last_journal_key = None
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning("Failed to complete journal entry: %s", e)
|
||||
return False
|
||||
|
||||
def reset(self, user_id: str) -> bool:
|
||||
"""Clear a user's session history. Returns True if it existed."""
|
||||
storage_key = self._storage_key(user_id)
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
"""
|
||||
Streaming response helpers for bot channels.
|
||||
|
||||
Provides live draft/streaming reply functionality that connects core SDK streaming
|
||||
events to bot channel message edits, enabling progressive response delivery.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional, Protocol, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents.streaming.events import StreamEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StreamingMode(Enum):
|
||||
"""Streaming reply modes for bot channels."""
|
||||
|
||||
OFF = "off" # Current behavior (single final message, chunked)
|
||||
DRAFT = "draft" # Send placeholder, edit in place with growing content
|
||||
PROGRESS = "progress" # Show compact status, then replace with final answer
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamingConfig:
|
||||
"""Configuration for streaming bot replies."""
|
||||
|
||||
mode: StreamingMode = StreamingMode.OFF
|
||||
min_interval: float = 1.5 # Minimum seconds between edits
|
||||
min_delta: int = 120 # Minimum characters before triggering edit
|
||||
placeholder_text: str = "🤔 Thinking..."
|
||||
progress_prefix: str = "🤔 "
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "StreamingConfig":
|
||||
"""Create config from dictionary (e.g., from YAML)."""
|
||||
mode_str = data.get("mode", "off")
|
||||
try:
|
||||
mode = StreamingMode(mode_str)
|
||||
except ValueError:
|
||||
logger.warning("Invalid streaming mode '%s', using 'off'", mode_str)
|
||||
mode = StreamingMode.OFF
|
||||
|
||||
return cls(
|
||||
mode=mode,
|
||||
min_interval=data.get("min_interval", 1.5),
|
||||
min_delta=data.get("min_delta", 120),
|
||||
placeholder_text=data.get("placeholder_text", "🤔 Thinking..."),
|
||||
progress_prefix=data.get("progress_prefix", "🤔 "),
|
||||
)
|
||||
|
||||
|
||||
class BotAdapter(Protocol):
|
||||
"""Protocol for bot adapters that support message editing."""
|
||||
|
||||
async def send_message(self, channel_id: str, content: str) -> Any:
|
||||
"""Send a new message and return message info."""
|
||||
...
|
||||
|
||||
async def edit_message(self, channel_id: str, message_id: str, content: str) -> Any:
|
||||
"""Edit an existing message."""
|
||||
...
|
||||
|
||||
|
||||
class DraftStreamer:
|
||||
"""
|
||||
Manages live draft updates for streaming bot replies.
|
||||
|
||||
Subscribes to agent streaming events and progressively updates a bot message
|
||||
with growing content, respecting rate limits and coalescing updates.
|
||||
|
||||
Usage:
|
||||
streamer = DraftStreamer(
|
||||
adapter=bot_adapter,
|
||||
channel_id="123456",
|
||||
config=StreamingConfig(mode=StreamingMode.DRAFT, min_interval=1.5)
|
||||
)
|
||||
|
||||
# Start streaming
|
||||
message_id = await streamer.start()
|
||||
|
||||
# Subscribe to agent events
|
||||
response = await agent.chat(
|
||||
prompt,
|
||||
stream_callback=streamer.on_event
|
||||
)
|
||||
|
||||
# Finalize with complete response
|
||||
await streamer.finalize(response)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
adapter: BotAdapter,
|
||||
channel_id: str,
|
||||
config: StreamingConfig,
|
||||
rate_limiter: Optional[Any] = None,
|
||||
):
|
||||
self._adapter = adapter
|
||||
self._channel_id = channel_id
|
||||
self._config = config
|
||||
self._rate_limiter = rate_limiter
|
||||
|
||||
# State
|
||||
self._message_id: Optional[str] = None
|
||||
self._content_buffer = ""
|
||||
self._last_edit_time = 0.0
|
||||
self._last_content_length = 0
|
||||
self._current_tool: Optional[str] = None
|
||||
self._is_finalized = False
|
||||
self._update_task: Optional[asyncio.Task] = None
|
||||
self._pending_update = False
|
||||
|
||||
logger.debug("DraftStreamer initialized for channel %s, mode=%s", channel_id, config.mode)
|
||||
|
||||
async def start(self) -> str:
|
||||
"""Start streaming by sending initial placeholder message.
|
||||
|
||||
Returns:
|
||||
Message ID of the placeholder message
|
||||
"""
|
||||
if self._config.mode == StreamingMode.OFF:
|
||||
raise ValueError("Cannot start streaming when mode is OFF")
|
||||
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
|
||||
# Send placeholder
|
||||
placeholder = self._config.placeholder_text
|
||||
message_info = await self._adapter.send_message(self._channel_id, placeholder)
|
||||
|
||||
# Extract message ID - adapt to different bot adapter formats
|
||||
if hasattr(message_info, 'message_id'):
|
||||
self._message_id = message_info.message_id
|
||||
elif isinstance(message_info, dict) and 'message_id' in message_info:
|
||||
self._message_id = message_info['message_id']
|
||||
elif hasattr(message_info, 'id'):
|
||||
self._message_id = str(message_info.id)
|
||||
else:
|
||||
# Fallback - try to extract from string representation
|
||||
self._message_id = str(message_info)
|
||||
|
||||
self._last_edit_time = time.monotonic()
|
||||
logger.debug("DraftStreamer started, placeholder sent with message_id=%s", self._message_id)
|
||||
return self._message_id
|
||||
|
||||
async def on_event(self, event: "StreamEvent") -> None:
|
||||
"""Stream event callback to handle progressive updates."""
|
||||
if self._is_finalized or not self._message_id:
|
||||
return
|
||||
|
||||
# Import here to avoid circular imports
|
||||
from praisonaiagents.streaming.events import StreamEventType
|
||||
|
||||
if event.type == StreamEventType.DELTA_TEXT and event.content:
|
||||
self._content_buffer += event.content
|
||||
await self._schedule_update()
|
||||
|
||||
elif event.type == StreamEventType.DELTA_TOOL_CALL and event.tool_call:
|
||||
tool_name = event.tool_call.get('name', 'unknown_tool')
|
||||
if self._config.mode == StreamingMode.PROGRESS:
|
||||
# Update current tool for progress display
|
||||
self._current_tool = tool_name
|
||||
await self._schedule_update()
|
||||
|
||||
elif event.type == StreamEventType.TOOL_CALL_START and event.tool_call:
|
||||
tool_name = event.tool_call.get('name', 'unknown_tool')
|
||||
if self._config.mode == StreamingMode.PROGRESS:
|
||||
self._current_tool = tool_name
|
||||
await self._schedule_update()
|
||||
|
||||
elif event.type == StreamEventType.STREAM_END:
|
||||
# Final update will be handled by finalize()
|
||||
pass
|
||||
|
||||
async def _schedule_update(self) -> None:
|
||||
"""Schedule a throttled update if conditions are met."""
|
||||
if self._pending_update:
|
||||
return # Update already scheduled
|
||||
|
||||
now = time.monotonic()
|
||||
elapsed = now - self._last_edit_time
|
||||
content_delta = len(self._content_buffer) - self._last_content_length
|
||||
|
||||
# Check if we should update based on time and content thresholds
|
||||
should_update = (
|
||||
elapsed >= self._config.min_interval and
|
||||
content_delta >= self._config.min_delta
|
||||
) or (
|
||||
# Always update for progress mode when tool changes
|
||||
self._config.mode == StreamingMode.PROGRESS and self._current_tool
|
||||
)
|
||||
|
||||
if not should_update:
|
||||
return
|
||||
|
||||
self._pending_update = True
|
||||
|
||||
# Calculate delay to respect min_interval
|
||||
delay = max(0, self._config.min_interval - elapsed)
|
||||
|
||||
# Schedule update as background task to avoid blocking token stream
|
||||
async def _delayed_update():
|
||||
if delay > 0:
|
||||
await asyncio.sleep(delay)
|
||||
await self._perform_update()
|
||||
self._pending_update = False
|
||||
|
||||
# Cancel existing update task if any
|
||||
if self._update_task and not self._update_task.done():
|
||||
self._update_task.cancel()
|
||||
|
||||
# Schedule new update task
|
||||
self._update_task = asyncio.create_task(_delayed_update())
|
||||
|
||||
async def _perform_update(self) -> None:
|
||||
"""Perform the actual message edit."""
|
||||
if not self._message_id or self._is_finalized:
|
||||
return
|
||||
|
||||
try:
|
||||
# Apply rate limiting
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
|
||||
# Prepare content based on mode
|
||||
if self._config.mode == StreamingMode.DRAFT:
|
||||
content = self._content_buffer or self._config.placeholder_text
|
||||
elif self._config.mode == StreamingMode.PROGRESS:
|
||||
if self._current_tool:
|
||||
content = f"{self._config.progress_prefix}Running {self._current_tool}..."
|
||||
elif self._content_buffer:
|
||||
content = self._content_buffer
|
||||
else:
|
||||
content = self._config.placeholder_text
|
||||
else:
|
||||
return # Should not happen
|
||||
|
||||
# Perform edit
|
||||
await self._adapter.edit_message(self._channel_id, self._message_id, content)
|
||||
|
||||
# Update state
|
||||
self._last_edit_time = time.monotonic()
|
||||
self._last_content_length = len(self._content_buffer)
|
||||
|
||||
logger.debug(
|
||||
"DraftStreamer updated message %s, content_length=%d",
|
||||
self._message_id, len(content)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("DraftStreamer update failed: %s", e)
|
||||
|
||||
async def finalize(self, final_content: str) -> None:
|
||||
"""Finalize with the complete response content."""
|
||||
if self._is_finalized or not self._message_id:
|
||||
return
|
||||
|
||||
self._is_finalized = True
|
||||
|
||||
# Cancel any pending updates
|
||||
if self._update_task and not self._update_task.done():
|
||||
self._update_task.cancel()
|
||||
try:
|
||||
await self._update_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
try:
|
||||
# Apply rate limiting
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
|
||||
# Final edit with complete content
|
||||
await self._adapter.edit_message(self._channel_id, self._message_id, final_content)
|
||||
|
||||
logger.debug(
|
||||
"DraftStreamer finalized message %s, final_length=%d",
|
||||
self._message_id, len(final_content)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("DraftStreamer finalization failed: %s", e)
|
||||
|
||||
@property
|
||||
def is_streaming(self) -> bool:
|
||||
"""Check if streaming is active."""
|
||||
return bool(self._message_id and not self._is_finalized)
|
||||
|
||||
@property
|
||||
def has_content(self) -> bool:
|
||||
"""Check if any content has been collected."""
|
||||
return bool(self._content_buffer.strip())
|
||||
@@ -661,6 +661,10 @@ class AgentMailBot(ChatCommandMixin, MessageHookMixin):
|
||||
self._agent,
|
||||
sender_id,
|
||||
body,
|
||||
chat_id=sender_id,
|
||||
user_name=message.sender.display_name if message.sender else "",
|
||||
message_id=message.message_id,
|
||||
account=self._config.get("account", "default"),
|
||||
)
|
||||
|
||||
if response:
|
||||
|
||||
@@ -251,6 +251,8 @@ class DiscordBot(ChatCommandMixin, MessageHookMixin):
|
||||
self._agent, user_id, text_to_send,
|
||||
chat_id=str(message.channel.id),
|
||||
user_name=str(getattr(message.author, "name", "")),
|
||||
message_id=str(message.id),
|
||||
account=self._config.get("account", "default"),
|
||||
)
|
||||
send_result = self.fire_message_sending(
|
||||
str(message.channel.id), str(response),
|
||||
|
||||
@@ -396,6 +396,10 @@ class EmailBot(ChatCommandMixin, MessageHookMixin):
|
||||
self._agent,
|
||||
message.sender.user_id,
|
||||
body,
|
||||
chat_id=message.sender.user_id,
|
||||
user_name=message.sender.display_name or "",
|
||||
message_id=message.message_id,
|
||||
account=self._config.get("account", "default"),
|
||||
)
|
||||
|
||||
if response:
|
||||
|
||||
@@ -240,12 +240,16 @@ class LinearBot(ChatCommandMixin, MessageHookMixin):
|
||||
# Read raw body for signature verification
|
||||
raw_body = await request.read()
|
||||
|
||||
# Verify signature if secret is configured
|
||||
from praisonai.bots.webhook_security import webhooks_require_verification
|
||||
|
||||
if self._signing_secret:
|
||||
signature = request.headers.get("Linear-Signature", "")
|
||||
if not self._verify_signature(raw_body, signature):
|
||||
logger.warning("Invalid webhook signature")
|
||||
return web.Response(status=401, text="Invalid signature")
|
||||
elif webhooks_require_verification():
|
||||
logger.warning("Webhook rejected: signing secret not configured")
|
||||
return web.Response(status=401, text="Webhook verification not configured")
|
||||
|
||||
# Parse JSON body
|
||||
try:
|
||||
|
||||
@@ -273,6 +273,8 @@ class SlackBot(ChatCommandMixin, MessageHookMixin):
|
||||
self._agent, user_id, text,
|
||||
chat_id=str(channel_id) if channel_id else "",
|
||||
thread_id=event.get("thread_ts", "") or "",
|
||||
message_id=event.get("ts", ""),
|
||||
account=self._config.get("account", "default"),
|
||||
)
|
||||
logger.info(f"Response sent: {response[:100]}...")
|
||||
|
||||
@@ -307,6 +309,21 @@ class SlackBot(ChatCommandMixin, MessageHookMixin):
|
||||
async def handle_mention(event, say):
|
||||
if event.get("bot_id"):
|
||||
return
|
||||
|
||||
bot_message = self._convert_event_to_message(event)
|
||||
bot_message._channel_type = "slack"
|
||||
|
||||
if not self.config.is_channel_allowed(
|
||||
bot_message.channel.channel_id if bot_message.channel else ""
|
||||
):
|
||||
return
|
||||
|
||||
user_id = bot_message.sender.user_id if bot_message.sender else ""
|
||||
is_explicitly_allowed = bool(self.config.allowed_users) and self.config.is_user_allowed(user_id)
|
||||
if not is_explicitly_allowed:
|
||||
user_allowed = await UnknownUserHandler.handle(bot_message, self._bot_context)
|
||||
if not user_allowed:
|
||||
return
|
||||
|
||||
text = event.get("text", "")
|
||||
if self._bot_user:
|
||||
@@ -320,6 +337,8 @@ class SlackBot(ChatCommandMixin, MessageHookMixin):
|
||||
self._agent, user_id, text,
|
||||
chat_id=str(event.get("channel", "")),
|
||||
thread_id=event.get("thread_ts", "") or "",
|
||||
message_id=event.get("ts", ""),
|
||||
account=self._config.get("account", "default"),
|
||||
)
|
||||
logger.info(f"Response sent: {response[:100]}...")
|
||||
|
||||
|
||||
@@ -33,6 +33,8 @@ from ._debounce import InboundDebouncer
|
||||
from ._ack import AckReactor
|
||||
from ._unknown_user import UnknownUserHandler, BotContext
|
||||
from ._pairing_ui import PairingUIBuilder, PairingCallbackHandler
|
||||
from ._streaming import StreamingConfig, StreamingMode, DraftStreamer
|
||||
from ._rate_limit import RateLimiter
|
||||
from ..gateway.unicode_utils import safe_error_message, safe_log_message, extract_root_cause_from_error
|
||||
from ..gateway.pairing import PairingStore
|
||||
|
||||
@@ -82,6 +84,10 @@ class TelegramBot(ChatCommandMixin, MessageHookMixin):
|
||||
self._agent = agent
|
||||
self.config = config or BotConfig(token=token)
|
||||
|
||||
# Initialize streaming config (None means streaming disabled)
|
||||
self._streaming_config = None
|
||||
self._rate_limiter = RateLimiter.for_platform("telegram")
|
||||
|
||||
self._is_running = False
|
||||
self._bot_user: Optional[BotUser] = None
|
||||
self._application = None
|
||||
@@ -94,9 +100,22 @@ class TelegramBot(ChatCommandMixin, MessageHookMixin):
|
||||
_store = get_default_session_store()
|
||||
except Exception:
|
||||
_store = None
|
||||
# Create run control if busy mode is configured
|
||||
run_control = None
|
||||
if hasattr(self.config, 'busy_mode') and self.config.busy_mode != "queue":
|
||||
try:
|
||||
from ._run_control import SessionRunControl
|
||||
run_control = SessionRunControl(
|
||||
busy_mode=self.config.busy_mode,
|
||||
busy_ack_template=getattr(self.config, 'busy_ack', "⏳ {action} — will be considered next")
|
||||
)
|
||||
except ImportError:
|
||||
logger.warning("Run control not available, falling back to basic session management")
|
||||
|
||||
self._session: BotSessionManager = BotSessionManager(
|
||||
store=_store,
|
||||
platform="telegram",
|
||||
run_control=run_control,
|
||||
)
|
||||
self._debouncer: InboundDebouncer = InboundDebouncer(
|
||||
debounce_ms=self.config.debounce_ms,
|
||||
@@ -120,6 +139,15 @@ class TelegramBot(ChatCommandMixin, MessageHookMixin):
|
||||
self._stop_event: Optional[asyncio.Event] = None
|
||||
|
||||
|
||||
def configure_streaming(self, config: StreamingConfig) -> None:
|
||||
"""Configure streaming reply mode.
|
||||
|
||||
Args:
|
||||
config: Streaming configuration. Set mode=StreamingMode.OFF to disable.
|
||||
"""
|
||||
self._streaming_config = config
|
||||
logger.debug("TelegramBot: streaming configured, mode=%s", config.mode)
|
||||
|
||||
def enable_stt(self, enabled: bool = True) -> None:
|
||||
"""Enable STT for voice message transcription."""
|
||||
self._stt_enabled = enabled
|
||||
@@ -227,41 +255,93 @@ class TelegramBot(ChatCommandMixin, MessageHookMixin):
|
||||
try:
|
||||
message_text = await self._debouncer.debounce(user_id, message.content)
|
||||
|
||||
# Show typing indicator with renewal during long operation
|
||||
if self.config.typing_indicator:
|
||||
from ._typing_indicator import with_typing_renewal
|
||||
|
||||
async def _typing_action():
|
||||
await update.message.chat.send_action("typing")
|
||||
|
||||
response = await with_typing_renewal(
|
||||
typing_func=_typing_action,
|
||||
operation_coro=self._session.chat(
|
||||
self._agent, user_id, message_text,
|
||||
chat_id=str(update.message.chat_id) if update.message.chat_id else "",
|
||||
user_name=user_name,
|
||||
)
|
||||
# Check if streaming is enabled and configured
|
||||
streaming_enabled = (
|
||||
self._streaming_config and
|
||||
self._streaming_config.mode != StreamingMode.OFF
|
||||
)
|
||||
|
||||
if streaming_enabled:
|
||||
# Streaming path
|
||||
streamer = DraftStreamer(
|
||||
adapter=self,
|
||||
channel_id=str(update.message.chat_id),
|
||||
config=self._streaming_config,
|
||||
rate_limiter=self._rate_limiter,
|
||||
)
|
||||
else:
|
||||
|
||||
# Start streaming (send placeholder)
|
||||
await streamer.start()
|
||||
|
||||
# Get response with streaming callback
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, message_text,
|
||||
chat_id=str(update.message.chat_id) if update.message.chat_id else "",
|
||||
user_name=user_name,
|
||||
message_id=str(update.message.message_id),
|
||||
account=self.config.get("account", "default"),
|
||||
stream_callback=streamer.on_event,
|
||||
)
|
||||
|
||||
# Apply message hooks to final response (same as non-streaming path)
|
||||
send_result = self.fire_message_sending(
|
||||
str(update.message.chat_id), str(response),
|
||||
reply_to=str(update.message.message_id),
|
||||
)
|
||||
if send_result["cancel"]:
|
||||
return
|
||||
|
||||
# Finalize with complete response (after hook processing)
|
||||
await streamer.finalize(send_result["content"])
|
||||
|
||||
# Fire sent hooks
|
||||
self.fire_message_sent(
|
||||
str(update.message.chat_id), send_result["content"],
|
||||
)
|
||||
|
||||
else:
|
||||
# Legacy non-streaming path
|
||||
# Show typing indicator with renewal during long operation
|
||||
if self.config.typing_indicator:
|
||||
from ._typing_indicator import with_typing_renewal
|
||||
|
||||
async def _typing_action():
|
||||
await update.message.chat.send_action("typing")
|
||||
|
||||
response = await with_typing_renewal(
|
||||
typing_func=_typing_action,
|
||||
operation_coro=self._session.chat(
|
||||
self._agent, user_id, message_text,
|
||||
chat_id=str(update.message.chat_id) if update.message.chat_id else "",
|
||||
user_name=user_name,
|
||||
message_id=str(update.message.message_id),
|
||||
account=self.config.get("account", "default"),
|
||||
)
|
||||
)
|
||||
else:
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, message_text,
|
||||
chat_id=str(update.message.chat_id) if update.message.chat_id else "",
|
||||
user_name=user_name,
|
||||
message_id=str(update.message.message_id),
|
||||
account=self.config.get("account", "default"),
|
||||
)
|
||||
|
||||
# Normal send flow for non-streaming
|
||||
send_result = self.fire_message_sending(
|
||||
str(update.message.chat_id), str(response),
|
||||
reply_to=str(update.message.message_id),
|
||||
)
|
||||
if send_result["cancel"]:
|
||||
return
|
||||
await self._send_response_with_media(
|
||||
update.message.chat_id,
|
||||
send_result["content"],
|
||||
reply_to=update.message.message_id,
|
||||
)
|
||||
self.fire_message_sent(
|
||||
str(update.message.chat_id), send_result["content"],
|
||||
)
|
||||
send_result = self.fire_message_sending(
|
||||
str(update.message.chat_id), str(response),
|
||||
reply_to=str(update.message.message_id),
|
||||
)
|
||||
if send_result["cancel"]:
|
||||
return
|
||||
await self._send_response_with_media(
|
||||
update.message.chat_id,
|
||||
send_result["content"],
|
||||
reply_to=update.message.message_id,
|
||||
)
|
||||
self.fire_message_sent(
|
||||
str(update.message.chat_id), send_result["content"],
|
||||
)
|
||||
# Done reaction
|
||||
if ack_ctx:
|
||||
await self._ack.done(ack_ctx, react_fn=_tg_react, unreact_fn=_tg_unreact)
|
||||
|
||||
Loaded 100 of 155 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user