mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 17:31:00 -10:00
Compare commits
28
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
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,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.58" \
|
||||
"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.58" \
|
||||
"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.58" \
|
||||
"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
|
||||
|
||||
@@ -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()
|
||||
@@ -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,
|
||||
|
||||
@@ -245,7 +245,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 +665,7 @@ class AgentTeam(SpawnAnnounceProtocol):
|
||||
_stream = _output_config.stream
|
||||
else:
|
||||
_verbose = 0
|
||||
_stream = True
|
||||
_stream = False
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Resolve EXECUTION param using canonical resolver
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -1089,8 +1089,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."""
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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."}
|
||||
|
||||
@@ -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.58"
|
||||
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
|
||||
|
||||
@@ -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,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
|
||||
Generated
+1
-1
@@ -3125,7 +3125,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.6.55"
|
||||
version = "1.6.58"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
{
|
||||
"version": "1.5.87",
|
||||
"lastUpdated": "2026-06-12",
|
||||
"lastUpdated": "2026-06-13",
|
||||
"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-13",
|
||||
"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-13
|
||||
> **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.58.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.6.58.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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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]}...")
|
||||
|
||||
@@ -320,6 +322,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)
|
||||
|
||||
@@ -445,7 +445,11 @@ class WhatsAppBot(ChatCommandMixin, MessageHookMixin):
|
||||
if self._agent and content:
|
||||
try:
|
||||
response = await self._session_mgr.chat(
|
||||
self._agent, sender_jid, content
|
||||
self._agent, sender_jid, content,
|
||||
chat_id=chat_jid,
|
||||
user_name=sender_name or "",
|
||||
message_id=msg_id,
|
||||
account=self._config.get("account", "default"),
|
||||
)
|
||||
if response:
|
||||
send_result = self.fire_message_sending(chat_jid, str(response))
|
||||
@@ -735,7 +739,11 @@ class WhatsAppBot(ChatCommandMixin, MessageHookMixin):
|
||||
|
||||
try:
|
||||
response = await self._session_mgr.chat(
|
||||
self._agent, sender_id, content
|
||||
self._agent, sender_id, content,
|
||||
chat_id=sender_id,
|
||||
user_name=sender_name or "",
|
||||
message_id=msg_id,
|
||||
account=self._config.get("account", "default"),
|
||||
)
|
||||
if response:
|
||||
send_result = self.fire_message_sending(sender_id, str(response))
|
||||
|
||||
@@ -120,7 +120,7 @@ class BrowserServer:
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=cors_origins,
|
||||
allow_origin_regex=r"chrome-extension://[a-z0-9]{32}",
|
||||
allow_origin_regex=r"chrome-extension://[a-z0-9]{32}$",
|
||||
allow_credentials=True,
|
||||
allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
|
||||
allow_headers=["Authorization", "Content-Type", "Origin", "Accept"],
|
||||
@@ -183,7 +183,7 @@ class BrowserServer:
|
||||
if origin in cors_origins:
|
||||
is_allowed = True
|
||||
# Check chrome extension regex pattern (same as CORS middleware)
|
||||
elif parsed_origin.scheme == "chrome-extension" and re.match(r"chrome-extension://[a-z0-9]{32}", origin):
|
||||
elif parsed_origin.scheme == "chrome-extension" and re.match(r"chrome-extension://[a-z0-9]{32}$", origin):
|
||||
is_allowed = True
|
||||
|
||||
if not is_allowed:
|
||||
|
||||
@@ -136,6 +136,7 @@ def chat_main(
|
||||
no_color: bool = typer.Option(False, "--no-color", help="Disable colors"),
|
||||
theme: str = typer.Option("default", "--theme", help="UI theme: default, dark, light, minimal"),
|
||||
compact: bool = typer.Option(False, "--compact", help="Compact output mode"),
|
||||
no_rules: bool = typer.Option(False, "--no-rules", help="Disable auto-injection of project instruction files"),
|
||||
):
|
||||
"""
|
||||
Start terminal-native interactive chat mode.
|
||||
@@ -209,6 +210,7 @@ def chat_main(
|
||||
workspace=workspace,
|
||||
debug=debug,
|
||||
autonomy_mode=autonomy,
|
||||
no_rules=no_rules,
|
||||
)
|
||||
|
||||
tui = AsyncTUI(config=tui_config)
|
||||
|
||||
@@ -192,6 +192,132 @@ def gateway_channels(
|
||||
print(f"{name:<20} {platform:<12} {has_token:<12}")
|
||||
|
||||
|
||||
@app.command("pause")
|
||||
def gateway_pause_channel(
|
||||
name: str = typer.Argument(help="Channel name to pause"),
|
||||
url: str = typer.Option("ws://127.0.0.1:8765", "--url", help="Gateway WebSocket URL"),
|
||||
):
|
||||
"""Pause a gateway channel.
|
||||
|
||||
Examples:
|
||||
praisonai gateway pause telegram
|
||||
praisonai gateway pause discord --url ws://localhost:8000
|
||||
"""
|
||||
import requests
|
||||
import sys
|
||||
from urllib.parse import urlparse, urlunparse
|
||||
|
||||
try:
|
||||
# Parse URL and convert WebSocket to HTTP
|
||||
parsed = urlparse(url)
|
||||
scheme = "https" if parsed.scheme == "wss" else "http"
|
||||
# Reconstruct base URL preserving path and query
|
||||
rest_url = urlunparse((
|
||||
scheme, parsed.netloc, parsed.path, parsed.params, parsed.query, parsed.fragment
|
||||
))
|
||||
if not rest_url.endswith("/"):
|
||||
rest_url += "/"
|
||||
|
||||
response = requests.post(f"{rest_url}api/channels/{name}/pause", timeout=10)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
if result.get("success"):
|
||||
print(f"✅ Channel '{name}' paused successfully")
|
||||
else:
|
||||
message = result.get("message", result.get("error", "Unknown error"))
|
||||
print(f"❌ Failed to pause channel '{name}': {message}")
|
||||
sys.exit(1)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error pausing channel '{name}': {str(e)}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@app.command("resume")
|
||||
def gateway_resume_channel(
|
||||
name: str = typer.Argument(help="Channel name to resume"),
|
||||
url: str = typer.Option("ws://127.0.0.1:8765", "--url", help="Gateway WebSocket URL"),
|
||||
):
|
||||
"""Resume a paused gateway channel.
|
||||
|
||||
Examples:
|
||||
praisonai gateway resume telegram
|
||||
praisonai gateway resume discord --url ws://localhost:8000
|
||||
"""
|
||||
import requests
|
||||
import sys
|
||||
from urllib.parse import urlparse, urlunparse
|
||||
|
||||
try:
|
||||
# Parse URL and convert WebSocket to HTTP
|
||||
parsed = urlparse(url)
|
||||
scheme = "https" if parsed.scheme == "wss" else "http"
|
||||
# Reconstruct base URL preserving path and query
|
||||
rest_url = urlunparse((
|
||||
scheme, parsed.netloc, parsed.path, parsed.params, parsed.query, parsed.fragment
|
||||
))
|
||||
if not rest_url.endswith("/"):
|
||||
rest_url += "/"
|
||||
|
||||
response = requests.post(f"{rest_url}api/channels/{name}/resume", timeout=10)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
if result.get("success"):
|
||||
print(f"✅ Channel '{name}' resumed successfully")
|
||||
else:
|
||||
message = result.get("message", result.get("error", "Unknown error"))
|
||||
print(f"❌ Failed to resume channel '{name}': {message}")
|
||||
sys.exit(1)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error resuming channel '{name}': {str(e)}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@app.command("reconnect")
|
||||
def gateway_reconnect_channel(
|
||||
name: str = typer.Argument(help="Channel name to reconnect"),
|
||||
url: str = typer.Option("ws://127.0.0.1:8765", "--url", help="Gateway WebSocket URL"),
|
||||
):
|
||||
"""Reconnect a gateway channel.
|
||||
|
||||
Examples:
|
||||
praisonai gateway reconnect telegram
|
||||
praisonai gateway reconnect discord --url ws://localhost:8000
|
||||
"""
|
||||
import requests
|
||||
import sys
|
||||
from urllib.parse import urlparse, urlunparse
|
||||
|
||||
try:
|
||||
# Parse URL and convert WebSocket to HTTP
|
||||
parsed = urlparse(url)
|
||||
scheme = "https" if parsed.scheme == "wss" else "http"
|
||||
# Reconstruct base URL preserving path and query
|
||||
rest_url = urlunparse((
|
||||
scheme, parsed.netloc, parsed.path, parsed.params, parsed.query, parsed.fragment
|
||||
))
|
||||
if not rest_url.endswith("/"):
|
||||
rest_url += "/"
|
||||
|
||||
response = requests.post(f"{rest_url}api/channels/{name}/reconnect", timeout=10)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
if result.get("success"):
|
||||
print(f"✅ Channel '{name}' reconnected successfully")
|
||||
else:
|
||||
message = result.get("message", result.get("error", "Unknown error"))
|
||||
print(f"❌ Failed to reconnect channel '{name}': {message}")
|
||||
sys.exit(1)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error reconnecting channel '{name}': {str(e)}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@app.command("install")
|
||||
def gateway_install(
|
||||
config: str = typer.Option(
|
||||
|
||||
@@ -34,6 +34,12 @@ def run_main(
|
||||
approval: Optional[str] = typer.Option(None, "--approval", help="Approval backend: console, slack, telegram, discord, webhook, http, agent, auto, none"),
|
||||
approve_all_tools: bool = typer.Option(False, "--approve-all-tools", help="Require approval for ALL tool calls, not just dangerous tools"),
|
||||
approval_timeout: Optional[str] = typer.Option(None, "--approval-timeout", help="Seconds to wait for approval. Use 'none' for indefinite wait"),
|
||||
no_rules: bool = typer.Option(False, "--no-rules", help="Disable auto-injection of project instruction files"),
|
||||
# Session continuity options
|
||||
continue_session: bool = typer.Option(False, "--continue", "-c", help="Continue the most recent session for this project"),
|
||||
session: Optional[str] = typer.Option(None, "--session", "-s", help="Resume a specific session ID"),
|
||||
fork: bool = typer.Option(False, "--fork", help="Fork from the specified session (requires --session)"),
|
||||
no_save: bool = typer.Option(False, "--no-save", help="Don't auto-save session after execution"),
|
||||
):
|
||||
"""
|
||||
Run agents from a file or prompt.
|
||||
@@ -41,18 +47,31 @@ def run_main(
|
||||
Examples:
|
||||
praisonai run agents.yaml
|
||||
praisonai run "What is the weather?"
|
||||
praisonai run "What is the weather?" --continue
|
||||
praisonai run "Add tests" --session abc123
|
||||
praisonai run agents.yaml --interactive
|
||||
praisonai run "What is 2+2?" --profile
|
||||
"""
|
||||
output = get_output_controller()
|
||||
_ = get_current_context() # Initialize context
|
||||
|
||||
# Validate session options
|
||||
if fork and not session:
|
||||
output.print_error("--fork requires --session to specify which session to fork from")
|
||||
raise typer.Exit(1)
|
||||
|
||||
if continue_session and session:
|
||||
output.print_error("Cannot use both --continue and --session together")
|
||||
raise typer.Exit(1)
|
||||
|
||||
if not target:
|
||||
output.print_panel(
|
||||
"Run agents from a file or prompt.\n\n"
|
||||
"Usage:\n"
|
||||
" praisonai run agents.yaml\n"
|
||||
" praisonai run \"What is the weather?\"\n"
|
||||
" praisonai run \"What is the weather?\" --continue\n"
|
||||
" praisonai run \"Add tests\" --session abc123\n"
|
||||
" praisonai run agents.yaml --interactive\n\n"
|
||||
"Options:\n"
|
||||
" --model, -m LLM model to use\n"
|
||||
@@ -60,7 +79,11 @@ def run_main(
|
||||
" --interactive, -i Interactive mode\n"
|
||||
" --verbose, -v Verbose output\n"
|
||||
" --trace Enable tracing\n"
|
||||
" --memory Enable memory",
|
||||
" --memory Enable memory\n"
|
||||
" --continue, -c Continue the most recent session\n"
|
||||
" --session, -s Resume a specific session ID\n"
|
||||
" --fork Fork from specified session\n"
|
||||
" --no-save Don't auto-save session",
|
||||
title="Run Command"
|
||||
)
|
||||
return
|
||||
@@ -89,6 +112,10 @@ def run_main(
|
||||
framework=framework,
|
||||
verbose=verbose,
|
||||
profile_deep=profile_deep,
|
||||
continue_session=continue_session,
|
||||
session=session,
|
||||
fork=fork,
|
||||
no_save=no_save,
|
||||
)
|
||||
else:
|
||||
# Profiling for direct prompt
|
||||
@@ -97,6 +124,10 @@ def run_main(
|
||||
model=model,
|
||||
verbose=verbose,
|
||||
profile_deep=profile_deep,
|
||||
continue_session=continue_session,
|
||||
session=session,
|
||||
fork=fork,
|
||||
no_save=no_save,
|
||||
)
|
||||
return
|
||||
|
||||
@@ -114,6 +145,10 @@ def run_main(
|
||||
tools=tools,
|
||||
max_tokens=max_tokens,
|
||||
output_mode=output_mode,
|
||||
continue_session=continue_session,
|
||||
session=session,
|
||||
fork=fork,
|
||||
no_save=no_save,
|
||||
)
|
||||
else:
|
||||
# Run as prompt
|
||||
@@ -131,6 +166,11 @@ def run_main(
|
||||
approval=approval,
|
||||
approve_all_tools=approve_all_tools,
|
||||
approval_timeout=approval_timeout,
|
||||
no_rules=no_rules,
|
||||
continue_session=continue_session,
|
||||
session=session,
|
||||
fork=fork,
|
||||
no_save=no_save,
|
||||
)
|
||||
|
||||
|
||||
@@ -146,6 +186,10 @@ def _run_from_file(
|
||||
tools: Optional[str] = None,
|
||||
max_tokens: int = 16000,
|
||||
output_mode: Optional[str] = None,
|
||||
continue_session: bool = False,
|
||||
session: Optional[str] = None,
|
||||
fork: bool = False,
|
||||
no_save: bool = False,
|
||||
):
|
||||
"""Run agents from a YAML file."""
|
||||
output = get_output_controller()
|
||||
@@ -163,6 +207,58 @@ def _run_from_file(
|
||||
if model:
|
||||
praison.config_list[0]['model'] = model
|
||||
|
||||
# Handle session continuity for YAML files
|
||||
session_id = None
|
||||
auto_save_name = None
|
||||
|
||||
if continue_session or session or fork:
|
||||
from ..state.project_sessions import get_project_session_store, find_last_session
|
||||
|
||||
if continue_session:
|
||||
# Find last session for this project
|
||||
session_id = find_last_session()
|
||||
if session_id:
|
||||
output.print_info(f"Continuing session: {session_id}")
|
||||
else:
|
||||
output.print_warning("No previous sessions found. Starting new session.")
|
||||
|
||||
elif session:
|
||||
# Use specific session
|
||||
project_store = get_project_session_store()
|
||||
if project_store.session_exists(session):
|
||||
session_id = session
|
||||
output.print_info(f"Resuming session: {session_id}")
|
||||
else:
|
||||
output.print_error(f"Session not found: {session}")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Handle forking
|
||||
if fork:
|
||||
from praisonaiagents.session.hierarchy import HierarchicalSessionStore
|
||||
from ..utils.project import get_project_sessions_dir
|
||||
|
||||
# Create hierarchical store for forking
|
||||
hierarchical_store = HierarchicalSessionStore(str(get_project_sessions_dir()))
|
||||
forked_session_id = hierarchical_store.fork_session(session_id)
|
||||
session_id = forked_session_id
|
||||
output.print_info(f"Forked session {session} -> {forked_session_id}")
|
||||
|
||||
# Enable auto-save if not disabled
|
||||
if not no_save:
|
||||
import uuid
|
||||
auto_save_name = session_id or "session-" + str(uuid.uuid4())[:8]
|
||||
|
||||
# Create args-like object for session configuration
|
||||
if session_id or auto_save_name:
|
||||
class Args:
|
||||
pass
|
||||
|
||||
args = Args()
|
||||
args.auto_save = auto_save_name
|
||||
args.resume_session = session_id
|
||||
|
||||
praison.args = args
|
||||
|
||||
# Run
|
||||
result = praison.run()
|
||||
|
||||
@@ -195,11 +291,65 @@ def _run_prompt(
|
||||
approval: Optional[str] = None,
|
||||
approve_all_tools: bool = False,
|
||||
approval_timeout: Optional[str] = None,
|
||||
no_rules: bool = False,
|
||||
continue_session: bool = False,
|
||||
session: Optional[str] = None,
|
||||
fork: bool = False,
|
||||
no_save: bool = False,
|
||||
):
|
||||
"""Run a direct prompt."""
|
||||
output = get_output_controller()
|
||||
|
||||
try:
|
||||
# Handle session continuity first (before any execution mode)
|
||||
from praisonai.cli.main import PraisonAI
|
||||
|
||||
praison = PraisonAI()
|
||||
|
||||
if model:
|
||||
praison.config_list[0]['model'] = model
|
||||
|
||||
# Handle session continuity
|
||||
session_id = None
|
||||
auto_save_name = None
|
||||
|
||||
if continue_session or session or fork:
|
||||
from ..state.project_sessions import get_project_session_store, find_last_session
|
||||
|
||||
if continue_session:
|
||||
# Find last session for this project
|
||||
session_id = find_last_session()
|
||||
if session_id:
|
||||
output.print_info(f"Continuing session: {session_id}")
|
||||
else:
|
||||
output.print_warning("No previous sessions found. Starting new session.")
|
||||
|
||||
elif session:
|
||||
# Use specific session
|
||||
project_store = get_project_session_store()
|
||||
if project_store.session_exists(session):
|
||||
session_id = session
|
||||
output.print_info(f"Resuming session: {session_id}")
|
||||
else:
|
||||
output.print_error(f"Session not found: {session}")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Handle forking
|
||||
if fork:
|
||||
from praisonaiagents.session.hierarchy import HierarchicalSessionStore
|
||||
from ..utils.project import get_project_sessions_dir
|
||||
|
||||
# Create hierarchical store for forking
|
||||
hierarchical_store = HierarchicalSessionStore(str(get_project_sessions_dir()))
|
||||
forked_session_id = hierarchical_store.fork_session(session_id)
|
||||
session_id = forked_session_id
|
||||
output.print_info(f"Forked session {session} -> {forked_session_id}")
|
||||
|
||||
# Enable auto-save if not disabled (for all runs, not just session continuity)
|
||||
if not no_save:
|
||||
import uuid
|
||||
auto_save_name = session_id or "session-" + str(uuid.uuid4())[:8]
|
||||
|
||||
# If output_mode is "actions", use direct Agent with actions preset
|
||||
if output_mode == "actions":
|
||||
from praisonaiagents import Agent
|
||||
@@ -220,6 +370,12 @@ def _run_prompt(
|
||||
approval, all_tools=approve_all_tools, timeout=approval_timeout,
|
||||
)
|
||||
|
||||
# Add session support to Agent if needed
|
||||
if session_id:
|
||||
agent_config["resume_session"] = session_id
|
||||
if auto_save_name:
|
||||
agent_config["auto_save"] = auto_save_name
|
||||
|
||||
agent = Agent(**agent_config)
|
||||
result = agent.start(prompt)
|
||||
|
||||
@@ -231,14 +387,8 @@ def _run_prompt(
|
||||
# Don't print result again - actions mode already shows output
|
||||
return
|
||||
|
||||
# Use existing handle_direct_prompt for other modes
|
||||
from praisonai.cli.main import PraisonAI
|
||||
|
||||
praison = PraisonAI()
|
||||
|
||||
if model:
|
||||
praison.config_list[0]['model'] = model
|
||||
|
||||
# Use handle_direct_prompt for other modes
|
||||
|
||||
# Create args-like object for handle_direct_prompt
|
||||
class Args:
|
||||
pass
|
||||
@@ -259,9 +409,12 @@ def _run_prompt(
|
||||
args.auto_approve_plan = False
|
||||
args.final_agent = None
|
||||
args.user_id = None
|
||||
args.auto_save = None
|
||||
# Enable session features based on flags
|
||||
args.auto_save = auto_save_name
|
||||
args.history = None
|
||||
args.include_rules = None
|
||||
args.resume_session = session_id
|
||||
args.include_rules = None if no_rules else "auto"
|
||||
args.no_rules = no_rules
|
||||
args.workflow = None
|
||||
args.workflow_var = None
|
||||
args.claude_memory = False
|
||||
@@ -286,6 +439,45 @@ def _run_prompt(
|
||||
args.approval = approval
|
||||
|
||||
praison.args = args
|
||||
|
||||
# If output_mode is "actions", use direct Agent with actions preset
|
||||
if output_mode == "actions":
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent_config = {
|
||||
"name": "RunAgent",
|
||||
"role": "Assistant",
|
||||
"goal": "Complete the task",
|
||||
"output": "actions", # Use actions preset
|
||||
}
|
||||
if model:
|
||||
agent_config["llm"] = model
|
||||
|
||||
# Resolve approval backend if specified
|
||||
if approval:
|
||||
from praisonai.cli.features.approval import resolve_approval_config
|
||||
agent_config["approval"] = resolve_approval_config(
|
||||
approval, all_tools=approve_all_tools, timeout=approval_timeout,
|
||||
)
|
||||
|
||||
# Add session support to Agent if needed
|
||||
if session_id:
|
||||
agent_config["resume_session"] = session_id
|
||||
if auto_save_name:
|
||||
agent_config["auto_save"] = auto_save_name
|
||||
|
||||
agent = Agent(**agent_config)
|
||||
result = agent.start(prompt)
|
||||
|
||||
output.emit_result(
|
||||
message="Prompt completed",
|
||||
data={"result": str(result) if result else None}
|
||||
)
|
||||
|
||||
# Don't print result again - actions mode already shows output
|
||||
return
|
||||
|
||||
# Use handle_direct_prompt for other modes
|
||||
result = praison.handle_direct_prompt(prompt)
|
||||
|
||||
output.emit_result(
|
||||
@@ -308,6 +500,10 @@ def _run_from_file_profiled(
|
||||
framework: Optional[str] = None,
|
||||
verbose: bool = False,
|
||||
profile_deep: bool = False,
|
||||
continue_session: bool = False,
|
||||
session: Optional[str] = None,
|
||||
fork: bool = False,
|
||||
no_save: bool = False,
|
||||
):
|
||||
"""Run agents from a YAML file with profiling enabled."""
|
||||
from praisonai.cli.features.cli_profiler import (
|
||||
@@ -362,6 +558,10 @@ def _run_prompt_profiled(
|
||||
model: Optional[str] = None,
|
||||
verbose: bool = False,
|
||||
profile_deep: bool = False,
|
||||
continue_session: bool = False,
|
||||
session: Optional[str] = None,
|
||||
fork: bool = False,
|
||||
no_save: bool = False,
|
||||
):
|
||||
"""Run a direct prompt with profiling enabled."""
|
||||
from praisonai.cli.features.cli_profiler import (
|
||||
|
||||
@@ -55,19 +55,75 @@ def session_list(
|
||||
"--storage-path",
|
||||
help="Path for storage backend",
|
||||
),
|
||||
all_projects: bool = typer.Option(
|
||||
False,
|
||||
"--all",
|
||||
help="Show sessions from all projects (default: current project only)",
|
||||
),
|
||||
project_id: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--project",
|
||||
help="Show sessions for specific project ID",
|
||||
),
|
||||
):
|
||||
"""List all sessions."""
|
||||
output = get_output_controller()
|
||||
|
||||
# Configure backend if specified
|
||||
if storage_backend:
|
||||
backend = _create_backend(storage_backend, storage_path)
|
||||
if backend:
|
||||
set_session_backend(backend)
|
||||
|
||||
manager = get_session_manager()
|
||||
|
||||
sessions = manager.list(limit=limit)
|
||||
# Handle project-scoped session listing
|
||||
if not all_projects and not storage_backend:
|
||||
# Use project-scoped session store by default
|
||||
from ..state.project_sessions import get_project_session_store
|
||||
from ..utils.project import get_project_id, get_project_name
|
||||
|
||||
# List sessions for specific or current project
|
||||
project_store = get_project_session_store(project_id=project_id)
|
||||
|
||||
sessions_data = project_store.list_sessions(limit=limit)
|
||||
|
||||
# Convert to expected format
|
||||
class SessionInfo:
|
||||
def __init__(self, data):
|
||||
from datetime import datetime
|
||||
|
||||
self.session_id = data.get("session_id", data.get("id", ""))
|
||||
self.name = data.get("agent_name", "")
|
||||
self.status = data.get("status") # Use actual status from data if available
|
||||
self.event_count = data.get("message_count", 0)
|
||||
|
||||
# Parse updated_at string
|
||||
updated_str = data.get("updated_at", "")
|
||||
try:
|
||||
self.updated_at = datetime.fromisoformat(updated_str.replace('Z', '+00:00')) if updated_str else datetime.now()
|
||||
except ValueError:
|
||||
self.updated_at = datetime.now()
|
||||
|
||||
def to_dict(self):
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"name": self.name,
|
||||
"status": self.status,
|
||||
"event_count": self.event_count,
|
||||
"updated_at": self.updated_at.isoformat(),
|
||||
}
|
||||
|
||||
sessions = [SessionInfo(data) for data in sessions_data]
|
||||
|
||||
# Add project info to output
|
||||
if not project_id:
|
||||
current_project = get_project_name()
|
||||
current_id = get_project_id()
|
||||
output.print_info(f"Project: {current_project} (ID: {current_id})")
|
||||
|
||||
else:
|
||||
# Use global session manager
|
||||
# Configure backend if specified
|
||||
if storage_backend:
|
||||
backend = _create_backend(storage_backend, storage_path)
|
||||
if backend:
|
||||
set_session_backend(backend)
|
||||
|
||||
manager = get_session_manager()
|
||||
sessions = manager.list(limit=limit)
|
||||
|
||||
if output.is_json_mode:
|
||||
output.print_json({
|
||||
|
||||
@@ -81,6 +81,13 @@ class SessionConfig:
|
||||
history_limit: int = 10
|
||||
|
||||
|
||||
@dataclass
|
||||
class RulesConfig:
|
||||
"""Rules/instructions configuration."""
|
||||
auto: bool = True # Auto-inject project instruction files
|
||||
max_chars: int = 32000 # Maximum total characters from rules
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConfigSchema:
|
||||
"""
|
||||
@@ -93,6 +100,7 @@ class ConfigSchema:
|
||||
mcp: MCPConfig = field(default_factory=MCPConfig)
|
||||
model: ModelConfig = field(default_factory=ModelConfig)
|
||||
session: SessionConfig = field(default_factory=SessionConfig)
|
||||
rules: RulesConfig = field(default_factory=RulesConfig)
|
||||
|
||||
# Additional settings stored as dict for flexibility
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
@@ -127,6 +135,10 @@ class ConfigSchema:
|
||||
"auto_save": self.session.auto_save,
|
||||
"history_limit": self.session.history_limit,
|
||||
},
|
||||
"rules": {
|
||||
"auto": self.rules.auto,
|
||||
"max_chars": self.rules.max_chars,
|
||||
},
|
||||
**self.extra,
|
||||
}
|
||||
|
||||
@@ -166,9 +178,10 @@ class ConfigSchema:
|
||||
mcp_data = data.get("mcp", {})
|
||||
model_data = data.get("model", {})
|
||||
session_data = data.get("session", {})
|
||||
rules_data = data.get("rules", {})
|
||||
|
||||
# Extract known keys
|
||||
known_keys = {"output", "traces", "mcp", "model", "session"}
|
||||
known_keys = {"output", "traces", "mcp", "model", "session", "rules"}
|
||||
extra = {k: v for k, v in data.items() if k not in known_keys}
|
||||
|
||||
# Parse MCP servers
|
||||
@@ -225,6 +238,10 @@ class ConfigSchema:
|
||||
auto_save=session_data.get("auto_save", False),
|
||||
history_limit=session_data.get("history_limit", 10),
|
||||
),
|
||||
rules=RulesConfig(
|
||||
auto=rules_data.get("auto", True),
|
||||
max_chars=rules_data.get("max_chars", 32000),
|
||||
),
|
||||
extra=extra,
|
||||
)
|
||||
|
||||
|
||||
@@ -144,6 +144,17 @@ class BaseFormatter:
|
||||
class TextFormatter(BaseFormatter):
|
||||
"""Text formatter with optional color support."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
no_color: bool = False,
|
||||
quiet: bool = False,
|
||||
redact: bool = True,
|
||||
show_prefix_suffix: bool = False,
|
||||
):
|
||||
super().__init__(no_color, quiet, redact, show_prefix_suffix)
|
||||
# Cache encoding detection result at initialization
|
||||
self._unicode_supported = self._can_encode_unicode()
|
||||
|
||||
# ANSI color codes
|
||||
COLORS = {
|
||||
"reset": "\033[0m",
|
||||
@@ -173,15 +184,49 @@ class TextFormatter(BaseFormatter):
|
||||
CheckStatus.ERROR: "✗",
|
||||
}
|
||||
|
||||
# ASCII fallback symbols for systems that don't support UTF-8
|
||||
STATUS_SYMBOLS_ASCII = {
|
||||
CheckStatus.PASS: "[OK]",
|
||||
CheckStatus.WARN: "[!]",
|
||||
CheckStatus.FAIL: "[X]",
|
||||
CheckStatus.SKIP: "[-]",
|
||||
CheckStatus.ERROR: "[X]",
|
||||
}
|
||||
|
||||
def _color(self, text: str, color: str) -> str:
|
||||
"""Apply color to text if colors are enabled."""
|
||||
if self.no_color:
|
||||
return text
|
||||
return f"{self.COLORS.get(color, '')}{text}{self.COLORS['reset']}"
|
||||
|
||||
def _can_encode_unicode(self) -> bool:
|
||||
"""Check if stdout can encode Unicode symbols."""
|
||||
try:
|
||||
# Get stdout encoding
|
||||
encoding = getattr(sys.stdout, 'encoding', 'ascii') or 'ascii'
|
||||
|
||||
# UTF-8 always supports Unicode
|
||||
if encoding.lower() in ('utf-8', 'utf8'):
|
||||
return True
|
||||
|
||||
# Test if we can encode our specific Unicode symbols
|
||||
test_symbols = ["✓", "⚠", "✗", "○"]
|
||||
for symbol in test_symbols:
|
||||
try:
|
||||
symbol.encode(encoding, errors='strict')
|
||||
except (UnicodeEncodeError, LookupError):
|
||||
return False
|
||||
return True
|
||||
except (AttributeError, TypeError):
|
||||
# If we can't determine encoding, assume ASCII-only is safer
|
||||
return False
|
||||
|
||||
def _status_symbol(self, status: CheckStatus) -> str:
|
||||
"""Get colored status symbol."""
|
||||
symbol = self.STATUS_SYMBOLS.get(status, "?")
|
||||
"""Get colored status symbol with encoding safety."""
|
||||
if self._unicode_supported:
|
||||
symbol = self.STATUS_SYMBOLS.get(status, "?")
|
||||
else:
|
||||
symbol = self.STATUS_SYMBOLS_ASCII.get(status, "[?]")
|
||||
color = self.STATUS_COLORS.get(status, "white")
|
||||
return self._color(symbol, color)
|
||||
|
||||
|
||||
@@ -19,6 +19,36 @@ import logging
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
def _install_api_key_middleware(
|
||||
app: Any,
|
||||
api_key: Optional[str],
|
||||
public_paths: Optional[set] = None,
|
||||
) -> None:
|
||||
"""Enforce --api-key on serve endpoints when a key is configured."""
|
||||
if not api_key:
|
||||
return
|
||||
|
||||
import hmac
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
public = public_paths or {"/health", "/", "/.well-known/agent.json"}
|
||||
|
||||
class APIKeyMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request, call_next):
|
||||
path = request.url.path
|
||||
if path in public or path.startswith("/__praisonai__/"):
|
||||
return await call_next(request)
|
||||
auth = request.headers.get("Authorization", "")
|
||||
header_key = request.headers.get("X-API-Key", "")
|
||||
token = auth[7:] if auth.startswith("Bearer ") else header_key
|
||||
if not token or not hmac.compare_digest(token, api_key):
|
||||
return JSONResponse({"error": "Unauthorized"}, status_code=401)
|
||||
return await call_next(request)
|
||||
|
||||
app.add_middleware(APIKeyMiddleware)
|
||||
|
||||
|
||||
class ServeHandler:
|
||||
"""
|
||||
CLI handler for serve operations.
|
||||
@@ -407,6 +437,8 @@ Launch PraisonAI servers with unified discovery support.
|
||||
"endpoints": [path, "/agents/{agent_name}", "/api/v1/agents/{agent_id}/invoke"],
|
||||
"discovery": "/__praisonai__/discovery",
|
||||
}
|
||||
|
||||
_install_api_key_middleware(app, config.get("api_key"))
|
||||
|
||||
return app
|
||||
|
||||
@@ -943,6 +975,8 @@ Launch PraisonAI servers with unified discovery support.
|
||||
"description": "PraisonAI Unified Server",
|
||||
"version": "1.0.0",
|
||||
}
|
||||
|
||||
_install_api_key_middleware(app, config.get("api_key"))
|
||||
|
||||
return app
|
||||
|
||||
|
||||
@@ -78,6 +78,7 @@ class AsyncTUIConfig:
|
||||
multiline_mode: bool = False
|
||||
autonomy_mode: bool = True # Enable autonomous task delegation (aligned with InteractiveConfig)
|
||||
debug: bool = False # Enable debug logging to file (~/.praisonai/async_tui_debug.log)
|
||||
no_rules: bool = False # Disable auto-injection of project instruction files
|
||||
|
||||
|
||||
# ============================================================================
|
||||
@@ -197,12 +198,42 @@ class AsyncTUI:
|
||||
tools = self._load_tools()
|
||||
logger.debug(f"Tools for agent: {len(tools) if tools else 0}")
|
||||
|
||||
# Auto-inject project instruction files unless disabled
|
||||
instructions = "You are Praison AI, a helpful assistant. Be concise and helpful. You have access to tools for file operations, code execution, and web search."
|
||||
max_chars = 32000 # Default cap
|
||||
|
||||
# Check config for rules settings with proper precedence
|
||||
try:
|
||||
from praisonai.cli.configuration.loader import load_config
|
||||
config = load_config()
|
||||
if self.config.no_rules:
|
||||
should_load_rules = False
|
||||
else:
|
||||
should_load_rules = config.rules.auto
|
||||
max_chars = config.rules.max_chars
|
||||
except Exception:
|
||||
# Config not available, use defaults
|
||||
should_load_rules = not self.config.no_rules
|
||||
|
||||
if should_load_rules:
|
||||
try:
|
||||
from praisonaiagents.memory import RulesManager
|
||||
rules_manager = RulesManager(workspace_path=self.config.workspace or os.getcwd(), verbose=0)
|
||||
rules_context = rules_manager.build_rules_context(max_chars=max_chars)
|
||||
if rules_context:
|
||||
instructions = f"{rules_context}\n\n{instructions}"
|
||||
logger.debug("Project instruction files loaded")
|
||||
except ImportError:
|
||||
pass # RulesManager not available
|
||||
except Exception as e:
|
||||
logger.debug(f"Error loading rules: {e}")
|
||||
|
||||
# Build agent config
|
||||
agent_config = {
|
||||
"name": "Praison",
|
||||
"role": "AI Assistant",
|
||||
"role": "AI Assistant",
|
||||
"goal": "Help the user with their requests",
|
||||
"instructions": "You are Praison AI, a helpful assistant. Be concise and helpful. You have access to tools for file operations, code execution, and web search.",
|
||||
"instructions": instructions,
|
||||
"llm": self.config.model,
|
||||
"tools": tools if tools else None,
|
||||
}
|
||||
|
||||
@@ -425,7 +425,6 @@ class PraisonAI:
|
||||
|
||||
# Handle backends command
|
||||
elif args.command == "backends":
|
||||
from rich import print
|
||||
subcommand = unknown_args[0] if unknown_args and not unknown_args[0].startswith('-') else None
|
||||
|
||||
if subcommand == "list" or subcommand is None:
|
||||
@@ -923,6 +922,8 @@ class PraisonAI:
|
||||
default_args.call = False
|
||||
default_args.public = False
|
||||
default_args.chat_mode = False
|
||||
default_args.include_rules = None
|
||||
default_args.no_rules = False
|
||||
return default_args
|
||||
|
||||
# Define special commands
|
||||
@@ -997,6 +998,7 @@ class PraisonAI:
|
||||
|
||||
# Rules arguments
|
||||
parser.add_argument("--include-rules", type=str, help="Include manual rules by name (comma-separated)")
|
||||
parser.add_argument("--no-rules", action="store_true", help="Disable automatic project rule injection")
|
||||
|
||||
# Workflow arguments (uses global --memory, --save, --verbose, --planning flags)
|
||||
parser.add_argument("--workflow", type=str, help="Run inline workflow steps (format: 'step1:action1,step2:action2')")
|
||||
@@ -4282,9 +4284,64 @@ Do NOT add any explanations or formatting."""
|
||||
# Apply prompt expansion if enabled
|
||||
prompt = self._expand_prompt_if_enabled(prompt)
|
||||
|
||||
# Prepend mention context to prompt
|
||||
if mention_context:
|
||||
# Auto-inject project instruction files unless disabled
|
||||
rules_context = ""
|
||||
try:
|
||||
should_load_rules = False
|
||||
max_chars = 32000 # Default cap
|
||||
include_manual = None
|
||||
|
||||
# First check for explicit --no-rules flag
|
||||
if hasattr(self, 'args') and getattr(self.args, 'no_rules', False):
|
||||
should_load_rules = False
|
||||
else:
|
||||
# Check for manual include_rules
|
||||
include_rules = getattr(self.args, 'include_rules', None) if hasattr(self, 'args') else None
|
||||
if include_rules:
|
||||
if include_rules == "auto":
|
||||
should_load_rules = True
|
||||
else:
|
||||
include_manual = [name.strip() for name in include_rules.split(",") if name.strip()]
|
||||
should_load_rules = bool(include_manual)
|
||||
else:
|
||||
# Check config for auto rules
|
||||
try:
|
||||
from praisonai.cli.configuration.loader import load_config
|
||||
config = load_config()
|
||||
should_load_rules = config.rules.auto
|
||||
max_chars = config.rules.max_chars
|
||||
except Exception:
|
||||
pass # Config not available, use defaults
|
||||
|
||||
if should_load_rules:
|
||||
from praisonaiagents.memory import RulesManager
|
||||
rules_manager = RulesManager(workspace_path=os.getcwd(), verbose=getattr(self.args, 'verbose', 0))
|
||||
rules_context = rules_manager.build_rules_context(
|
||||
include_manual=include_manual,
|
||||
max_chars=max_chars
|
||||
)
|
||||
if rules_context and getattr(self.args, 'verbose', 0):
|
||||
# Show loaded files in verbose mode
|
||||
active_rules = rules_manager.get_active_rules()
|
||||
loaded_files = []
|
||||
for rule in active_rules:
|
||||
if rule.file_path and rule.priority >= 500: # Root instruction files
|
||||
loaded_files.append(os.path.basename(rule.file_path))
|
||||
if loaded_files:
|
||||
print(f"[cyan]Loaded project instructions: {', '.join(loaded_files)}[/cyan]")
|
||||
except ImportError:
|
||||
pass # RulesManager not available
|
||||
except Exception as e:
|
||||
if hasattr(self, 'args') and getattr(self.args, 'verbose', 0):
|
||||
logging.debug(f"Error loading rules: {e}")
|
||||
|
||||
# Prepend mention context and rules context to prompt
|
||||
if mention_context and rules_context:
|
||||
prompt = f"{mention_context}{rules_context}\n\n# Task:\n{prompt}"
|
||||
elif mention_context:
|
||||
prompt = f"{mention_context}# Task:\n{prompt}"
|
||||
elif rules_context:
|
||||
prompt = f"{rules_context}\n\n# Task:\n{prompt}"
|
||||
|
||||
if PRAISONAI_AVAILABLE:
|
||||
from praisonaiagents import Agent as PraisonAgent
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
"""
|
||||
Project-scoped session management for CLI.
|
||||
|
||||
Provides session continuity within project boundaries.
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
from ..utils.project import get_project_id, get_project_name, get_project_sessions_dir
|
||||
|
||||
|
||||
class ProjectSessionStore(DefaultSessionStore):
|
||||
"""
|
||||
Project-scoped session store.
|
||||
|
||||
Extends DefaultSessionStore to scope sessions to the current project.
|
||||
"""
|
||||
|
||||
def __init__(self, project_path: Optional[str] = None, **kwargs):
|
||||
"""
|
||||
Initialize project-scoped session store.
|
||||
|
||||
Args:
|
||||
project_path: Project root path (defaults to cwd)
|
||||
**kwargs: Additional arguments for DefaultSessionStore
|
||||
"""
|
||||
self.project_path = project_path
|
||||
self.project_id = get_project_id(project_path)
|
||||
self.project_name = get_project_name(project_path)
|
||||
|
||||
# Use project-specific session directory
|
||||
project_session_dir = get_project_sessions_dir(project_path)
|
||||
|
||||
super().__init__(
|
||||
session_dir=str(project_session_dir),
|
||||
**kwargs
|
||||
)
|
||||
|
||||
def get_last_session_id(self) -> Optional[str]:
|
||||
"""
|
||||
Get the most recent session ID for this project.
|
||||
|
||||
Returns:
|
||||
Session ID of the most recent session, or None if no sessions exist
|
||||
"""
|
||||
sessions = self.list_sessions(limit=1)
|
||||
return sessions[0].get("session_id") if sessions else None
|
||||
|
||||
def get_project_info(self) -> dict:
|
||||
"""
|
||||
Get project information.
|
||||
|
||||
Returns:
|
||||
Dictionary with project details
|
||||
"""
|
||||
return {
|
||||
"project_id": self.project_id,
|
||||
"project_name": self.project_name,
|
||||
"project_path": self.project_path or str(Path.cwd()),
|
||||
"session_dir": self.session_dir,
|
||||
}
|
||||
|
||||
|
||||
def get_project_session_store(project_path: Optional[str] = None, project_id: Optional[str] = None) -> ProjectSessionStore:
|
||||
"""
|
||||
Get a project-scoped session store.
|
||||
|
||||
Args:
|
||||
project_path: Project root path (defaults to cwd)
|
||||
project_id: Specific project ID to use (if provided, creates store for that project)
|
||||
|
||||
Returns:
|
||||
ProjectSessionStore instance
|
||||
"""
|
||||
if project_id:
|
||||
# Create store for specific project ID
|
||||
from praisonaiagents.paths import get_sessions_dir
|
||||
project_session_dir = get_sessions_dir() / f"projects/{project_id}"
|
||||
# Use DefaultSessionStore directly with the specific directory
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
return DefaultSessionStore(session_dir=str(project_session_dir))
|
||||
else:
|
||||
# Use current project
|
||||
return ProjectSessionStore(project_path)
|
||||
|
||||
|
||||
def find_last_session(project_path: Optional[str] = None) -> Optional[str]:
|
||||
"""
|
||||
Find the last session ID for the current project.
|
||||
|
||||
Args:
|
||||
project_path: Project root path (defaults to cwd)
|
||||
|
||||
Returns:
|
||||
Session ID or None if no sessions exist
|
||||
"""
|
||||
store = get_project_session_store(project_path)
|
||||
return store.get_last_session_id()
|
||||
@@ -0,0 +1,96 @@
|
||||
"""
|
||||
Project identity utilities for CLI session scoping.
|
||||
|
||||
Provides functions to identify the current project context for scoped sessions.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def get_git_root(path: Optional[str] = None) -> Optional[Path]:
|
||||
"""
|
||||
Find the git repository root from the given path.
|
||||
|
||||
Args:
|
||||
path: Starting path (defaults to cwd)
|
||||
|
||||
Returns:
|
||||
Path to git root, or None if not in a git repo
|
||||
"""
|
||||
start_path = Path(path) if path else Path.cwd()
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['git', 'rev-parse', '--show-toplevel'],
|
||||
cwd=start_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
timeout=5
|
||||
)
|
||||
return Path(result.stdout.strip())
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
|
||||
return None
|
||||
|
||||
|
||||
def get_project_id(path: Optional[str] = None) -> str:
|
||||
"""
|
||||
Generate a stable project identifier.
|
||||
|
||||
Uses git root path if available, falls back to current directory.
|
||||
Returns a short hash for scoped session storage.
|
||||
|
||||
Args:
|
||||
path: Starting path (defaults to cwd)
|
||||
|
||||
Returns:
|
||||
Short project hash (8 chars)
|
||||
"""
|
||||
# Try git root first
|
||||
git_root = get_git_root(path)
|
||||
if git_root:
|
||||
project_path = str(git_root.resolve())
|
||||
else:
|
||||
# Fall back to current directory
|
||||
project_path = str(Path(path).resolve() if path else Path.cwd().resolve())
|
||||
|
||||
# Generate short hash
|
||||
return hashlib.sha256(project_path.encode()).hexdigest()[:8]
|
||||
|
||||
|
||||
def get_project_name(path: Optional[str] = None) -> str:
|
||||
"""
|
||||
Get a human-readable project name.
|
||||
|
||||
Args:
|
||||
path: Starting path (defaults to cwd)
|
||||
|
||||
Returns:
|
||||
Project directory name
|
||||
"""
|
||||
git_root = get_git_root(path)
|
||||
if git_root:
|
||||
return git_root.name
|
||||
|
||||
current_path = Path(path) if path else Path.cwd()
|
||||
return current_path.name
|
||||
|
||||
|
||||
def get_project_sessions_dir(path: Optional[str] = None) -> Path:
|
||||
"""
|
||||
Get project-scoped session directory.
|
||||
|
||||
Args:
|
||||
path: Starting path (defaults to cwd)
|
||||
|
||||
Returns:
|
||||
Path to project sessions directory
|
||||
"""
|
||||
from praisonaiagents.paths import get_sessions_dir
|
||||
|
||||
project_id = get_project_id(path)
|
||||
return get_sessions_dir() / f"projects/{project_id}"
|
||||
@@ -248,9 +248,9 @@ def is_path_within_directory(file_path: str, directory: str) -> bool:
|
||||
Returns:
|
||||
True if file_path is within directory
|
||||
"""
|
||||
# Resolve to absolute paths
|
||||
abs_file = os.path.abspath(file_path)
|
||||
abs_dir = os.path.abspath(directory)
|
||||
# Resolve to real paths so symlinks cannot escape the workspace
|
||||
abs_file = os.path.realpath(file_path)
|
||||
abs_dir = os.path.realpath(directory)
|
||||
|
||||
# Ensure directory ends with separator for proper prefix matching
|
||||
if not abs_dir.endswith(os.sep):
|
||||
|
||||
@@ -1,223 +0,0 @@
|
||||
import subprocess
|
||||
import os
|
||||
import platform
|
||||
from dotenv import load_dotenv
|
||||
|
||||
class CloudDeployer:
|
||||
"""
|
||||
A class for deploying a cloud-based application.
|
||||
|
||||
Attributes:
|
||||
None
|
||||
|
||||
Methods:
|
||||
__init__(self):
|
||||
Loads environment variables from .env file or system and sets them.
|
||||
|
||||
"""
|
||||
def __init__(self):
|
||||
"""
|
||||
Loads environment variables from .env file or system and sets them.
|
||||
|
||||
Parameters:
|
||||
self: An instance of the CloudDeployer class.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
None
|
||||
|
||||
"""
|
||||
# Load environment variables from .env file or system
|
||||
load_dotenv()
|
||||
self.set_environment_variables()
|
||||
|
||||
def create_dockerfile(self):
|
||||
"""
|
||||
Creates a Dockerfile for the application.
|
||||
|
||||
Parameters:
|
||||
self: An instance of the CloudDeployer class.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
None
|
||||
|
||||
This method creates a Dockerfile in the current directory with the specified content.
|
||||
The Dockerfile is used to build a Docker image for the application.
|
||||
The content of the Dockerfile includes instructions to use the Python 3.11-slim base image,
|
||||
set the working directory to /app, copy the current directory contents into the container,
|
||||
install the required Python packages (flask, praisonai, gunicorn, and markdown),
|
||||
expose port 8080, and run the application using Gunicorn.
|
||||
"""
|
||||
with open("Dockerfile", "w") as file:
|
||||
file.write("FROM python:3.11-slim\n")
|
||||
file.write("WORKDIR /app\n")
|
||||
file.write("COPY . .\n")
|
||||
file.write("RUN pip install flask praisonai==4.6.55 gunicorn markdown\n")
|
||||
file.write("EXPOSE 8080\n")
|
||||
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
|
||||
|
||||
def create_api_file(self):
|
||||
"""
|
||||
Creates an API file for the application.
|
||||
|
||||
Parameters:
|
||||
self (CloudDeployer): An instance of the CloudDeployer class.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
This method creates an API file named "api.py" in the current directory. The file contains a basic Flask application that uses the PraisonAI library to run a simple agent and returns the output as an HTML page. The application listens on the root path ("/") and uses the Markdown library to format the output.
|
||||
"""
|
||||
with open("api.py", "w") as file:
|
||||
file.write("from flask import Flask\n")
|
||||
file.write("from praisonai import PraisonAI\n")
|
||||
file.write("import markdown\n")
|
||||
file.write("import bleach\n\n")
|
||||
file.write("app = Flask(__name__)\n\n")
|
||||
file.write("def basic():\n")
|
||||
file.write(" praisonai = PraisonAI(agent_file=\"agents.yaml\")\n")
|
||||
file.write(" return praisonai.run()\n\n")
|
||||
file.write("@app.route('/')\n")
|
||||
file.write("def home():\n")
|
||||
file.write(" output = basic()\n")
|
||||
file.write(" rendered = markdown.markdown(str(output))\n")
|
||||
file.write(" safe_html = bleach.clean(rendered, tags=bleach.sanitizer.ALLOWED_TAGS, attributes=bleach.sanitizer.ALLOWED_ATTRIBUTES)\n")
|
||||
file.write(" return f'<html><body>{safe_html}</body></html>'\n\n")
|
||||
file.write("if __name__ == \"__main__\":\n")
|
||||
file.write(" import os\n")
|
||||
file.write(" app.run(debug=os.environ.get('DEBUG', 'false').lower() == 'true')\n")
|
||||
|
||||
def set_environment_variables(self):
|
||||
"""Sets environment variables with fallback to .env values or defaults."""
|
||||
from praisonai.llm.env import resolve_llm_endpoint
|
||||
ep = resolve_llm_endpoint()
|
||||
|
||||
os.environ["OPENAI_MODEL_NAME"] = ep.model
|
||||
os.environ["OPENAI_API_KEY"] = ep.api_key or "Enter your API key"
|
||||
os.environ["OPENAI_API_BASE"] = ep.base_url
|
||||
|
||||
def run_commands(self):
|
||||
"""
|
||||
Sets environment variables with fallback to .env values or defaults.
|
||||
|
||||
Parameters:
|
||||
None
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
None
|
||||
|
||||
This method sets environment variables for the application. It uses the `os.environ` dictionary to set the following environment variables:
|
||||
|
||||
- `OPENAI_MODEL_NAME`: The name of the OpenAI model to use. If not specified in the .env file, it defaults to "gpt-4o-mini".
|
||||
- `OPENAI_API_KEY`: The API key for accessing the OpenAI API. If not specified in the .env file, it defaults to "Enter your API key".
|
||||
- `OPENAI_API_BASE`: The base URL for the OpenAI API. If not specified in the .env file, it defaults to "https://api.openai.com/v1".
|
||||
"""
|
||||
self.create_api_file()
|
||||
self.create_dockerfile()
|
||||
"""Runs a sequence of shell commands for deployment, continues on error."""
|
||||
|
||||
# Get project ID upfront for Windows compatibility
|
||||
try:
|
||||
result = subprocess.run(['gcloud', 'config', 'get-value', 'project'],
|
||||
capture_output=True, text=True, check=True)
|
||||
project_id = result.stdout.strip()
|
||||
except subprocess.CalledProcessError:
|
||||
print("ERROR: Failed to get GCP project ID. Ensure gcloud is configured.")
|
||||
return
|
||||
|
||||
# Get environment variables
|
||||
from praisonai.llm.env import resolve_llm_endpoint
|
||||
ep = resolve_llm_endpoint()
|
||||
openai_model = ep.model
|
||||
openai_key = ep.api_key or 'Enter your API key'
|
||||
openai_base = ep.base_url
|
||||
|
||||
# Create temporary env vars file to avoid exposing secrets in argv
|
||||
import tempfile
|
||||
import yaml
|
||||
import os
|
||||
|
||||
env_vars_file = None
|
||||
try:
|
||||
# Create secure temp file for environment variables
|
||||
fd, env_vars_file = tempfile.mkstemp(suffix=".yaml", prefix="praisonai-deploy-")
|
||||
os.close(fd)
|
||||
os.chmod(env_vars_file, 0o600) # Secure file permissions
|
||||
|
||||
# Write env vars to file instead of passing in argv
|
||||
env_vars = {
|
||||
"OPENAI_MODEL_NAME": openai_model,
|
||||
"OPENAI_API_KEY": openai_key,
|
||||
"OPENAI_API_BASE": openai_base
|
||||
}
|
||||
|
||||
with open(env_vars_file, "w") as f:
|
||||
yaml.safe_dump(env_vars, f)
|
||||
|
||||
# Build commands with secure env vars file
|
||||
commands = [
|
||||
['gcloud', 'auth', 'configure-docker', 'us-central1-docker.pkg.dev'],
|
||||
['gcloud', 'artifacts', 'repositories', 'create', 'praisonai-repository',
|
||||
'--repository-format=docker', '--location=us-central1'],
|
||||
['docker', 'build', '--platform', 'linux/amd64', '-t',
|
||||
f'gcr.io/{project_id}/praisonai-app:latest', '.'],
|
||||
['docker', 'tag', f'gcr.io/{project_id}/praisonai-app:latest',
|
||||
f'us-central1-docker.pkg.dev/{project_id}/praisonai-repository/praisonai-app:latest'],
|
||||
['docker', 'push',
|
||||
f'us-central1-docker.pkg.dev/{project_id}/praisonai-repository/praisonai-app:latest'],
|
||||
['gcloud', 'run', 'deploy', 'praisonai-service',
|
||||
'--image', f'us-central1-docker.pkg.dev/{project_id}/praisonai-repository/praisonai-app:latest',
|
||||
'--platform', 'managed', '--region', 'us-central1', '--allow-unauthenticated',
|
||||
'--env-vars-file', env_vars_file]
|
||||
]
|
||||
|
||||
# Run commands with appropriate handling for each platform
|
||||
for i, cmd in enumerate(commands):
|
||||
try:
|
||||
if i == 0: # First command (gcloud auth configure-docker)
|
||||
if platform.system() != 'Windows':
|
||||
# On Unix, pipe 'yes' to auto-confirm
|
||||
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
|
||||
proc.communicate(input=b'Y\n')
|
||||
if proc.returncode != 0:
|
||||
raise subprocess.CalledProcessError(proc.returncode, cmd)
|
||||
else:
|
||||
# On Windows, try with --quiet flag to avoid prompts
|
||||
cmd_with_quiet = cmd + ['--quiet']
|
||||
try:
|
||||
subprocess.run(cmd_with_quiet, check=True)
|
||||
except subprocess.CalledProcessError:
|
||||
# If --quiet fails, try without it
|
||||
print("Note: You may need to manually confirm the authentication prompt")
|
||||
subprocess.run(cmd, check=True)
|
||||
else:
|
||||
# Run other commands normally
|
||||
subprocess.run(cmd, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"ERROR: Command failed with exit status {e.returncode}")
|
||||
# Commands 2 (build) and 4 (push) and 5 (deploy) are critical
|
||||
if i in [2, 4, 5]:
|
||||
print("Critical command failed. Aborting deployment.")
|
||||
return
|
||||
print(f"Continuing with the next command...")
|
||||
|
||||
finally:
|
||||
# Always cleanup the temporary env vars file
|
||||
if env_vars_file:
|
||||
try:
|
||||
os.remove(env_vars_file)
|
||||
except Exception:
|
||||
pass # Don't fail deployment if cleanup fails
|
||||
|
||||
# Usage
|
||||
if __name__ == "__main__":
|
||||
deployer = CloudDeployer()
|
||||
deployer.run_commands()
|
||||
@@ -110,8 +110,13 @@ class KanbanDispatcher:
|
||||
# Release claim on error
|
||||
try:
|
||||
store.release_claim(task.id, self.worker_id)
|
||||
except:
|
||||
pass
|
||||
except Exception as release_err:
|
||||
logger.error(
|
||||
"Failed to release claim for task %s (worker %s); task may be stuck "
|
||||
"in 'claimed' state until manually released: %s",
|
||||
task.id, self.worker_id, release_err,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
if spawned > 0:
|
||||
logger.info(f"Spawned {spawned} kanban task workers")
|
||||
@@ -149,13 +154,15 @@ class KanbanDispatcher:
|
||||
with tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix='.log') as temp_log:
|
||||
temp_log_path = temp_log.name
|
||||
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
stdout=open(temp_log_path, 'w'),
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True
|
||||
)
|
||||
with open(temp_log_path, 'w') as log_handle:
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
stdout=log_handle,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True
|
||||
)
|
||||
# parent FD closed here; child still has its duped copy
|
||||
|
||||
# Store log path for later cleanup
|
||||
if not hasattr(self, '_temp_logs'):
|
||||
@@ -275,8 +282,12 @@ class KanbanDispatcher:
|
||||
# Release claim as fallback
|
||||
try:
|
||||
store.release_claim(task_id, self.worker_id)
|
||||
except:
|
||||
pass
|
||||
except Exception as release_err:
|
||||
logger.error(
|
||||
"Failed to release claim during cleanup for task %s: %s",
|
||||
task_id, release_err,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
# Remove completed tasks from tracking and clean up temp logs
|
||||
for task_id in completed:
|
||||
@@ -330,12 +341,23 @@ class KanbanDispatcher:
|
||||
|
||||
# Force terminate remaining tasks
|
||||
for task_id, process in self.running_tasks.items():
|
||||
logger.warning(f"Force terminating task {task_id}")
|
||||
try:
|
||||
logger.warning(f"Force terminating task {task_id}")
|
||||
process.terminate()
|
||||
process.wait(timeout=5)
|
||||
except:
|
||||
process.kill()
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(f"Task {task_id} did not terminate in 5s; sending SIGKILL")
|
||||
try:
|
||||
process.kill()
|
||||
except OSError as kill_err:
|
||||
logger.error(f"Failed to kill task {task_id}: {kill_err}")
|
||||
except OSError as os_err:
|
||||
logger.error(f"OS error while terminating task {task_id}: {os_err}")
|
||||
try:
|
||||
process.kill()
|
||||
except OSError as kill_err:
|
||||
logger.error(f"Failed to kill task {task_id}: {kill_err}")
|
||||
# KeyboardInterrupt, SystemExit, CancelledError now propagate as intended
|
||||
|
||||
self.running_tasks.clear()
|
||||
|
||||
|
||||
@@ -30,6 +30,7 @@ from praisonaiagents.gateway import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .unicode_utils import safe_error_message, safe_log_message, extract_root_cause_from_error
|
||||
from .supervisor import ChannelSupervisor
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -278,6 +279,9 @@ class WebSocketGateway:
|
||||
|
||||
# PID lock for single-instance enforcement
|
||||
self._pid_lock: Optional[Any] = None
|
||||
|
||||
# Channel supervisor for resilient bot management
|
||||
self._channel_supervisor = ChannelSupervisor()
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
@@ -718,6 +722,43 @@ class WebSocketGateway:
|
||||
status_code=500
|
||||
)
|
||||
|
||||
# Channel control endpoints
|
||||
async def pause_channel_handler(request) -> JSONResponse:
|
||||
"""POST /api/channels/{name}/pause — pause a channel."""
|
||||
auth_err = _check_auth(request)
|
||||
if auth_err:
|
||||
return auth_err
|
||||
channel_name = request.path_params["name"]
|
||||
success = self.pause_channel(channel_name)
|
||||
return JSONResponse({
|
||||
"success": success,
|
||||
"message": f"Channel '{channel_name}' {'paused' if success else 'could not be paused'}"
|
||||
})
|
||||
|
||||
async def resume_channel_handler(request) -> JSONResponse:
|
||||
"""POST /api/channels/{name}/resume — resume a paused channel."""
|
||||
auth_err = _check_auth(request)
|
||||
if auth_err:
|
||||
return auth_err
|
||||
channel_name = request.path_params["name"]
|
||||
success = self.resume_channel(channel_name)
|
||||
return JSONResponse({
|
||||
"success": success,
|
||||
"message": f"Channel '{channel_name}' {'resumed' if success else 'could not be resumed'}"
|
||||
})
|
||||
|
||||
async def reconnect_channel_handler(request) -> JSONResponse:
|
||||
"""POST /api/channels/{name}/reconnect — reconnect a channel."""
|
||||
auth_err = _check_auth(request)
|
||||
if auth_err:
|
||||
return auth_err
|
||||
channel_name = request.path_params["name"]
|
||||
success = self.reconnect_channel(channel_name)
|
||||
return JSONResponse({
|
||||
"success": success,
|
||||
"message": f"Channel '{channel_name}' {'reconnected' if success else 'could not be reconnected'}"
|
||||
})
|
||||
|
||||
routes = [
|
||||
Route("/", magic_link_handler, methods=["GET"]),
|
||||
Route("/health", health, methods=["GET"]),
|
||||
@@ -728,6 +769,9 @@ class WebSocketGateway:
|
||||
Route("/api/pairing/pending", _pairing_routes["pending"], methods=["GET"]),
|
||||
Route("/api/pairing/approve", _pairing_routes["approve"], methods=["POST"]),
|
||||
Route("/api/pairing/revoke", _pairing_routes["revoke"], methods=["POST"]),
|
||||
Route("/api/channels/{name}/pause", pause_channel_handler, methods=["POST"]),
|
||||
Route("/api/channels/{name}/resume", resume_channel_handler, methods=["POST"]),
|
||||
Route("/api/channels/{name}/reconnect", reconnect_channel_handler, methods=["POST"]),
|
||||
WebSocketRoute("/ws", websocket_endpoint),
|
||||
]
|
||||
|
||||
@@ -1189,16 +1233,36 @@ class WebSocketGateway:
|
||||
logger.error(f"Broadcast error to {client_id}: {e}")
|
||||
|
||||
def health(self) -> Dict[str, Any]:
|
||||
"""Get gateway health status including per-channel bot status."""
|
||||
"""Get gateway health status including per-channel bot status and supervision state."""
|
||||
uptime = time.time() - self._started_at if self._started_at else 0
|
||||
channel_status = {}
|
||||
supervision_status = self._channel_supervisor.get_all_status()
|
||||
|
||||
for name, bot in self._channel_bots.items():
|
||||
running = getattr(bot, "is_running", False)
|
||||
platform = getattr(bot, "platform", "unknown")
|
||||
channel_status[name] = {
|
||||
"platform": platform,
|
||||
"running": running,
|
||||
}
|
||||
|
||||
# Get supervision state
|
||||
sup_status = supervision_status.get(name)
|
||||
if sup_status:
|
||||
channel_status[name] = {
|
||||
"platform": platform,
|
||||
"running": running,
|
||||
"supervision": {
|
||||
"state": sup_status.state.value,
|
||||
"last_error": sup_status.last_error,
|
||||
"last_error_time": sup_status.last_error_time,
|
||||
"next_retry_at": sup_status.next_retry_at,
|
||||
"total_recoveries": sup_status.total_recoveries,
|
||||
"manual_pause": sup_status.manual_pause,
|
||||
}
|
||||
}
|
||||
else:
|
||||
channel_status[name] = {
|
||||
"platform": platform,
|
||||
"running": running,
|
||||
}
|
||||
|
||||
result = {
|
||||
"status": "healthy" if self._is_running else "stopped",
|
||||
"uptime": uptime,
|
||||
@@ -1769,44 +1833,29 @@ class WebSocketGateway:
|
||||
return None
|
||||
|
||||
async def _run_bot_safe(self, name: str, bot: Any) -> None:
|
||||
"""Run a single bot, catching errors so other bots stay alive.
|
||||
"""Run a single bot with supervision and resilient error handling.
|
||||
|
||||
For TelegramBot, uses the low-level initialize/start/start_polling
|
||||
API to avoid event-loop conflicts with ``run_polling()``.
|
||||
For Discord/Slack, injects a routing-aware message handler before
|
||||
starting so the gateway's routing rules are respected.
|
||||
Uses the channel supervisor for unlimited retries with error classification,
|
||||
exponential backoff, and operator controls.
|
||||
"""
|
||||
max_retries = 5
|
||||
base_delay = 5 # seconds
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
logger.info(f"Starting bot for channel '{name}'..." + (f" (retry {attempt})" if attempt else ""))
|
||||
|
||||
# TelegramBot special handling: run_polling() tries to manage
|
||||
# its own event loop which conflicts with our gateway loop.
|
||||
# Use the lower-level API instead.
|
||||
if self._is_telegram_bot(bot):
|
||||
await self._start_telegram_bot_polling(name, bot)
|
||||
elif type(bot).__name__ in ("WhatsAppBot", "LinearBot"):
|
||||
# WhatsApp/Linear run their own aiohttp webhook servers
|
||||
self._inject_routing_handler(name, bot)
|
||||
await bot.start()
|
||||
else:
|
||||
# Inject routing-aware handler for Discord/Slack
|
||||
self._inject_routing_handler(name, bot)
|
||||
await bot.start()
|
||||
break # clean exit
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"Bot '{name}' cancelled")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Bot '{name}' crashed: {e}")
|
||||
if attempt < max_retries:
|
||||
delay = base_delay * (2 ** attempt)
|
||||
logger.info(f"Reconnecting '{name}' in {delay}s...")
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Bot '{name}' failed after {max_retries} retries, giving up")
|
||||
async def start_bot(name: str, bot: Any) -> None:
|
||||
"""Start function for the supervisor."""
|
||||
# TelegramBot special handling: run_polling() tries to manage
|
||||
# its own event loop which conflicts with our gateway loop.
|
||||
# Use the lower-level API instead.
|
||||
if self._is_telegram_bot(bot):
|
||||
await self._start_telegram_bot_polling(name, bot)
|
||||
elif type(bot).__name__ in ("WhatsAppBot", "LinearBot"):
|
||||
# WhatsApp/Linear run their own aiohttp webhook servers
|
||||
self._inject_routing_handler(name, bot)
|
||||
await bot.start()
|
||||
else:
|
||||
# Inject routing-aware handler for Discord/Slack
|
||||
self._inject_routing_handler(name, bot)
|
||||
await bot.start()
|
||||
|
||||
# Use supervisor for resilient channel management
|
||||
await self._channel_supervisor.run(name, bot, start_bot)
|
||||
|
||||
@staticmethod
|
||||
def _is_telegram_bot(bot: Any) -> bool:
|
||||
@@ -2028,6 +2077,10 @@ class WebSocketGateway:
|
||||
except Exception as e:
|
||||
logger.error(f"Error stopping bot '{name}': {e}")
|
||||
|
||||
# Clean up supervisor state for all channels
|
||||
for name in list(self._channel_bots.keys()):
|
||||
self._channel_supervisor.cleanup(name)
|
||||
|
||||
self._channel_bots.clear()
|
||||
self._routing_rules.clear()
|
||||
|
||||
@@ -2087,6 +2140,59 @@ class WebSocketGateway:
|
||||
except Exception as e:
|
||||
logger.warning(f"Config watch error: {e}")
|
||||
|
||||
# Channel supervision control methods
|
||||
|
||||
def pause_channel(self, name: str) -> bool:
|
||||
"""Pause a channel.
|
||||
|
||||
Args:
|
||||
name: Channel name
|
||||
|
||||
Returns:
|
||||
True if channel was paused, False if not found or not running
|
||||
"""
|
||||
return self._channel_supervisor.pause(name)
|
||||
|
||||
def resume_channel(self, name: str) -> bool:
|
||||
"""Resume a paused channel.
|
||||
|
||||
Args:
|
||||
name: Channel name
|
||||
|
||||
Returns:
|
||||
True if channel was resumed, False if not found or not paused
|
||||
"""
|
||||
return self._channel_supervisor.resume(name)
|
||||
|
||||
def reconnect_channel(self, name: str) -> bool:
|
||||
"""Force reconnect a channel.
|
||||
|
||||
Args:
|
||||
name: Channel name
|
||||
|
||||
Returns:
|
||||
True if channel exists, False otherwise
|
||||
"""
|
||||
return self._channel_supervisor.reconnect(name)
|
||||
|
||||
def get_channel_supervision_status(self) -> Dict[str, Any]:
|
||||
"""Get supervision status for all channels.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping channel names to their supervision status
|
||||
"""
|
||||
status_dict = {}
|
||||
for name, status in self._channel_supervisor.get_all_status().items():
|
||||
status_dict[name] = {
|
||||
"state": status.state.value,
|
||||
"last_error": status.last_error,
|
||||
"last_error_time": status.last_error_time,
|
||||
"next_retry_at": status.next_retry_at,
|
||||
"total_recoveries": status.total_recoveries,
|
||||
"manual_pause": status.manual_pause,
|
||||
}
|
||||
return status_dict
|
||||
|
||||
async def start_with_config(self, config_path: str) -> None:
|
||||
"""Start the gateway with a gateway.yaml configuration.
|
||||
|
||||
|
||||
@@ -0,0 +1,277 @@
|
||||
"""
|
||||
Channel supervision for WebSocket Gateway.
|
||||
|
||||
Provides resilient channel management with error classification,
|
||||
unlimited retries for recoverable errors, and operator controls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
from ..bots._resilience import (
|
||||
BackoffPolicy,
|
||||
ConnectionMonitor,
|
||||
is_recoverable_error,
|
||||
is_conflict_error,
|
||||
sleep_with_abort,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChannelState(Enum):
|
||||
"""Channel supervision states."""
|
||||
RUNNING = "running"
|
||||
FAILED = "failed"
|
||||
PAUSED = "paused"
|
||||
STOPPED = "stopped"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChannelStatus:
|
||||
"""Channel supervision status."""
|
||||
state: ChannelState = ChannelState.STOPPED
|
||||
last_error: Optional[str] = None
|
||||
last_error_time: Optional[float] = None
|
||||
next_retry_at: Optional[float] = None
|
||||
total_recoveries: int = 0
|
||||
manual_pause: bool = False
|
||||
|
||||
|
||||
class ChannelSupervisor:
|
||||
"""Supervises channel lifecycle with resilient error handling.
|
||||
|
||||
Provides:
|
||||
- Error classification (fatal vs recoverable)
|
||||
- Unlimited retries for recoverable errors with capped exponential backoff
|
||||
- Manual pause/resume/reconnect controls
|
||||
- Status tracking for health reporting
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
policy: Optional[BackoffPolicy] = None,
|
||||
classify_fn: Optional[Callable[[BaseException, str], bool]] = None
|
||||
):
|
||||
"""Initialize channel supervisor.
|
||||
|
||||
Args:
|
||||
policy: Backoff policy for retries (uses default if None)
|
||||
classify_fn: Error classification function (uses is_recoverable_error if None)
|
||||
"""
|
||||
self._policy = policy or BackoffPolicy(max_attempts=0) # Unlimited retries
|
||||
self._classify_fn = classify_fn or is_recoverable_error
|
||||
self._channels: Dict[str, ChannelStatus] = {}
|
||||
self._monitors: Dict[str, ConnectionMonitor] = {}
|
||||
self._abort_signals: Dict[str, asyncio.Event] = {}
|
||||
self._tasks: Dict[str, asyncio.Task] = {}
|
||||
|
||||
def get_status(self, name: str) -> ChannelStatus:
|
||||
"""Get current status of a channel."""
|
||||
return self._channels.get(name, ChannelStatus())
|
||||
|
||||
def get_all_status(self) -> Dict[str, ChannelStatus]:
|
||||
"""Get status of all supervised channels."""
|
||||
return dict(self._channels)
|
||||
|
||||
def pause(self, name: str) -> bool:
|
||||
"""Manually pause a channel.
|
||||
|
||||
Args:
|
||||
name: Channel name
|
||||
|
||||
Returns:
|
||||
True if channel was running and paused, False otherwise
|
||||
"""
|
||||
if name not in self._channels:
|
||||
return False
|
||||
|
||||
status = self._channels[name]
|
||||
if status.state == ChannelState.RUNNING:
|
||||
status.state = ChannelState.PAUSED
|
||||
status.manual_pause = True
|
||||
|
||||
# Signal abort to stop current operations
|
||||
if name in self._abort_signals:
|
||||
self._abort_signals[name].set()
|
||||
|
||||
logger.info(f"Channel '{name}' manually paused")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def resume(self, name: str) -> bool:
|
||||
"""Resume a manually paused channel.
|
||||
|
||||
Args:
|
||||
name: Channel name
|
||||
|
||||
Returns:
|
||||
True if channel was paused and resumed, False otherwise
|
||||
"""
|
||||
if name not in self._channels:
|
||||
return False
|
||||
|
||||
status = self._channels[name]
|
||||
if status.state == ChannelState.PAUSED and status.manual_pause:
|
||||
status.state = ChannelState.STOPPED # Will be restarted by supervision
|
||||
status.manual_pause = False
|
||||
|
||||
# SET abort signal to wake the paused supervision loop
|
||||
if name in self._abort_signals:
|
||||
self._abort_signals[name].set()
|
||||
|
||||
logger.info(f"Channel '{name}' manually resumed")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def reconnect(self, name: str) -> bool:
|
||||
"""Force reconnect of a channel.
|
||||
|
||||
Args:
|
||||
name: Channel name
|
||||
|
||||
Returns:
|
||||
True if channel exists, False otherwise
|
||||
"""
|
||||
if name not in self._channels:
|
||||
return False
|
||||
|
||||
# Reset monitor state and force restart
|
||||
if name in self._monitors:
|
||||
self._monitors[name].attempt = 0
|
||||
self._monitors[name].last_error = None
|
||||
self._monitors[name].last_error_time = None
|
||||
|
||||
status = self._channels[name]
|
||||
status.state = ChannelState.STOPPED
|
||||
status.manual_pause = False # Clear manual pause flag
|
||||
status.last_error = None
|
||||
status.last_error_time = None
|
||||
status.next_retry_at = None
|
||||
|
||||
# Signal abort to stop current operations
|
||||
if name in self._abort_signals:
|
||||
self._abort_signals[name].set()
|
||||
|
||||
logger.info(f"Channel '{name}' manually reconnected")
|
||||
return True
|
||||
|
||||
async def run(self, name: str, bot: Any, start_fn: Callable[[str, Any], Any]) -> None:
|
||||
"""Run a channel with supervision.
|
||||
|
||||
Args:
|
||||
name: Channel name
|
||||
bot: Bot instance
|
||||
start_fn: Function to start the bot (should be async)
|
||||
"""
|
||||
# Initialize supervision state
|
||||
if name not in self._channels:
|
||||
self._channels[name] = ChannelStatus()
|
||||
if name not in self._monitors:
|
||||
# Determine platform from bot for error classification
|
||||
platform = getattr(bot, 'platform', type(bot).__name__.lower())
|
||||
self._monitors[name] = ConnectionMonitor(platform=platform, policy=self._policy)
|
||||
if name not in self._abort_signals:
|
||||
self._abort_signals[name] = asyncio.Event()
|
||||
|
||||
status = self._channels[name]
|
||||
monitor = self._monitors[name]
|
||||
abort_signal = self._abort_signals[name]
|
||||
|
||||
logger.info(f"Starting supervision for channel '{name}'")
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Check if manually paused
|
||||
if status.manual_pause:
|
||||
status.state = ChannelState.PAUSED
|
||||
await abort_signal.wait() # Wait until resumed
|
||||
abort_signal.clear()
|
||||
continue
|
||||
|
||||
# Check if task was cancelled
|
||||
if abort_signal.is_set():
|
||||
abort_signal.clear()
|
||||
|
||||
status.state = ChannelState.RUNNING
|
||||
logger.info(f"Starting channel '{name}'..." +
|
||||
(f" (attempt {monitor.attempt + 1})" if monitor.attempt > 0 else ""))
|
||||
|
||||
# Start the bot
|
||||
await start_fn(name, bot)
|
||||
|
||||
# If we get here, the bot exited cleanly
|
||||
monitor.record_success()
|
||||
status.total_recoveries = monitor.total_recoveries # Sync recovery count
|
||||
status.state = ChannelState.STOPPED
|
||||
logger.info(f"Channel '{name}' stopped cleanly")
|
||||
break
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"Channel '{name}' supervision cancelled")
|
||||
status.state = ChannelState.STOPPED
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
# Classify the error
|
||||
is_recoverable = self._classify_fn(e, monitor.platform)
|
||||
is_conflict = is_conflict_error(e)
|
||||
|
||||
if is_conflict:
|
||||
# Conflict errors are fatal - another bot instance using same token
|
||||
status.state = ChannelState.FAILED
|
||||
status.last_error = f"Conflict error (fatal): {str(e)}"
|
||||
status.last_error_time = time.time()
|
||||
status.next_retry_at = None
|
||||
logger.error(f"Channel '{name}' failed with conflict error: {e}")
|
||||
break
|
||||
|
||||
elif not is_recoverable:
|
||||
# Non-recoverable error - treat as fatal
|
||||
status.state = ChannelState.FAILED
|
||||
status.last_error = f"Fatal error: {str(e)}"
|
||||
status.last_error_time = time.time()
|
||||
status.next_retry_at = None
|
||||
logger.error(f"Channel '{name}' failed with fatal error: {e}")
|
||||
break
|
||||
|
||||
else:
|
||||
# Recoverable error - retry with backoff
|
||||
delay = monitor.record_error(e)
|
||||
status.last_error = str(e)
|
||||
status.last_error_time = time.time()
|
||||
status.next_retry_at = time.time() + delay
|
||||
|
||||
logger.warning(f"Channel '{name}' error (recoverable): {e}")
|
||||
logger.info(f"Retrying channel '{name}' in {delay:.1f}s...")
|
||||
|
||||
# Sleep with abort signal support
|
||||
completed = await sleep_with_abort(delay, abort_signal)
|
||||
if not completed:
|
||||
# Aborted - check if paused or reconnect requested
|
||||
continue
|
||||
|
||||
# Cleanup - don't overwrite terminal failure states
|
||||
if status.state != ChannelState.FAILED:
|
||||
status.state = ChannelState.STOPPED
|
||||
logger.info(f"Supervision ended for channel '{name}'")
|
||||
|
||||
def cleanup(self, name: str) -> None:
|
||||
"""Clean up supervision state for a channel."""
|
||||
self._channels.pop(name, None)
|
||||
self._monitors.pop(name, None)
|
||||
if name in self._abort_signals:
|
||||
self._abort_signals[name].set()
|
||||
self._abort_signals.pop(name, None)
|
||||
if name in self._tasks:
|
||||
task = self._tasks.pop(name)
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
@@ -7,9 +7,64 @@ Set ``PRAISONAI_HOST_LEGACY=1`` to skip provider wiring (callback-only mode).
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import contextvars
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
|
||||
_CONFIGURED = False
|
||||
# Use context variable instead of module-level global for multi-agent safety
|
||||
_configured_context: contextvars.ContextVar[bool] = contextvars.ContextVar('host_configured', default=False)
|
||||
|
||||
# For testing: also maintain a module-level flag that can be accessed across contexts
|
||||
_configured_global = False
|
||||
|
||||
# Backward compatibility for tests
|
||||
def reset_configuration() -> None:
|
||||
"""Reset configuration state for testing. Use instead of host_app._CONFIGURED = False."""
|
||||
global _configured_global
|
||||
_configured_context.set(False)
|
||||
_configured_global = False
|
||||
|
||||
# Also clear any cached state that might be lingering
|
||||
try:
|
||||
import praisonaiui.server as srv
|
||||
if hasattr(srv, '_provider'):
|
||||
srv._provider = None
|
||||
if hasattr(srv, '_datastore'):
|
||||
srv._datastore = None
|
||||
# Clear any other potentially cached state
|
||||
if hasattr(srv, '_app'):
|
||||
srv._app = None
|
||||
# Clear backends registry if it exists
|
||||
try:
|
||||
import praisonaiui.backends as backends
|
||||
if hasattr(backends, 'clear_backends'):
|
||||
backends.clear_backends()
|
||||
elif hasattr(backends, '_backends'):
|
||||
backends._backends.clear()
|
||||
except (ImportError, AttributeError):
|
||||
pass
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
def is_configured() -> bool:
|
||||
"""Check if configuration has been applied in current context."""
|
||||
return _configured_context.get() or _configured_global
|
||||
|
||||
# Backward compatibility shim for tests that assign host_app._CONFIGURED = False
|
||||
class _ConfiguredShim:
|
||||
"""Backward compatibility shim that proxies to both ContextVar and global."""
|
||||
def __get__(self, obj, objtype=None):
|
||||
return _configured_context.get() or _configured_global
|
||||
|
||||
def __set__(self, obj, value):
|
||||
global _configured_global
|
||||
_configured_context.set(bool(value))
|
||||
_configured_global = bool(value)
|
||||
|
||||
def __bool__(self):
|
||||
return _configured_context.get() or _configured_global
|
||||
|
||||
# Expose the shim so tests can still use host_app._CONFIGURED = False
|
||||
_CONFIGURED = _ConfiguredShim()
|
||||
|
||||
|
||||
def is_legacy_host() -> bool:
|
||||
@@ -38,7 +93,11 @@ def configure_host(
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Apply PraisonAIUI host settings and wire L1 backends (unless legacy mode)."""
|
||||
global _CONFIGURED
|
||||
global _configured_global
|
||||
|
||||
# Check if already configured in this context to avoid duplicate configuration
|
||||
if _configured_context.get() or _configured_global:
|
||||
return
|
||||
|
||||
import praisonaiui as aiui
|
||||
from praisonai.ui._aiui_datastore import PraisonAISessionDataStore
|
||||
@@ -78,7 +137,9 @@ def configure_host(
|
||||
if not is_legacy_host():
|
||||
from praisonaiui.providers import PraisonAIProvider
|
||||
from praisonaiui.server import set_provider
|
||||
|
||||
|
||||
# Always proceed with provider setup unless explicitly skipped by legacy mode
|
||||
# Tests that need to override providers should do so after configure_host() completes
|
||||
kwargs = dict(agent_kwargs or {})
|
||||
if agents:
|
||||
set_provider(PraisonAIProvider(agents=list(agents), **kwargs))
|
||||
@@ -112,7 +173,8 @@ def configure_host(
|
||||
except ImportError:
|
||||
pass # L3 pages are optional
|
||||
|
||||
_CONFIGURED = True
|
||||
_configured_context.set(True)
|
||||
_configured_global = True
|
||||
|
||||
|
||||
def setup_bridges() -> None:
|
||||
@@ -170,8 +232,6 @@ def setup_bridges() -> None:
|
||||
from praisonai.integration.bridges.kanban_bridge import register_kanban_backends
|
||||
register_kanban_backends()
|
||||
|
||||
except Exception as exc:
|
||||
log.debug("aiui backend injection failed: %s", exc)
|
||||
except Exception as exc:
|
||||
log.warning("aiui backend injection failed: %s", exc)
|
||||
|
||||
@@ -180,7 +240,7 @@ def create_host_app():
|
||||
"""Return the Starlette app from PraisonAIUI (call after ``configure_host``)."""
|
||||
from praisonaiui.server import create_app
|
||||
|
||||
if not _CONFIGURED:
|
||||
if not (_configured_context.get() or _configured_global):
|
||||
configure_host()
|
||||
return create_app()
|
||||
|
||||
|
||||
@@ -102,13 +102,25 @@ class BaseCLIIntegration(ABC):
|
||||
"""
|
||||
cmd = self.cli_command
|
||||
cache = BaseCLIIntegration._availability_cache
|
||||
if cmd in cache:
|
||||
return cache[cmd]
|
||||
with BaseCLIIntegration._availability_cache_lock:
|
||||
if cmd not in cache:
|
||||
cache[cmd] = shutil.which(cmd) is not None
|
||||
return cache[cmd]
|
||||
|
||||
@classmethod
|
||||
def invalidate_availability(cls, command: Optional[str] = None) -> None:
|
||||
"""
|
||||
Invalidate availability cache for CLI tools.
|
||||
|
||||
Args:
|
||||
command: Specific CLI command to invalidate, or None to clear all
|
||||
"""
|
||||
with cls._availability_cache_lock:
|
||||
if command is None:
|
||||
cls._availability_cache.clear()
|
||||
else:
|
||||
cls._availability_cache.pop(command, None)
|
||||
|
||||
@abstractmethod
|
||||
async def execute(self, prompt: str, **options) -> str:
|
||||
"""
|
||||
|
||||
@@ -19,8 +19,10 @@ Usage:
|
||||
# Register custom integration
|
||||
registry.register('my-agent', MyCustomIntegration)
|
||||
|
||||
# Create integration
|
||||
agent = registry.create('claude', workspace="/path/to/project")
|
||||
# Create integration (use try_create for optional behavior)
|
||||
agent = registry.try_create('claude', workspace="/path/to/project")
|
||||
if agent is None:
|
||||
print("Integration 'claude' not found or unavailable")
|
||||
|
||||
# List available integrations
|
||||
available = await registry.get_available()
|
||||
@@ -104,9 +106,9 @@ class ExternalAgentRegistry(PluginRegistry[BaseCLIIntegration]):
|
||||
"""
|
||||
return self.list_names()
|
||||
|
||||
def create(self, name: str, **kwargs: Any) -> Optional[BaseCLIIntegration]:
|
||||
def try_create(self, name: str, **kwargs: Any) -> Optional[BaseCLIIntegration]:
|
||||
"""
|
||||
Create an instance of the specified integration.
|
||||
Try to create an instance of the specified integration.
|
||||
|
||||
Args:
|
||||
name: Name of the integration
|
||||
@@ -217,7 +219,7 @@ def create_integration(name: str, **kwargs: Any) -> Optional[BaseCLIIntegration]
|
||||
BaseCLIIntegration: Instance of the integration, or None if not found
|
||||
"""
|
||||
registry = get_default_registry()
|
||||
return registry.create(name, **kwargs)
|
||||
return registry.try_create(name, **kwargs)
|
||||
|
||||
|
||||
def get_available_integrations() -> Dict[str, bool]:
|
||||
|
||||
@@ -61,7 +61,7 @@ class JobSubmitRequest(BaseModel):
|
||||
if ip_obj.is_private or ip_obj.is_loopback or ip_obj.is_link_local or ip_obj.is_multicast:
|
||||
raise ValueError("Webhook URL resolves to a private or restricted network address")
|
||||
except socket.gaierror:
|
||||
pass
|
||||
raise ValueError("Webhook URL hostname could not be resolved")
|
||||
|
||||
return v
|
||||
|
||||
@@ -162,7 +162,7 @@ class Job(BaseModel):
|
||||
if ip_obj.is_private or ip_obj.is_loopback or ip_obj.is_link_local or ip_obj.is_multicast:
|
||||
raise ValueError("Webhook URL resolves to a private or restricted network address")
|
||||
except socket.gaierror:
|
||||
pass
|
||||
raise ValueError("Webhook URL hostname could not be resolved")
|
||||
|
||||
return v
|
||||
|
||||
|
||||
@@ -123,6 +123,25 @@ def create_app(
|
||||
allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
|
||||
allow_headers=["Authorization", "Content-Type", "Origin", "Accept", "Idempotency-Key"],
|
||||
)
|
||||
|
||||
jobs_api_key = os.environ.get("PRAISONAI_JOBS_API_KEY")
|
||||
if jobs_api_key:
|
||||
import hmac
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
class JobsAPIKeyMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request, call_next):
|
||||
if request.url.path == "/health":
|
||||
return await call_next(request)
|
||||
auth = request.headers.get("Authorization", "")
|
||||
header_key = request.headers.get("X-API-Key", "")
|
||||
token = auth[7:] if auth.startswith("Bearer ") else header_key
|
||||
if not token or not hmac.compare_digest(token, jobs_api_key):
|
||||
return JSONResponse({"error": "Unauthorized"}, status_code=401)
|
||||
return await call_next(request)
|
||||
|
||||
app.add_middleware(JobsAPIKeyMiddleware)
|
||||
|
||||
# Add jobs router
|
||||
jobs_router = create_router(get_store(), get_executor())
|
||||
|
||||
@@ -81,6 +81,7 @@ class MCPServer:
|
||||
# Cancellation support
|
||||
self._active_requests: Dict[str, asyncio.Task] = {}
|
||||
self._cancelled_requests: Set[str] = set()
|
||||
self._max_cancelled_requests = 10000
|
||||
|
||||
# Progress notification callback
|
||||
self._progress_callback: Optional[Callable] = None
|
||||
@@ -170,6 +171,10 @@ class MCPServer:
|
||||
request_id = params.get("requestId")
|
||||
if request_id:
|
||||
self._cancelled_requests.add(str(request_id))
|
||||
if len(self._cancelled_requests) > self._max_cancelled_requests:
|
||||
overflow = len(self._cancelled_requests) - self._max_cancelled_requests
|
||||
for _ in range(overflow):
|
||||
self._cancelled_requests.pop()
|
||||
# Cancel active task if exists
|
||||
if str(request_id) in self._active_requests:
|
||||
task = self._active_requests[str(request_id)]
|
||||
|
||||
@@ -14,6 +14,7 @@ Implements the MCP Streamable HTTP transport (MCP 2025-11-25 spec):
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
@@ -90,6 +91,7 @@ class HTTPStreamTransport:
|
||||
self.cors_origins = [origin for origin in cors_origins if origin != "*"]
|
||||
self.api_key = api_key
|
||||
self.session_ttl = session_ttl
|
||||
self.max_sessions = int(os.getenv("PRAISONAI_MCP_MAX_SESSIONS", "1000"))
|
||||
self.allow_client_termination = allow_client_termination
|
||||
self.response_mode = response_mode
|
||||
self.resumability_enabled = resumability_enabled
|
||||
@@ -142,13 +144,17 @@ class HTTPStreamTransport:
|
||||
if self.allowed_origins is None:
|
||||
# No allowed origins configured - reject all Origin headers
|
||||
return False
|
||||
|
||||
# Check if origin matches any allowed origin
|
||||
for allowed in self.allowed_origins:
|
||||
if request_origin == allowed or request_origin.startswith(allowed):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
from urllib.parse import urlparse
|
||||
from praisonaiagents.mcp.mcp_security import is_valid_origin
|
||||
|
||||
allowed_hosts = []
|
||||
for origin in self.allowed_origins:
|
||||
host = urlparse(origin).hostname
|
||||
if host:
|
||||
allowed_hosts.append(host.lower())
|
||||
|
||||
return is_valid_origin(request_origin, allowed_hosts, allow_missing=False)
|
||||
|
||||
def _validate_protocol_version(self, version: Optional[str]) -> bool:
|
||||
"""Validate MCP-Protocol-Version header."""
|
||||
@@ -171,6 +177,7 @@ class HTTPStreamTransport:
|
||||
|
||||
async def mcp_post(request: Request) -> Response:
|
||||
"""Handle POST requests (client→server messages)."""
|
||||
self._cleanup_sessions()
|
||||
# Validate Origin header (MCP 2025-11-25 security requirement)
|
||||
origin = request.headers.get("Origin")
|
||||
if not self._validate_origin(origin):
|
||||
@@ -220,6 +227,11 @@ class HTTPStreamTransport:
|
||||
|
||||
# Check if this is an initialize request
|
||||
if body.get("method") == "initialize":
|
||||
if len(self._sessions) >= self.max_sessions:
|
||||
return JSONResponse(
|
||||
{"jsonrpc": "2.0", "id": body.get("id"), "error": {"code": -32000, "message": "Too many active sessions"}},
|
||||
status_code=503,
|
||||
)
|
||||
# Create new session
|
||||
new_session_id = str(uuid.uuid4())
|
||||
self._sessions[new_session_id] = {
|
||||
|
||||
@@ -330,8 +330,10 @@ def create_auth_middleware(auth_type: str, api_key: Optional[str] = None, jwt_se
|
||||
expected_key = api_key or os.environ.get("PRAISONAI_API_KEY")
|
||||
|
||||
if not expected_key:
|
||||
# No key configured, allow request
|
||||
return await call_next(request)
|
||||
return JSONResponse(
|
||||
{"error": {"code": "misconfigured", "message": "API key authentication is enabled but no key is configured"}},
|
||||
status_code=503,
|
||||
)
|
||||
|
||||
if not provided_key or not hmac.compare_digest(provided_key, expected_key):
|
||||
return JSONResponse(
|
||||
@@ -352,7 +354,10 @@ def create_auth_middleware(auth_type: str, api_key: Optional[str] = None, jwt_se
|
||||
# Get JWT secret
|
||||
secret = jwt_secret or os.environ.get("PRAISONAI_JWT_SECRET")
|
||||
if not secret:
|
||||
return await call_next(request)
|
||||
return JSONResponse(
|
||||
{"error": {"code": "misconfigured", "message": "JWT authentication is enabled but no secret is configured"}},
|
||||
status_code=503,
|
||||
)
|
||||
|
||||
# Check Authorization header
|
||||
auth_header = request.headers.get("Authorization", "")
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "4.6.55"
|
||||
__version__ = "4.6.58"
|
||||
@@ -12,7 +12,7 @@ dependencies = [
|
||||
"rich>=13.7",
|
||||
"markdown>=3.5",
|
||||
"pyparsing>=3.0.0",
|
||||
"praisonaiagents>=1.6.55",
|
||||
"praisonaiagents>=1.6.58",
|
||||
"python-dotenv>=0.19.0",
|
||||
"litellm>=1.83.14,<2",
|
||||
"PyYAML>=6.0",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -139,13 +139,17 @@ def bump_version(new_version: str, agents_version: Optional[str] = None):
|
||||
root
|
||||
)
|
||||
|
||||
# 2. Update deploy.py (Dockerfile template)
|
||||
print("\n🐳 Deploy Script:")
|
||||
update_file(
|
||||
praisonai_dir / "praisonai/deploy.py",
|
||||
[(r'praisonai==[0-9.]+', f'praisonai=={new_version}')],
|
||||
root
|
||||
)
|
||||
# 2. Update deploy/docker.py (Docker deployment scripts)
|
||||
print("\n🐳 Deploy Scripts:")
|
||||
docker_deploy_file = praisonai_dir / "praisonai/deploy/docker.py"
|
||||
if docker_deploy_file.exists():
|
||||
update_file(
|
||||
docker_deploy_file,
|
||||
[(r'praisonai==[0-9.]+', f'praisonai=={new_version}')],
|
||||
root
|
||||
)
|
||||
else:
|
||||
print(f" ⏭️ Docker deploy script not found: {docker_deploy_file.relative_to(root)}")
|
||||
|
||||
# 3. Update Dockerfiles
|
||||
print("\n🐳 Dockerfiles:")
|
||||
@@ -240,7 +244,7 @@ def release(version: str, use_frozen_lock: bool = False, no_add_all: bool = Fals
|
||||
|
||||
release_files = [
|
||||
"src/praisonai/praisonai/version.py",
|
||||
"src/praisonai/praisonai/deploy.py",
|
||||
"src/praisonai/praisonai/deploy/docker.py",
|
||||
"docker/Dockerfile",
|
||||
"docker/Dockerfile.chat",
|
||||
"docker/Dockerfile.dev",
|
||||
|
||||
@@ -4,7 +4,7 @@ Version bump script for PraisonAI package.
|
||||
|
||||
This script updates the version number in all required locations:
|
||||
- praisonai/version.py (single source of truth for Python package)
|
||||
- praisonai/deploy.py (Dockerfile template)
|
||||
- praisonai/deploy/docker.py (Docker deployment scripts)
|
||||
- ../../docker/Dockerfile, Dockerfile.chat, Dockerfile.dev, Dockerfile.ui
|
||||
- praisonai.rb (Homebrew formula)
|
||||
|
||||
@@ -63,12 +63,16 @@ def bump_version(new_version: str, agents_version: str | None = None):
|
||||
[(r'__version__ = "[^"]+"', f'__version__ = "{new_version}"')]
|
||||
)
|
||||
|
||||
# 2. Update deploy.py (Dockerfile template)
|
||||
print("\n🐳 Deploy Script:")
|
||||
update_file(
|
||||
praisonai_dir / "praisonai/deploy.py",
|
||||
[(r'praisonai==[0-9.]+', f'praisonai=={new_version}')]
|
||||
)
|
||||
# 2. Update deploy/docker.py (Docker deployment scripts)
|
||||
print("\n🐳 Deploy Scripts:")
|
||||
docker_deploy_file = praisonai_dir / "praisonai/deploy/docker.py"
|
||||
if docker_deploy_file.exists():
|
||||
update_file(
|
||||
docker_deploy_file,
|
||||
[(r'praisonai==[0-9.]+', f'praisonai=={new_version}')]
|
||||
)
|
||||
else:
|
||||
print(f" ⏭️ Docker deploy script not found: {docker_deploy_file.relative_to(get_project_root())}")
|
||||
|
||||
# 3. Update Dockerfiles
|
||||
print("\n🐳 Dockerfiles:")
|
||||
|
||||
@@ -13,10 +13,10 @@ def test_gateway_start_wires_datastore(monkeypatch):
|
||||
import praisonaiui.server as srv
|
||||
from praisonai.integration import host_app
|
||||
|
||||
host_app._CONFIGURED = False
|
||||
host_app.reset_configuration()
|
||||
backends.clear_backends()
|
||||
srv._provider = None
|
||||
|
||||
host_app.configure_host(pages=["chat"])
|
||||
assert srv.get_datastore() is not None
|
||||
assert "hooks" in backends.list_backends() or host_app._CONFIGURED
|
||||
assert "hooks" in backends.list_backends() or host_app.is_configured()
|
||||
Loaded 100 of 109 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user