mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-12 01:41:26 -10:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 praisonai-triage-agent[bot]andMervin Praison
|
0891b29e9e | ||
|
老憨
|
4a4b7a3208 |
No files matched your search
@@ -1,29 +0,0 @@
|
||||
# Contributing to PraisonAI
|
||||
|
||||
Thank you for your interest! Here's how to get started:
|
||||
|
||||
## Getting Started
|
||||
|
||||
1. Fork the repository
|
||||
2. Clone your fork and install dependencies:
|
||||
```bash
|
||||
git clone https://github.com/your-username/PraisonAI.git
|
||||
cd PraisonAI
|
||||
pip install -e src/praisonai-agents/
|
||||
pip install -e src/praisonai/
|
||||
```
|
||||
3. Create a feature branch: `git checkout -b feature/your-feature`
|
||||
4. Make your changes
|
||||
5. Submit a pull request
|
||||
|
||||
## Guidelines
|
||||
|
||||
- Follow existing code style (PEP 8)
|
||||
- Add tests for new features and run them locally using pytest:
|
||||
```bash
|
||||
cd src/praisonai-agents && pytest tests/
|
||||
```
|
||||
- Update documentation as needed
|
||||
- Check existing issues or open a new one to discuss your idea before implementing large changes
|
||||
|
||||
For questions or support, please visit our [GitHub Issues](https://github.com/MervinPraison/PraisonAI/issues).
|
||||
@@ -786,7 +786,6 @@ Methods:
|
||||
* <code title="cli">praisonai bot discord <a href="./src/praisonai/praisonai/cli/commands/bot.py">--help</a></code>
|
||||
* <code title="cli">praisonai bot email <a href="./src/praisonai/praisonai/cli/commands/bot.py">--help</a></code>
|
||||
* <code title="cli">praisonai bot install-daemon <a href="./src/praisonai/praisonai/cli/commands/bot.py">--help</a></code>
|
||||
* <code title="cli">praisonai bot linear <a href="./src/praisonai/praisonai/cli/commands/bot.py">--help</a></code>
|
||||
* <code title="cli">praisonai bot slack <a href="./src/praisonai/praisonai/cli/commands/bot.py">--help</a></code>
|
||||
* <code title="cli">praisonai bot start <a href="./src/praisonai/praisonai/cli/commands/bot.py">--help</a></code>
|
||||
* <code title="cli">praisonai bot telegram <a href="./src/praisonai/praisonai/cli/commands/bot.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.37" \
|
||||
"praisonai>=4.6.35" \
|
||||
"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.37" \
|
||||
"praisonai>=4.6.35" \
|
||||
"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.37" \
|
||||
"praisonai>=4.6.35" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -1,68 +0,0 @@
|
||||
"""
|
||||
Linear Agent Example for PraisonAI.
|
||||
|
||||
Demonstrates how to create a Linear bot that responds to issue mentions
|
||||
and assignments with full coding capabilities.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from praisonaiagents import Agent
|
||||
from praisonai.bots import LinearBot
|
||||
|
||||
# Create an agent with Linear tools and coding capabilities
|
||||
agent = Agent(
|
||||
name="PraisonAI Coder",
|
||||
instructions="""
|
||||
You are an autonomous coding agent integrated with Linear.
|
||||
|
||||
When you are mentioned or assigned an issue:
|
||||
1. Analyze the issue description and requirements
|
||||
2. Break down complex tasks into smaller steps
|
||||
3. Use available tools to implement solutions
|
||||
4. Update the issue with progress comments
|
||||
5. Create GitHub pull requests when code changes are needed
|
||||
|
||||
Focus on being helpful, thorough, and providing clear updates.
|
||||
""",
|
||||
llm="gpt-4o-mini",
|
||||
tools=[
|
||||
"linear_search_issues",
|
||||
"linear_get_issue",
|
||||
"linear_add_comment",
|
||||
"linear_update_issue",
|
||||
"linear_list_teams",
|
||||
"linear_list_issue_states",
|
||||
"github_create_branch",
|
||||
"github_commit_and_push",
|
||||
"github_create_pull_request",
|
||||
"read_file",
|
||||
"write_file",
|
||||
"execute_command"
|
||||
],
|
||||
memory=True,
|
||||
web_search=True,
|
||||
auto_approve_tools=True,
|
||||
)
|
||||
|
||||
# Create Linear bot
|
||||
bot = LinearBot(
|
||||
token="your-linear-oauth-token", # or set LINEAR_OAUTH_TOKEN env var
|
||||
signing_secret="your-webhook-secret", # or set LINEAR_WEBHOOK_SECRET env var
|
||||
agent=agent,
|
||||
webhook_port=8080,
|
||||
)
|
||||
|
||||
async def main():
|
||||
"""Start the Linear bot."""
|
||||
print("Starting Linear bot...")
|
||||
print("Set up your Linear webhook to point to: http://your-host:8080/webhook")
|
||||
print("Configure AgentSession and Comment events in Linear webhook settings")
|
||||
|
||||
try:
|
||||
await bot.start()
|
||||
except KeyboardInterrupt:
|
||||
print("\nStopping bot...")
|
||||
await bot.stop()
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -1,35 +0,0 @@
|
||||
platform: linear
|
||||
token: ${LINEAR_OAUTH_TOKEN}
|
||||
signing_secret: ${LINEAR_WEBHOOK_SECRET}
|
||||
webhook_port: 8080
|
||||
|
||||
agent:
|
||||
name: "PraisonAI Coder"
|
||||
instructions: |
|
||||
You are an autonomous coding agent integrated with Linear.
|
||||
|
||||
When you are mentioned or assigned an issue:
|
||||
1. Analyze the issue description and requirements
|
||||
2. Break down complex tasks into smaller steps
|
||||
3. Use available tools to implement solutions
|
||||
4. Update the issue with progress comments
|
||||
5. Create GitHub pull requests when code changes are needed
|
||||
|
||||
Focus on being helpful, thorough, and providing clear updates.
|
||||
llm: gpt-4o-mini
|
||||
tools:
|
||||
- linear_search_issues
|
||||
- linear_get_issue
|
||||
- linear_add_comment
|
||||
- linear_update_issue
|
||||
- linear_list_teams
|
||||
- linear_list_issue_states
|
||||
- github_create_branch
|
||||
- github_commit_and_push
|
||||
- github_create_pull_request
|
||||
- read_file
|
||||
- write_file
|
||||
- execute_command
|
||||
memory: true
|
||||
web_search: true
|
||||
auto_approve_tools: true
|
||||
@@ -1,99 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""W1 real agentic smoke test (high-level BotOS API, no manual mgr).
|
||||
|
||||
Uses the user-facing API only — ``Bot``, ``BotOS``, ``InMemoryIdentityResolver`` —
|
||||
and exercises the same cross-platform-continuity scenario as
|
||||
``smoke_w1_real.py`` but without touching ``BotSessionManager`` directly.
|
||||
|
||||
The mock-platform pattern: we instantiate two ``Bot``s but reach into
|
||||
their adapter ``_session`` to drive the chat manually — this proves the
|
||||
wire-up from ``BotOS(identity_resolver=...)`` → ``Bot._identity_resolver``
|
||||
→ adapter splice works end-to-end.
|
||||
|
||||
Run::
|
||||
|
||||
# From the repository root:
|
||||
PYTHONPATH=src/praisonai-agents:src/praisonai python scripts/smoke_w1_botos_real.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
|
||||
async def main() -> int:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.session import (
|
||||
FileIdentityResolver,
|
||||
DefaultSessionStore,
|
||||
)
|
||||
from praisonai.bots import Bot, BotOS
|
||||
|
||||
if os.getenv("ANTHROPIC_API_KEY"):
|
||||
model = "claude-3-5-haiku-latest"
|
||||
elif os.getenv("GOOGLE_API_KEY"):
|
||||
model = "gemini-1.5-flash"
|
||||
else:
|
||||
model = "gpt-4o-mini"
|
||||
print(f"Using model: {model}")
|
||||
|
||||
agent = Agent(
|
||||
name="W1BotOSTester",
|
||||
instructions=(
|
||||
"You are a helpful assistant. Always remember facts the user "
|
||||
"tells you and recall them precisely when asked."
|
||||
),
|
||||
llm=model,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
# Persistent FileIdentityResolver — survives restart.
|
||||
resolver = FileIdentityResolver(path=f"{tmp}/identity.json")
|
||||
resolver.link("telegram", "tg-aaa", "alice-prod")
|
||||
resolver.link("discord", "dc-bbb", "alice-prod")
|
||||
store = DefaultSessionStore(session_dir=f"{tmp}/sessions")
|
||||
|
||||
# Create two Bots; BotOS wires the resolver to each.
|
||||
bot_t = Bot("telegram", agent=agent, token="dummy")
|
||||
bot_d = Bot("discord", agent=agent, token="dummy")
|
||||
|
||||
os_ = BotOS(bots=[bot_t, bot_d], identity_resolver=resolver)
|
||||
assert bot_t._identity_resolver is resolver
|
||||
assert bot_d._identity_resolver is resolver
|
||||
print(f"BotOS wired resolver to {os_.list_bots()}")
|
||||
|
||||
# We don't need the network; build adapters lazily and use their
|
||||
# _session managers directly. Splicing logic is the same code
|
||||
# path real adapters take during start().
|
||||
adapter_t = bot_t._build_adapter()
|
||||
adapter_d = bot_d._build_adapter()
|
||||
# Inject the shared store into both managers so unification works.
|
||||
adapter_t._session._store = store
|
||||
adapter_d._session._store = store
|
||||
|
||||
out_t = await adapter_t._session.chat(
|
||||
agent, "tg-aaa",
|
||||
"Remember this: my favourite music genre is darkwave.",
|
||||
chat_id="100", user_name="Alice",
|
||||
)
|
||||
print(f"\n[Telegram] {out_t}\n")
|
||||
|
||||
out_d = await adapter_d._session.chat(
|
||||
agent, "dc-bbb",
|
||||
"Hey, what was my favourite music genre I just told you?",
|
||||
chat_id="200", user_name="Alice",
|
||||
)
|
||||
print(f"[Discord] {out_d}\n")
|
||||
|
||||
if out_d and "darkwave" in out_d.lower():
|
||||
print("PASS: BotOS-wired cross-platform continuity verified.")
|
||||
return 0
|
||||
print("FAIL: agent did not recall via BotOS-wired path.")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(asyncio.run(main()))
|
||||
@@ -1,76 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""W1 real agentic smoke test (bypasses pytest conftest gating).
|
||||
|
||||
Runs::
|
||||
|
||||
cd src/praisonai
|
||||
python ../../scripts/smoke_w1_real.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
|
||||
async def main() -> int:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.session.identity import InMemoryIdentityResolver
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
if os.getenv("ANTHROPIC_API_KEY"):
|
||||
model = "claude-3-5-haiku-latest"
|
||||
elif os.getenv("GOOGLE_API_KEY"):
|
||||
model = "gemini-1.5-flash"
|
||||
else:
|
||||
model = "gpt-4o-mini"
|
||||
print(f"Using model: {model}")
|
||||
|
||||
agent = Agent(
|
||||
name="W1Tester",
|
||||
instructions=(
|
||||
"You are a helpful assistant. Always remember facts the user "
|
||||
"tells you and recall them precisely when asked."
|
||||
),
|
||||
llm=model,
|
||||
)
|
||||
|
||||
resolver = InMemoryIdentityResolver()
|
||||
resolver.link("telegram", "tg-12345", "alice-global")
|
||||
resolver.link("discord", "dc-67890", "alice-global")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_path:
|
||||
store = DefaultSessionStore(session_dir=tmp_path)
|
||||
mgr_t = BotSessionManager(
|
||||
platform="telegram", identity_resolver=resolver, store=store
|
||||
)
|
||||
mgr_d = BotSessionManager(
|
||||
platform="discord", identity_resolver=resolver, store=store
|
||||
)
|
||||
|
||||
out_t = await mgr_t.chat(
|
||||
agent, "tg-12345",
|
||||
"Remember this: my favourite colour is octarine.",
|
||||
)
|
||||
print("\n[Telegram] in: Remember this: my favourite colour is octarine.")
|
||||
print(f"[Telegram] out: {out_t}\n")
|
||||
|
||||
out_d = await mgr_d.chat(
|
||||
agent, "dc-67890",
|
||||
"What did I just tell you my favourite colour was?",
|
||||
)
|
||||
print("[Discord] in: What did I just tell you my favourite colour was?")
|
||||
print(f"[Discord] out: {out_d}\n")
|
||||
|
||||
if out_d and "octarine" in out_d.lower():
|
||||
print("PASS: Cross-platform context recalled.")
|
||||
return 0
|
||||
print("FAIL: Agent did not recall cross-platform fact.")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(asyncio.run(main()))
|
||||
@@ -1,304 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""W1 deep-validation smoke test.
|
||||
|
||||
Exercises every robustness dimension of W1 with REAL LLM calls:
|
||||
|
||||
T1. 3-platform unification (telegram → discord → slack, one user)
|
||||
T2. Persistence across restart (FileIdentityResolver)
|
||||
T3. Concurrent users do NOT leak history
|
||||
T4. SessionContext is visible to a tool the agent calls
|
||||
T5. Self-improving probe — ask the agent to "learn a new skill"
|
||||
(this surfaces the W1↔W2 boundary honestly)
|
||||
|
||||
Run::
|
||||
|
||||
cd /Users/praison/worktrees/hermes-parity
|
||||
PYTHONPATH=src/praisonai-agents:src/praisonai python scripts/smoke_w1_robust.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _pick_model() -> str:
|
||||
if os.getenv("ANTHROPIC_API_KEY"):
|
||||
return "anthropic/claude-haiku-4-5"
|
||||
if os.getenv("GOOGLE_API_KEY"):
|
||||
return "gemini/gemini-2.5-flash"
|
||||
return "gpt-4o-mini"
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# T1 — 3-platform unification
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
async def test_three_platform_unification(tmp_path: Path, model: str) -> bool:
|
||||
print("\n" + "=" * 72)
|
||||
print("T1: 3-platform unification (telegram → discord → slack, same user)")
|
||||
print("=" * 72)
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.session import (
|
||||
FileIdentityResolver,
|
||||
DefaultSessionStore,
|
||||
)
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
resolver = FileIdentityResolver(path=tmp_path / "identity.json")
|
||||
resolver.link("telegram", "tg-1", "alice")
|
||||
resolver.link("discord", "dc-1", "alice")
|
||||
resolver.link("slack", "sl-1", "alice")
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path / "sessions"))
|
||||
agent = Agent(
|
||||
name="T1",
|
||||
instructions="Remember facts the user tells you and recall them precisely.",
|
||||
llm=model,
|
||||
)
|
||||
|
||||
mgr = lambda p: BotSessionManager(platform=p, identity_resolver=resolver, store=store)
|
||||
mgr_t, mgr_d, mgr_s = mgr("telegram"), mgr("discord"), mgr("slack")
|
||||
|
||||
out_t = await mgr_t.chat(agent, "tg-1", "My name is Vimes and my badge number is 177.")
|
||||
print(f"[Telegram] {out_t}")
|
||||
|
||||
out_d = await mgr_d.chat(agent, "dc-1", "What did I tell you my badge number was?")
|
||||
print(f"[Discord] {out_d}")
|
||||
|
||||
out_s = await mgr_s.chat(agent, "sl-1", "And what did I say my name was?")
|
||||
print(f"[Slack] {out_s}")
|
||||
|
||||
ok = (
|
||||
out_d and "177" in out_d
|
||||
and out_s and "vimes" in out_s.lower()
|
||||
)
|
||||
print(" RESULT:", "PASS" if ok else "FAIL")
|
||||
return bool(ok)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# T2 — Persistence across restart
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
async def test_persistence_across_restart(tmp_path: Path, model: str) -> bool:
|
||||
print("\n" + "=" * 72)
|
||||
print("T2: Persistence — links + history survive 'restart' (new resolver+store)")
|
||||
print("=" * 72)
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.session import (
|
||||
FileIdentityResolver,
|
||||
DefaultSessionStore,
|
||||
)
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
id_path = tmp_path / "identity.json"
|
||||
sess_dir = tmp_path / "sessions"
|
||||
|
||||
# ── First "process" ─────────────────────────────────────────────
|
||||
r1 = FileIdentityResolver(path=id_path)
|
||||
r1.link("telegram", "tg-2", "bob")
|
||||
s1 = DefaultSessionStore(session_dir=str(sess_dir))
|
||||
agent1 = Agent(name="T2", instructions="You remember facts.", llm=model)
|
||||
mgr1 = BotSessionManager(platform="telegram", identity_resolver=r1, store=s1)
|
||||
out1 = await mgr1.chat(
|
||||
agent1, "tg-2",
|
||||
"Remember: project name is 'pyramid-of-tsort'.",
|
||||
)
|
||||
print(f"[Run1 Telegram] {out1}")
|
||||
|
||||
# ── Second "process" — fresh objects, same files ────────────────
|
||||
r2 = FileIdentityResolver(path=id_path)
|
||||
s2 = DefaultSessionStore(session_dir=str(sess_dir))
|
||||
agent2 = Agent(name="T2", instructions="You remember facts.", llm=model)
|
||||
mgr2 = BotSessionManager(platform="telegram", identity_resolver=r2, store=s2)
|
||||
|
||||
# Link must have survived
|
||||
assert r2.resolve("telegram", "tg-2") == "bob", "FileIdentityResolver lost link!"
|
||||
|
||||
out2 = await mgr2.chat(agent2, "tg-2", "What was the project name I told you?")
|
||||
print(f"[Run2 Telegram] {out2}")
|
||||
|
||||
ok = bool(out2 and "pyramid" in out2.lower())
|
||||
print(" RESULT:", "PASS" if ok else "FAIL")
|
||||
return ok
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# T3 — Concurrent users do NOT leak history
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
async def test_concurrent_users_isolated(tmp_path: Path, model: str) -> bool:
|
||||
print("\n" + "=" * 72)
|
||||
print("T3: Concurrent users — no history leakage")
|
||||
print("=" * 72)
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.session import (
|
||||
FileIdentityResolver,
|
||||
DefaultSessionStore,
|
||||
)
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
resolver = FileIdentityResolver(path=tmp_path / "identity.json")
|
||||
resolver.link("telegram", "u-alice", "alice")
|
||||
resolver.link("telegram", "u-bob", "bob")
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path / "sessions"))
|
||||
agent = Agent(name="T3", instructions="You remember facts per user.", llm=model)
|
||||
mgr = BotSessionManager(
|
||||
platform="telegram", identity_resolver=resolver, store=store
|
||||
)
|
||||
|
||||
await asyncio.gather(
|
||||
mgr.chat(agent, "u-alice", "I love mangoes."),
|
||||
mgr.chat(agent, "u-bob", "I love durian."),
|
||||
)
|
||||
out_alice = await mgr.chat(agent, "u-alice", "What fruit did I say I love?")
|
||||
out_bob = await mgr.chat(agent, "u-bob", "What fruit did I say I love?")
|
||||
print(f"[Alice] {out_alice}")
|
||||
print(f"[Bob] {out_bob}")
|
||||
|
||||
a_ok = out_alice and "mango" in out_alice.lower() and "durian" not in out_alice.lower()
|
||||
b_ok = out_bob and "durian" in out_bob.lower() and "mango" not in out_bob.lower()
|
||||
ok = bool(a_ok and b_ok)
|
||||
print(" RESULT:", "PASS" if ok else "FAIL")
|
||||
return ok
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# T4 — SessionContext visible to a tool
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
async def test_session_context_visible_to_tool(tmp_path: Path, model: str) -> bool:
|
||||
print("\n" + "=" * 72)
|
||||
print("T4: SessionContext is visible to a tool the agent invokes")
|
||||
print("=" * 72)
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.session import (
|
||||
FileIdentityResolver,
|
||||
DefaultSessionStore,
|
||||
get_session_context,
|
||||
)
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
captured = {}
|
||||
|
||||
def whoami() -> str:
|
||||
"""Return the platform and user id of the caller."""
|
||||
ctx = get_session_context()
|
||||
captured["platform"] = ctx.platform
|
||||
captured["unified_user_id"] = ctx.unified_user_id
|
||||
return f"platform={ctx.platform} user={ctx.unified_user_id}"
|
||||
|
||||
resolver = FileIdentityResolver(path=tmp_path / "identity.json")
|
||||
resolver.link("telegram", "tg-99", "alice")
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path / "sessions"))
|
||||
|
||||
agent = Agent(
|
||||
name="T4",
|
||||
instructions=(
|
||||
"When the user asks 'who am I' or about platform info, "
|
||||
"ALWAYS call the whoami tool and report its result verbatim."
|
||||
),
|
||||
tools=[whoami],
|
||||
llm=model,
|
||||
)
|
||||
mgr = BotSessionManager(
|
||||
platform="telegram", identity_resolver=resolver, store=store
|
||||
)
|
||||
|
||||
out = await mgr.chat(
|
||||
agent, "tg-99",
|
||||
"Please call the whoami tool and tell me what you see.",
|
||||
chat_id="100", user_name="Alice",
|
||||
)
|
||||
print(f"[Agent] {out}")
|
||||
print(f"[Tool captured] platform={captured.get('platform')!r} "
|
||||
f"unified_user_id={captured.get('unified_user_id')!r}")
|
||||
|
||||
ok = (
|
||||
captured.get("platform") == "telegram"
|
||||
and captured.get("unified_user_id") == "alice"
|
||||
)
|
||||
print(" RESULT:", "PASS" if ok else "FAIL")
|
||||
return ok
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# T5 — Self-improving probe (W1 ↔ W2 boundary)
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
async def test_self_improving_probe(tmp_path: Path, model: str) -> bool:
|
||||
print("\n" + "=" * 72)
|
||||
print("T5: Self-improving probe — ask the agent to 'learn a new skill'")
|
||||
print(" (Expectation: W1 does NOT enable this — it is W2's job.)")
|
||||
print("=" * 72)
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.session import FileIdentityResolver, DefaultSessionStore
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
resolver = FileIdentityResolver(path=tmp_path / "identity.json")
|
||||
resolver.link("telegram", "tg-77", "alice")
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path / "sessions"))
|
||||
agent = Agent(name="T5", instructions="You are a helpful assistant.", llm=model)
|
||||
mgr = BotSessionManager(platform="telegram", identity_resolver=resolver, store=store)
|
||||
|
||||
out = await mgr.chat(
|
||||
agent, "tg-77",
|
||||
"I want you to learn a new skill called 'fizzbuzz_writer' that "
|
||||
"writes Python FizzBuzz. Add it to your permanent skills so a "
|
||||
"future session can use it. Then describe what you did.",
|
||||
)
|
||||
print(f"[Agent] {out}\n")
|
||||
|
||||
# Probe whether anything actually got persisted to disk as a skill.
|
||||
candidate_paths = [
|
||||
Path.home() / ".praisonai" / "skills",
|
||||
Path.cwd() / ".praison" / "skills",
|
||||
Path.cwd() / ".claude" / "skills",
|
||||
]
|
||||
found = []
|
||||
for p in candidate_paths:
|
||||
if p.exists():
|
||||
for child in p.glob("**/SKILL.md"):
|
||||
if "fizzbuzz" in child.read_text(errors="ignore").lower():
|
||||
found.append(str(child))
|
||||
|
||||
print(f"[Disk] SKILL.md files matching 'fizzbuzz': {found or 'none'}")
|
||||
print()
|
||||
print(" EXPECTED: W1 does NOT persist new skills — that's W2's curator loop.")
|
||||
print(" This test PASSES by demonstrating the gap, not by creating a file.")
|
||||
return True # Always pass — this is a documentation/probe test.
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
async def main() -> int:
|
||||
model = _pick_model()
|
||||
print(f"\n{'#' * 72}\n# W1 ROBUSTNESS SUITE — model: {model}\n{'#' * 72}")
|
||||
|
||||
results = {}
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
for name, coro_factory in [
|
||||
("T1", lambda p: test_three_platform_unification(p, model)),
|
||||
("T2", lambda p: test_persistence_across_restart(p, model)),
|
||||
("T3", lambda p: test_concurrent_users_isolated(p, model)),
|
||||
("T4", lambda p: test_session_context_visible_to_tool(p, model)),
|
||||
("T5", lambda p: test_self_improving_probe(p, model)),
|
||||
]:
|
||||
sub = Path(td) / name
|
||||
sub.mkdir()
|
||||
try:
|
||||
results[name] = await coro_factory(sub)
|
||||
except Exception as e:
|
||||
print(f"\n EXCEPTION in {name}: {type(e).__name__}: {e}")
|
||||
results[name] = False
|
||||
|
||||
print(f"\n{'#' * 72}\n# SUMMARY\n{'#' * 72}")
|
||||
for name, ok in results.items():
|
||||
print(f" {name}: {'PASS' if ok else 'FAIL'}")
|
||||
return 0 if all(results.values()) else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(asyncio.run(main()))
|
||||
@@ -519,7 +519,6 @@ class Agent(UnifiedExecutionMixin, ToolExecutionMixin, ChatHandlerMixin, Session
|
||||
model: Optional[Union[str, Any]] = None, # Alias for llm=
|
||||
base_url: Optional[str] = None, # Kept separate (connection/auth)
|
||||
api_key: Optional[str] = None, # Kept separate (connection/auth)
|
||||
auth: Optional[str] = None, # Subscription auth provider: "claude-code", "codex", etc.
|
||||
# Tools
|
||||
tools: Optional[List[Any]] = None,
|
||||
allow_delegation: bool = False, # Deprecated: use handoffs= instead
|
||||
@@ -1522,8 +1521,6 @@ class Agent(UnifiedExecutionMixin, ToolExecutionMixin, ChatHandlerMixin, Session
|
||||
llm_config['base_url'] = base_url
|
||||
if api_key:
|
||||
llm_config['api_key'] = api_key
|
||||
if auth:
|
||||
llm_config['auth'] = auth
|
||||
llm_config['metrics'] = metrics
|
||||
self.llm_instance = LLM(**llm_config)
|
||||
self.llm = llm.get('model', Agent._get_default_model())
|
||||
@@ -1534,7 +1531,6 @@ class Agent(UnifiedExecutionMixin, ToolExecutionMixin, ChatHandlerMixin, Session
|
||||
model=model_name,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
auth=auth,
|
||||
metrics=metrics,
|
||||
web_search=web_search,
|
||||
web_fetch=web_fetch,
|
||||
@@ -1556,10 +1552,6 @@ class Agent(UnifiedExecutionMixin, ToolExecutionMixin, ChatHandlerMixin, Session
|
||||
if api_key and 'api_key' not in llm:
|
||||
llm = llm.copy()
|
||||
llm['api_key'] = api_key
|
||||
# Add auth if provided and not in dict
|
||||
if auth and 'auth' not in llm:
|
||||
llm = llm.copy()
|
||||
llm['auth'] = auth
|
||||
# Add metrics parameter
|
||||
llm = llm.copy()
|
||||
llm['metrics'] = metrics
|
||||
@@ -1579,8 +1571,6 @@ class Agent(UnifiedExecutionMixin, ToolExecutionMixin, ChatHandlerMixin, Session
|
||||
llm_params = {'model': llm}
|
||||
if api_key:
|
||||
llm_params['api_key'] = api_key
|
||||
if auth:
|
||||
llm_params['auth'] = auth
|
||||
llm_params['metrics'] = metrics
|
||||
llm_params['web_search'] = web_search
|
||||
llm_params['web_fetch'] = web_fetch
|
||||
|
||||
@@ -11,6 +11,8 @@ import time
|
||||
import json
|
||||
import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
import asyncio
|
||||
import threading
|
||||
from ..errors import BudgetExceededError
|
||||
|
||||
# Fallback helpers to avoid circular imports
|
||||
@@ -43,7 +45,9 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
||||
from typing import List, Optional, Any, Dict, Union, Generator, TYPE_CHECKING
|
||||
import traceback
|
||||
from typing import List, Optional, Any, Dict, Union, Callable, Generator, TYPE_CHECKING
|
||||
from collections import OrderedDict
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
@@ -556,11 +560,7 @@ Your Goal: {self.goal}"""
|
||||
temperature=temperature
|
||||
)
|
||||
self._hook_runner.execute_sync(HookEvent.BEFORE_LLM, before_llm_input)
|
||||
# C7 — honour any BEFORE_LLM hook that mutated the message stream
|
||||
# (e.g. PII redactor). The runner applies modified_input in-place on
|
||||
# before_llm_input.messages; adopt that value for the actual LLM call.
|
||||
messages = before_llm_input.messages
|
||||
|
||||
|
||||
logging.debug(f"{self.name} sending messages to LLM: {messages}")
|
||||
|
||||
# Emit LLM request trace event (zero overhead when not set)
|
||||
@@ -1149,7 +1149,7 @@ Your Goal: {self.goal}"""
|
||||
)
|
||||
return rendered
|
||||
|
||||
def chat(self, prompt: str, temperature: float = 1.0, tools: Optional[List[Any]] = None, output_json: Optional[Any] = None, output_pydantic: Optional[Any] = None, reasoning_steps: bool = False, stream: Optional[bool] = None, task_name: Optional[str] = None, task_description: Optional[str] = None, task_id: Optional[str] = None, config: Optional[Dict[str, Any]] = None, force_retrieval: bool = False, skip_retrieval: bool = False, attachments: Optional[List[str]] = None, tool_choice: Optional[str] = None, seed: Optional[int] = None, cancel_token: Optional[Any] = None) -> Optional[str]:
|
||||
def chat(self, prompt: str, temperature: float = 1.0, tools: Optional[List[Any]] = None, output_json: Optional[Any] = None, output_pydantic: Optional[Any] = None, reasoning_steps: bool = False, stream: Optional[bool] = None, task_name: Optional[str] = None, task_description: Optional[str] = None, task_id: Optional[str] = None, config: Optional[Dict[str, Any]] = None, force_retrieval: bool = False, skip_retrieval: bool = False, attachments: Optional[List[str]] = None, tool_choice: Optional[str] = None) -> Optional[str]:
|
||||
"""
|
||||
Chat with the agent.
|
||||
|
||||
@@ -1193,17 +1193,11 @@ Your Goal: {self.goal}"""
|
||||
_trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
|
||||
|
||||
try:
|
||||
# C2 — cooperative cancellation: abort early if a pre-set token is given
|
||||
_cancel = cancel_token if cancel_token is not None else getattr(self, "interrupt_controller", None)
|
||||
if _cancel is not None and getattr(_cancel, "is_set", lambda: False)():
|
||||
reason = getattr(_cancel, "reason", None) or "cancelled before LLM call"
|
||||
raise InterruptedError(f"Agent chat cancelled: {reason}")
|
||||
|
||||
return self._chat_impl(prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice, seed=seed, cancel_token=_cancel)
|
||||
return self._chat_impl(prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice)
|
||||
finally:
|
||||
_trace_emitter.agent_end(self.name)
|
||||
|
||||
def _chat_impl(self, prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice=None, seed=None, cancel_token=None):
|
||||
def _chat_impl(self, prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice=None):
|
||||
"""Internal chat implementation (extracted for trace wrapping)."""
|
||||
# Apply rate limiter if configured (before any LLM call)
|
||||
if self._rate_limiter is not None:
|
||||
@@ -1424,16 +1418,7 @@ Your Goal: {self.goal}"""
|
||||
effective_tool_choice = tool_choice or getattr(self, '_yaml_tool_choice', None)
|
||||
if effective_tool_choice:
|
||||
llm_kwargs['tool_choice'] = effective_tool_choice
|
||||
|
||||
# C1 — per-call seed overrides llm_instance.seed for determinism
|
||||
if seed is not None:
|
||||
llm_kwargs['seed'] = seed
|
||||
|
||||
# C2 — last-chance cancel check before handing to the LLM
|
||||
if cancel_token is not None and getattr(cancel_token, 'is_set', lambda: False)():
|
||||
reason = getattr(cancel_token, 'reason', None) or 'cancelled'
|
||||
raise InterruptedError(f"Agent chat cancelled: {reason}")
|
||||
|
||||
|
||||
response_text = self.llm_instance.get_response(**llm_kwargs)
|
||||
|
||||
self._add_to_chat_history("assistant", response_text)
|
||||
@@ -1729,7 +1714,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
cleaned = cleaned[:-3].strip()
|
||||
return cleaned
|
||||
|
||||
async def achat(self, prompt: str, temperature: float = 1.0, tools: Optional[List[Any]] = None, output_json: Optional[Any] = None, output_pydantic: Optional[Any] = None, reasoning_steps: bool = False, stream: Optional[bool] = None, task_name: Optional[str] = None, task_description: Optional[str] = None, task_id: Optional[str] = None, config: Optional[Dict[str, Any]] = None, force_retrieval: bool = False, skip_retrieval: bool = False, attachments: Optional[List[str]] = None, tool_choice: Optional[str] = None, seed: Optional[int] = None, cancel_token: Optional[Any] = None):
|
||||
async def achat(self, prompt: str, temperature: float = 1.0, tools: Optional[List[Any]] = None, output_json: Optional[Any] = None, output_pydantic: Optional[Any] = None, reasoning_steps: bool = False, stream: Optional[bool] = None, task_name: Optional[str] = None, task_description: Optional[str] = None, task_id: Optional[str] = None, config: Optional[Dict[str, Any]] = None, force_retrieval: bool = False, skip_retrieval: bool = False, attachments: Optional[List[str]] = None, tool_choice: Optional[str] = None):
|
||||
"""Async version of chat method with self-reflection support.
|
||||
|
||||
Args:
|
||||
@@ -1752,20 +1737,13 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
reasoning_steps=reasoning_steps, stream=stream,
|
||||
task_name=task_name, task_description=task_description, task_id=task_id,
|
||||
config=config, force_retrieval=force_retrieval, skip_retrieval=skip_retrieval,
|
||||
attachments=attachments, _trace_emitter=_trace_emitter, tool_choice=tool_choice,
|
||||
seed=seed, cancel_token=cancel_token
|
||||
attachments=attachments, _trace_emitter=_trace_emitter, tool_choice=tool_choice
|
||||
)
|
||||
finally:
|
||||
_trace_emitter.agent_end(self.name)
|
||||
|
||||
async def _achat_impl(self, prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice=None, seed=None, cancel_token=None):
|
||||
async def _achat_impl(self, prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice=None):
|
||||
"""Internal async chat implementation (extracted for trace wrapping)."""
|
||||
# C2 — cooperative cancellation: abort early if a pre-set token is given
|
||||
_cancel = cancel_token if cancel_token is not None else getattr(self, "interrupt_controller", None)
|
||||
if _cancel is not None and getattr(_cancel, "is_set", lambda: False)():
|
||||
reason = getattr(_cancel, "reason", None) or "cancelled before LLM call"
|
||||
raise InterruptedError(f"Agent chat cancelled: {reason}")
|
||||
|
||||
# Use agent's stream setting if not explicitly provided
|
||||
if stream is None:
|
||||
stream = self.stream
|
||||
@@ -1888,43 +1866,31 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
logging.debug(f"[compaction] skipped (non-fatal): {_ce}")
|
||||
|
||||
try:
|
||||
# C1 — per-call seed forwarding (async path)
|
||||
llm_kwargs = {
|
||||
'prompt': prompt,
|
||||
'system_prompt': self._build_system_prompt(tools),
|
||||
'chat_history': self.chat_history,
|
||||
'temperature': temperature,
|
||||
'tools': tools,
|
||||
'output_json': output_json,
|
||||
'output_pydantic': output_pydantic,
|
||||
'verbose': self.verbose,
|
||||
'markdown': self.markdown,
|
||||
'reflection': self.self_reflect,
|
||||
'max_reflect': self.max_reflect,
|
||||
'min_reflect': self.min_reflect,
|
||||
'console': self.console,
|
||||
'agent_name': self.name,
|
||||
'agent_role': self.role,
|
||||
'agent_tools': [t.__name__ if hasattr(t, '__name__') else str(t) for t in (tools if tools is not None else self.tools)],
|
||||
'task_name': task_name,
|
||||
'task_description': task_description,
|
||||
'task_id': task_id,
|
||||
'execute_tool_fn': self.execute_tool_async,
|
||||
'parallel_tool_calls': getattr(getattr(self, "execution", None), "parallel_tool_calls", False),
|
||||
'reasoning_steps': reasoning_steps,
|
||||
'stream': stream
|
||||
}
|
||||
|
||||
# C1 — per-call seed overrides llm_instance.seed for determinism
|
||||
if seed is not None:
|
||||
llm_kwargs['seed'] = seed
|
||||
|
||||
# C2 — last-chance cancel check before handing to the LLM
|
||||
if _cancel is not None and getattr(_cancel, 'is_set', lambda: False)():
|
||||
reason = getattr(_cancel, 'reason', None) or 'cancelled'
|
||||
raise InterruptedError(f"Agent chat cancelled: {reason}")
|
||||
|
||||
response_text = await self.llm_instance.get_response_async(**llm_kwargs)
|
||||
response_text = await self.llm_instance.get_response_async(
|
||||
prompt=prompt,
|
||||
system_prompt=self._build_system_prompt(tools),
|
||||
chat_history=self.chat_history,
|
||||
temperature=temperature,
|
||||
tools=tools,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
verbose=self.verbose,
|
||||
markdown=self.markdown,
|
||||
reflection=self.self_reflect,
|
||||
max_reflect=self.max_reflect,
|
||||
min_reflect=self.min_reflect,
|
||||
console=self.console,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in (tools if tools is not None else self.tools)],
|
||||
task_name=task_name,
|
||||
task_description=task_description,
|
||||
task_id=task_id,
|
||||
execute_tool_fn=self.execute_tool_async,
|
||||
parallel_tool_calls=getattr(getattr(self, "execution", None), "parallel_tool_calls", False),
|
||||
reasoning_steps=reasoning_steps,
|
||||
stream=stream
|
||||
)
|
||||
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
|
||||
|
||||
@@ -5,10 +5,12 @@ Contains all methods for run/start/launch lifecycle, planning,
|
||||
and server endpoints. Extracted from agent.py for maintainability.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import json
|
||||
import logging
|
||||
import inspect
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
@@ -124,7 +126,7 @@ class ExecutionMixin:
|
||||
config=HandoffConfig(context_policy=ContextPolicy.NONE),
|
||||
)
|
||||
|
||||
async def arun(self, prompt: str, **kwargs) -> Optional[str]:
|
||||
async def arun(self, prompt: str, **kwargs):
|
||||
"""Async version of run() - silent, non-streaming, returns structured result.
|
||||
|
||||
Production-friendly async execution. Does not stream or display output.
|
||||
|
||||
@@ -9,6 +9,7 @@ for maintainability.
|
||||
import os
|
||||
import logging
|
||||
from typing import Optional
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
# Fallback helpers to avoid circular imports
|
||||
def _get_console():
|
||||
@@ -41,7 +42,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
import contextlib
|
||||
from typing import Optional, Dict, Generator, TYPE_CHECKING
|
||||
from typing import List, Optional, Any, Dict, Generator, TYPE_CHECKING
|
||||
from collections import OrderedDict
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -81,7 +82,7 @@ class MemoryMixin:
|
||||
with self._history_lock:
|
||||
self.chat_history.append({"role": role, "content": content})
|
||||
|
||||
def _add_to_chat_history_if_not_duplicate(self, role, content) -> bool:
|
||||
def _add_to_chat_history_if_not_duplicate(self, role, content):
|
||||
"""Thread-safe method to add messages to chat history only if not duplicate.
|
||||
|
||||
Atomically checks for duplicate and adds message under the same lock to prevent TOCTOU races.
|
||||
|
||||
@@ -194,21 +194,6 @@ class ToolExecutionMixin:
|
||||
if res.output and res.output.modified_data:
|
||||
arguments.update(res.output.modified_data)
|
||||
|
||||
# C4 — optional tool-argument validation via ToolValidatorProtocol.
|
||||
# Zero overhead when not set. Users wire via `agent._tool_validator = MyValidator()`.
|
||||
_validator = getattr(self, '_tool_validator', None)
|
||||
if _validator is not None:
|
||||
try:
|
||||
_vres = _validator.validate_args(function_name, arguments)
|
||||
if _vres is not None and not getattr(_vres, 'valid', True):
|
||||
_errs = "; ".join(getattr(_vres, 'errors', []) or ["validation failed"])
|
||||
logging.warning(
|
||||
f"Tool {function_name} args rejected by validator: {_errs}"
|
||||
)
|
||||
return f"Tool arguments rejected: {_errs}"
|
||||
except Exception as _ve: # noqa: BLE001 — never break tool exec on validator bug
|
||||
logging.debug(f"Tool validator raised; skipping validation: {_ve}")
|
||||
|
||||
# P8/G11: Apply tool timeout if configured
|
||||
tool_timeout = getattr(self, '_tool_timeout', None)
|
||||
if tool_timeout and tool_timeout > 0:
|
||||
@@ -773,9 +758,11 @@ class ToolExecutionMixin:
|
||||
pass
|
||||
|
||||
if func is None:
|
||||
# Tool not found in declared tools or registry — do not fall back to
|
||||
# globals() or __main__ as that allows undeclared callables to execute.
|
||||
pass
|
||||
# If not found in tools or registry, try globals and main
|
||||
func = globals().get(function_name)
|
||||
if not func:
|
||||
import __main__
|
||||
func = getattr(__main__, function_name, None)
|
||||
|
||||
if func:
|
||||
try:
|
||||
|
||||
@@ -14,13 +14,6 @@ Quick start::
|
||||
|
||||
async def authorize(self, identity, resource, action):
|
||||
return "owner" in identity.roles
|
||||
|
||||
Subscription auth (lazy-loaded)::
|
||||
|
||||
from praisonaiagents.auth import resolve_subscription_credentials
|
||||
|
||||
creds = resolve_subscription_credentials("claude-code")
|
||||
print(creds.api_key, creds.headers)
|
||||
"""
|
||||
|
||||
from .protocols import (
|
||||
@@ -32,36 +25,7 @@ from .protocols import (
|
||||
|
||||
__all__ = [
|
||||
"AuthIdentity",
|
||||
"AuthBackendProtocol",
|
||||
"AuthBackendProtocol",
|
||||
"WorkspaceContextProtocol",
|
||||
"AuthConfig",
|
||||
# Subscription auth (lazy-loaded)
|
||||
"SubscriptionAuthProtocol",
|
||||
"SubscriptionCredentials",
|
||||
"register_subscription_provider",
|
||||
"list_subscription_providers",
|
||||
"resolve_subscription_credentials",
|
||||
"get_subscription_provider",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
"""Lazy-load subscription auth to avoid import overhead."""
|
||||
if name in (
|
||||
"SubscriptionAuthProtocol",
|
||||
"SubscriptionCredentials",
|
||||
"register_subscription_provider",
|
||||
"list_subscription_providers",
|
||||
"resolve_subscription_credentials",
|
||||
"get_subscription_provider",
|
||||
):
|
||||
from .subscription import (
|
||||
SubscriptionAuthProtocol,
|
||||
SubscriptionCredentials,
|
||||
register_subscription_provider,
|
||||
list_subscription_providers,
|
||||
resolve_subscription_credentials,
|
||||
get_subscription_provider,
|
||||
)
|
||||
return locals()[name]
|
||||
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
|
||||
@@ -1,35 +0,0 @@
|
||||
"""Subscription / OAuth auth for LLM providers.
|
||||
|
||||
Opt-in, zero-overhead. Lazy-loaded via __getattr__ so importing
|
||||
praisonaiagents stays fast.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = [
|
||||
"SubscriptionAuthProtocol",
|
||||
"SubscriptionCredentials",
|
||||
"register_subscription_provider",
|
||||
"list_subscription_providers",
|
||||
"resolve_subscription_credentials",
|
||||
"get_subscription_provider",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
if name in ("SubscriptionAuthProtocol", "SubscriptionCredentials"):
|
||||
from .protocols import SubscriptionAuthProtocol, SubscriptionCredentials
|
||||
return locals()[name]
|
||||
if name in (
|
||||
"register_subscription_provider",
|
||||
"list_subscription_providers",
|
||||
"resolve_subscription_credentials",
|
||||
"get_subscription_provider",
|
||||
):
|
||||
from .registry import (
|
||||
register_subscription_provider,
|
||||
list_subscription_providers,
|
||||
resolve_subscription_credentials,
|
||||
get_subscription_provider,
|
||||
)
|
||||
return locals()[name]
|
||||
raise AttributeError(name)
|
||||
@@ -1,202 +0,0 @@
|
||||
"""Claude Code OAuth — read ~/.claude/.credentials.json or macOS Keychain."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from .protocols import AuthError, SubscriptionAuthProtocol, SubscriptionCredentials
|
||||
from .registry import register_subscription_provider
|
||||
|
||||
_OAUTH_CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
|
||||
_TOKEN_ENDPOINTS = (
|
||||
"https://platform.claude.com/v1/oauth/token",
|
||||
"https://console.anthropic.com/v1/oauth/token",
|
||||
)
|
||||
_CLAUDE_CODE_VERSION_FALLBACK = "2.1.74"
|
||||
|
||||
|
||||
def _detect_claude_code_version() -> str:
|
||||
"""Detect the installed Claude Code version, fall back to a static constant.
|
||||
|
||||
Anthropic's OAuth infrastructure validates the user-agent version and may
|
||||
reject requests with a version that's too old. Detecting dynamically means
|
||||
users who keep Claude Code updated never hit stale-version 400s.
|
||||
"""
|
||||
cached = getattr(_detect_claude_code_version, "_cache", None)
|
||||
if cached:
|
||||
return cached
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["claude", "--version"], capture_output=True, text=True, timeout=2, check=False
|
||||
)
|
||||
match = re.match(r"(\d+\.\d+\.\d+)", out.stdout.strip())
|
||||
version = match.group(1) if match else _CLAUDE_CODE_VERSION_FALLBACK
|
||||
except Exception:
|
||||
version = _CLAUDE_CODE_VERSION_FALLBACK
|
||||
_detect_claude_code_version._cache = version
|
||||
return version
|
||||
|
||||
|
||||
_CLI_USER_AGENT_FALLBACK = f"claude-cli/{_CLAUDE_CODE_VERSION_FALLBACK} (external, cli)"
|
||||
|
||||
|
||||
def _read_keychain_credentials() -> Optional[Dict[str, Any]]:
|
||||
"""Read 'Claude Code-credentials' entry from macOS Keychain.
|
||||
|
||||
Returns dict with accessToken/refreshToken/expiresAt or None.
|
||||
"""
|
||||
if sys.platform != "darwin":
|
||||
return None
|
||||
import subprocess
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["security", "find-generic-password", "-s", "Claude Code-credentials", "-w"],
|
||||
capture_output=True, text=True, timeout=5, check=False,
|
||||
)
|
||||
if out.returncode != 0 or not out.stdout.strip():
|
||||
return None
|
||||
data = json.loads(out.stdout)
|
||||
oauth = data.get("claudeAiOauth") or {}
|
||||
if oauth.get("accessToken"):
|
||||
return {
|
||||
"accessToken": oauth["accessToken"],
|
||||
"refreshToken": oauth.get("refreshToken", ""),
|
||||
"expiresAt": int(oauth.get("expiresAt", 0)),
|
||||
"source": "claude-code-keychain",
|
||||
}
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _read_file_credentials() -> Optional[Dict[str, Any]]:
|
||||
"""Read ~/.claude/.credentials.json."""
|
||||
path = Path.home() / ".claude" / ".credentials.json"
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
oauth = data.get("claudeAiOauth") or {}
|
||||
if not oauth.get("accessToken"):
|
||||
return None
|
||||
return {
|
||||
"accessToken": oauth["accessToken"],
|
||||
"refreshToken": oauth.get("refreshToken", ""),
|
||||
"expiresAt": int(oauth.get("expiresAt", 0)),
|
||||
"source": "claude-code-file",
|
||||
}
|
||||
|
||||
|
||||
def _is_expiring(expires_at_ms: Optional[int], skew_ms: int = 60_000) -> bool:
|
||||
if not expires_at_ms:
|
||||
return False
|
||||
return int(time.time() * 1000) >= (expires_at_ms - skew_ms)
|
||||
|
||||
|
||||
def _refresh(refresh_token: str) -> Dict[str, Any]:
|
||||
"""Pure refresh — does not touch local files."""
|
||||
if not refresh_token:
|
||||
raise AuthError("no refresh_token; please re-run 'claude /login'")
|
||||
body = urllib.parse.urlencode({
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
"client_id": _OAUTH_CLIENT_ID,
|
||||
}).encode()
|
||||
last = None
|
||||
for endpoint in _TOKEN_ENDPOINTS:
|
||||
req = urllib.request.Request(
|
||||
endpoint, data=body, method="POST",
|
||||
headers={
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
"User-Agent": f"claude-cli/{_detect_claude_code_version()} (external, cli)",
|
||||
},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
payload = json.loads(resp.read().decode())
|
||||
return {
|
||||
"accessToken": payload["access_token"],
|
||||
"refreshToken": payload.get("refresh_token", refresh_token),
|
||||
"expiresAt": int(time.time() * 1000) + int(payload.get("expires_in", 3600)) * 1000,
|
||||
}
|
||||
except Exception as exc:
|
||||
last = exc
|
||||
raise AuthError(f"Anthropic refresh failed: {last}")
|
||||
|
||||
|
||||
class ClaudeCodeAuth:
|
||||
"""Claude Code OAuth subscription auth."""
|
||||
|
||||
name = "claude-code"
|
||||
|
||||
def resolve_credentials(self) -> SubscriptionCredentials:
|
||||
# 1. ANTHROPIC_TOKEN / CLAUDE_CODE_OAUTH_TOKEN env (Hermes-style)
|
||||
for env_key in ("ANTHROPIC_TOKEN", "CLAUDE_CODE_OAUTH_TOKEN"):
|
||||
tok = os.environ.get(env_key, "").strip()
|
||||
if tok:
|
||||
return SubscriptionCredentials(
|
||||
api_key=tok,
|
||||
base_url="https://api.anthropic.com",
|
||||
headers=self.headers_for("https://api.anthropic.com", ""),
|
||||
auth_scheme="bearer",
|
||||
source=f"env:{env_key}",
|
||||
)
|
||||
|
||||
# 2. Keychain or file
|
||||
creds = _read_keychain_credentials() or _read_file_credentials()
|
||||
if not creds:
|
||||
raise AuthError(
|
||||
"No Claude Code credentials found. Install Claude Code and run "
|
||||
"'claude /login', or set ANTHROPIC_TOKEN."
|
||||
)
|
||||
if _is_expiring(creds["expiresAt"]) and creds.get("refreshToken"):
|
||||
creds.update(_refresh(creds["refreshToken"]))
|
||||
return SubscriptionCredentials(
|
||||
api_key=creds["accessToken"],
|
||||
base_url="https://api.anthropic.com",
|
||||
headers=self.headers_for("https://api.anthropic.com", ""),
|
||||
auth_scheme="bearer",
|
||||
expires_at_ms=creds.get("expiresAt"),
|
||||
source=creds["source"],
|
||||
)
|
||||
|
||||
def refresh(self) -> SubscriptionCredentials:
|
||||
creds = _read_keychain_credentials() or _read_file_credentials()
|
||||
if not creds or not creds.get("refreshToken"):
|
||||
raise AuthError("cannot refresh: no refresh token available")
|
||||
new = _refresh(creds["refreshToken"])
|
||||
return SubscriptionCredentials(
|
||||
api_key=new["accessToken"],
|
||||
base_url="https://api.anthropic.com",
|
||||
headers=self.headers_for("https://api.anthropic.com", ""),
|
||||
auth_scheme="bearer",
|
||||
expires_at_ms=new["expiresAt"],
|
||||
source="claude-code-refreshed",
|
||||
)
|
||||
|
||||
def headers_for(self, base_url: str, model: str) -> Dict[str, str]:
|
||||
# CRITICAL: without these headers Anthropic returns 500s on OAuth tokens.
|
||||
# Match Hermes' OAuth header set exactly. litellm auto-adds oauth-2025-04-20.
|
||||
return {
|
||||
"anthropic-beta": ",".join([
|
||||
"interleaved-thinking-2025-05-14",
|
||||
"fine-grained-tool-streaming-2025-05-14",
|
||||
"context-1m-2025-08-07",
|
||||
"claude-code-20250219",
|
||||
]),
|
||||
"user-agent": f"claude-cli/{_detect_claude_code_version()} (external, cli)",
|
||||
"x-app": "cli",
|
||||
}
|
||||
|
||||
|
||||
register_subscription_provider("claude-code", lambda: ClaudeCodeAuth())
|
||||
@@ -1,69 +0,0 @@
|
||||
"""Codex / ChatGPT subscription — read ~/.codex/auth.json (or CODEX_HOME)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from .protocols import AuthError, SubscriptionCredentials
|
||||
from .registry import register_subscription_provider
|
||||
|
||||
|
||||
_CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"
|
||||
|
||||
|
||||
def _read_codex_tokens() -> Optional[Dict[str, Any]]:
|
||||
home = os.environ.get("CODEX_HOME", "").strip() or str(Path.home() / ".codex")
|
||||
auth_path = Path(home).expanduser() / "auth.json"
|
||||
if not auth_path.is_file():
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(auth_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
tokens = payload.get("tokens") or {}
|
||||
if not tokens.get("access_token") or not tokens.get("refresh_token"):
|
||||
return None
|
||||
return tokens
|
||||
|
||||
|
||||
def _is_jwt_expiring(jwt: str, skew_seconds: int = 60) -> bool:
|
||||
"""Decode JWT exp claim without verifying signature."""
|
||||
import base64
|
||||
try:
|
||||
_hdr, payload_b64, _sig = jwt.split(".")
|
||||
pad = "=" * (-len(payload_b64) % 4)
|
||||
payload = json.loads(base64.urlsafe_b64decode(payload_b64 + pad))
|
||||
exp = int(payload.get("exp", 0))
|
||||
return exp and (time.time() + skew_seconds) >= exp
|
||||
except Exception:
|
||||
return False # treat unknown as not-expiring; refresh on actual 401
|
||||
|
||||
|
||||
class CodexAuth:
|
||||
name = "codex"
|
||||
|
||||
def resolve_credentials(self) -> SubscriptionCredentials:
|
||||
# EXPERIMENTAL: Codex uses Responses API, not Chat Completions
|
||||
# OpenAI chat-completions requests will 404 against this endpoint
|
||||
raise AuthError(
|
||||
"Codex auth is experimental and requires custom Responses API transport. "
|
||||
"Use a standard OpenAI API key or wait for proper Codex support in a future release."
|
||||
)
|
||||
|
||||
def refresh(self) -> SubscriptionCredentials:
|
||||
# OpenAI Codex tokens are refreshed by the Codex CLI itself; we rely
|
||||
# on the CLI being run periodically. A direct refresh would require
|
||||
# replicating their device-flow client — out of scope for v1.
|
||||
return self.resolve_credentials()
|
||||
|
||||
def headers_for(self, base_url: str, model: str) -> Dict[str, str]:
|
||||
return {
|
||||
"openai-beta": "responses=v1",
|
||||
"user-agent": "codex-cli/0.1.0 (external)",
|
||||
}
|
||||
|
||||
|
||||
register_subscription_provider("codex", lambda: CodexAuth())
|
||||
@@ -1,51 +0,0 @@
|
||||
"""Gemini CLI OAuth — read ~/.gemini/oauth_creds.json."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from .protocols import AuthError, SubscriptionCredentials
|
||||
from .registry import register_subscription_provider
|
||||
|
||||
|
||||
_GEMINI_BASE_URL = "https://cloudcode-pa.googleapis.com/v1internal"
|
||||
|
||||
|
||||
def _read_gemini_tokens() -> Optional[Dict[str, Any]]:
|
||||
"""Read Gemini CLI OAuth credentials."""
|
||||
auth_path = Path.home() / ".gemini" / "oauth_creds.json"
|
||||
if not auth_path.is_file():
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(auth_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
if not payload.get("access_token"):
|
||||
return None
|
||||
return payload
|
||||
|
||||
|
||||
class GeminiCliAuth:
|
||||
name = "gemini-cli"
|
||||
|
||||
def resolve_credentials(self) -> SubscriptionCredentials:
|
||||
# EXPERIMENTAL: Gemini CLI uses Code Assist API, not OpenAI chat format
|
||||
# Standard OpenAI chat-completions requests will not work with cloudcode-pa endpoints
|
||||
raise AuthError(
|
||||
"Gemini CLI auth is experimental and requires custom Code Assist transport. "
|
||||
"Use standard Gemini API keys or wait for proper Code Assist support in a future release."
|
||||
)
|
||||
|
||||
def refresh(self) -> SubscriptionCredentials:
|
||||
# Gemini CLI handles refresh internally
|
||||
return self.resolve_credentials()
|
||||
|
||||
def headers_for(self, base_url: str, model: str) -> Dict[str, str]:
|
||||
return {
|
||||
"user-agent": "gemini-cli/1.0.0 (external)",
|
||||
}
|
||||
|
||||
|
||||
register_subscription_provider("gemini-cli", lambda: GeminiCliAuth())
|
||||
@@ -1,41 +0,0 @@
|
||||
"""Subscription-auth protocol — what every provider must implement."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Optional, Protocol, runtime_checkable
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubscriptionCredentials:
|
||||
"""Resolved credentials ready to inject into an LLM client."""
|
||||
api_key: str # access_token (or "" for SDK-default chains like AWS)
|
||||
base_url: str = "" # provider endpoint
|
||||
headers: Dict[str, str] = field(default_factory=dict) # extra headers (UA, x-app, betas)
|
||||
auth_scheme: str = "bearer" # "bearer" | "x-api-key" | "sdk-default"
|
||||
expires_at_ms: Optional[int] = None # epoch ms; None = no expiry tracked
|
||||
source: str = "" # "claude-code-keychain", "codex-cli-file", etc.
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SubscriptionAuthProtocol(Protocol):
|
||||
"""Discover, refresh, and inject subscription credentials."""
|
||||
|
||||
def resolve_credentials(self) -> SubscriptionCredentials:
|
||||
"""Return current valid credentials, refreshing if expiring soon.
|
||||
|
||||
Raises:
|
||||
AuthError: if no usable credentials can be resolved.
|
||||
"""
|
||||
...
|
||||
|
||||
def refresh(self) -> SubscriptionCredentials:
|
||||
"""Force a refresh and return the new credentials."""
|
||||
...
|
||||
|
||||
def headers_for(self, base_url: str, model: str) -> Dict[str, str]:
|
||||
"""Provider-specific headers (UA, beta flags) keyed off endpoint/model."""
|
||||
...
|
||||
|
||||
|
||||
class AuthError(RuntimeError):
|
||||
"""Raised when subscription credentials cannot be resolved."""
|
||||
@@ -1,58 +0,0 @@
|
||||
"""Qwen CLI OAuth — read ~/.qwen/oauth_creds.json."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from .protocols import AuthError, SubscriptionCredentials
|
||||
from .registry import register_subscription_provider
|
||||
|
||||
|
||||
_QWEN_BASE_URL = "https://portal.qwen.ai/v1"
|
||||
|
||||
|
||||
def _read_qwen_tokens() -> Optional[Dict[str, Any]]:
|
||||
"""Read Qwen CLI OAuth credentials."""
|
||||
auth_path = Path.home() / ".qwen" / "oauth_creds.json"
|
||||
if not auth_path.is_file():
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(auth_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
if not payload.get("access_token"):
|
||||
return None
|
||||
return payload
|
||||
|
||||
|
||||
class QwenCliAuth:
|
||||
name = "qwen-cli"
|
||||
|
||||
def resolve_credentials(self) -> SubscriptionCredentials:
|
||||
tokens = _read_qwen_tokens()
|
||||
if not tokens:
|
||||
raise AuthError(
|
||||
"No Qwen CLI credentials at ~/.qwen/oauth_creds.json. "
|
||||
"Install Qwen CLI and run 'qwen login'."
|
||||
)
|
||||
return SubscriptionCredentials(
|
||||
api_key=tokens["access_token"],
|
||||
base_url=os.environ.get("PRAISONAI_QWEN_BASE_URL", "").strip() or _QWEN_BASE_URL,
|
||||
headers=self.headers_for(_QWEN_BASE_URL, ""),
|
||||
auth_scheme="bearer",
|
||||
source="qwen-cli-file",
|
||||
)
|
||||
|
||||
def refresh(self) -> SubscriptionCredentials:
|
||||
# Qwen CLI handles refresh internally
|
||||
return self.resolve_credentials()
|
||||
|
||||
def headers_for(self, base_url: str, model: str) -> Dict[str, str]:
|
||||
return {
|
||||
"user-agent": "qwen-cli/1.0.0 (external)",
|
||||
}
|
||||
|
||||
|
||||
register_subscription_provider("qwen-cli", lambda: QwenCliAuth())
|
||||
@@ -1,64 +0,0 @@
|
||||
"""Thread-safe provider registry for subscription auth."""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
from .protocols import AuthError, SubscriptionAuthProtocol, SubscriptionCredentials
|
||||
|
||||
_REGISTRY: Dict[str, Callable[[], SubscriptionAuthProtocol]] = {}
|
||||
_LOCK = threading.RLock()
|
||||
_BUILTIN_REGISTERED = False
|
||||
|
||||
|
||||
def register_subscription_provider(
|
||||
provider_id: str,
|
||||
factory: Callable[[], SubscriptionAuthProtocol],
|
||||
) -> None:
|
||||
"""Register a new subscription-auth provider.
|
||||
|
||||
Args:
|
||||
provider_id: short id e.g. "claude-code", "codex", "gemini-cli".
|
||||
factory: zero-arg callable returning a SubscriptionAuthProtocol impl.
|
||||
"""
|
||||
with _LOCK:
|
||||
_REGISTRY[provider_id] = factory
|
||||
|
||||
|
||||
def list_subscription_providers() -> List[str]:
|
||||
_ensure_builtins()
|
||||
with _LOCK:
|
||||
return sorted(_REGISTRY)
|
||||
|
||||
|
||||
def resolve_subscription_credentials(provider_id: str) -> SubscriptionCredentials:
|
||||
"""Look up provider by id and resolve its credentials."""
|
||||
_ensure_builtins()
|
||||
with _LOCK:
|
||||
factory = _REGISTRY.get(provider_id)
|
||||
if factory is None:
|
||||
raise AuthError(
|
||||
f"Unknown subscription provider {provider_id!r}. "
|
||||
f"Registered: {list_subscription_providers()}"
|
||||
)
|
||||
return factory().resolve_credentials()
|
||||
|
||||
|
||||
def get_subscription_provider(provider_id: str) -> Optional[SubscriptionAuthProtocol]:
|
||||
"""Return a new provider instance for *provider_id*, or None if not registered."""
|
||||
_ensure_builtins()
|
||||
with _LOCK:
|
||||
factory = _REGISTRY.get(provider_id)
|
||||
return factory() if factory is not None else None
|
||||
|
||||
|
||||
def _ensure_builtins() -> None:
|
||||
"""Lazy register built-in providers on first use."""
|
||||
global _BUILTIN_REGISTERED
|
||||
with _LOCK:
|
||||
if _BUILTIN_REGISTERED:
|
||||
return
|
||||
# Import modules inside the lock to prevent race conditions
|
||||
from . import claude_code, codex, gemini_cli, qwen_cli # noqa: F401
|
||||
# Each module registers itself at import time.
|
||||
_BUILTIN_REGISTERED = True
|
||||
@@ -71,8 +71,7 @@ class Mem0Adapter:
|
||||
from mem0 import Memory
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"mem0ai is not installed. "
|
||||
"Run: pip install 'praisonaiagents[memory]'"
|
||||
"mem0 is not installed. Install with: pip install mem0ai"
|
||||
)
|
||||
|
||||
if self._config:
|
||||
|
||||
@@ -361,7 +361,6 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
prompt_caching: Optional[bool] = None,
|
||||
claude_memory: Optional[Union[bool, Any]] = None,
|
||||
failover_manager: Optional[FailoverManagerProtocol] = None,
|
||||
auth: Optional[str] = None, # NEW: "claude-code", "codex", "gemini-cli", "qwen-cli"
|
||||
**extra_settings
|
||||
):
|
||||
# Configure logging only once at the class level
|
||||
@@ -438,10 +437,6 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
self._max_retries = extra_settings.get('max_retries', 3)
|
||||
self._retry_delay = extra_settings.get('retry_delay', 60) # Default 60 seconds
|
||||
|
||||
# Subscription auth
|
||||
self._auth_provider_id = auth
|
||||
self._cached_subscription_creds = None # SubscriptionCredentials | None
|
||||
|
||||
# Failover management
|
||||
self._failover_manager = failover_manager
|
||||
self._current_profile = None # Track current auth profile for failover
|
||||
@@ -720,35 +715,6 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
logging.info(f"Failover: switching from {self.model} to {profile.model}")
|
||||
self.model = profile.model
|
||||
|
||||
def _resolve_subscription_creds(self):
|
||||
"""Lazy resolve + cache subscription credentials, checking expiration."""
|
||||
if not self._auth_provider_id:
|
||||
return None
|
||||
|
||||
# Check if cached credentials are expired
|
||||
if (self._cached_subscription_creds and
|
||||
self._cached_subscription_creds.expires_at_ms and
|
||||
self._cached_subscription_creds.expires_at_ms <= int(time.time() * 1000)):
|
||||
self._cached_subscription_creds = None
|
||||
|
||||
if self._cached_subscription_creds is None:
|
||||
from ..auth import resolve_subscription_credentials
|
||||
self._cached_subscription_creds = resolve_subscription_credentials(
|
||||
self._auth_provider_id,
|
||||
)
|
||||
return self._cached_subscription_creds
|
||||
|
||||
def _refresh_subscription_creds(self):
|
||||
"""Force refresh of subscription credentials."""
|
||||
if not self._auth_provider_id:
|
||||
return None
|
||||
from ..auth.subscription.registry import get_subscription_provider
|
||||
provider = get_subscription_provider(self._auth_provider_id)
|
||||
if provider is None:
|
||||
return None
|
||||
self._cached_subscription_creds = provider.refresh()
|
||||
return self._cached_subscription_creds
|
||||
|
||||
def _call_with_retry(self, func, *args, **kwargs):
|
||||
"""Call a function with automatic retry on rate limit errors and failover support.
|
||||
|
||||
@@ -785,21 +751,6 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
last_error = e
|
||||
error_str = str(e)
|
||||
|
||||
# Check for auth errors and try refreshing subscription credentials
|
||||
if category == "auth" and self._auth_provider_id and attempt == 0:
|
||||
try:
|
||||
logging.info("Authentication error detected - attempting credential refresh")
|
||||
refreshed_creds = self._refresh_subscription_creds()
|
||||
if refreshed_creds:
|
||||
# Update parameters with refreshed credentials (don't clear cache)
|
||||
kwargs = self._build_completion_params(**kwargs)
|
||||
# Retry immediately with refreshed credentials
|
||||
can_retry = True
|
||||
retry_delay = 0.0
|
||||
logging.info("Subscription credentials refreshed, retrying...")
|
||||
except Exception as refresh_error:
|
||||
logging.warning(f"Failed to refresh subscription credentials: {refresh_error}")
|
||||
|
||||
# Failover: mark failure and try next profile (do this before early exit)
|
||||
if self._failover_manager and self._current_profile:
|
||||
is_rate_limit = (category == "rate_limit")
|
||||
@@ -889,21 +840,6 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
last_error = e
|
||||
error_str = str(e)
|
||||
|
||||
# Check for auth errors and try refreshing subscription credentials
|
||||
if category == "auth" and self._auth_provider_id and attempt == 0:
|
||||
try:
|
||||
logging.info("Authentication error detected - attempting credential refresh")
|
||||
refreshed_creds = self._refresh_subscription_creds()
|
||||
if refreshed_creds:
|
||||
# Update parameters with refreshed credentials (don't clear cache)
|
||||
kwargs = self._build_completion_params(**kwargs)
|
||||
# Retry immediately with refreshed credentials
|
||||
can_retry = True
|
||||
retry_delay = 0.0
|
||||
logging.info("Subscription credentials refreshed, retrying...")
|
||||
except Exception as refresh_error:
|
||||
logging.warning(f"Failed to refresh subscription credentials: {refresh_error}")
|
||||
|
||||
# Failover: mark failure and try next profile (do this before early exit)
|
||||
if self._failover_manager and self._current_profile:
|
||||
is_rate_limit = (category == "rate_limit")
|
||||
@@ -4657,18 +4593,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
filtered_extra_settings = {k: v for k, v in self.extra_settings.items() if k != 'metrics'}
|
||||
params.update(filtered_extra_settings)
|
||||
|
||||
# Inject subscription credentials if auth provider is set
|
||||
creds = self._resolve_subscription_creds()
|
||||
if creds:
|
||||
# Use litellm's native OAuth detection (auto-detects sk-ant-oat-* and switches to Bearer)
|
||||
params["api_key"] = creds.api_key
|
||||
if creds.base_url:
|
||||
params["base_url"] = creds.base_url
|
||||
if creds.headers:
|
||||
extra_headers = dict(params.get("extra_headers") or {})
|
||||
extra_headers.update(creds.headers)
|
||||
params["extra_headers"] = extra_headers
|
||||
|
||||
# Override with any provided parameters
|
||||
params.update(override_params)
|
||||
|
||||
|
||||
@@ -40,8 +40,7 @@ def create_mem0_memory_adapter(**kwargs) -> MemoryProtocol:
|
||||
import mem0
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"mem0ai is not installed. "
|
||||
"Run: pip install 'praisonaiagents[memory]'"
|
||||
"mem0 is required for mem0 adapter. Install with: pip install mem0ai"
|
||||
)
|
||||
|
||||
return Mem0MemoryAdapter(mem0_config=kwargs)
|
||||
@@ -364,7 +363,7 @@ class MongoDBMemoryAdapter:
|
||||
|
||||
def store_short_term(self, text: str, metadata: Optional[Dict[str, Any]] = None, **kwargs) -> str:
|
||||
"""Store in MongoDB short-term collection."""
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime
|
||||
import time
|
||||
|
||||
doc_id = str(time.time_ns())
|
||||
@@ -372,7 +371,7 @@ class MongoDBMemoryAdapter:
|
||||
"_id": doc_id,
|
||||
"content": text,
|
||||
"metadata": metadata or {},
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
"created_at": datetime.now(datetime.timezone.utc),
|
||||
"memory_type": "short_term"
|
||||
}
|
||||
|
||||
@@ -396,7 +395,7 @@ class MongoDBMemoryAdapter:
|
||||
|
||||
def store_long_term(self, text: str, metadata: Optional[Dict[str, Any]] = None, **kwargs) -> str:
|
||||
"""Store in MongoDB long-term collection."""
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime
|
||||
import time
|
||||
|
||||
doc_id = str(time.time_ns())
|
||||
@@ -404,7 +403,7 @@ class MongoDBMemoryAdapter:
|
||||
"_id": doc_id,
|
||||
"content": text,
|
||||
"metadata": metadata or {},
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
"created_at": datetime.now(datetime.timezone.utc),
|
||||
"memory_type": "long_term"
|
||||
}
|
||||
|
||||
|
||||
@@ -143,9 +143,6 @@ class Memory(StorageMixin, SearchMixin, MemoryCoreMixin):
|
||||
adapter_name = provider_mapping.get(provider, provider)
|
||||
|
||||
# Try to get preferred adapter, fallback to available ones
|
||||
adapter = None
|
||||
provider_explicitly_requested = provider != "rag" # "rag" is default, others are explicit
|
||||
|
||||
try:
|
||||
adapter_config = self._get_adapter_config_for_provider(adapter_name)
|
||||
adapter = get_memory_adapter(adapter_name, **adapter_config)
|
||||
@@ -157,37 +154,6 @@ class Memory(StorageMixin, SearchMixin, MemoryCoreMixin):
|
||||
adapter = None
|
||||
|
||||
if adapter is None:
|
||||
# If the provider was explicitly requested, try to give a helpful error message
|
||||
if provider_explicitly_requested:
|
||||
try:
|
||||
# Try the factory function directly to get the specific ImportError message
|
||||
from .adapters.factories import (
|
||||
create_mem0_memory_adapter,
|
||||
create_chroma_memory_adapter,
|
||||
create_mongodb_memory_adapter,
|
||||
)
|
||||
factory_map = {
|
||||
"mem0": create_mem0_memory_adapter,
|
||||
"chroma": create_chroma_memory_adapter,
|
||||
"mongodb": create_mongodb_memory_adapter,
|
||||
}
|
||||
if adapter_name in factory_map:
|
||||
# Call the factory to get ImportError with installation hint
|
||||
# or capture successful adapter creation
|
||||
adapter_config = self._get_adapter_config_for_provider(adapter_name)
|
||||
adapter = factory_map[adapter_name](**adapter_config)
|
||||
except ImportError:
|
||||
# Re-raise ImportError with installation instructions
|
||||
raise
|
||||
# All other exceptions (e.g., ValueError for missing API key)
|
||||
# propagate naturally, preserving backward compatibility
|
||||
|
||||
# If factory succeeded, use that adapter
|
||||
if adapter is not None:
|
||||
self.memory_adapter = adapter
|
||||
self.provider = adapter_name
|
||||
return
|
||||
|
||||
# Fallback to first available adapter
|
||||
self._log_verbose(f"Provider '{adapter_name}' not available, trying fallbacks")
|
||||
# Try each fallback preference individually
|
||||
|
||||
@@ -90,43 +90,7 @@ def __getattr__(name: str):
|
||||
from .hierarchy import ExtendedSessionData
|
||||
_module_cache[name] = ExtendedSessionData
|
||||
return ExtendedSessionData
|
||||
|
||||
# W1 — Identity resolver
|
||||
if name == "IdentityResolverProtocol":
|
||||
from .identity import IdentityResolverProtocol
|
||||
_module_cache[name] = IdentityResolverProtocol
|
||||
return IdentityResolverProtocol
|
||||
if name == "IdentityLink":
|
||||
from .identity import IdentityLink
|
||||
_module_cache[name] = IdentityLink
|
||||
return IdentityLink
|
||||
if name == "InMemoryIdentityResolver":
|
||||
from .identity import InMemoryIdentityResolver
|
||||
_module_cache[name] = InMemoryIdentityResolver
|
||||
return InMemoryIdentityResolver
|
||||
if name == "FileIdentityResolver":
|
||||
from .identity import FileIdentityResolver
|
||||
_module_cache[name] = FileIdentityResolver
|
||||
return FileIdentityResolver
|
||||
|
||||
# W1 — Task-local session context
|
||||
if name == "SessionContext":
|
||||
from .context import SessionContext
|
||||
_module_cache[name] = SessionContext
|
||||
return SessionContext
|
||||
if name == "set_session_context":
|
||||
from .context import set_session_context
|
||||
_module_cache[name] = set_session_context
|
||||
return set_session_context
|
||||
if name == "get_session_context":
|
||||
from .context import get_session_context
|
||||
_module_cache[name] = get_session_context
|
||||
return get_session_context
|
||||
if name == "clear_session_context":
|
||||
from .context import clear_session_context
|
||||
_module_cache[name] = clear_session_context
|
||||
return clear_session_context
|
||||
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
@@ -141,13 +105,4 @@ __all__ = [
|
||||
"get_hierarchical_session_store",
|
||||
"SessionSnapshot",
|
||||
"ExtendedSessionData",
|
||||
# W1
|
||||
"IdentityResolverProtocol",
|
||||
"IdentityLink",
|
||||
"InMemoryIdentityResolver",
|
||||
"FileIdentityResolver",
|
||||
"SessionContext",
|
||||
"set_session_context",
|
||||
"get_session_context",
|
||||
"clear_session_context",
|
||||
]
|
||||
@@ -1,112 +0,0 @@
|
||||
"""Task-local session context for concurrent message handlers.
|
||||
|
||||
W1 — Replaces process-global ``os.environ`` based session metadata
|
||||
with ``contextvars.ContextVar`` so that two messages handled
|
||||
concurrently never overwrite each other's platform / chat / user IDs.
|
||||
|
||||
Each ``set_session_context`` call returns a token; pass it to
|
||||
``clear_session_context`` in a ``finally`` block to restore the
|
||||
previous state.
|
||||
|
||||
Usage in a bot handler::
|
||||
|
||||
from praisonaiagents.session.context import (
|
||||
set_session_context,
|
||||
clear_session_context,
|
||||
get_session_context,
|
||||
)
|
||||
|
||||
token = set_session_context(
|
||||
platform="telegram",
|
||||
chat_id="100",
|
||||
user_id="alice",
|
||||
unified_user_id="alice-global",
|
||||
)
|
||||
try:
|
||||
# ... agent handles message; tools can call get_session_context()
|
||||
ctx = get_session_context()
|
||||
finally:
|
||||
clear_session_context(token)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextvars import ContextVar, Token
|
||||
from dataclasses import asdict, dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SessionContext:
|
||||
"""Snapshot of the current message-handling context."""
|
||||
|
||||
platform: str = ""
|
||||
chat_id: str = ""
|
||||
chat_name: str = ""
|
||||
thread_id: str = ""
|
||||
user_id: str = ""
|
||||
user_name: str = ""
|
||||
unified_user_id: str = ""
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return asdict(self)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict) -> "SessionContext":
|
||||
return cls(
|
||||
platform=d.get("platform", ""),
|
||||
chat_id=d.get("chat_id", ""),
|
||||
chat_name=d.get("chat_name", ""),
|
||||
thread_id=d.get("thread_id", ""),
|
||||
user_id=d.get("user_id", ""),
|
||||
user_name=d.get("user_name", ""),
|
||||
unified_user_id=d.get("unified_user_id", ""),
|
||||
)
|
||||
|
||||
|
||||
_CTX: ContextVar[SessionContext] = ContextVar(
|
||||
"praisonai_session_context", default=SessionContext()
|
||||
)
|
||||
|
||||
|
||||
def set_session_context(
|
||||
platform: str = "",
|
||||
chat_id: str = "",
|
||||
chat_name: str = "",
|
||||
thread_id: str = "",
|
||||
user_id: str = "",
|
||||
user_name: str = "",
|
||||
unified_user_id: str = "",
|
||||
) -> Token:
|
||||
"""Set the task-local session context. Returns a token for ``clear``."""
|
||||
ctx = SessionContext(
|
||||
platform=platform,
|
||||
chat_id=chat_id,
|
||||
chat_name=chat_name,
|
||||
thread_id=thread_id,
|
||||
user_id=user_id,
|
||||
user_name=user_name,
|
||||
unified_user_id=unified_user_id,
|
||||
)
|
||||
return _CTX.set(ctx)
|
||||
|
||||
|
||||
def get_session_context() -> SessionContext:
|
||||
"""Return the current task's session context (empty if unset)."""
|
||||
return _CTX.get()
|
||||
|
||||
|
||||
def clear_session_context(token: Token) -> None:
|
||||
"""Restore the previous session context using the token from ``set``."""
|
||||
try:
|
||||
_CTX.reset(token)
|
||||
except (LookupError, ValueError):
|
||||
# Token from a different context; reset to empty
|
||||
_CTX.set(SessionContext())
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SessionContext",
|
||||
"set_session_context",
|
||||
"get_session_context",
|
||||
"clear_session_context",
|
||||
]
|
||||
@@ -1,231 +0,0 @@
|
||||
"""Cross-platform identity resolution for agents.
|
||||
|
||||
W1 — Optional, opt-in mapping from ``(platform, platform_user_id)``
|
||||
to a single ``unified_user_id`` so the same human is recognised across
|
||||
Telegram, Discord, Slack, etc.
|
||||
|
||||
By default (no links registered) ``resolve()`` returns the
|
||||
platform-prefixed string unchanged — no surprises, no privacy leak.
|
||||
|
||||
Linking is performed explicitly, typically after a DM-based pairing
|
||||
flow has cryptographically verified ownership of both accounts.
|
||||
|
||||
Usage::
|
||||
|
||||
from praisonaiagents.session.identity import InMemoryIdentityResolver
|
||||
|
||||
resolver = InMemoryIdentityResolver()
|
||||
resolver.link("telegram", "12345", "alice")
|
||||
resolver.link("discord", "snowflake-1", "alice")
|
||||
assert resolver.resolve("telegram", "12345") == "alice"
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Protocol, runtime_checkable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IdentityLink:
|
||||
"""Immutable mapping from a platform-scoped user ID to a unified ID."""
|
||||
|
||||
platform: str
|
||||
platform_user_id: str
|
||||
unified_user_id: str
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class IdentityResolverProtocol(Protocol):
|
||||
"""Maps platform-scoped user IDs to a unified identity.
|
||||
|
||||
Implementations MUST be thread-safe. Implementations SHOULD return
|
||||
``f"{platform}:{platform_user_id}"`` for unlinked users so callers
|
||||
never need to special-case the unlinked path.
|
||||
"""
|
||||
|
||||
def resolve(self, platform: str, platform_user_id: str) -> str:
|
||||
"""Return ``unified_user_id`` for the given platform user.
|
||||
|
||||
If no link is registered, return ``f"{platform}:{platform_user_id}"``.
|
||||
"""
|
||||
...
|
||||
|
||||
def link(
|
||||
self, platform: str, platform_user_id: str, unified_user_id: str
|
||||
) -> None:
|
||||
"""Register a mapping. Overwrites any existing mapping for the same
|
||||
``(platform, platform_user_id)`` pair.
|
||||
"""
|
||||
...
|
||||
|
||||
def unlink(self, platform: str, platform_user_id: str) -> None:
|
||||
"""Remove a mapping. No-op when no link exists."""
|
||||
...
|
||||
|
||||
def links_for(self, unified_user_id: str) -> List[IdentityLink]:
|
||||
"""Return all links pointing to the given unified ID."""
|
||||
...
|
||||
|
||||
|
||||
class InMemoryIdentityResolver:
|
||||
"""Thread-safe in-memory ``IdentityResolverProtocol`` implementation.
|
||||
|
||||
Suitable for tests and single-process deployments. Production
|
||||
deployments with multiple processes should use a persistent
|
||||
implementation backed by SQLite, Redis, or a database.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._links: dict[tuple[str, str], str] = {}
|
||||
self._lock = threading.RLock()
|
||||
|
||||
def resolve(self, platform: str, platform_user_id: str) -> str:
|
||||
with self._lock:
|
||||
return self._links.get(
|
||||
(platform, platform_user_id),
|
||||
f"{platform}:{platform_user_id}",
|
||||
)
|
||||
|
||||
def link(
|
||||
self, platform: str, platform_user_id: str, unified_user_id: str
|
||||
) -> None:
|
||||
with self._lock:
|
||||
self._links[(platform, platform_user_id)] = unified_user_id
|
||||
|
||||
def unlink(self, platform: str, platform_user_id: str) -> None:
|
||||
with self._lock:
|
||||
self._links.pop((platform, platform_user_id), None)
|
||||
|
||||
def links_for(self, unified_user_id: str) -> List[IdentityLink]:
|
||||
with self._lock:
|
||||
return [
|
||||
IdentityLink(p, u, unified_user_id)
|
||||
for (p, u), uid in self._links.items()
|
||||
if uid == unified_user_id
|
||||
]
|
||||
|
||||
|
||||
_DEFAULT_IDENTITY_PATH = Path(
|
||||
os.environ.get(
|
||||
"PRAISONAI_IDENTITY_PATH",
|
||||
os.path.expanduser("~/.praisonai/identity.json"),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class FileIdentityResolver(InMemoryIdentityResolver):
|
||||
"""JSON-file-backed ``IdentityResolverProtocol`` implementation.
|
||||
|
||||
Loads links from disk on construction; persists on every mutation
|
||||
via atomic temp-file + ``os.replace``. Thread-safe.
|
||||
|
||||
Default path: ``$PRAISONAI_IDENTITY_PATH`` or ``~/.praisonai/identity.json``.
|
||||
|
||||
Storage format::
|
||||
|
||||
{
|
||||
"links": {
|
||||
"telegram::12345": "alice",
|
||||
"discord::snowflake-1": "alice"
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(self, path: Optional[Path | str] = None) -> None:
|
||||
super().__init__()
|
||||
self._path: Path = Path(path) if path else _DEFAULT_IDENTITY_PATH
|
||||
# Serialises concurrent disk writes so that two simultaneous
|
||||
# link()/unlink() calls cannot interleave their os.replace() and
|
||||
# lose the second writer's data.
|
||||
self._flush_lock = threading.Lock()
|
||||
self._load()
|
||||
|
||||
@staticmethod
|
||||
def _encode_key(platform: str, platform_user_id: str) -> str:
|
||||
return f"{platform}::{platform_user_id}"
|
||||
|
||||
@staticmethod
|
||||
def _decode_key(encoded: str) -> tuple[str, str]:
|
||||
platform, _, user = encoded.partition("::")
|
||||
return platform, user
|
||||
|
||||
def _load(self) -> None:
|
||||
try:
|
||||
if not self._path.exists():
|
||||
return
|
||||
data = json.loads(self._path.read_text(encoding="utf-8"))
|
||||
links = data.get("links", {}) if isinstance(data, dict) else {}
|
||||
with self._lock:
|
||||
self._links = {
|
||||
self._decode_key(k): v
|
||||
for k, v in links.items()
|
||||
if isinstance(v, str)
|
||||
}
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
logger.warning("FileIdentityResolver: failed to load %s: %s", self._path, e)
|
||||
|
||||
def _flush(self) -> None:
|
||||
try:
|
||||
self._path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Hold _flush_lock for the entire write (snapshot + disk I/O) so
|
||||
# that concurrent calls always write the latest in-memory state and
|
||||
# the last os.replace() wins rather than an older snapshot.
|
||||
with self._flush_lock:
|
||||
with self._lock:
|
||||
payload = {
|
||||
"links": {
|
||||
self._encode_key(p, u): uid
|
||||
for (p, u), uid in self._links.items()
|
||||
}
|
||||
}
|
||||
# Keep the _lock held across the entire write operation to prevent
|
||||
# TOCTOU race where a concurrent link() could modify _links between
|
||||
# payload snapshot and os.replace() completion.
|
||||
fd, tmp = tempfile.mkstemp(
|
||||
dir=str(self._path.parent), prefix=".identity-", suffix=".tmp"
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
json.dump(payload, f, indent=2, ensure_ascii=False)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp, self._path)
|
||||
try:
|
||||
os.chmod(self._path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
except OSError as e:
|
||||
logger.warning("FileIdentityResolver: failed to flush %s: %s", self._path, e)
|
||||
|
||||
def link(
|
||||
self, platform: str, platform_user_id: str, unified_user_id: str
|
||||
) -> None:
|
||||
super().link(platform, platform_user_id, unified_user_id)
|
||||
self._flush()
|
||||
|
||||
def unlink(self, platform: str, platform_user_id: str) -> None:
|
||||
super().unlink(platform, platform_user_id)
|
||||
self._flush()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"IdentityLink",
|
||||
"IdentityResolverProtocol",
|
||||
"InMemoryIdentityResolver",
|
||||
"FileIdentityResolver",
|
||||
]
|
||||
@@ -45,10 +45,6 @@ __all__ = [
|
||||
# Redaction
|
||||
"redact_dict",
|
||||
"REDACT_KEYS",
|
||||
# C7 — PII redaction for LLM egress
|
||||
"scrub_pii_text",
|
||||
"enable_pii_redaction",
|
||||
"disable_pii_redaction",
|
||||
# Context events (for replay)
|
||||
"ContextEvent",
|
||||
"ContextEventType",
|
||||
@@ -83,10 +79,6 @@ def __getattr__(name: str):
|
||||
if name in ("redact_dict", "REDACT_KEYS"):
|
||||
from .redact import redact_dict, REDACT_KEYS
|
||||
return locals()[name]
|
||||
|
||||
if name in ("scrub_pii_text", "enable_pii_redaction", "disable_pii_redaction"):
|
||||
from .redact import scrub_pii_text, enable_pii_redaction, disable_pii_redaction
|
||||
return locals()[name]
|
||||
|
||||
# Context events (for replay)
|
||||
if name in ("ContextEvent", "ContextEventType"):
|
||||
|
||||
@@ -5,9 +5,8 @@ Provides deterministic redaction of sensitive data in tool arguments
|
||||
and results. Uses only stdlib for zero dependencies.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Optional, Set
|
||||
from typing import Any, Dict, Set
|
||||
import re
|
||||
import threading
|
||||
|
||||
# Keys that should always be redacted (case-insensitive matching)
|
||||
REDACT_KEYS: Set[str] = {
|
||||
@@ -120,109 +119,6 @@ def _redact_key_value(key: str, value: Any) -> Any:
|
||||
return _redact_value(value)
|
||||
|
||||
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
# C7 — PII redaction for LLM egress (opt-in BEFORE_LLM middleware)
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Additional value-pattern rules for strings that look like secrets
|
||||
# even when no key=value pair surrounds them.
|
||||
_VALUE_PATTERNS = (
|
||||
# OpenAI-style API keys
|
||||
(re.compile(r"\bsk-[A-Za-z0-9]{12,}\b"), "[REDACTED]"),
|
||||
# US SSN
|
||||
(re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), "[REDACTED-SSN]"),
|
||||
# Credit card — canonical 4-group format only (e.g. "4111 1111 1111 1111").
|
||||
# @claude: The bare \b\d{16}\b branch was removed — it matches ANY 16-digit sequence:
|
||||
# UNIX microsecond timestamps, distributed trace IDs, order/invoice numbers, phone
|
||||
# numbers — causing false-positive redactions in LLM tool outputs and RAG results.
|
||||
# Unformatted 16-digit strings are indistinguishable from non-card IDs at regex level.
|
||||
# The 4-group pattern covers real-world formatted card numbers with near-zero FP rate.
|
||||
(re.compile(r"\b(?:\d{4}[ -]){3}\d{4}\b"), "[REDACTED-CC]"),
|
||||
# Email (optional — often safe, but default-scrub for compliance)
|
||||
(re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b"), "[REDACTED-EMAIL]"),
|
||||
)
|
||||
|
||||
|
||||
def scrub_pii_text(text: str) -> str:
|
||||
"""Redact API keys, passwords, SSNs, credit cards, and emails from free text.
|
||||
|
||||
Reuses REDACT_KEYS for ``key=value`` pairs and adds value-pattern rules
|
||||
for naked secrets. Zero dependency (stdlib re only).
|
||||
|
||||
Args:
|
||||
text: Input string — returned unchanged if empty or None.
|
||||
|
||||
Returns:
|
||||
Scrubbed string with literal secrets replaced by ``[REDACTED*]``.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
result = redact_string(text, enabled=True)
|
||||
for pat, repl in _VALUE_PATTERNS:
|
||||
result = pat.sub(repl, result)
|
||||
return result
|
||||
|
||||
|
||||
# Module-level state for idempotent enable/disable
|
||||
_PII_HOOK_ID: Optional[str] = None
|
||||
_PII_HOOK_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def _pii_before_llm_hook(event_data):
|
||||
"""BEFORE_LLM hook that scrubs every message's content in-place."""
|
||||
try:
|
||||
messages = getattr(event_data, "messages", None) or []
|
||||
modified = []
|
||||
for m in messages:
|
||||
if isinstance(m, dict) and isinstance(m.get("content"), str):
|
||||
m = {**m, "content": scrub_pii_text(m["content"])}
|
||||
modified.append(m)
|
||||
# Import locally to avoid module-level hooks dep (keeps redact.py lightweight)
|
||||
from ..hooks.types import HookResult
|
||||
return HookResult(decision="allow", modified_input={"messages": modified})
|
||||
except Exception:
|
||||
# Never block the LLM call on a scrubber bug
|
||||
from ..hooks.types import HookResult
|
||||
return HookResult(decision="allow")
|
||||
|
||||
|
||||
def enable_pii_redaction() -> str:
|
||||
"""Register a BEFORE_LLM hook that scrubs secrets from every message.
|
||||
|
||||
Idempotent — calling twice leaves exactly one hook registered.
|
||||
|
||||
Returns:
|
||||
The hook id (useful for :func:`disable_pii_redaction`).
|
||||
"""
|
||||
global _PII_HOOK_ID
|
||||
with _PII_HOOK_LOCK:
|
||||
if _PII_HOOK_ID is not None:
|
||||
return _PII_HOOK_ID
|
||||
from ..hooks.registry import get_default_registry
|
||||
from ..hooks.types import HookEvent
|
||||
reg = get_default_registry()
|
||||
_PII_HOOK_ID = reg.register_function(
|
||||
event=HookEvent.BEFORE_LLM,
|
||||
func=_pii_before_llm_hook,
|
||||
name="praisonaiagents.pii_redactor",
|
||||
)
|
||||
return _PII_HOOK_ID
|
||||
|
||||
|
||||
def disable_pii_redaction() -> bool:
|
||||
"""Unregister the PII-redaction hook. No-op if never enabled."""
|
||||
global _PII_HOOK_ID
|
||||
if _PII_HOOK_ID is None:
|
||||
return False
|
||||
from ..hooks.registry import get_default_registry
|
||||
reg = get_default_registry()
|
||||
try:
|
||||
reg.unregister(_PII_HOOK_ID)
|
||||
finally:
|
||||
_PII_HOOK_ID = None
|
||||
return True
|
||||
|
||||
|
||||
def redact_string(text: str, enabled: bool = True) -> str:
|
||||
"""
|
||||
Redact potential secrets from a string.
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "1.6.37"
|
||||
version = "1.6.35"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
@@ -30,7 +30,7 @@ mcp = [
|
||||
|
||||
memory = [
|
||||
"chromadb>=1.0.0",
|
||||
"litellm>=1.83.14",
|
||||
"litellm>=1.81.0,<=1.82.6",
|
||||
]
|
||||
|
||||
knowledge = [
|
||||
@@ -48,7 +48,7 @@ graph = [
|
||||
|
||||
# Add LLM dependencies
|
||||
llm = [
|
||||
"litellm>=1.83.14",
|
||||
"litellm>=1.81.0,<=1.82.6",
|
||||
"pydantic>=2.10.0"
|
||||
]
|
||||
|
||||
@@ -113,9 +113,7 @@ all = [
|
||||
"praisonaiagents[mongodb]",
|
||||
"praisonaiagents[auth]",
|
||||
"praisonaiagents[search]",
|
||||
"praisonaiagents[crawl]",
|
||||
"praisonaiagents[autonomy]",
|
||||
"praisonaiagents[os]"
|
||||
"praisonaiagents[crawl]"
|
||||
]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
|
||||
@@ -1,158 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Validate subscription auth fixes without pytest dependency.
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add current directory to Python path
|
||||
sys.path.insert(0, '.')
|
||||
sys.path.insert(0, os.path.join(os.getcwd()))
|
||||
|
||||
def test_imports():
|
||||
"""Test that imports work correctly."""
|
||||
print("Testing imports...")
|
||||
try:
|
||||
from praisonaiagents.auth.subscription.protocols import SubscriptionCredentials, AuthError
|
||||
print("✅ Protocol imports work")
|
||||
|
||||
from praisonaiagents.auth.subscription.claude_code import ClaudeCodeAuth
|
||||
print("✅ Claude Code auth imports work")
|
||||
|
||||
from praisonaiagents.auth.subscription.registry import resolve_subscription_credentials
|
||||
print("✅ Registry imports work")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Import test failed: {e}")
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def test_codex_experimental():
|
||||
"""Test that Codex raises experimental error."""
|
||||
print("\nTesting Codex experimental error...")
|
||||
try:
|
||||
from praisonaiagents.auth.subscription.codex import CodexAuth
|
||||
auth = CodexAuth()
|
||||
try:
|
||||
auth.resolve_credentials()
|
||||
print("❌ Codex should have raised experimental error")
|
||||
return False
|
||||
except Exception as e:
|
||||
if "experimental" in str(e).lower():
|
||||
print("✅ Codex correctly raises experimental error")
|
||||
return True
|
||||
else:
|
||||
print(f"❌ Codex raised wrong error: {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"❌ Codex test failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_gemini_experimental():
|
||||
"""Test that Gemini CLI raises experimental error."""
|
||||
print("\nTesting Gemini CLI experimental error...")
|
||||
try:
|
||||
from praisonaiagents.auth.subscription.gemini_cli import GeminiCliAuth
|
||||
auth = GeminiCliAuth()
|
||||
try:
|
||||
auth.resolve_credentials()
|
||||
print("❌ Gemini CLI should have raised experimental error")
|
||||
return False
|
||||
except Exception as e:
|
||||
if "experimental" in str(e).lower():
|
||||
print("✅ Gemini CLI correctly raises experimental error")
|
||||
return True
|
||||
else:
|
||||
print(f"❌ Gemini CLI raised wrong error: {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"❌ Gemini CLI test failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_header_simplification():
|
||||
"""Test that Claude Code headers are simplified."""
|
||||
print("\nTesting Claude Code header simplification...")
|
||||
try:
|
||||
from praisonaiagents.auth.subscription.claude_code import ClaudeCodeAuth
|
||||
auth = ClaudeCodeAuth()
|
||||
headers = auth.headers_for("https://api.anthropic.com", "claude-3-haiku")
|
||||
|
||||
# Should only have interleaved-thinking, not oauth header (litellm handles that)
|
||||
if "anthropic-beta" in headers:
|
||||
beta_header = headers["anthropic-beta"]
|
||||
if beta_header == "interleaved-thinking-2025-05-14":
|
||||
print("✅ Claude Code headers correctly simplified")
|
||||
return True
|
||||
else:
|
||||
print(f"❌ Unexpected beta header value: {beta_header}")
|
||||
return False
|
||||
else:
|
||||
print("❌ anthropic-beta header missing")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Header test failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_agent_construction():
|
||||
"""Test basic agent construction with auth parameter."""
|
||||
print("\nTesting Agent construction with auth...")
|
||||
try:
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# This should work (though credential resolution will fail without actual creds)
|
||||
agent = Agent(name="test", auth="claude-code")
|
||||
print("✅ Agent construction with auth parameter works")
|
||||
|
||||
# Test invalid provider
|
||||
try:
|
||||
agent = Agent(name="test", auth="invalid-provider")
|
||||
print("⚠️ Invalid provider didn't fail at construction (will fail at runtime)")
|
||||
except Exception:
|
||||
print("✅ Invalid provider handling works")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Agent construction test failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all validation tests."""
|
||||
print("=== Validating Subscription Auth Fixes ===\n")
|
||||
|
||||
tests = [
|
||||
test_imports,
|
||||
test_codex_experimental,
|
||||
test_gemini_experimental,
|
||||
test_header_simplification,
|
||||
test_agent_construction,
|
||||
]
|
||||
|
||||
passed = 0
|
||||
total = len(tests)
|
||||
|
||||
for test in tests:
|
||||
if test():
|
||||
passed += 1
|
||||
else:
|
||||
print()
|
||||
|
||||
print(f"\n=== Results: {passed}/{total} tests passed ===")
|
||||
|
||||
if passed == total:
|
||||
print("🎉 All validation tests passed!")
|
||||
return True
|
||||
else:
|
||||
print("💥 Some tests failed!")
|
||||
return False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
success = main()
|
||||
sys.exit(0 if success else 1)
|
||||
@@ -1,69 +0,0 @@
|
||||
"""Integration tests for Agent with subscription auth."""
|
||||
import pytest
|
||||
|
||||
# Skip entire module when LLM dependencies (e.g. pydantic) are not installed.
|
||||
# CI runs without the [llm] extra, so these are skipped there.
|
||||
pytest.importorskip("pydantic", reason="requires praisonaiagents[llm]")
|
||||
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.auth.subscription.protocols import SubscriptionCredentials
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_subscription_auth():
|
||||
"""Mock subscription auth that returns test credentials."""
|
||||
def mock_resolve_credentials(provider_id):
|
||||
if provider_id == "claude-code":
|
||||
return SubscriptionCredentials(
|
||||
api_key="sk-ant-oat-stub",
|
||||
base_url="https://api.anthropic.com",
|
||||
headers={
|
||||
"anthropic-beta": "oauth-2025-04-20,interleaved-thinking-2025-05-14",
|
||||
"user-agent": "claude-cli/2.1.0 (external, cli)",
|
||||
"x-app": "cli",
|
||||
},
|
||||
auth_scheme="bearer",
|
||||
source="test-mock"
|
||||
)
|
||||
else:
|
||||
raise Exception(f"Unknown provider: {provider_id}")
|
||||
|
||||
return mock_resolve_credentials
|
||||
|
||||
|
||||
def test_agent_accepts_auth_kwarg(mock_subscription_auth):
|
||||
"""Test that Agent accepts auth parameter and passes it to LLM."""
|
||||
with patch('praisonaiagents.auth.resolve_subscription_credentials', mock_subscription_auth):
|
||||
agent = Agent(
|
||||
name="test-agent",
|
||||
llm="anthropic/claude-3-haiku-20240307",
|
||||
auth="claude-code"
|
||||
)
|
||||
|
||||
assert hasattr(agent, 'llm_instance')
|
||||
assert agent.llm_instance._auth_provider_id == "claude-code"
|
||||
|
||||
|
||||
def test_llm_kwargs_carry_oauth_headers(mock_subscription_auth):
|
||||
"""Test that LLM completion params include OAuth headers."""
|
||||
with patch('praisonaiagents.auth.resolve_subscription_credentials', mock_subscription_auth):
|
||||
from praisonaiagents.llm import LLM
|
||||
|
||||
llm = LLM(model="anthropic/claude-3-haiku-20240307", auth="claude-code")
|
||||
params = llm._build_completion_params()
|
||||
|
||||
assert params["api_key"] == "sk-ant-oat-stub"
|
||||
assert params["base_url"] == "https://api.anthropic.com"
|
||||
assert params["extra_headers"]["x-app"] == "cli"
|
||||
assert params["extra_headers"]["user-agent"] == "claude-cli/2.1.0 (external, cli)"
|
||||
|
||||
|
||||
def test_agent_without_auth_works_normally():
|
||||
"""Test that Agent without auth parameter works as before."""
|
||||
agent = Agent(name="test-agent", llm="anthropic/claude-3-haiku-20240307")
|
||||
|
||||
if hasattr(agent, 'llm_instance'):
|
||||
assert agent.llm_instance._auth_provider_id is None
|
||||
assert agent.llm_instance._cached_subscription_creds is None
|
||||
@@ -1,167 +0,0 @@
|
||||
"""
|
||||
Smoke test for subscription authentication - exercises real auth flows.
|
||||
"""
|
||||
import pytest
|
||||
from unittest.mock import patch, Mock
|
||||
|
||||
try:
|
||||
import litellm
|
||||
except ImportError:
|
||||
pytest.skip("litellm not available", allow_module_level=True)
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.auth.subscription.protocols import SubscriptionCredentials
|
||||
|
||||
|
||||
def test_claude_code_auth_flow():
|
||||
"""Test that Claude Code auth properly injects OAuth credentials."""
|
||||
# Mock litellm completion to capture the actual params sent
|
||||
mock_response = Mock()
|
||||
mock_response.choices = [Mock()]
|
||||
mock_response.choices[0].message = Mock()
|
||||
mock_response.choices[0].message.content = "Hello! This is a test response."
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
mock_completion.return_value = mock_response
|
||||
|
||||
# Mock the credential resolution to return OAuth token
|
||||
mock_creds = SubscriptionCredentials(
|
||||
api_key="sk-ant-oat-test-oauth-token", # OAuth token format
|
||||
base_url="https://api.anthropic.com",
|
||||
headers={"user-agent": "claude-cli/2.1.0", "x-app": "cli"},
|
||||
auth_scheme="bearer",
|
||||
source="claude-code-test"
|
||||
)
|
||||
|
||||
with patch('praisonaiagents.auth.resolve_subscription_credentials') as mock_resolve:
|
||||
mock_resolve.return_value = mock_creds
|
||||
|
||||
# Create agent with Claude Code auth
|
||||
agent = Agent(
|
||||
name="test-agent",
|
||||
instructions="You are a test assistant.",
|
||||
auth="claude-code"
|
||||
)
|
||||
|
||||
# Make a real request
|
||||
response = agent.start("Say hello")
|
||||
|
||||
# Verify the call was made
|
||||
assert mock_completion.called
|
||||
call_args = mock_completion.call_args
|
||||
params = call_args.kwargs if call_args.kwargs else {}
|
||||
|
||||
# Verify OAuth token was passed as api_key (litellm will auto-detect and use Bearer)
|
||||
assert params.get("api_key") == "sk-ant-oat-test-oauth-token"
|
||||
assert params.get("base_url") == "https://api.anthropic.com"
|
||||
|
||||
# Verify provider headers are included
|
||||
extra_headers = params.get("extra_headers", {})
|
||||
assert "user-agent" in extra_headers
|
||||
assert "x-app" in extra_headers
|
||||
|
||||
# Verify we got a response (just check that completion was called with right params)
|
||||
assert params.get("api_key") == "sk-ant-oat-test-oauth-token"
|
||||
|
||||
|
||||
def test_auth_refresh_on_401():
|
||||
"""Test that auth errors trigger credential refresh and retry."""
|
||||
# First call: return 401, second call: success
|
||||
mock_response = Mock()
|
||||
mock_response.choices = [Mock()]
|
||||
mock_response.choices[0].message = Mock()
|
||||
mock_response.choices[0].message.content = "Success after refresh"
|
||||
|
||||
auth_error = Exception("AuthenticationError: Invalid API key")
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
# First call raises auth error, second succeeds
|
||||
mock_completion.side_effect = [auth_error, mock_response]
|
||||
|
||||
# Mock credential resolution and refresh
|
||||
initial_creds = SubscriptionCredentials(
|
||||
api_key="sk-ant-oat-expired-token",
|
||||
base_url="https://api.anthropic.com",
|
||||
headers={},
|
||||
auth_scheme="bearer",
|
||||
source="claude-code-test"
|
||||
)
|
||||
|
||||
refreshed_creds = SubscriptionCredentials(
|
||||
api_key="sk-ant-oat-fresh-token",
|
||||
base_url="https://api.anthropic.com",
|
||||
headers={},
|
||||
auth_scheme="bearer",
|
||||
source="claude-code-refreshed"
|
||||
)
|
||||
|
||||
with patch('praisonaiagents.auth.resolve_subscription_credentials') as mock_resolve:
|
||||
mock_resolve.return_value = initial_creds
|
||||
|
||||
with patch('praisonaiagents.auth.subscription.registry.get_subscription_provider') as mock_provider:
|
||||
mock_auth_provider = Mock()
|
||||
mock_auth_provider.refresh.return_value = refreshed_creds
|
||||
mock_provider.return_value = mock_auth_provider
|
||||
|
||||
# Mock error classification to detect auth error
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
with patch.object(LLM, '_classify_error_and_should_retry') as mock_classify:
|
||||
# First call: auth error, can retry
|
||||
# Second call: shouldn't be called since retry succeeds
|
||||
mock_classify.return_value = ("auth", True, 0.0)
|
||||
|
||||
agent = Agent(
|
||||
name="test-agent",
|
||||
instructions="Test agent",
|
||||
auth="claude-code"
|
||||
)
|
||||
|
||||
# This should trigger refresh on first 401 and succeed on retry
|
||||
response = agent.start("Test refresh flow")
|
||||
|
||||
# Verify refresh was called
|
||||
assert mock_auth_provider.refresh.called
|
||||
|
||||
# Verify we got success response
|
||||
assert "Success after refresh" in response
|
||||
|
||||
|
||||
def test_invalid_auth_provider():
|
||||
"""Test that invalid auth providers raise clear errors."""
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
agent = Agent(
|
||||
name="test-agent",
|
||||
auth="nonexistent-provider"
|
||||
)
|
||||
agent.start("This should fail")
|
||||
|
||||
error_msg = str(exc_info.value).lower()
|
||||
assert "unknown" in error_msg or "provider" in error_msg
|
||||
|
||||
|
||||
def test_codex_experimental_error():
|
||||
"""Test that Codex auth raises experimental error."""
|
||||
from praisonaiagents.auth.subscription.protocols import AuthError
|
||||
with pytest.raises(AuthError) as exc_info:
|
||||
agent = Agent(name="test", auth="codex")
|
||||
agent.start("test")
|
||||
assert "experimental" in str(exc_info.value).lower()
|
||||
|
||||
|
||||
def test_gemini_experimental_error():
|
||||
"""Test that Gemini CLI auth raises experimental error."""
|
||||
from praisonaiagents.auth.subscription.protocols import AuthError
|
||||
with pytest.raises(AuthError) as exc_info:
|
||||
agent = Agent(name="test", auth="gemini-cli")
|
||||
agent.start("test")
|
||||
assert "experimental" in str(exc_info.value).lower()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Manual test runner for development
|
||||
test_claude_code_auth_flow()
|
||||
test_auth_refresh_on_401()
|
||||
test_invalid_auth_provider()
|
||||
test_codex_experimental_error()
|
||||
test_gemini_experimental_error()
|
||||
print("✅ All subscription auth smoke tests passed!")
|
||||
@@ -1,131 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Real agentic E2E — Phase A (C1/C2/C4/C7).
|
||||
|
||||
Exercises all four gap closures in a single live LLM run:
|
||||
|
||||
* C1 seed — request deterministic output
|
||||
* C2 cancel — pass a non-tripped InterruptController
|
||||
* C4 validator — install a passthrough ToolValidator (asserts it was invoked)
|
||||
* C7 PII — enable_pii_redaction() so secrets in the prompt are scrubbed
|
||||
|
||||
Run::
|
||||
|
||||
ANTHROPIC_API_KEY=... python tests/smoke_core_phase_a_real.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def _pick_model() -> str:
|
||||
if os.getenv("ANTHROPIC_API_KEY"):
|
||||
return "anthropic/claude-haiku-4-5"
|
||||
if os.getenv("GOOGLE_API_KEY"):
|
||||
return "gemini/gemini-2.5-flash"
|
||||
return "gpt-4o-mini"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.agent.interrupt import InterruptController
|
||||
from praisonaiagents.tools.validators import ValidationResult
|
||||
from praisonaiagents.trace.redact import (
|
||||
enable_pii_redaction,
|
||||
disable_pii_redaction,
|
||||
)
|
||||
|
||||
model = _pick_model()
|
||||
print(f"Model: {model}\n")
|
||||
|
||||
# C7 — register the PII scrubber for the whole run
|
||||
enable_pii_redaction()
|
||||
|
||||
# A trivial tool the LLM may or may not call — mainly to exercise the
|
||||
# validator wiring. We track how often it is asked to validate.
|
||||
validated: list[str] = []
|
||||
|
||||
class LoggingValidator:
|
||||
def validate_args(self, tool_name, args, context=None):
|
||||
validated.append(tool_name)
|
||||
return ValidationResult(valid=True)
|
||||
|
||||
def validate_result(self, tool_name, result, context=None):
|
||||
return ValidationResult(valid=True)
|
||||
|
||||
def echo(text: str) -> str:
|
||||
"""Echo a string back."""
|
||||
return f"echoed: {text}"
|
||||
|
||||
agent = Agent(
|
||||
name="PhaseATester",
|
||||
instructions=(
|
||||
"You are a concise assistant. Answer with one short sentence. "
|
||||
"Feel free to call the echo tool if asked."
|
||||
),
|
||||
llm={"model": model},
|
||||
tools=[echo],
|
||||
)
|
||||
agent._tool_validator = LoggingValidator() # C4 wiring
|
||||
|
||||
tok = InterruptController() # C2 — not set, must not affect anything
|
||||
|
||||
# Prompt contains a fake secret — C7 must scrub it before LLM egress.
|
||||
prompt = (
|
||||
"My api_key=sk-XYZSECRET12345 — please answer: what is 2 + 2? "
|
||||
"Reply with a single digit."
|
||||
)
|
||||
|
||||
print(f"[prompt raw] {prompt}")
|
||||
reply = agent.chat(prompt, seed=42, cancel_token=tok)
|
||||
print(f"[reply] {reply}\n")
|
||||
|
||||
# Validate outcomes
|
||||
ok = True
|
||||
|
||||
if not reply:
|
||||
print("FAIL: LLM returned empty reply")
|
||||
ok = False
|
||||
elif "4" not in reply:
|
||||
print("FAIL: reply does not contain expected '4'")
|
||||
ok = False
|
||||
else:
|
||||
print("PASS C1 — seed threaded and LLM returned a valid reply")
|
||||
|
||||
# Confirm the scrubber registered its hook (C7)
|
||||
# (chat_history stores the user's raw prompt by design — scrubbing happens
|
||||
# only on egress to the LLM, verified here via hook registry state.)
|
||||
from praisonaiagents.hooks import get_default_registry
|
||||
reg = get_default_registry()
|
||||
has_pii = any(
|
||||
(getattr(h, "name", "") or "").startswith("praisonaiagents.pii_redactor")
|
||||
for ev in getattr(reg, "_hooks", {}).values()
|
||||
for h in (ev if isinstance(ev, list) else [])
|
||||
)
|
||||
if has_pii:
|
||||
print("PASS C7 — PII redactor hook active during LLM call")
|
||||
else:
|
||||
print("FAIL C7 — PII redactor hook missing from registry")
|
||||
ok = False
|
||||
|
||||
# Cancel-token path did not trip (C2)
|
||||
if tok.is_set():
|
||||
print("FAIL C2 — cancel token unexpectedly set")
|
||||
ok = False
|
||||
else:
|
||||
print("PASS C2 — cancel_token honoured without tripping")
|
||||
|
||||
# Validator invocation (C4) is exercised only if LLM actually called echo;
|
||||
# don't fail the run if it chose not to call a tool.
|
||||
if validated:
|
||||
print(f"PASS C4 — validator invoked for: {validated}")
|
||||
else:
|
||||
print("SKIP C4 — LLM chose not to call a tool this run (non-fatal)")
|
||||
|
||||
disable_pii_redaction()
|
||||
return 0 if ok else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,184 +0,0 @@
|
||||
"""Phase A gap-closure tests: C1 seed, C2 cancel_token, C4 tool validator, C7 PII.
|
||||
|
||||
These tests exercise the four gap closures identified in the core-features gap
|
||||
analysis. They target the existing infrastructure (InterruptController,
|
||||
ToolValidatorProtocol, trace/redact rules) which was defined but not wired into
|
||||
`Agent.chat()` / tool execution paths.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.tools.validators import (
|
||||
ToolValidatorProtocol,
|
||||
ValidationResult,
|
||||
)
|
||||
|
||||
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
# C1 — seed threaded through Agent.chat()
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
class TestC1SeedPassthrough:
|
||||
def test_chat_signature_accepts_seed(self):
|
||||
"""`Agent.chat` must accept a per-call `seed=` kwarg."""
|
||||
import inspect
|
||||
sig = inspect.signature(Agent.chat)
|
||||
assert "seed" in sig.parameters, (
|
||||
"Agent.chat() must expose seed= to let callers request determinism"
|
||||
)
|
||||
|
||||
def test_seed_reaches_llm_build_params(self, monkeypatch):
|
||||
"""A `seed=42` passed to chat() must appear in the LLM kwargs."""
|
||||
captured = {}
|
||||
|
||||
# Use dict-form llm= so the custom-LLM path (with llm_instance) is taken.
|
||||
agent = Agent(name="C1", instructions="say hi", llm={"model": "gpt-4o-mini"})
|
||||
assert getattr(agent, "llm_instance", None) is not None, \
|
||||
"dict llm= must produce a materialised llm_instance"
|
||||
|
||||
# Short-circuit the LLM call so nothing hits the network
|
||||
monkeypatch.setattr(
|
||||
agent.llm_instance,
|
||||
"get_response",
|
||||
lambda **kw: (captured.update(kw) or "ok"),
|
||||
)
|
||||
|
||||
agent.chat("hello", seed=42)
|
||||
assert captured.get("seed") == 42, (
|
||||
f"seed=42 did not reach LLM call; kwargs were {list(captured)}"
|
||||
)
|
||||
|
||||
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
# C2 — cancel_token wired into chat()
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
class TestC2CancelToken:
|
||||
def test_chat_signature_accepts_cancel_token(self):
|
||||
import inspect
|
||||
sig = inspect.signature(Agent.chat)
|
||||
assert "cancel_token" in sig.parameters, (
|
||||
"Agent.chat() must expose cancel_token= for cooperative cancellation"
|
||||
)
|
||||
|
||||
def test_preset_cancelled_token_aborts_chat_before_llm(self, monkeypatch):
|
||||
"""If the cancel token is already set when chat() is called,
|
||||
the LLM is NEVER invoked and the call returns cleanly (None)."""
|
||||
from praisonaiagents.agent.interrupt import InterruptController
|
||||
|
||||
agent = Agent(name="C2", instructions="irrelevant", llm={"model": "gpt-4o-mini"})
|
||||
|
||||
called = {"n": 0}
|
||||
|
||||
def must_not_be_called(**kwargs):
|
||||
called["n"] += 1
|
||||
return "should not be reached"
|
||||
|
||||
monkeypatch.setattr(agent.llm_instance, "get_response", must_not_be_called)
|
||||
|
||||
tok = InterruptController()
|
||||
tok.request("user-cancel")
|
||||
|
||||
try:
|
||||
agent.chat("anything", cancel_token=tok)
|
||||
except InterruptedError:
|
||||
pass # acceptable outcome
|
||||
assert called["n"] == 0, (
|
||||
"LLM was invoked despite pre-set cancel_token — cancellation not wired"
|
||||
)
|
||||
|
||||
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
# C4 — ToolValidatorProtocol wired into tool execution
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
class _RejectingValidator:
|
||||
"""Minimal ToolValidatorProtocol impl that always rejects."""
|
||||
|
||||
def validate_args(self, tool_name, args, context=None):
|
||||
return ValidationResult(
|
||||
valid=False,
|
||||
errors=[f"arg-validator rejected {tool_name}"],
|
||||
remediation="Fix your tool-call schema.",
|
||||
)
|
||||
|
||||
def validate_result(self, tool_name, result, context=None):
|
||||
return ValidationResult(valid=True)
|
||||
|
||||
|
||||
class TestC4ToolValidator:
|
||||
def test_validator_implements_protocol(self):
|
||||
"""Sanity: our fixture validator satisfies the protocol."""
|
||||
v = _RejectingValidator()
|
||||
assert isinstance(v, ToolValidatorProtocol)
|
||||
|
||||
def test_validator_rejects_before_tool_runs(self):
|
||||
"""When agent._tool_validator rejects, the tool function is NEVER called
|
||||
and the returned string contains the validator's error."""
|
||||
|
||||
calls = {"n": 0}
|
||||
|
||||
def fake_tool(x: int) -> int:
|
||||
calls["n"] += 1
|
||||
return x * 2
|
||||
|
||||
agent = Agent(name="C4", instructions="…", tools=[fake_tool], llm="gpt-4o-mini")
|
||||
agent._tool_validator = _RejectingValidator()
|
||||
|
||||
result = agent.execute_tool("fake_tool", {"x": 5})
|
||||
|
||||
assert calls["n"] == 0, "tool function was executed despite validator rejection"
|
||||
assert "rejected" in str(result).lower(), (
|
||||
f"validator error not surfaced: {result!r}"
|
||||
)
|
||||
|
||||
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
# C7 — enable_pii_redaction() registers a BEFORE_LLM middleware
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
class TestC7PIIRedaction:
|
||||
def test_enable_pii_redaction_is_importable(self):
|
||||
from praisonaiagents.trace.redact import enable_pii_redaction
|
||||
assert callable(enable_pii_redaction)
|
||||
|
||||
def test_enable_pii_redaction_is_idempotent(self):
|
||||
"""Calling enable_pii_redaction() twice must not double-register."""
|
||||
from praisonaiagents.trace.redact import (
|
||||
enable_pii_redaction,
|
||||
disable_pii_redaction,
|
||||
)
|
||||
disable_pii_redaction() # ensure clean baseline
|
||||
enable_pii_redaction()
|
||||
first_state = _pii_state_snapshot()
|
||||
enable_pii_redaction()
|
||||
second_state = _pii_state_snapshot()
|
||||
assert first_state == second_state, "enable_pii_redaction is not idempotent"
|
||||
disable_pii_redaction()
|
||||
|
||||
def test_pii_scrubber_scrubs_api_keys(self):
|
||||
"""The scrubber function must redact API-key-shaped strings."""
|
||||
from praisonaiagents.trace.redact import scrub_pii_text
|
||||
|
||||
raw = "my api_key=sk-ABCDEF12345 and password=hunter2"
|
||||
scrubbed = scrub_pii_text(raw)
|
||||
assert "sk-ABCDEF12345" not in scrubbed
|
||||
assert "hunter2" not in scrubbed
|
||||
assert "[REDACTED]" in scrubbed
|
||||
|
||||
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
# Helpers
|
||||
# ────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _pii_state_snapshot():
|
||||
"""Snapshot of hook registry state used by C7 idempotency test."""
|
||||
from praisonaiagents.hooks import get_default_registry
|
||||
reg = get_default_registry()
|
||||
# Count hooks registered for BEFORE_LLM — we just want stable identity
|
||||
hooks = getattr(reg, "_hooks", None) or getattr(reg, "hooks", None) or {}
|
||||
if isinstance(hooks, dict):
|
||||
bl = hooks.get("before_llm") or hooks.get("BEFORE_LLM") or []
|
||||
return len(bl)
|
||||
return None
|
||||
@@ -1,101 +0,0 @@
|
||||
"""Regression tests for GHSA-gmjg-hv98-qggq:
|
||||
Undeclared __main__ callables must NOT be executed via execute_tool.
|
||||
"""
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from praisonaiagents.agent.tool_execution import ToolExecutionMixin
|
||||
|
||||
|
||||
class _HookRunner:
|
||||
def execute_sync(self, *args, **kwargs):
|
||||
return []
|
||||
|
||||
def is_blocked(self, results):
|
||||
return False
|
||||
|
||||
|
||||
class _DummyAgent(ToolExecutionMixin):
|
||||
def __init__(self, tools=None):
|
||||
self.name = "test_agent"
|
||||
self.tools = tools or []
|
||||
self.chat_history = []
|
||||
self._hook_runner = _HookRunner()
|
||||
self.context_manager = None
|
||||
self._doom_loop_tracker = None
|
||||
self._perm_deny = frozenset()
|
||||
self._perm_allow = None
|
||||
self._approval_backend = None
|
||||
|
||||
|
||||
def _make_approved_registry():
|
||||
reg = MagicMock()
|
||||
reg.approve_sync.return_value = MagicMock(
|
||||
approved=True, reason="mock", modified_args=None
|
||||
)
|
||||
reg.mark_approved = MagicMock()
|
||||
return reg
|
||||
|
||||
|
||||
def _sneaky_undeclared(msg="pwned"):
|
||||
"""Callable present only in __main__, never declared as a tool."""
|
||||
return {"ran": msg}
|
||||
|
||||
|
||||
def test_undeclared_main_callable_not_executed():
|
||||
"""An undeclared callable in __main__ must NOT be invoked."""
|
||||
# Inject the function into __main__ to simulate the exploit scenario.
|
||||
original = getattr(sys.modules["__main__"], "_sneaky_undeclared", None)
|
||||
sys.modules["__main__"]._sneaky_undeclared = _sneaky_undeclared
|
||||
try:
|
||||
reg = _make_approved_registry()
|
||||
with patch("praisonaiagents.approval.get_approval_registry", return_value=reg):
|
||||
agent = _DummyAgent(tools=[])
|
||||
result = agent.execute_tool("_sneaky_undeclared", {"msg": "hello"})
|
||||
# Should NOT have run the function — result must indicate not-found, not success.
|
||||
assert result != {"ran": "hello"}, (
|
||||
"Undeclared __main__ callable was executed — vulnerability still present"
|
||||
)
|
||||
finally:
|
||||
if original is None:
|
||||
delattr(sys.modules["__main__"], "_sneaky_undeclared")
|
||||
else:
|
||||
sys.modules["__main__"]._sneaky_undeclared = original
|
||||
|
||||
|
||||
def test_declared_tool_still_executes():
|
||||
"""A properly declared tool must still be callable after the fix."""
|
||||
def my_tool(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
reg = _make_approved_registry()
|
||||
with patch("praisonaiagents.approval.get_approval_registry", return_value=reg):
|
||||
agent = _DummyAgent(tools=[my_tool])
|
||||
result = agent.execute_tool("my_tool", {"x": 3})
|
||||
assert result == 6, f"Declared tool did not execute correctly: {result}"
|
||||
|
||||
|
||||
def test_globals_callable_not_executed():
|
||||
"""A callable only in tool_execution module globals must NOT be invoked."""
|
||||
reg = _make_approved_registry()
|
||||
with patch("praisonaiagents.approval.get_approval_registry", return_value=reg):
|
||||
agent = _DummyAgent(tools=[])
|
||||
# 'logging' is a module-level name in tool_execution; it must not be invocable.
|
||||
result = agent.execute_tool("logging", {})
|
||||
assert not isinstance(result, dict) or result.get("ran") is None, (
|
||||
"globals() fallback executed an unintended name"
|
||||
)
|
||||
|
||||
|
||||
def test_undeclared_returns_error_or_none():
|
||||
"""execute_tool with an unknown name should return None or an error dict, not raise."""
|
||||
reg = _make_approved_registry()
|
||||
with patch("praisonaiagents.approval.get_approval_registry", return_value=reg):
|
||||
agent = _DummyAgent(tools=[])
|
||||
result = agent.execute_tool("completely_nonexistent_xyz", {})
|
||||
# Result should be None or an error dict — not a successful execution.
|
||||
assert result is None or isinstance(result, dict), (
|
||||
f"Unexpected return type: {type(result)}"
|
||||
)
|
||||
if isinstance(result, dict):
|
||||
assert result.get("ran") is None
|
||||
@@ -1,107 +0,0 @@
|
||||
"""Tests for Claude Code subscription auth."""
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import tempfile
|
||||
import pytest
|
||||
from unittest.mock import patch, mock_open
|
||||
from pathlib import Path
|
||||
|
||||
from praisonaiagents.auth.subscription.claude_code import (
|
||||
ClaudeCodeAuth,
|
||||
_read_keychain_credentials,
|
||||
_read_file_credentials,
|
||||
_is_expiring,
|
||||
AuthError
|
||||
)
|
||||
|
||||
|
||||
def test_keychain_path_returns_none_off_darwin():
|
||||
"""Test keychain reading returns None on non-macOS platforms."""
|
||||
with patch.object(sys, 'platform', 'linux'):
|
||||
assert _read_keychain_credentials() is None
|
||||
|
||||
|
||||
def test_file_path_reads_credentials(monkeypatch):
|
||||
"""Test file credential reading."""
|
||||
# Create a temporary file with Claude Code credentials
|
||||
credentials_data = {
|
||||
"claudeAiOauth": {
|
||||
"accessToken": "sk-ant-oat-test",
|
||||
"refreshToken": "rt-test",
|
||||
"expiresAt": 9999999999000,
|
||||
}
|
||||
}
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tmp:
|
||||
json.dump(credentials_data, tmp)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
with patch('praisonaiagents.auth.subscription.claude_code.Path') as mock_path:
|
||||
mock_home_path = mock_path.home.return_value
|
||||
mock_claude_path = mock_home_path / ".claude" / ".credentials.json"
|
||||
mock_claude_path.exists.return_value = True
|
||||
mock_claude_path.read_text.return_value = json.dumps(credentials_data)
|
||||
|
||||
creds = _read_file_credentials()
|
||||
assert creds["accessToken"] == "sk-ant-oat-test"
|
||||
assert creds["refreshToken"] == "rt-test"
|
||||
assert creds["source"] == "claude-code-file"
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
|
||||
def test_is_expiring():
|
||||
"""Test token expiry detection."""
|
||||
import time
|
||||
current_ms = int(time.time() * 1000)
|
||||
|
||||
# Token expires in 30 seconds (should be considered expiring with default 60s skew)
|
||||
soon_expired = current_ms + 30_000
|
||||
assert _is_expiring(soon_expired) is True
|
||||
|
||||
# Token expires in 2 minutes (should not be expiring)
|
||||
not_expired = current_ms + 120_000
|
||||
assert _is_expiring(not_expired) is False
|
||||
|
||||
# No expiry time
|
||||
assert _is_expiring(0) is False
|
||||
assert _is_expiring(None) is False
|
||||
|
||||
|
||||
def test_resolve_credentials_uses_env_first(monkeypatch):
|
||||
"""Test that environment variables take precedence."""
|
||||
monkeypatch.setenv("ANTHROPIC_TOKEN", "sk-ant-oat-from-env")
|
||||
|
||||
auth = ClaudeCodeAuth()
|
||||
creds = auth.resolve_credentials()
|
||||
|
||||
assert creds.api_key == "sk-ant-oat-from-env"
|
||||
assert creds.base_url == "https://api.anthropic.com"
|
||||
assert creds.headers["x-app"] == "cli"
|
||||
assert creds.headers["user-agent"].startswith("claude-cli/")
|
||||
assert creds.source == "env:ANTHROPIC_TOKEN"
|
||||
|
||||
|
||||
def test_resolve_credentials_no_auth_raises_error():
|
||||
"""Test that missing credentials raises AuthError."""
|
||||
auth = ClaudeCodeAuth()
|
||||
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
with patch('praisonaiagents.auth.subscription.claude_code._read_keychain_credentials', return_value=None):
|
||||
with patch('praisonaiagents.auth.subscription.claude_code._read_file_credentials', return_value=None):
|
||||
with pytest.raises(AuthError) as exc_info:
|
||||
auth.resolve_credentials()
|
||||
|
||||
assert "No Claude Code credentials found" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_headers_for_includes_required_oauth_headers():
|
||||
"""Test that OAuth-specific headers are included."""
|
||||
auth = ClaudeCodeAuth()
|
||||
headers = auth.headers_for("https://api.anthropic.com", "claude-3-haiku")
|
||||
|
||||
assert headers["anthropic-beta"] == "oauth-2025-04-20,interleaved-thinking-2025-05-14"
|
||||
assert headers["user-agent"] == "claude-cli/2.1.0 (external, cli)"
|
||||
assert headers["x-app"] == "cli"
|
||||
@@ -1,38 +0,0 @@
|
||||
"""Tests for subscription auth protocols."""
|
||||
import pytest
|
||||
from praisonaiagents.auth.subscription.protocols import (
|
||||
SubscriptionCredentials,
|
||||
SubscriptionAuthProtocol,
|
||||
AuthError
|
||||
)
|
||||
|
||||
|
||||
def test_credentials_dataclass_defaults():
|
||||
"""Test SubscriptionCredentials dataclass defaults."""
|
||||
c = SubscriptionCredentials(api_key="test-key")
|
||||
assert c.api_key == "test-key"
|
||||
assert c.base_url == ""
|
||||
assert c.headers == {}
|
||||
assert c.auth_scheme == "bearer"
|
||||
assert c.expires_at_ms is None
|
||||
assert c.source == ""
|
||||
|
||||
|
||||
def test_protocol_runtime_check():
|
||||
"""Test that classes implementing the protocol are recognized."""
|
||||
class MockAuth:
|
||||
def resolve_credentials(self):
|
||||
return SubscriptionCredentials(api_key="test")
|
||||
def refresh(self):
|
||||
return SubscriptionCredentials(api_key="refreshed")
|
||||
def headers_for(self, base_url, model):
|
||||
return {"x-test": "header"}
|
||||
|
||||
assert isinstance(MockAuth(), SubscriptionAuthProtocol)
|
||||
|
||||
|
||||
def test_auth_error():
|
||||
"""Test AuthError exception."""
|
||||
with pytest.raises(AuthError) as exc_info:
|
||||
raise AuthError("test error")
|
||||
assert str(exc_info.value) == "test error"
|
||||
@@ -1,54 +0,0 @@
|
||||
"""Tests for subscription auth registry."""
|
||||
import pytest
|
||||
from praisonaiagents.auth.subscription.protocols import SubscriptionCredentials, AuthError
|
||||
from praisonaiagents.auth.subscription import registry as _reg
|
||||
from praisonaiagents.auth.subscription.registry import (
|
||||
register_subscription_provider,
|
||||
list_subscription_providers,
|
||||
resolve_subscription_credentials,
|
||||
_REGISTRY
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clean_registry():
|
||||
"""Clean the registry before and after each test, resetting the builtin flag.
|
||||
|
||||
Direct access to _REGISTRY and _BUILTIN_REGISTERED is intentional here:
|
||||
we are testing the registry implementation itself and need full state control
|
||||
for test isolation.
|
||||
"""
|
||||
original_registry = _REGISTRY.copy()
|
||||
original_builtin_flag = _reg._BUILTIN_REGISTERED
|
||||
_REGISTRY.clear()
|
||||
_reg._BUILTIN_REGISTERED = False
|
||||
yield
|
||||
_REGISTRY.clear()
|
||||
_REGISTRY.update(original_registry)
|
||||
_reg._BUILTIN_REGISTERED = original_builtin_flag
|
||||
|
||||
|
||||
def test_register_and_resolve():
|
||||
"""Test registering and resolving a provider."""
|
||||
class MockAuth:
|
||||
def resolve_credentials(self):
|
||||
return SubscriptionCredentials(api_key="test-key", source="mock")
|
||||
def refresh(self):
|
||||
return self.resolve_credentials()
|
||||
def headers_for(self, *_):
|
||||
return {}
|
||||
|
||||
register_subscription_provider("test-mock", lambda: MockAuth())
|
||||
assert "test-mock" in list_subscription_providers()
|
||||
|
||||
creds = resolve_subscription_credentials("test-mock")
|
||||
assert creds.api_key == "test-key"
|
||||
assert creds.source == "mock"
|
||||
|
||||
|
||||
def test_resolve_unknown_provider():
|
||||
"""Test resolving an unknown provider raises AuthError."""
|
||||
with pytest.raises(AuthError) as exc_info:
|
||||
resolve_subscription_credentials("unknown-provider")
|
||||
|
||||
assert "Unknown subscription provider 'unknown-provider'" in str(exc_info.value)
|
||||
@@ -1,164 +0,0 @@
|
||||
"""Tests for IdentityResolverProtocol and InMemoryIdentityResolver.
|
||||
|
||||
W1 — Cross-platform identity linking. An IdentityResolver maps
|
||||
``(platform, platform_user_id)`` → ``unified_user_id`` so that the same
|
||||
human can be recognised across Telegram, Discord, Slack, etc.
|
||||
|
||||
Linking is OPT-IN. By default the resolver returns the platform-prefixed
|
||||
ID unchanged (no surprises, no privacy leak).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.identity import (
|
||||
IdentityLink,
|
||||
IdentityResolverProtocol,
|
||||
InMemoryIdentityResolver,
|
||||
)
|
||||
|
||||
|
||||
class TestIdentityLink:
|
||||
def test_link_holds_platform_user_unified(self):
|
||||
link = IdentityLink(
|
||||
platform="telegram",
|
||||
platform_user_id="12345",
|
||||
unified_user_id="alice",
|
||||
)
|
||||
assert link.platform == "telegram"
|
||||
assert link.platform_user_id == "12345"
|
||||
assert link.unified_user_id == "alice"
|
||||
|
||||
def test_link_is_immutable(self):
|
||||
link = IdentityLink("telegram", "12345", "alice")
|
||||
with pytest.raises((AttributeError, Exception)):
|
||||
link.unified_user_id = "bob" # type: ignore[misc]
|
||||
|
||||
|
||||
class TestProtocolConformance:
|
||||
def test_in_memory_satisfies_protocol(self):
|
||||
resolver = InMemoryIdentityResolver()
|
||||
assert isinstance(resolver, IdentityResolverProtocol)
|
||||
|
||||
def test_protocol_has_required_methods(self):
|
||||
proto_methods = {"resolve", "link", "unlink", "links_for"}
|
||||
assert proto_methods.issubset(set(dir(IdentityResolverProtocol)))
|
||||
|
||||
|
||||
class TestInMemoryResolverDefaults:
|
||||
def test_unlinked_returns_platform_prefixed_id(self):
|
||||
"""When no link exists, return ``{platform}:{user_id}`` unchanged."""
|
||||
r = InMemoryIdentityResolver()
|
||||
assert r.resolve("telegram", "12345") == "telegram:12345"
|
||||
|
||||
def test_unlinked_does_not_create_link(self):
|
||||
r = InMemoryIdentityResolver()
|
||||
r.resolve("telegram", "12345")
|
||||
assert r.links_for("telegram:12345") == []
|
||||
|
||||
|
||||
class TestLinking:
|
||||
def test_link_two_platforms_resolves_to_same_unified(self):
|
||||
r = InMemoryIdentityResolver()
|
||||
r.link("telegram", "12345", "alice")
|
||||
r.link("discord", "snowflake-1", "alice")
|
||||
assert r.resolve("telegram", "12345") == "alice"
|
||||
assert r.resolve("discord", "snowflake-1") == "alice"
|
||||
|
||||
def test_unlink_reverts_to_default(self):
|
||||
r = InMemoryIdentityResolver()
|
||||
r.link("telegram", "12345", "alice")
|
||||
r.unlink("telegram", "12345")
|
||||
assert r.resolve("telegram", "12345") == "telegram:12345"
|
||||
|
||||
def test_links_for_returns_all_platforms(self):
|
||||
r = InMemoryIdentityResolver()
|
||||
r.link("telegram", "12345", "alice")
|
||||
r.link("discord", "snowflake-1", "alice")
|
||||
links = r.links_for("alice")
|
||||
assert len(links) == 2
|
||||
platforms = {link.platform for link in links}
|
||||
assert platforms == {"telegram", "discord"}
|
||||
|
||||
def test_link_overwrites_previous(self):
|
||||
r = InMemoryIdentityResolver()
|
||||
r.link("telegram", "12345", "alice")
|
||||
r.link("telegram", "12345", "bob")
|
||||
assert r.resolve("telegram", "12345") == "bob"
|
||||
|
||||
|
||||
class TestThreadSafety:
|
||||
def test_concurrent_links_dont_corrupt_state(self):
|
||||
import threading
|
||||
|
||||
r = InMemoryIdentityResolver()
|
||||
|
||||
def worker(i):
|
||||
r.link("telegram", f"u{i}", f"unified-{i}")
|
||||
|
||||
threads = [threading.Thread(target=worker, args=(i,)) for i in range(50)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
for i in range(50):
|
||||
assert r.resolve("telegram", f"u{i}") == f"unified-{i}"
|
||||
|
||||
|
||||
class TestFileIdentityResolver:
|
||||
def test_persists_across_instances(self, tmp_path):
|
||||
from praisonaiagents.session.identity import FileIdentityResolver
|
||||
|
||||
path = tmp_path / "identity.json"
|
||||
r1 = FileIdentityResolver(path=path)
|
||||
r1.link("telegram", "12345", "alice")
|
||||
r1.link("discord", "snowflake-1", "alice")
|
||||
|
||||
# New instance, same file
|
||||
r2 = FileIdentityResolver(path=path)
|
||||
assert r2.resolve("telegram", "12345") == "alice"
|
||||
assert r2.resolve("discord", "snowflake-1") == "alice"
|
||||
|
||||
def test_unlink_persists(self, tmp_path):
|
||||
from praisonaiagents.session.identity import FileIdentityResolver
|
||||
|
||||
path = tmp_path / "identity.json"
|
||||
r1 = FileIdentityResolver(path=path)
|
||||
r1.link("telegram", "12345", "alice")
|
||||
r1.unlink("telegram", "12345")
|
||||
|
||||
r2 = FileIdentityResolver(path=path)
|
||||
assert r2.resolve("telegram", "12345") == "telegram:12345"
|
||||
|
||||
def test_corrupt_file_does_not_crash(self, tmp_path):
|
||||
from praisonaiagents.session.identity import FileIdentityResolver
|
||||
|
||||
path = tmp_path / "identity.json"
|
||||
path.write_text("{not valid json", encoding="utf-8")
|
||||
r = FileIdentityResolver(path=path)
|
||||
# Falls back to empty state
|
||||
assert r.resolve("telegram", "12345") == "telegram:12345"
|
||||
|
||||
def test_satisfies_protocol(self, tmp_path):
|
||||
from praisonaiagents.session.identity import (
|
||||
FileIdentityResolver,
|
||||
IdentityResolverProtocol,
|
||||
)
|
||||
r = FileIdentityResolver(path=tmp_path / "identity.json")
|
||||
assert isinstance(r, IdentityResolverProtocol)
|
||||
|
||||
def test_file_permissions_restricted(self, tmp_path):
|
||||
import sys
|
||||
if sys.platform.startswith("win"):
|
||||
pytest.skip("chmod semantics differ on Windows")
|
||||
|
||||
from praisonaiagents.session.identity import FileIdentityResolver
|
||||
|
||||
path = tmp_path / "identity.json"
|
||||
r = FileIdentityResolver(path=path)
|
||||
r.link("telegram", "12345", "alice")
|
||||
# 0o600 = owner-only read/write
|
||||
mode = path.stat().st_mode & 0o777
|
||||
assert mode == 0o600
|
||||
@@ -1,88 +0,0 @@
|
||||
"""Tests for task-local SessionContextVars.
|
||||
|
||||
W1 — Concurrent message handlers must NOT clobber each other's
|
||||
session metadata. Use ``contextvars.ContextVar`` so each asyncio task
|
||||
has its own copy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.context import (
|
||||
SessionContext,
|
||||
clear_session_context,
|
||||
get_session_context,
|
||||
set_session_context,
|
||||
)
|
||||
|
||||
|
||||
class TestSetGet:
|
||||
def test_set_and_get_roundtrips(self):
|
||||
token = set_session_context(
|
||||
platform="telegram",
|
||||
chat_id="100",
|
||||
user_id="alice",
|
||||
)
|
||||
try:
|
||||
ctx = get_session_context()
|
||||
assert ctx.platform == "telegram"
|
||||
assert ctx.chat_id == "100"
|
||||
assert ctx.user_id == "alice"
|
||||
finally:
|
||||
clear_session_context(token)
|
||||
|
||||
def test_default_context_is_empty(self):
|
||||
ctx = get_session_context()
|
||||
assert ctx.platform == ""
|
||||
assert ctx.chat_id == ""
|
||||
assert ctx.user_id == ""
|
||||
|
||||
def test_clear_restores_previous(self):
|
||||
token = set_session_context(platform="telegram", user_id="alice")
|
||||
clear_session_context(token)
|
||||
assert get_session_context().platform == ""
|
||||
|
||||
|
||||
class TestAsyncTaskIsolation:
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_tasks_dont_clobber(self):
|
||||
seen = []
|
||||
|
||||
async def handler(name, delay):
|
||||
token = set_session_context(platform=name, user_id=name)
|
||||
try:
|
||||
await asyncio.sleep(delay)
|
||||
seen.append(get_session_context().platform)
|
||||
finally:
|
||||
clear_session_context(token)
|
||||
|
||||
await asyncio.gather(
|
||||
handler("telegram", 0.02),
|
||||
handler("discord", 0.01),
|
||||
handler("slack", 0.03),
|
||||
)
|
||||
assert sorted(seen) == ["discord", "slack", "telegram"]
|
||||
|
||||
|
||||
class TestSessionContextDataclass:
|
||||
def test_context_has_all_fields(self):
|
||||
ctx = SessionContext(
|
||||
platform="telegram",
|
||||
chat_id="100",
|
||||
chat_name="DM",
|
||||
thread_id="t1",
|
||||
user_id="alice",
|
||||
user_name="Alice",
|
||||
unified_user_id="alice-global",
|
||||
)
|
||||
assert ctx.unified_user_id == "alice-global"
|
||||
|
||||
def test_context_dict_roundtrip(self):
|
||||
ctx = SessionContext(platform="telegram", user_id="alice")
|
||||
d = ctx.to_dict()
|
||||
ctx2 = SessionContext.from_dict(d)
|
||||
assert ctx2.platform == "telegram"
|
||||
assert ctx2.user_id == "alice"
|
||||
Generated
+1
-1
@@ -2951,7 +2951,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.6.37"
|
||||
version = "1.6.35"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"version": "1.5.87",
|
||||
"lastUpdated": "2026-05-08",
|
||||
"lastUpdated": "2026-05-01",
|
||||
"generatedBy": "praisonai._dev.parity.generator",
|
||||
"sourceOfTruth": "Python SDK (praisonaiagents)",
|
||||
"status": "PARITY_ACHIEVED",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"version": "1.5.87",
|
||||
"lastUpdated": "2026-05-08",
|
||||
"lastUpdated": "2026-05-01",
|
||||
"generatedBy": "praisonai._dev.parity.generator",
|
||||
"sourceOfTruth": "Python SDK (praisonaiagents)",
|
||||
"summary": {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# Feature Parity Tracker
|
||||
|
||||
> **Version:** 1.5.87 | **Last Updated:** 2026-05-08
|
||||
> **Version:** 1.5.87 | **Last Updated:** 2026-05-01
|
||||
> **Source of Truth:** Python SDK (praisonaiagents)
|
||||
|
||||
## Summary
|
||||
|
||||
@@ -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.37.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.6.37.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.6.35.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.6.35.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
# Suppress crewai.cli.config logger BEFORE any imports to prevent INFO log
|
||||
import logging
|
||||
import threading
|
||||
logging.getLogger('crewai.cli.config').setLevel(logging.ERROR)
|
||||
|
||||
# Version is lightweight, import directly
|
||||
from .version import __version__
|
||||
|
||||
@@ -15,7 +15,6 @@ import os
|
||||
import logging
|
||||
import re
|
||||
import keyword
|
||||
import difflib
|
||||
|
||||
# Import new architecture components
|
||||
from .framework_adapters.base import FrameworkAdapter
|
||||
@@ -233,8 +232,10 @@ class AgentsGenerator:
|
||||
self.tool_registry = ToolRegistry()
|
||||
self.tool_registry.register_builtin_autogen_adapters()
|
||||
|
||||
# Get framework adapter (availability already validated at CLI entry)
|
||||
# Get framework adapter and validate availability
|
||||
self.framework_adapter = self._get_framework_adapter(framework)
|
||||
if not self.framework_adapter.is_available():
|
||||
raise ImportError(f"Framework '{framework}' is not available. Please install the required dependencies.")
|
||||
|
||||
def _get_framework_adapter(self, framework: str) -> FrameworkAdapter:
|
||||
"""
|
||||
@@ -320,47 +321,6 @@ class AgentsGenerator:
|
||||
agent_config[field] = value
|
||||
self.logger.debug(f"CLI override for agent {agent_name}: {field} = {value}")
|
||||
|
||||
def _validate_agents_config(self, config):
|
||||
"""
|
||||
Validate agent configuration for typos in field names and provide suggestions.
|
||||
|
||||
Args:
|
||||
config (dict): The parsed YAML configuration
|
||||
"""
|
||||
known_fields = {
|
||||
'role', 'goal', 'instructions', 'backstory', 'tools', 'tasks', 'llm',
|
||||
'function_calling_llm', 'allow_delegation', 'max_iter', 'max_rpm',
|
||||
'max_execution_time', 'verbose', 'cache', 'system_template',
|
||||
'prompt_template', 'response_template', 'tool_timeout', 'planning_tools',
|
||||
'planning', 'autonomy', 'guardrails', 'streaming', 'stream',
|
||||
'approval', 'skills', 'cli_backend', 'reflection'
|
||||
}
|
||||
|
||||
for section_name in ('agents', 'roles'):
|
||||
section = config.get(section_name, {})
|
||||
if not isinstance(section, dict):
|
||||
continue
|
||||
|
||||
entity_name = 'agent' if section_name == 'agents' else 'role'
|
||||
for name, section_config in section.items():
|
||||
if not isinstance(section_config, dict):
|
||||
continue
|
||||
|
||||
for field_name in section_config:
|
||||
if field_name in known_fields:
|
||||
continue
|
||||
|
||||
close_matches = difflib.get_close_matches(
|
||||
field_name,
|
||||
known_fields,
|
||||
n=1,
|
||||
cutoff=0.6
|
||||
)
|
||||
suggestion = f" Did you mean '{close_matches[0]}'?" if close_matches else ""
|
||||
self.logger.warning(
|
||||
f"Unknown field '{field_name}' in {entity_name} '{name}'.{suggestion}"
|
||||
)
|
||||
|
||||
def is_function_or_decorated(self, obj):
|
||||
"""
|
||||
Checks if the given object is a function or has a __call__ method.
|
||||
@@ -546,10 +506,6 @@ class AgentsGenerator:
|
||||
|
||||
# Get workflow input: 'input' is canonical, 'topic' is alias for backward compatibility
|
||||
topic = config.get('input', config.get('topic', ''))
|
||||
|
||||
# Validate agents configuration for typos in field names
|
||||
self._validate_agents_config(config)
|
||||
|
||||
tools_dict = {}
|
||||
|
||||
# Use ToolResolver to get available tools (consistent tool resolution)
|
||||
@@ -641,20 +597,12 @@ class AgentsGenerator:
|
||||
self.framework = framework
|
||||
self.framework_adapter = self._get_framework_adapter(framework)
|
||||
|
||||
# Validate framework availability for non-CLI callers
|
||||
from .framework_adapters.validators import assert_framework_available
|
||||
assert_framework_available(framework)
|
||||
|
||||
# Final availability check
|
||||
if not self.framework_adapter.is_available():
|
||||
raise ImportError(f"Framework '{framework}' is not available. Please install the required dependencies.")
|
||||
|
||||
self.logger.info(f"Using framework: {framework}")
|
||||
return self.framework_adapter.run(
|
||||
config,
|
||||
self.config_list,
|
||||
topic,
|
||||
tools_dict=tools_dict,
|
||||
agent_callback=getattr(self, 'agent_callback', None),
|
||||
task_callback=getattr(self, 'task_callback', None),
|
||||
cli_config=getattr(self, 'cli_config', None),
|
||||
)
|
||||
return self.framework_adapter.run(config, self.config_list, topic)
|
||||
|
||||
def _run_yaml_workflow(self, config):
|
||||
"""
|
||||
@@ -936,13 +884,10 @@ class AgentsGenerator:
|
||||
api_type = _resolve("api_type", default="openai").lower()
|
||||
model_name = _resolve("model", default="gpt-4o-mini")
|
||||
api_key = _resolve("api_key", env_var="OPENAI_API_KEY")
|
||||
# Use resolver for consistent env-var precedence as fallback
|
||||
from praisonai.llm.env import resolve_llm_endpoint
|
||||
ep = resolve_llm_endpoint()
|
||||
|
||||
base_url = (model_config.get("base_url")
|
||||
or yaml_llm.get("base_url")
|
||||
or ep.base_url)
|
||||
or os.environ.get("OPENAI_BASE_URL")
|
||||
or os.environ.get("OPENAI_API_BASE"))
|
||||
|
||||
# Build LLMConfig — Bedrock needs no api_key
|
||||
if api_type == "bedrock":
|
||||
|
||||
@@ -255,11 +255,7 @@ class AsyncAgentScheduler:
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get current execution statistics (synchronous, best-effort).
|
||||
|
||||
Warning: This method provides a best-effort view of stats without
|
||||
guaranteeing atomicity. For consistent snapshots in async context,
|
||||
use get_stats_async() instead.
|
||||
Get execution statistics.
|
||||
|
||||
Returns:
|
||||
Dictionary with execution stats
|
||||
@@ -272,40 +268,6 @@ class AsyncAgentScheduler:
|
||||
"success_rate": (self._success_count / self._execution_count * 100) if self._execution_count > 0 else 0
|
||||
}
|
||||
|
||||
async def get_stats_async(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get current execution statistics with atomic snapshot (async).
|
||||
|
||||
Returns:
|
||||
Dictionary with execution stats
|
||||
"""
|
||||
if self._stats_lock is None:
|
||||
# Not yet started: stats are all zero, no lock needed
|
||||
execs, success, failed = 0, 0, 0
|
||||
else:
|
||||
# Take atomic snapshot of all counters
|
||||
async with self._stats_lock:
|
||||
execs = self._execution_count
|
||||
success = self._success_count
|
||||
failed = self._failure_count
|
||||
|
||||
return {
|
||||
"is_running": self.is_running,
|
||||
"total_executions": execs,
|
||||
"successful_executions": success,
|
||||
"failed_executions": failed,
|
||||
"success_rate": (success / execs * 100) if execs > 0 else 0
|
||||
}
|
||||
|
||||
def get_stats_sync(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Alias for get_stats() for clarity.
|
||||
|
||||
Returns:
|
||||
Dictionary with execution stats
|
||||
"""
|
||||
return self.get_stats()
|
||||
|
||||
async def _run_schedule(self, interval: int, max_retries: int):
|
||||
"""Internal method to run scheduled agent executions."""
|
||||
try:
|
||||
@@ -327,13 +289,7 @@ class AsyncAgentScheduler:
|
||||
|
||||
async def _execute_with_retry(self, max_retries: int):
|
||||
"""Execute agent with retry logic."""
|
||||
# Ensure async primitives are available
|
||||
if self._stats_lock is None:
|
||||
self._ensure_async_primitives()
|
||||
|
||||
# Atomically increment execution count
|
||||
async with self._stats_lock:
|
||||
self._execution_count += 1
|
||||
self._execution_count += 1
|
||||
|
||||
last_exc: Optional[Exception] = None
|
||||
for attempt in range(max_retries):
|
||||
@@ -344,9 +300,7 @@ class AsyncAgentScheduler:
|
||||
logger.info(f"Async agent execution successful on attempt {attempt + 1}")
|
||||
logger.info(f"Result: {result}")
|
||||
|
||||
# Atomically increment success count
|
||||
async with self._stats_lock:
|
||||
self._success_count += 1
|
||||
self._success_count += 1
|
||||
safe_call(self.on_success, result)
|
||||
return
|
||||
|
||||
@@ -359,9 +313,7 @@ class AsyncAgentScheduler:
|
||||
logger.info(f"Waiting {wait_time}s before async retry...")
|
||||
await asyncio.sleep(wait_time)
|
||||
|
||||
# Atomically increment failure count
|
||||
async with self._stats_lock:
|
||||
self._failure_count += 1
|
||||
self._failure_count += 1
|
||||
logger.error(f"Async agent execution failed after {max_retries} attempts")
|
||||
safe_call(
|
||||
self.on_failure,
|
||||
|
||||
@@ -467,15 +467,19 @@ class BaseAutoGenerator:
|
||||
Args:
|
||||
config_list: Optional LLM configuration list
|
||||
"""
|
||||
# Resolve LLM endpoint configuration from environment variables
|
||||
from praisonai.llm.env import resolve_llm_endpoint
|
||||
ep = resolve_llm_endpoint()
|
||||
# Support multiple environment variable patterns for better compatibility
|
||||
model_name = os.environ.get("MODEL_NAME") or os.environ.get("OPENAI_MODEL_NAME", "gpt-4o-mini")
|
||||
base_url = (
|
||||
os.environ.get("OPENAI_BASE_URL") or
|
||||
os.environ.get("OPENAI_API_BASE") or
|
||||
os.environ.get("OLLAMA_API_BASE", "https://api.openai.com/v1")
|
||||
)
|
||||
|
||||
self.config_list = config_list or [
|
||||
{
|
||||
'model': ep.model,
|
||||
'base_url': ep.base_url,
|
||||
'api_key': ep.api_key
|
||||
'model': model_name,
|
||||
'base_url': base_url,
|
||||
'api_key': os.environ.get("OPENAI_API_KEY")
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@@ -12,7 +12,6 @@ if TYPE_CHECKING:
|
||||
from .discord import DiscordBot
|
||||
from .slack import SlackBot
|
||||
from .whatsapp import WhatsAppBot
|
||||
from .linear import LinearBot
|
||||
from .email import EmailBot
|
||||
from .agentmail import AgentMailBot
|
||||
from .bot import Bot
|
||||
@@ -37,9 +36,6 @@ def __getattr__(name: str):
|
||||
if name == "WhatsAppBot":
|
||||
from .whatsapp import WhatsAppBot
|
||||
return WhatsAppBot
|
||||
if name == "LinearBot":
|
||||
from .linear import LinearBot
|
||||
return LinearBot
|
||||
if name == "EmailBot":
|
||||
from .email import EmailBot
|
||||
return EmailBot
|
||||
@@ -67,20 +63,11 @@ def __getattr__(name: str):
|
||||
if name == "HTTPApproval":
|
||||
from ._http_approval import HTTPApproval
|
||||
return HTTPApproval
|
||||
# W1 — cross-platform mirror + identity
|
||||
if name == "mirror_to_session":
|
||||
from ._mirror import mirror_to_session
|
||||
return mirror_to_session
|
||||
if name == "BotSessionManager":
|
||||
from ._session import BotSessionManager
|
||||
return BotSessionManager
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
__all__ = [
|
||||
"TelegramBot", "DiscordBot", "SlackBot", "WhatsAppBot", "LinearBot", "EmailBot", "AgentMailBot",
|
||||
"TelegramBot", "DiscordBot", "SlackBot", "WhatsAppBot", "EmailBot", "AgentMailBot",
|
||||
"Bot", "BotOS",
|
||||
"SlackApproval", "TelegramApproval", "DiscordApproval",
|
||||
"WebhookApproval", "HTTPApproval",
|
||||
# W1
|
||||
"mirror_to_session", "BotSessionManager",
|
||||
]
|
||||
@@ -94,15 +94,12 @@ async def classify_with_llm(
|
||||
)
|
||||
|
||||
try:
|
||||
from praisonai.llm.env import resolve_llm_endpoint
|
||||
ep = resolve_llm_endpoint()
|
||||
|
||||
client = OpenAI(
|
||||
api_key=ep.api_key or "",
|
||||
base_url=ep.base_url,
|
||||
api_key=_os.environ.get("OPENAI_API_KEY", ""),
|
||||
base_url=_os.environ.get("OPENAI_BASE_URL"),
|
||||
)
|
||||
response = client.chat.completions.create(
|
||||
model=_os.environ.get("APPROVAL_LLM_MODEL", ep.model),
|
||||
model=_os.environ.get("APPROVAL_LLM_MODEL", "gpt-4o-mini"),
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
|
||||
@@ -1,134 +0,0 @@
|
||||
"""Outbound delivery mirror for bot sessions.
|
||||
|
||||
W1 — When a bot pushes a message via ``Bot.send_message()`` (notifications,
|
||||
scheduled deliveries, cross-platform replies), this module appends an
|
||||
``{"role": "assistant", "mirror": True}`` entry to the user's session
|
||||
so the next inbound turn carries that context.
|
||||
|
||||
Standalone helper — designed to be safe to call from any context (sync
|
||||
handler, async coroutine, cron job, scheduled delivery). Errors are
|
||||
swallowed and logged: a mirror failure must never break the outbound
|
||||
delivery itself.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# @claude: No module-level lock here by design.
|
||||
# A process-wide _mirror_lock = threading.RLock() would serialise mirror_to_session()
|
||||
# calls across ALL BotSessionManager instances — slow disk I/O for one bot would
|
||||
# block mirrors for every other bot in the process.
|
||||
# Lock ownership belongs to the caller (passed via the optional ``lock`` param),
|
||||
# or a fresh per-call RLock is created as fallback — scoped to one call only.
|
||||
|
||||
|
||||
def mirror_to_session(
|
||||
session_mgr: Any,
|
||||
user_id: str,
|
||||
message_text: str,
|
||||
source_label: str = "delivery",
|
||||
metadata: Optional[dict] = None,
|
||||
lock: Optional[threading.RLock] = None,
|
||||
) -> bool:
|
||||
"""Append a mirror entry to ``user_id``'s session history.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
session_mgr
|
||||
A ``BotSessionManager``-shaped object that exposes
|
||||
``_storage_key(user_id)``, ``_load_history(user_id)``, and
|
||||
``_save_history(user_id, history)``. Any object with these three
|
||||
methods is acceptable (duck-typed, no Protocol required).
|
||||
user_id
|
||||
Raw platform user id. Will be passed through the manager's
|
||||
identity resolver if one is configured.
|
||||
message_text
|
||||
The outbound message text to mirror.
|
||||
source_label
|
||||
Free-form tag identifying the source of the delivery
|
||||
(e.g. ``"cron"``, ``"web"``, ``"cross_platform"``).
|
||||
metadata
|
||||
Optional extra metadata to merge into the mirror entry.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
``True`` on success, ``False`` if anything went wrong. Errors
|
||||
are logged at WARN level and never raised.
|
||||
"""
|
||||
if not message_text:
|
||||
return False
|
||||
|
||||
try:
|
||||
# Verify that session_mgr is a compatible BotSessionManager-shaped
|
||||
# object by probing its _storage_key method (duck-typing check).
|
||||
session_mgr._storage_key(user_id)
|
||||
except Exception as e:
|
||||
logger.warning("mirror: storage_key failed: %s", e)
|
||||
return False
|
||||
|
||||
# Check if the session manager has a thread-safe mirror method
|
||||
# This allows the session manager to handle synchronization properly
|
||||
# with its internal asyncio locks.
|
||||
if hasattr(session_mgr, '_add_mirror_entry_sync'):
|
||||
try:
|
||||
entry: dict = {
|
||||
"role": "assistant",
|
||||
"content": message_text,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"mirror": True,
|
||||
"mirror_source": source_label,
|
||||
}
|
||||
if metadata:
|
||||
# Filter out reserved fields to preserve entry invariants
|
||||
_reserved = {"role", "content", "timestamp", "mirror", "mirror_source"}
|
||||
entry.update({k: v for k, v in metadata.items() if k not in _reserved})
|
||||
|
||||
return session_mgr._add_mirror_entry_sync(user_id, entry)
|
||||
except Exception as e:
|
||||
logger.warning("mirror: _add_mirror_entry_sync failed: %s", e)
|
||||
return False
|
||||
|
||||
# Fallback: session manager has no _add_mirror_entry_sync.
|
||||
# Use the caller-supplied lock, or a fresh per-call RLock — never a global.
|
||||
logger.warning(
|
||||
"mirror_to_session using fallback synchronization - "
|
||||
"session manager should implement _add_mirror_entry_sync for better safety"
|
||||
)
|
||||
_call_lock = lock if lock is not None else threading.RLock()
|
||||
with _call_lock:
|
||||
try:
|
||||
# Load current history
|
||||
history = list(session_mgr._load_history(user_id))
|
||||
|
||||
# Create mirror entry
|
||||
entry: dict = {
|
||||
"role": "assistant",
|
||||
"content": message_text,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"mirror": True,
|
||||
"mirror_source": source_label,
|
||||
}
|
||||
if metadata:
|
||||
# Filter out reserved fields to preserve entry invariants
|
||||
_reserved = {"role", "content", "timestamp", "mirror", "mirror_source"}
|
||||
entry.update({k: v for k, v in metadata.items() if k not in _reserved})
|
||||
history.append(entry)
|
||||
|
||||
# Save updated history atomically within the lock
|
||||
session_mgr._save_history(user_id, history)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("mirror: save_history failed: %s", e)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
__all__ = ["mirror_to_session"]
|
||||
@@ -15,7 +15,6 @@ _BUILTIN_PLATFORMS: Dict[str, tuple] = {
|
||||
"discord": ("praisonai.bots.discord", "DiscordBot"),
|
||||
"slack": ("praisonai.bots.slack", "SlackBot"),
|
||||
"whatsapp": ("praisonai.bots.whatsapp", "WhatsAppBot"),
|
||||
"linear": ("praisonai.bots.linear", "LinearBot"),
|
||||
"email": ("praisonai.bots.email", "EmailBot"),
|
||||
"agentmail": ("praisonai.bots.agentmail", "AgentMailBot"),
|
||||
}
|
||||
|
||||
@@ -62,7 +62,6 @@ class BotSessionManager:
|
||||
max_history: int = 100,
|
||||
store: Optional[Any] = None,
|
||||
platform: str = "",
|
||||
identity_resolver: Optional[Any] = None,
|
||||
) -> None:
|
||||
self._histories: Dict[str, List[Dict[str, Any]]] = {}
|
||||
self._locks: Dict[str, asyncio.Lock] = {}
|
||||
@@ -71,46 +70,17 @@ class BotSessionManager:
|
||||
self._store = store
|
||||
self._platform = platform
|
||||
self._last_active: Dict[str, float] = {}
|
||||
# W1: optional cross-platform identity resolver. When set, the
|
||||
# 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
|
||||
|
||||
def _storage_key(self, user_id: str) -> str:
|
||||
"""Resolve a raw platform user id to the in-memory/store key.
|
||||
|
||||
With an identity resolver this is the unified user id; without
|
||||
one, behaviour is unchanged (raw ``user_id``).
|
||||
"""
|
||||
if self._identity_resolver is not None and self._platform:
|
||||
try:
|
||||
return self._identity_resolver.resolve(self._platform, user_id)
|
||||
except Exception as e: # pragma: no cover — defensive
|
||||
logger.warning("identity resolver failed: %s", e)
|
||||
return user_id
|
||||
|
||||
def _persist_key(self, storage_key: str) -> str:
|
||||
"""Derive the persistent-store key from an in-memory storage key.
|
||||
|
||||
With a resolver, ``storage_key`` is already the unified user id
|
||||
and is used directly. Without a resolver, the legacy
|
||||
``bot_{platform}_{user_id}`` prefix is preserved for back-compat.
|
||||
"""
|
||||
if self._identity_resolver is not None and self._platform:
|
||||
return storage_key
|
||||
prefix = f"bot_{self._platform}" if self._platform else "bot"
|
||||
return f"{prefix}_{storage_key}"
|
||||
|
||||
def _session_key(self, user_id: str) -> str:
|
||||
"""Persistent-store key for a raw platform user id (back-compat API)."""
|
||||
return self._persist_key(self._storage_key(user_id))
|
||||
"""Generate a deterministic session key for persistent storage."""
|
||||
prefix = f"bot_{self._platform}" if self._platform else "bot"
|
||||
return f"{prefix}_{user_id}"
|
||||
|
||||
def _get_lock(self, user_id: str) -> asyncio.Lock:
|
||||
"""Get or create an asyncio.Lock for *user_id* (storage-keyed)."""
|
||||
key = self._storage_key(user_id)
|
||||
if key not in self._locks:
|
||||
self._locks[key] = asyncio.Lock()
|
||||
return self._locks[key]
|
||||
"""Get or create an asyncio.Lock for *user_id*."""
|
||||
if user_id not in self._locks:
|
||||
self._locks[user_id] = asyncio.Lock()
|
||||
return self._locks[user_id]
|
||||
|
||||
def _get_agent_lock(self, agent: "Agent") -> asyncio.Lock:
|
||||
"""Get or create a lock for the *agent* instance (by id)."""
|
||||
@@ -129,7 +99,7 @@ class BotSessionManager:
|
||||
return list(history)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to load session from store: %s", e)
|
||||
return list(self._histories.get(self._storage_key(user_id), []))
|
||||
return list(self._histories.get(user_id, []))
|
||||
|
||||
def _save_history(
|
||||
self, user_id: str, history: List[Dict[str, Any]]
|
||||
@@ -138,8 +108,8 @@ class BotSessionManager:
|
||||
if self._max_history > 0 and len(history) > self._max_history:
|
||||
history = history[-self._max_history:]
|
||||
|
||||
# Always update in-memory cache (keyed by storage key)
|
||||
self._histories[self._storage_key(user_id)] = history
|
||||
# Always update in-memory cache
|
||||
self._histories[user_id] = history
|
||||
|
||||
if self._store is not None:
|
||||
key = self._session_key(user_id)
|
||||
@@ -161,9 +131,6 @@ class BotSessionManager:
|
||||
agent: "Agent",
|
||||
user_id: str,
|
||||
prompt: str,
|
||||
chat_id: str = "",
|
||||
thread_id: str = "",
|
||||
user_name: str = "",
|
||||
) -> str:
|
||||
"""Run ``agent.chat(prompt)`` with *user_id*-scoped history.
|
||||
|
||||
@@ -173,76 +140,35 @@ class BotSessionManager:
|
||||
Uses both a per-user lock (serialise same user) and a per-agent
|
||||
lock (prevent concurrent history swaps on a shared Agent).
|
||||
"""
|
||||
self._last_active[self._storage_key(user_id)] = time.monotonic()
|
||||
self._last_active[user_id] = time.monotonic()
|
||||
user_lock = self._get_lock(user_id)
|
||||
agent_lock = self._get_agent_lock(agent)
|
||||
|
||||
# W1: set task-local session context so any tool the agent
|
||||
# invokes can read platform / chat / user metadata without
|
||||
# relying on os.environ globals.
|
||||
ctx_token = None
|
||||
_clear_ctx = None
|
||||
try:
|
||||
from praisonaiagents.session.context import (
|
||||
set_session_context as _set_ctx,
|
||||
clear_session_context as _clear_ctx,
|
||||
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
|
||||
)
|
||||
ctx_token = _set_ctx(
|
||||
platform=self._platform,
|
||||
chat_id=chat_id,
|
||||
thread_id=thread_id,
|
||||
user_id=user_id,
|
||||
user_name=user_name,
|
||||
unified_user_id=self._storage_key(user_id),
|
||||
)
|
||||
except Exception: # pragma: no cover — defensive
|
||||
_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
|
||||
)
|
||||
async with agent_lock:
|
||||
# Swap histories
|
||||
saved_history = agent.chat_history
|
||||
agent.chat_history = user_history
|
||||
|
||||
# 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.
|
||||
# Throughput on a shared agent is then bounded by the LLM's
|
||||
# serial latency — this is correct: a single Agent's
|
||||
# ``chat_history`` is mutable and cannot be safely interleaved.
|
||||
try:
|
||||
response = await loop.run_in_executor(None, agent.chat, prompt)
|
||||
finally:
|
||||
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
|
||||
finally:
|
||||
agent.chat_history = saved_history
|
||||
# Capture updated history and restore agent's original
|
||||
updated_history = agent.chat_history
|
||||
agent.chat_history = saved_history
|
||||
|
||||
# Persist outside the agent_lock — it's per-user and the agent
|
||||
# is no longer touched.
|
||||
# Persist (may hit disk via run_in_executor)
|
||||
await loop.run_in_executor(
|
||||
None, self._save_history, user_id, updated_history
|
||||
)
|
||||
|
||||
return response
|
||||
finally:
|
||||
# Always clear task-local session context, even if an exception occurred.
|
||||
if ctx_token is not None and _clear_ctx is not None:
|
||||
try:
|
||||
_clear_ctx(ctx_token)
|
||||
except Exception as e:
|
||||
logger.debug("Failed to clear task-local session context for ctx_token=%r: %s", ctx_token, e)
|
||||
return response
|
||||
|
||||
def reap_stale(self, max_age_seconds: int) -> int:
|
||||
"""Remove sessions older than *max_age_seconds*. Returns count reaped.
|
||||
@@ -257,12 +183,12 @@ class BotSessionManager:
|
||||
uid for uid, ts in self._last_active.items()
|
||||
if (now - ts) > max_age_seconds
|
||||
]
|
||||
for storage_key in stale:
|
||||
self._histories.pop(storage_key, None)
|
||||
self._last_active.pop(storage_key, None)
|
||||
self._locks.pop(storage_key, None)
|
||||
for uid in stale:
|
||||
self._histories.pop(uid, None)
|
||||
self._last_active.pop(uid, None)
|
||||
self._locks.pop(uid, None)
|
||||
if self._store is not None:
|
||||
key = self._persist_key(storage_key)
|
||||
key = self._session_key(uid)
|
||||
try:
|
||||
self._store.clear_session(key)
|
||||
except Exception as e:
|
||||
@@ -273,109 +199,26 @@ class BotSessionManager:
|
||||
|
||||
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)
|
||||
existed = storage_key in self._histories
|
||||
self._histories.pop(storage_key, None)
|
||||
self._last_active.pop(storage_key, None)
|
||||
self._locks.pop(storage_key, None)
|
||||
existed = user_id in self._histories
|
||||
self._histories.pop(user_id, None)
|
||||
self._last_active.pop(user_id, None)
|
||||
|
||||
if self._store is not None:
|
||||
persist_key = self._session_key(user_id)
|
||||
key = self._session_key(user_id)
|
||||
try:
|
||||
self._store.clear_session(persist_key)
|
||||
self._store.clear_session(key)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to clear session in store: %s", e)
|
||||
|
||||
return existed
|
||||
|
||||
def _add_mirror_entry_sync(self, user_id: str, entry: dict) -> bool:
|
||||
"""Thread-safe method to add a mirror entry to user's history.
|
||||
|
||||
This method coordinates with the asyncio locks used by chat()
|
||||
to prevent race conditions. Called by mirror_to_session().
|
||||
|
||||
Args:
|
||||
user_id: Raw platform user id
|
||||
entry: Mirror entry dict to append to history
|
||||
|
||||
Returns:
|
||||
bool: True on success, False on failure
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
def _sync_add_entry():
|
||||
# Get the event loop that owns the asyncio locks
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
# No event loop running - we need to use threading synchronization
|
||||
# This happens when called from sync contexts like cron jobs
|
||||
storage_key = self._storage_key(user_id)
|
||||
self._last_active[storage_key] = time.monotonic()
|
||||
|
||||
# Create a temporary threading lock for this user to prevent concurrent
|
||||
# access to the same user's history from multiple sync threads
|
||||
import threading
|
||||
user_sync_lock = getattr(self, '_user_sync_locks', None)
|
||||
if user_sync_lock is None:
|
||||
self._user_sync_locks = {}
|
||||
user_sync_lock = self._user_sync_locks
|
||||
|
||||
if storage_key not in user_sync_lock:
|
||||
user_sync_lock[storage_key] = threading.Lock()
|
||||
|
||||
with user_sync_lock[storage_key]:
|
||||
history = list(self._load_history(user_id))
|
||||
history.append(entry)
|
||||
self._save_history(user_id, history)
|
||||
return True
|
||||
|
||||
# There's an event loop - we need to coordinate with asyncio locks
|
||||
# This is more complex but necessary for thread safety
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
self._add_mirror_entry_async(user_id, entry), loop
|
||||
)
|
||||
return future.result(timeout=10.0) # 10 second timeout
|
||||
|
||||
try:
|
||||
return _sync_add_entry()
|
||||
except Exception as e:
|
||||
logger.warning("_add_mirror_entry_sync failed: %s", e)
|
||||
return False
|
||||
|
||||
async def _add_mirror_entry_async(self, user_id: str, entry: dict) -> bool:
|
||||
"""Async helper for _add_mirror_entry_sync."""
|
||||
try:
|
||||
storage_key = self._storage_key(user_id)
|
||||
self._last_active[storage_key] = time.monotonic()
|
||||
|
||||
# Use the same per-user lock that chat() uses
|
||||
user_lock = self._get_lock(user_id)
|
||||
async with user_lock:
|
||||
# Load history (may hit disk via run_in_executor for async safety)
|
||||
loop = asyncio.get_running_loop()
|
||||
history = await loop.run_in_executor(
|
||||
None, self._load_history, user_id
|
||||
)
|
||||
history = list(history) # Ensure it's mutable
|
||||
history.append(entry)
|
||||
|
||||
# Save updated history
|
||||
await loop.run_in_executor(
|
||||
None, self._save_history, user_id, history
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning("_add_mirror_entry_async failed: %s", e)
|
||||
return False
|
||||
|
||||
def reset_all(self) -> int:
|
||||
"""Clear all user sessions. Returns the count cleared."""
|
||||
count = len(self._histories)
|
||||
|
||||
if self._store is not None:
|
||||
for storage_key in list(self._histories.keys()):
|
||||
key = self._persist_key(storage_key)
|
||||
for user_id in list(self._histories.keys()):
|
||||
key = self._session_key(user_id)
|
||||
try:
|
||||
self._store.clear_session(key)
|
||||
except Exception as e:
|
||||
|
||||
@@ -45,7 +45,6 @@ _TOKEN_ENV_MAP = {
|
||||
"discord": "DISCORD_BOT_TOKEN",
|
||||
"slack": "SLACK_BOT_TOKEN",
|
||||
"whatsapp": "WHATSAPP_ACCESS_TOKEN",
|
||||
"linear": "LINEAR_OAUTH_TOKEN",
|
||||
"email": "EMAIL_APP_PASSWORD",
|
||||
"agentmail": "AGENTMAIL_API_KEY",
|
||||
}
|
||||
@@ -54,7 +53,6 @@ _TOKEN_ENV_MAP = {
|
||||
_EXTRA_ENV_MAP = {
|
||||
"slack": {"app_token": "SLACK_APP_TOKEN"},
|
||||
"whatsapp": {"phone_number_id": "WHATSAPP_PHONE_NUMBER_ID"},
|
||||
"linear": {"signing_secret": "LINEAR_WEBHOOK_SECRET"},
|
||||
"email": {
|
||||
"email_address": "EMAIL_ADDRESS",
|
||||
"imap_server": "EMAIL_IMAP_SERVER",
|
||||
@@ -87,18 +85,12 @@ class Bot:
|
||||
agent: Optional[Any] = None,
|
||||
token: Optional[str] = None,
|
||||
config: Optional[Any] = None,
|
||||
identity_resolver: Optional[Any] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
self._platform = platform.lower().strip()
|
||||
self._agent = agent
|
||||
self._explicit_token = token
|
||||
self._config = config
|
||||
# W1: optional cross-platform identity resolver. Applied to the
|
||||
# adapter's ``_session`` after construction (duck-typed; works
|
||||
# with any adapter that exposes a BotSessionManager-compatible
|
||||
# ``_session`` attribute).
|
||||
self._identity_resolver = identity_resolver
|
||||
self._kwargs = kwargs
|
||||
|
||||
self._adapter: Optional[Any] = None
|
||||
@@ -186,24 +178,7 @@ class Bot:
|
||||
# Merge user kwargs (override defaults)
|
||||
init_kwargs.update(self._kwargs)
|
||||
|
||||
adapter = adapter_cls(**init_kwargs)
|
||||
|
||||
# W1: post-construction wire-up for the identity resolver.
|
||||
# Adapters create their own BotSessionManager during __init__;
|
||||
# we splice the resolver in here so existing adapters need no
|
||||
# signature change.
|
||||
if self._identity_resolver is not None:
|
||||
session = getattr(adapter, "_session", None)
|
||||
if session is not None and hasattr(session, "_identity_resolver"):
|
||||
session._identity_resolver = self._identity_resolver
|
||||
else:
|
||||
logger.warning(
|
||||
"Bot(%s): adapter has no BotSessionManager-compatible "
|
||||
"_session; identity_resolver ignored.",
|
||||
self._platform,
|
||||
)
|
||||
|
||||
return adapter
|
||||
return adapter_cls(**init_kwargs)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Build the adapter and start the bot."""
|
||||
|
||||
@@ -69,14 +69,10 @@ class BotOS:
|
||||
agent: Optional[Any] = None,
|
||||
platforms: Optional[List[str]] = None,
|
||||
config: Optional[Any] = None,
|
||||
identity_resolver: Optional[Any] = None,
|
||||
):
|
||||
self._bots: Dict[str, Bot] = {}
|
||||
self._is_running = False
|
||||
self._config = config
|
||||
# W1: shared identity resolver applied to every managed bot —
|
||||
# gives cross-platform unified-user sessions out of the box.
|
||||
self._identity_resolver = identity_resolver
|
||||
self._tasks: List[asyncio.Task] = []
|
||||
|
||||
# Register explicit bots
|
||||
@@ -88,13 +84,7 @@ class BotOS:
|
||||
if agent and platforms:
|
||||
for plat in platforms:
|
||||
if plat not in self._bots:
|
||||
self.add_bot(
|
||||
Bot(
|
||||
plat,
|
||||
agent=agent,
|
||||
identity_resolver=self._identity_resolver,
|
||||
)
|
||||
)
|
||||
self.add_bot(Bot(plat, agent=agent))
|
||||
|
||||
# ── Properties ──────────────────────────────────────────────────
|
||||
|
||||
@@ -110,14 +100,6 @@ class BotOS:
|
||||
Args:
|
||||
bot: A Bot instance.
|
||||
"""
|
||||
# W1: propagate the BotOS-level resolver to bots that don't
|
||||
# already have one, so cross-platform unification works whether
|
||||
# the user uses the shortcut API or constructs Bots manually.
|
||||
if (
|
||||
self._identity_resolver is not None
|
||||
and getattr(bot, "_identity_resolver", None) is None
|
||||
):
|
||||
bot._identity_resolver = self._identity_resolver
|
||||
self._bots[bot.platform] = bot
|
||||
|
||||
def list_bots(self) -> List[str]:
|
||||
|
||||
@@ -247,11 +247,7 @@ class DiscordBot(ChatCommandMixin, MessageHookMixin):
|
||||
|
||||
async def _send_agent_response():
|
||||
text_to_send = await self._debouncer.debounce(user_id, bot_message.text)
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, text_to_send,
|
||||
chat_id=str(message.channel.id),
|
||||
user_name=str(getattr(message.author, "name", "")),
|
||||
)
|
||||
response = await self._session.chat(self._agent, user_id, text_to_send)
|
||||
send_result = self.fire_message_sending(
|
||||
str(message.channel.id), str(response),
|
||||
)
|
||||
|
||||
@@ -1,534 +0,0 @@
|
||||
"""
|
||||
Linear Bot implementation for PraisonAI.
|
||||
|
||||
Supports Linear agent integration via AgentSession webhooks.
|
||||
|
||||
Usage:
|
||||
bot = LinearBot(token="linear-oauth-token", agent=agent,
|
||||
signing_secret="webhook-secret")
|
||||
await bot.start()
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent
|
||||
|
||||
from praisonai.bots._protocol_mixin import ChatCommandMixin, MessageHookMixin
|
||||
from praisonaiagents.bots import (
|
||||
BotConfig,
|
||||
BotMessage,
|
||||
BotUser,
|
||||
BotChannel,
|
||||
MessageType,
|
||||
)
|
||||
|
||||
from ._commands import format_status, format_help
|
||||
from ._session import BotSessionManager
|
||||
from ._rate_limit import RateLimiter
|
||||
from ._ack import AckReactor
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Linear GraphQL endpoint
|
||||
LINEAR_API_BASE = "https://api.linear.app/graphql"
|
||||
|
||||
|
||||
class LinearBot(ChatCommandMixin, MessageHookMixin):
|
||||
"""Linear bot runtime for PraisonAI agents.
|
||||
|
||||
Connects an agent to Linear via AgentSession webhooks, handling
|
||||
mentions, assignments, and providing full bot functionality.
|
||||
|
||||
Example:
|
||||
from praisonai.bots import LinearBot
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="assistant")
|
||||
bot = LinearBot(
|
||||
token="YOUR_OAUTH_TOKEN",
|
||||
signing_secret="YOUR_WEBHOOK_SECRET",
|
||||
agent=agent,
|
||||
)
|
||||
|
||||
await bot.start()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
token: str = "",
|
||||
agent: Optional["Agent"] = None,
|
||||
config: Optional[BotConfig] = None,
|
||||
signing_secret: str = "",
|
||||
webhook_port: int = 8080,
|
||||
webhook_path: str = "/webhook",
|
||||
**kwargs,
|
||||
):
|
||||
# Store extra kwargs for forward compatibility
|
||||
self._extra_kwargs = kwargs
|
||||
|
||||
# Determine token source for proper authorization format
|
||||
self._oauth_token = token or os.environ.get("LINEAR_OAUTH_TOKEN", "")
|
||||
self._api_key = os.environ.get("LINEAR_API_KEY", "") if not self._oauth_token else ""
|
||||
self._token = self._oauth_token or self._api_key
|
||||
self._is_oauth = bool(self._oauth_token)
|
||||
self._agent = agent
|
||||
self.config = config or BotConfig(token=self._token, mode="webhook")
|
||||
self._signing_secret = signing_secret or os.environ.get("LINEAR_WEBHOOK_SECRET", "")
|
||||
self._webhook_port = webhook_port
|
||||
self._webhook_path = webhook_path
|
||||
|
||||
self._is_running = False
|
||||
self._started_at: Optional[float] = None
|
||||
self._bot_user: Optional[BotUser] = None
|
||||
|
||||
try:
|
||||
from praisonaiagents.session import get_default_session_store
|
||||
_store = get_default_session_store()
|
||||
except Exception:
|
||||
_store = None
|
||||
|
||||
self._session_mgr = BotSessionManager(
|
||||
store=_store,
|
||||
platform="linear",
|
||||
)
|
||||
self._message_handlers: List[Callable] = []
|
||||
self._runner: Any = None
|
||||
self._site: Any = None
|
||||
self._http_session: Any = None
|
||||
self._background_tasks: set = set()
|
||||
self._rate_limiter = RateLimiter.for_platform("linear")
|
||||
self._ack: AckReactor = AckReactor(
|
||||
ack_emoji=self.config.ack_emoji,
|
||||
done_emoji=self.config.done_emoji,
|
||||
)
|
||||
|
||||
# ChatCommandMixin setup
|
||||
self._command_handlers: Dict[str, Callable] = {}
|
||||
self._command_info: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
# Register built-in commands
|
||||
self._register_builtins()
|
||||
|
||||
def _register_builtins(self) -> None:
|
||||
"""Register built-in /status, /new, /help commands."""
|
||||
|
||||
async def _status(msg):
|
||||
return format_status(self._agent, "linear", self._started_at, self._is_running)
|
||||
|
||||
async def _new(msg):
|
||||
user_id = msg.sender.user_id if msg.sender else "unknown"
|
||||
self._session_mgr.reset(user_id)
|
||||
return "Session reset. Send a message to start a new conversation."
|
||||
|
||||
async def _help(msg):
|
||||
extra = {
|
||||
name: info.get("description", "")
|
||||
for name, info in self._command_info.items()
|
||||
if name not in ("status", "new", "help")
|
||||
}
|
||||
return format_help(self._agent, "linear", extra or None)
|
||||
|
||||
self.register_command("status", _status, description="Show bot status and info")
|
||||
self.register_command("new", _new, description="Reset conversation session")
|
||||
self.register_command("help", _help, description="Show this help message")
|
||||
|
||||
# ── Properties ──────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return self._is_running
|
||||
|
||||
@property
|
||||
def platform(self) -> str:
|
||||
return "linear"
|
||||
|
||||
@property
|
||||
def bot_user(self) -> Optional[BotUser]:
|
||||
return self._bot_user
|
||||
|
||||
# ── Lifecycle ───────────────────────────────────────────────────
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the Linear webhook server."""
|
||||
if self._is_running:
|
||||
return
|
||||
|
||||
if not self._token:
|
||||
raise ValueError("LINEAR_OAUTH_TOKEN or LINEAR_API_KEY required")
|
||||
|
||||
if not self._signing_secret:
|
||||
logger.warning("LINEAR_WEBHOOK_SECRET not set - webhook signatures will not be verified")
|
||||
|
||||
logger.info(f"Starting Linear bot on port {self._webhook_port}")
|
||||
|
||||
try:
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
except ImportError:
|
||||
raise ImportError("aiohttp required: pip install aiohttp")
|
||||
|
||||
# Create HTTP session for API calls
|
||||
self._http_session = aiohttp.ClientSession()
|
||||
|
||||
# Create webhook server
|
||||
app = web.Application()
|
||||
app.router.add_post(self._webhook_path, self._handle_webhook)
|
||||
app.router.add_get(self._webhook_path, self._handle_webhook_verify)
|
||||
|
||||
self._runner = web.AppRunner(app)
|
||||
await self._runner.setup()
|
||||
|
||||
self._site = web.TCPSite(self._runner, "0.0.0.0", self._webhook_port)
|
||||
await self._site.start()
|
||||
|
||||
self._is_running = True
|
||||
self._started_at = time.time()
|
||||
|
||||
# Get bot user info
|
||||
await self._fetch_bot_user()
|
||||
|
||||
logger.info(f"Linear bot started on http://0.0.0.0:{self._webhook_port}{self._webhook_path}")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the Linear webhook server."""
|
||||
if not self._is_running:
|
||||
return
|
||||
|
||||
logger.info("Stopping Linear bot...")
|
||||
|
||||
# Wait for background tasks to complete
|
||||
if self._background_tasks:
|
||||
logger.info(f"Waiting for {len(self._background_tasks)} background tasks...")
|
||||
await asyncio.gather(*self._background_tasks, return_exceptions=True)
|
||||
|
||||
# Stop HTTP server
|
||||
if self._site:
|
||||
await self._site.stop()
|
||||
if self._runner:
|
||||
await self._runner.cleanup()
|
||||
|
||||
# Close HTTP session
|
||||
if self._http_session:
|
||||
await self._http_session.close()
|
||||
|
||||
self._is_running = False
|
||||
self._started_at = None
|
||||
logger.info("Linear bot stopped")
|
||||
|
||||
# ── Webhook Handling ─────────────────────────────────────────────
|
||||
|
||||
async def _handle_webhook_verify(self, request) -> Any:
|
||||
"""Handle webhook verification (GET request)."""
|
||||
from aiohttp import web
|
||||
return web.Response(status=200, text="Linear webhook endpoint")
|
||||
|
||||
async def _handle_webhook(self, request) -> Any:
|
||||
"""Handle incoming Linear webhooks."""
|
||||
from aiohttp import web
|
||||
|
||||
try:
|
||||
# Read raw body for signature verification
|
||||
raw_body = await request.read()
|
||||
|
||||
# Verify signature if secret is configured
|
||||
if self._signing_secret:
|
||||
signature = request.headers.get("Linear-Signature", "")
|
||||
if not self._verify_signature(raw_body, signature):
|
||||
logger.warning("Invalid webhook signature")
|
||||
return web.Response(status=401, text="Invalid signature")
|
||||
|
||||
# Parse JSON body
|
||||
try:
|
||||
body = json.loads(raw_body.decode('utf-8'))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
||||
logger.error(f"Invalid JSON payload: {e}")
|
||||
return web.Response(status=400, text="Invalid JSON")
|
||||
|
||||
# Check timestamp for replay protection (Linear recommendation: 60s window)
|
||||
webhook_timestamp = body.get("webhookTimestamp", 0)
|
||||
if webhook_timestamp:
|
||||
time_diff = abs(time.time() * 1000 - webhook_timestamp)
|
||||
if time_diff > 60_000: # 60 seconds in milliseconds
|
||||
logger.warning(f"Webhook timestamp too old: {time_diff}ms")
|
||||
return web.Response(status=401, text="Timestamp too old")
|
||||
|
||||
# Get event type
|
||||
event_type = request.headers.get("Linear-Event", "")
|
||||
|
||||
# Process webhook in background to avoid blocking
|
||||
task = asyncio.create_task(self._process_webhook(event_type, body))
|
||||
self._background_tasks.add(task)
|
||||
task.add_done_callback(self._background_tasks.discard)
|
||||
|
||||
return web.Response(status=200, text="OK")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Webhook handling error: {e}")
|
||||
return web.Response(status=500, text="Internal error")
|
||||
|
||||
def _verify_signature(self, body: bytes, signature: str) -> bool:
|
||||
"""Verify Linear webhook signature using HMAC-SHA256."""
|
||||
if not self._signing_secret or not signature:
|
||||
return False
|
||||
|
||||
expected = hmac.new(
|
||||
self._signing_secret.encode(),
|
||||
body,
|
||||
hashlib.sha256
|
||||
).hexdigest()
|
||||
|
||||
return hmac.compare_digest(expected, signature)
|
||||
|
||||
async def _process_webhook(self, event_type: str, body: Dict[str, Any]) -> None:
|
||||
"""Process webhook events."""
|
||||
try:
|
||||
logger.debug(f"Processing {event_type} event")
|
||||
|
||||
if event_type == "AgentSession":
|
||||
await self._handle_agent_session(body)
|
||||
elif event_type == "Comment":
|
||||
await self._handle_comment(body)
|
||||
elif event_type == "Issue":
|
||||
await self._handle_issue(body)
|
||||
else:
|
||||
logger.debug(f"Ignoring {event_type} event")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing {event_type} webhook: {e}")
|
||||
|
||||
async def _handle_agent_session(self, body: Dict[str, Any]) -> None:
|
||||
"""Handle AgentSession events (mentions, assignments)."""
|
||||
action = body.get("action")
|
||||
data = body.get("data", {})
|
||||
|
||||
if action == "create":
|
||||
# New agent session created (mention or assignment)
|
||||
session_id = data.get("id")
|
||||
issue_data = data.get("issue", {})
|
||||
|
||||
if not session_id or not issue_data:
|
||||
logger.warning("Missing session or issue data")
|
||||
return
|
||||
|
||||
# Create bot message from issue
|
||||
issue_id = issue_data.get("id", "")
|
||||
issue_title = issue_data.get("title", "")
|
||||
issue_description = issue_data.get("description", "")
|
||||
|
||||
message_text = f"Issue: {issue_title}"
|
||||
if issue_description:
|
||||
message_text += f"\n\n{issue_description}"
|
||||
|
||||
bot_message = BotMessage(
|
||||
message_id=session_id,
|
||||
content=message_text,
|
||||
message_type=MessageType.TEXT,
|
||||
channel=BotChannel(channel_id=issue_id, name=f"Issue {issue_data.get('identifier', '')}"),
|
||||
sender=BotUser(user_id="linear-system", display_name="Linear"),
|
||||
timestamp=time.time(),
|
||||
metadata={"issue": issue_data, "session_id": session_id}
|
||||
)
|
||||
|
||||
# Process with agent
|
||||
await self._handle_agent_message(bot_message)
|
||||
|
||||
async def _handle_comment(self, body: Dict[str, Any]) -> None:
|
||||
"""Handle Comment events."""
|
||||
# For future implementation - comment threads
|
||||
logger.debug("Comment event received")
|
||||
|
||||
async def _handle_issue(self, body: Dict[str, Any]) -> None:
|
||||
"""Handle Issue events."""
|
||||
# For future implementation - issue updates
|
||||
logger.debug("Issue event received")
|
||||
|
||||
async def _handle_agent_message(self, message: BotMessage) -> None:
|
||||
"""Process message with the agent."""
|
||||
if not self._agent:
|
||||
logger.warning("No agent configured")
|
||||
return
|
||||
|
||||
try:
|
||||
user_id = message.sender.user_id if message.sender else "unknown"
|
||||
session_id = message.metadata.get("session_id") if message.metadata else None
|
||||
|
||||
# Use BotSessionManager.chat() which handles history isolation and run_in_executor
|
||||
response = await self._session_mgr.chat(self._agent, user_id, message.content)
|
||||
|
||||
# Send response back to Linear
|
||||
if response and message.metadata:
|
||||
await self._send_comment(
|
||||
issue_id=message.metadata.get("issue", {}).get("id", ""),
|
||||
comment=response,
|
||||
session_id=session_id,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing agent message: {e}")
|
||||
|
||||
# ── Linear API Methods ──────────────────────────────────────────
|
||||
|
||||
async def _fetch_bot_user(self) -> None:
|
||||
"""Fetch bot user information from Linear API."""
|
||||
try:
|
||||
query = """query {
|
||||
viewer { id name email }
|
||||
}"""
|
||||
|
||||
data = await self._gql_query(query)
|
||||
viewer = data.get("viewer", {})
|
||||
|
||||
if viewer:
|
||||
self._bot_user = BotUser(
|
||||
user_id=viewer.get("id", ""),
|
||||
display_name=viewer.get("name", "Linear Bot"),
|
||||
metadata=viewer
|
||||
)
|
||||
logger.info(f"Connected as: {self._bot_user.display_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not fetch bot user: {e}")
|
||||
|
||||
async def _send_comment(self, issue_id: str, comment: str, session_id: Optional[str] = None) -> None:
|
||||
"""Send a comment to a Linear issue."""
|
||||
try:
|
||||
mutation = """mutation($input: CommentCreateInput!) {
|
||||
commentCreate(input: $input) {
|
||||
success
|
||||
comment { id url }
|
||||
}
|
||||
}"""
|
||||
|
||||
variables = {
|
||||
"input": {
|
||||
"issueId": issue_id,
|
||||
"body": comment
|
||||
}
|
||||
}
|
||||
|
||||
data = await self._gql_query(mutation, variables)
|
||||
|
||||
if data.get("commentCreate", {}).get("success"):
|
||||
logger.info(f"Comment posted to issue {issue_id}")
|
||||
|
||||
# Optionally update agent session status
|
||||
if session_id:
|
||||
await self._update_agent_session(session_id, "response", comment)
|
||||
else:
|
||||
logger.error("Failed to post comment to Linear")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending comment: {e}")
|
||||
|
||||
async def _update_agent_session(self, session_id: str, activity_type: str, content: str) -> None:
|
||||
"""Update agent session with activity."""
|
||||
try:
|
||||
mutation = """mutation($input: AgentActivityCreateInput!) {
|
||||
agentActivityCreate(input: $input) {
|
||||
success
|
||||
}
|
||||
}"""
|
||||
|
||||
variables = {
|
||||
"input": {
|
||||
"agentSessionId": session_id,
|
||||
"type": activity_type,
|
||||
"content": content
|
||||
}
|
||||
}
|
||||
|
||||
await self._gql_query(mutation, variables)
|
||||
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not update agent session: {e}")
|
||||
|
||||
async def _gql_query(self, query: str, variables: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Execute GraphQL query against Linear API."""
|
||||
if not self._http_session:
|
||||
raise RuntimeError("HTTP session not initialized")
|
||||
|
||||
# OAuth tokens require Bearer prefix, API keys are sent raw
|
||||
auth_header = f"Bearer {self._token}" if self._is_oauth else self._token
|
||||
headers = {
|
||||
"Authorization": auth_header,
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"query": query,
|
||||
"variables": variables or {}
|
||||
}
|
||||
|
||||
import aiohttp
|
||||
async with self._http_session.post(
|
||||
LINEAR_API_BASE,
|
||||
headers=headers,
|
||||
data=json.dumps(payload),
|
||||
timeout=aiohttp.ClientTimeout(total=30),
|
||||
) as response:
|
||||
data = await response.json()
|
||||
|
||||
if "errors" in data:
|
||||
raise RuntimeError(f"Linear API error: {data['errors']}")
|
||||
|
||||
return data.get("data", {})
|
||||
|
||||
# ── Agent Integration ───────────────────────────────────────────
|
||||
|
||||
def set_agent(self, agent: "Agent") -> None:
|
||||
"""Set the agent that handles messages."""
|
||||
self._agent = agent
|
||||
|
||||
def get_agent(self) -> Optional["Agent"]:
|
||||
"""Get the current agent."""
|
||||
return self._agent
|
||||
|
||||
# ── Message Sending ─────────────────────────────────────────────
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
channel_id: str,
|
||||
content: Union[str, Dict[str, Any]],
|
||||
reply_to: Optional[str] = None,
|
||||
thread_id: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> BotMessage:
|
||||
"""Send a message (comment) to a Linear issue."""
|
||||
text = content if isinstance(content, str) else str(content)
|
||||
try:
|
||||
await self._send_comment(channel_id, text)
|
||||
return BotMessage(
|
||||
message_id=f"linear-{channel_id}-{int(time.time())}",
|
||||
content=text,
|
||||
message_type=MessageType.TEXT,
|
||||
channel=BotChannel(channel_id=channel_id, name="Linear Issue"),
|
||||
sender=self._bot_user or BotUser(user_id="linear-bot", display_name="Linear Bot"),
|
||||
timestamp=time.time(),
|
||||
reply_to=reply_to,
|
||||
thread_id=thread_id,
|
||||
metadata={"linear_channel": channel_id}
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send message: {e}")
|
||||
# Return a failed message rather than None to maintain protocol
|
||||
return BotMessage(
|
||||
message_id=f"linear-failed-{int(time.time())}",
|
||||
content=text,
|
||||
message_type=MessageType.TEXT,
|
||||
channel=BotChannel(channel_id=channel_id, name="Linear Issue"),
|
||||
sender=self._bot_user or BotUser(user_id="linear-bot", display_name="Linear Bot"),
|
||||
timestamp=time.time(),
|
||||
reply_to=reply_to,
|
||||
thread_id=thread_id,
|
||||
metadata={"error": str(e), "linear_channel": channel_id}
|
||||
)
|
||||
@@ -269,11 +269,7 @@ class SlackBot(ChatCommandMixin, MessageHookMixin):
|
||||
user_id = event.get("user", "unknown")
|
||||
logger.info(f"Message received: {text[:100]}...")
|
||||
text = await self._debouncer.debounce(user_id, text)
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, text,
|
||||
chat_id=str(channel_id) if channel_id else "",
|
||||
thread_id=event.get("thread_ts", "") or "",
|
||||
)
|
||||
response = await self._session.chat(self._agent, user_id, text)
|
||||
logger.info(f"Response sent: {response[:100]}...")
|
||||
|
||||
send_result = self.fire_message_sending(channel_id, str(response))
|
||||
@@ -316,11 +312,7 @@ class SlackBot(ChatCommandMixin, MessageHookMixin):
|
||||
try:
|
||||
user_id = event.get("user", "unknown")
|
||||
logger.info(f"@mention received: {text[:100]}...")
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, text,
|
||||
chat_id=str(event.get("channel", "")),
|
||||
thread_id=event.get("thread_ts", "") or "",
|
||||
)
|
||||
response = await self._session.chat(self._agent, user_id, text)
|
||||
logger.info(f"Response sent: {response[:100]}...")
|
||||
|
||||
# Determine if we should reply in thread
|
||||
|
||||
@@ -250,16 +250,9 @@ class TelegramBot(ChatCommandMixin, MessageHookMixin):
|
||||
ack_ctx = await self._ack.ack(react_fn=_tg_react)
|
||||
|
||||
user_id = str(update.message.from_user.id) if update.message.from_user else "unknown"
|
||||
user_name = (
|
||||
update.message.from_user.username or update.message.from_user.first_name or ""
|
||||
) if update.message.from_user else ""
|
||||
try:
|
||||
message_text = await self._debouncer.debounce(user_id, message_text)
|
||||
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,
|
||||
)
|
||||
response = await self._session.chat(self._agent, user_id, message_text)
|
||||
send_result = self.fire_message_sending(
|
||||
str(update.message.chat_id), str(response),
|
||||
reply_to=str(update.message.message_id),
|
||||
|
||||
@@ -2707,10 +2707,9 @@ def launch_browser(
|
||||
if debug:
|
||||
console.print("[dim] [DEBUG] Page loaded, content script should be active[/dim]")
|
||||
break
|
||||
except (asyncio.TimeoutError, Exception) as e:
|
||||
except:
|
||||
if debug:
|
||||
console.print(f"[dim] [DEBUG] Page load event timeout/error ({type(e).__name__}), continuing anyway[/dim]")
|
||||
logging.debug("Page load event timeout/error: %s", e)
|
||||
console.print("[dim] [DEBUG] Page load event timeout, continuing anyway[/dim]")
|
||||
|
||||
# Wait for content script to inject
|
||||
await asyncio.sleep(2)
|
||||
|
||||
@@ -62,22 +62,11 @@ def realtime_connect(
|
||||
session_id = f"realtime-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
# Build WebSocket URL
|
||||
if api_base:
|
||||
base = api_base
|
||||
else:
|
||||
from praisonai.llm.env import resolve_llm_endpoint
|
||||
ep = resolve_llm_endpoint(default_base="wss://api.openai.com")
|
||||
base = ep.base_url
|
||||
|
||||
base = api_base or os.environ.get("OPENAI_API_BASE", "wss://api.openai.com")
|
||||
if base.startswith("http"):
|
||||
base = base.replace("https://", "wss://").replace("http://", "ws://")
|
||||
|
||||
# Strip any existing /v1 path to avoid double /v1/v1/realtime
|
||||
base = base.rstrip('/')
|
||||
if base.endswith('/v1'):
|
||||
base = base[:-3]
|
||||
|
||||
url = f"{base}/v1/realtime?model={model}"
|
||||
url = f"{base.rstrip('/')}/v1/realtime?model={model}"
|
||||
|
||||
return RealtimeSession(
|
||||
id=session_id,
|
||||
|
||||
@@ -363,74 +363,6 @@ def bot_whatsapp(
|
||||
)
|
||||
|
||||
|
||||
@app.command("linear")
|
||||
def bot_linear(
|
||||
token: Optional[str] = typer.Option(None, "--token", "-t", help="Linear OAuth token", envvar="LINEAR_OAUTH_TOKEN"),
|
||||
signing_secret: Optional[str] = typer.Option(None, "--signing-secret", help="Linear webhook signing secret", envvar="LINEAR_WEBHOOK_SECRET"),
|
||||
port: int = typer.Option(8080, "--port", "-p", help="Webhook server port"),
|
||||
agent: Optional[str] = typer.Option(None, "--agent", "-a", help="Agent YAML configuration file"),
|
||||
model: Optional[str] = typer.Option(None, "--model", "-m", help="LLM model to use"),
|
||||
browser: bool = typer.Option(False, "--browser", help="Enable browser control"),
|
||||
browser_profile: str = typer.Option("default", "--browser-profile", help="Browser profile name"),
|
||||
browser_headless: bool = typer.Option(False, "--browser-headless", help="Run browser headless"),
|
||||
tools: Optional[List[str]] = typer.Option(None, "--tools", help="Tools to enable"),
|
||||
skills: Optional[List[str]] = typer.Option(None, "--skills", help="Skills to enable"),
|
||||
skills_dir: Optional[str] = typer.Option(None, "--skills-dir", help="Custom skills directory"),
|
||||
memory: bool = typer.Option(False, "--memory", help="Enable memory"),
|
||||
memory_provider: str = typer.Option("default", "--memory-provider", help="Memory provider"),
|
||||
knowledge: bool = typer.Option(False, "--knowledge", help="Enable knowledge/RAG"),
|
||||
knowledge_sources: Optional[List[str]] = typer.Option(None, "--knowledge-sources", help="Knowledge sources"),
|
||||
web_search: bool = typer.Option(False, "--web", "--web-search", help="Enable web search"),
|
||||
web_provider: str = typer.Option("duckduckgo", "--web-provider", help="Web search provider"),
|
||||
sandbox: bool = typer.Option(False, "--sandbox", help="Enable sandbox mode"),
|
||||
exec_enabled: bool = typer.Option(False, "--exec", help="Enable exec tool"),
|
||||
auto_approve: bool = typer.Option(False, "--auto-approve", help="Auto-approve all tool executions"),
|
||||
session_id: Optional[str] = typer.Option(None, "--session-id", help="Session ID"),
|
||||
user_id: Optional[str] = typer.Option(None, "--user-id", help="User ID for memory isolation"),
|
||||
thinking: Optional[str] = typer.Option(None, "--thinking", help="Thinking mode (off, minimal, low, medium, high)"),
|
||||
):
|
||||
"""Start a Linear bot with full agent capabilities.
|
||||
|
||||
Handles Linear AgentSession webhooks for issue mentions and assignments.
|
||||
|
||||
Examples:
|
||||
praisonai bot linear --token $LINEAR_OAUTH_TOKEN --signing-secret $LINEAR_WEBHOOK_SECRET
|
||||
praisonai bot linear --agent agents.yaml --memory --web --tools linear_tools
|
||||
"""
|
||||
from ..features.bots_cli import BotHandler, BotCapabilities
|
||||
|
||||
capabilities = BotCapabilities(
|
||||
model=model,
|
||||
browser=browser,
|
||||
browser_profile=browser_profile,
|
||||
browser_headless=browser_headless,
|
||||
tools=tools or [],
|
||||
skills=skills or [],
|
||||
skills_dir=skills_dir,
|
||||
memory=memory,
|
||||
memory_provider=memory_provider,
|
||||
knowledge=knowledge,
|
||||
knowledge_sources=knowledge_sources or [],
|
||||
web_search=web_search,
|
||||
web_search_provider=web_provider,
|
||||
sandbox=sandbox,
|
||||
exec_enabled=exec_enabled,
|
||||
auto_approve=auto_approve,
|
||||
session_id=session_id,
|
||||
user_id=user_id,
|
||||
thinking=thinking,
|
||||
)
|
||||
|
||||
handler = BotHandler()
|
||||
handler.start_linear(
|
||||
token=token,
|
||||
signing_secret=signing_secret,
|
||||
webhook_port=port,
|
||||
agent_file=agent,
|
||||
capabilities=capabilities,
|
||||
)
|
||||
|
||||
|
||||
@app.command("email")
|
||||
def bot_email(
|
||||
token: Optional[str] = typer.Option(None, "--token", "-t", help="Email app password", envvar="EMAIL_APP_PASSWORD"),
|
||||
@@ -570,7 +502,6 @@ Start a bot on any platform with: praisonai bot <platform>
|
||||
[green]discord[/green] Discord Bot API
|
||||
[green]slack[/green] Slack Socket Mode
|
||||
[green]whatsapp[/green] WhatsApp Cloud API or Web mode (QR scan)
|
||||
[green]linear[/green] Linear AgentSession webhooks
|
||||
[green]email[/green] Email via IMAP/SMTP
|
||||
[green]agentmail[/green] AgentMail API (API-first email for AI agents)
|
||||
|
||||
@@ -600,7 +531,6 @@ Start a bot on any platform with: praisonai bot <platform>
|
||||
praisonai bot discord --tools DuckDuckGoTool --memory --model gpt-4o
|
||||
praisonai bot whatsapp --token $WHATSAPP_ACCESS_TOKEN --phone-id $WHATSAPP_PHONE_NUMBER_ID
|
||||
praisonai bot whatsapp --mode web
|
||||
praisonai bot linear --token $LINEAR_OAUTH_TOKEN --signing-secret $LINEAR_WEBHOOK_SECRET
|
||||
"""
|
||||
try:
|
||||
from rich import print as rprint
|
||||
|
||||
@@ -434,8 +434,8 @@ class AgentSchedulerHandler:
|
||||
hours = int(uptime_delta.total_seconds() // 3600)
|
||||
minutes = int((uptime_delta.total_seconds() % 3600) // 60)
|
||||
uptime = f"{hours}h {minutes}m"
|
||||
except (ValueError, TypeError) as e:
|
||||
logger.debug("Failed to parse started_at timestamp '%s': %s", started_at, e)
|
||||
except:
|
||||
pass
|
||||
|
||||
# Display detailed info
|
||||
print(f"\n{'='*60}")
|
||||
|
||||
@@ -557,74 +557,6 @@ class BotHandler:
|
||||
except Exception:
|
||||
pass # neonize Go threads may already be gone
|
||||
|
||||
def start_linear(
|
||||
self,
|
||||
token: Optional[str] = None,
|
||||
signing_secret: Optional[str] = None,
|
||||
webhook_port: int = 8080,
|
||||
agent_file: Optional[str] = None,
|
||||
capabilities: Optional[BotCapabilities] = None,
|
||||
agent_config_dict: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Start a Linear bot.
|
||||
|
||||
Args:
|
||||
token: Linear OAuth token (or LINEAR_OAUTH_TOKEN env var)
|
||||
signing_secret: Webhook signing secret (or LINEAR_WEBHOOK_SECRET env var)
|
||||
webhook_port: Port for webhook server (default 8080)
|
||||
agent_file: Optional path to agent configuration file
|
||||
capabilities: Optional capabilities configuration
|
||||
agent_config_dict: Optional agent config dict from bot YAML
|
||||
"""
|
||||
self._load_dotenv()
|
||||
|
||||
# Environment variable fallbacks
|
||||
token = token or os.environ.get("LINEAR_OAUTH_TOKEN") or os.environ.get("LINEAR_API_KEY")
|
||||
signing_secret = signing_secret or os.environ.get("LINEAR_WEBHOOK_SECRET", "")
|
||||
|
||||
if not token:
|
||||
print("Error: Linear token required")
|
||||
print("Provide via --token or LINEAR_OAUTH_TOKEN environment variable")
|
||||
print("For personal API key, set LINEAR_API_KEY instead")
|
||||
return
|
||||
|
||||
if not signing_secret:
|
||||
print("Warning: LINEAR_WEBHOOK_SECRET not set - webhook signatures will not be verified")
|
||||
|
||||
# Set auto-approve if enabled
|
||||
if capabilities and capabilities.auto_approve:
|
||||
os.environ["PRAISONAI_AUTO_APPROVE"] = "true"
|
||||
logger.info("Auto-approve enabled for all tool executions")
|
||||
|
||||
try:
|
||||
from praisonai.bots import LinearBot
|
||||
except ImportError as e:
|
||||
print(f"Error: LinearBot import failed. {e}")
|
||||
print("Install with: pip install aiohttp")
|
||||
return
|
||||
|
||||
agent = self._load_agent(agent_file, capabilities, agent_config_dict=agent_config_dict)
|
||||
bot = LinearBot(
|
||||
token=token,
|
||||
agent=agent,
|
||||
signing_secret=signing_secret,
|
||||
webhook_port=webhook_port,
|
||||
)
|
||||
|
||||
self._print_startup_info("Linear", capabilities)
|
||||
print(f"Webhook server on port {webhook_port}")
|
||||
print(f"Webhook endpoint: http://0.0.0.0:{webhook_port}/webhook")
|
||||
if signing_secret:
|
||||
print("Webhook signature verification: enabled")
|
||||
else:
|
||||
print("Webhook signature verification: disabled (set LINEAR_WEBHOOK_SECRET)")
|
||||
|
||||
try:
|
||||
asyncio.run(bot.start())
|
||||
except KeyboardInterrupt:
|
||||
print("\nStopping bot...")
|
||||
asyncio.run(bot.stop())
|
||||
|
||||
def start_email(
|
||||
self,
|
||||
token: Optional[str] = None,
|
||||
|
||||
@@ -5,12 +5,9 @@ Provides token usage and cost tracking for agent executions.
|
||||
Usage: praisonai "prompt" --metrics
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, Tuple
|
||||
from .base import FlagHandler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MetricsHandler(FlagHandler):
|
||||
"""
|
||||
@@ -136,29 +133,33 @@ class MetricsHandler(FlagHandler):
|
||||
ctx = litellm._thread_context
|
||||
if hasattr(ctx, 'usage'):
|
||||
usage = ctx.usage
|
||||
metrics['prompt_tokens'] = getattr(usage, 'prompt_tokens', None)
|
||||
metrics['completion_tokens'] = getattr(usage, 'completion_tokens', None)
|
||||
metrics['total_tokens'] = getattr(usage, 'total_tokens', None)
|
||||
except ImportError:
|
||||
logger.debug("litellm not installed; skipping token metrics")
|
||||
except AttributeError as e:
|
||||
logger.debug("litellm internals changed (%s); skipping token metrics", e)
|
||||
if hasattr(usage, 'prompt_tokens'):
|
||||
metrics['prompt_tokens'] = usage.prompt_tokens
|
||||
if hasattr(usage, 'completion_tokens'):
|
||||
metrics['completion_tokens'] = usage.completion_tokens
|
||||
if hasattr(usage, 'total_tokens'):
|
||||
metrics['total_tokens'] = usage.total_tokens
|
||||
except:
|
||||
pass
|
||||
|
||||
# Try to get cost from litellm cost tracking
|
||||
try:
|
||||
import litellm
|
||||
from litellm import completion_cost
|
||||
if metrics.get('prompt_tokens') and metrics.get('completion_tokens'):
|
||||
# If we have token counts, estimate cost
|
||||
if 'prompt_tokens' in metrics and 'completion_tokens' in metrics:
|
||||
model = metrics.get('model', 'gpt-4o-mini')
|
||||
try:
|
||||
metrics['cost'] = completion_cost(
|
||||
cost = completion_cost(
|
||||
model=model,
|
||||
prompt_tokens=metrics['prompt_tokens'],
|
||||
completion_tokens=metrics['completion_tokens'],
|
||||
completion_tokens=metrics['completion_tokens']
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("cost calc failed for model=%s: %s", model, e)
|
||||
except ImportError:
|
||||
logger.debug("litellm.completion_cost unavailable; skipping cost metrics")
|
||||
metrics['cost'] = cost
|
||||
except:
|
||||
pass
|
||||
except:
|
||||
pass
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
@@ -70,6 +70,7 @@ import yaml
|
||||
import time
|
||||
from rich import print
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
import shutil
|
||||
import subprocess
|
||||
import logging
|
||||
@@ -109,11 +110,6 @@ _BLOCKED_ENV_KEYS = frozenset({
|
||||
_BLOCKED_ENV_KEYS_UPPER = frozenset(k.upper() for k in _BLOCKED_ENV_KEYS)
|
||||
|
||||
|
||||
def _load_env_once():
|
||||
"""Load environment variables from .env file once at CLI startup."""
|
||||
load_dotenv()
|
||||
|
||||
|
||||
def _validate_env_key(key) -> None:
|
||||
"""Raise ``ValueError`` if *key* is a blocked environment variable name.
|
||||
|
||||
@@ -272,26 +268,29 @@ class PraisonAI:
|
||||
self.agent_yaml = agent_yaml
|
||||
self._interactive_mode = False # Flag for interactive TUI mode
|
||||
# Create config_list with AutoGen compatibility
|
||||
# Resolve LLM endpoint configuration from environment variables
|
||||
from praisonai.llm.env import resolve_llm_endpoint
|
||||
ep = resolve_llm_endpoint()
|
||||
# Support multiple environment variable patterns for better compatibility
|
||||
# Priority order: MODEL_NAME > OPENAI_MODEL_NAME for model selection
|
||||
model_name = os.environ.get("MODEL_NAME") or os.environ.get("OPENAI_MODEL_NAME", "gpt-4o-mini")
|
||||
|
||||
# Priority order for base_url: OPENAI_BASE_URL > OPENAI_API_BASE > OLLAMA_API_BASE
|
||||
# OPENAI_BASE_URL is the standard OpenAI SDK environment variable
|
||||
base_url = (
|
||||
os.environ.get("OPENAI_BASE_URL") or
|
||||
os.environ.get("OPENAI_API_BASE") or
|
||||
os.environ.get("OLLAMA_API_BASE", "https://api.openai.com/v1")
|
||||
)
|
||||
|
||||
api_key = os.environ.get("OPENAI_API_KEY")
|
||||
self.config_list = [
|
||||
{
|
||||
'model': ep.model,
|
||||
'base_url': ep.base_url,
|
||||
'api_key': ep.api_key,
|
||||
'model': model_name,
|
||||
'base_url': base_url,
|
||||
'api_key': api_key,
|
||||
'api_type': 'openai' # AutoGen expects this field
|
||||
}
|
||||
]
|
||||
self.agent_file = agent_file
|
||||
self.framework = framework
|
||||
|
||||
# Validate framework availability early to fail fast
|
||||
if self.framework:
|
||||
from praisonai.framework_adapters.validators import assert_framework_available
|
||||
assert_framework_available(self.framework)
|
||||
|
||||
self.auto = auto
|
||||
self.init = init
|
||||
self.tools = tools or [] # Store tool class names as a list
|
||||
@@ -349,9 +348,6 @@ class PraisonAI:
|
||||
initializes the necessary attributes, and then calls the appropriate methods based on the
|
||||
provided arguments.
|
||||
"""
|
||||
# Load environment variables from .env file
|
||||
_load_env_once()
|
||||
|
||||
# Warning filters now installed via Typer callback for CLI-only usage
|
||||
|
||||
# Telemetry defaults now handled in PraisonAI.__init__ with Langfuse awareness
|
||||
@@ -380,11 +376,6 @@ class PraisonAI:
|
||||
|
||||
self.framework = args.framework or self.framework
|
||||
|
||||
# Validate framework availability early to fail fast
|
||||
if self.framework:
|
||||
from praisonai.framework_adapters.validators import assert_framework_available
|
||||
assert_framework_available(self.framework)
|
||||
|
||||
# Update config_list model if --model flag is provided
|
||||
if getattr(args, 'model', None):
|
||||
self.config_list[0]['model'] = args.model
|
||||
|
||||
@@ -97,4 +97,5 @@ def __getattr__(name: str):
|
||||
_register_builtin_backends()
|
||||
raise AttributeError(f"No attribute {name}")
|
||||
|
||||
# Registration happens lazily via __getattr__ on first attribute access
|
||||
# Trigger registration on import
|
||||
_register_builtin_backends()
|
||||
@@ -290,7 +290,10 @@ class PraisonAIDB:
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Called when a new run (turn) starts."""
|
||||
# Store run start in state if available (even without conversation store)
|
||||
if not self._conversation_store:
|
||||
return
|
||||
|
||||
# Store run start in state if available
|
||||
if self._state_store:
|
||||
run_key = f"run:{session_id}:{run_id}"
|
||||
self._state_store.set(run_key, {
|
||||
@@ -301,10 +304,6 @@ class PraisonAIDB:
|
||||
"status": "running",
|
||||
"metadata": metadata or {}
|
||||
})
|
||||
|
||||
# Early return only applies to conversation store operations
|
||||
if not self._conversation_store:
|
||||
return
|
||||
|
||||
def on_run_end(
|
||||
self,
|
||||
@@ -333,33 +332,10 @@ class PraisonAIDB:
|
||||
session_id: str,
|
||||
limit: Optional[int] = None,
|
||||
) -> List:
|
||||
"""Return persisted runs for a session, newest first."""
|
||||
self._init_stores()
|
||||
if not self._state_store:
|
||||
logger.warning("get_runs() called but no state_url configured; returning []")
|
||||
return []
|
||||
|
||||
prefix = f"run:{session_id}:"
|
||||
# State stores expose scan_prefix() per the protocol; fall back to keys() for legacy.
|
||||
if hasattr(self._state_store, "scan_prefix"):
|
||||
keys = self._state_store.scan_prefix(prefix)
|
||||
else:
|
||||
keys = [k for k in self._state_store.keys() if k.startswith(prefix)]
|
||||
|
||||
runs = []
|
||||
for k in keys:
|
||||
try:
|
||||
run_data = self._state_store.get(k)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to load run key '%s': %s", k, exc)
|
||||
continue
|
||||
if isinstance(run_data, dict):
|
||||
runs.append(run_data)
|
||||
|
||||
runs.sort(key=lambda r: r.get("started_at", 0), reverse=True)
|
||||
if limit is not None:
|
||||
return runs[: max(limit, 0)]
|
||||
return runs
|
||||
"""Get runs for a session."""
|
||||
# For now, runs are stored in state store
|
||||
# Full implementation would query from conversation store
|
||||
return []
|
||||
|
||||
def export_session(
|
||||
self,
|
||||
@@ -523,36 +499,9 @@ class PraisonAIDB:
|
||||
user_id: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> List:
|
||||
"""Return persisted traces, optionally filtered by session/user."""
|
||||
self._init_stores()
|
||||
if not self._state_store:
|
||||
logger.warning("get_traces() called but no state_url configured; returning []")
|
||||
return []
|
||||
|
||||
prefix = "trace:"
|
||||
if hasattr(self._state_store, "scan_prefix"):
|
||||
keys = self._state_store.scan_prefix(prefix)
|
||||
else:
|
||||
keys = [k for k in self._state_store.keys() if k.startswith(prefix)]
|
||||
|
||||
traces = []
|
||||
for k in keys:
|
||||
try:
|
||||
trace_data = self._state_store.get(k)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to load trace key '%s': %s", k, exc)
|
||||
continue
|
||||
if isinstance(trace_data, dict):
|
||||
traces.append(trace_data)
|
||||
|
||||
if session_id is not None:
|
||||
traces = [t for t in traces if t.get("session_id") == session_id]
|
||||
if user_id is not None:
|
||||
traces = [t for t in traces if t.get("user_id") == user_id]
|
||||
traces.sort(key=lambda t: t.get("started_at", 0), reverse=True)
|
||||
if limit is not None:
|
||||
return traces[: max(limit, 0)]
|
||||
return traces
|
||||
"""Get traces with optional filters."""
|
||||
# For now, return empty list - full implementation would scan state store
|
||||
return []
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close all database connections."""
|
||||
|
||||
@@ -57,7 +57,7 @@ class CloudDeployer:
|
||||
file.write("FROM python:3.11-slim\n")
|
||||
file.write("WORKDIR /app\n")
|
||||
file.write("COPY . .\n")
|
||||
file.write("RUN pip install flask praisonai==4.6.37 gunicorn markdown\n")
|
||||
file.write("RUN pip install flask praisonai==4.6.35 gunicorn markdown\n")
|
||||
file.write("EXPOSE 8080\n")
|
||||
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
|
||||
|
||||
@@ -92,12 +92,9 @@ class CloudDeployer:
|
||||
|
||||
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
|
||||
os.environ["OPENAI_MODEL_NAME"] = os.getenv("OPENAI_MODEL_NAME", "gpt-4o-mini")
|
||||
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "Enter your API key")
|
||||
os.environ["OPENAI_API_BASE"] = os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1")
|
||||
|
||||
def run_commands(self):
|
||||
"""
|
||||
@@ -132,11 +129,9 @@ class CloudDeployer:
|
||||
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
|
||||
openai_model = os.environ.get('OPENAI_MODEL_NAME', 'gpt-5-nano')
|
||||
openai_key = os.environ.get('OPENAI_API_KEY', 'Enter your API key')
|
||||
openai_base = os.environ.get('OPENAI_API_BASE', 'https://api.openai.com/v1')
|
||||
|
||||
# Build commands with actual values
|
||||
commands = [
|
||||
|
||||
@@ -38,8 +38,9 @@ class AutoGenAdapter(BaseFrameworkAdapter):
|
||||
Returns:
|
||||
Execution result as string
|
||||
"""
|
||||
# Availability already validated at CLI entry
|
||||
|
||||
if not self.is_available():
|
||||
raise ImportError("AutoGen v0.2 is not available. Install with: pip install autogen")
|
||||
|
||||
# Import AutoGen only when needed
|
||||
import autogen
|
||||
|
||||
@@ -122,8 +123,9 @@ class AutoGenV4Adapter(BaseFrameworkAdapter):
|
||||
Returns:
|
||||
Execution result as string
|
||||
"""
|
||||
# Availability already validated at CLI entry
|
||||
|
||||
if not self.is_available():
|
||||
raise ImportError("AutoGen v0.4 is not available. Install with: pip install autogen-agentchat autogen-ext")
|
||||
|
||||
logger.info("Starting AutoGen v0.4 execution...")
|
||||
# For now, return a proper error message instead of delegating
|
||||
# TODO: Implement full AutoGen v0.4 adapter logic
|
||||
@@ -160,7 +162,8 @@ class AG2Adapter(BaseFrameworkAdapter):
|
||||
Returns:
|
||||
Execution result as string
|
||||
"""
|
||||
# Availability already validated at CLI entry
|
||||
if not self.is_available():
|
||||
raise ImportError("AG2 is not available. Install with: pip install ag2")
|
||||
|
||||
logger.info("Starting AG2 execution...")
|
||||
# For now, return a proper error message instead of delegating
|
||||
|
||||
@@ -5,7 +5,7 @@ This module defines the protocol that all framework adapters must implement,
|
||||
enabling lazy-loaded, protocol-driven framework support.
|
||||
"""
|
||||
|
||||
from typing import Protocol, Dict, List, Any, Optional, Callable
|
||||
from typing import Protocol, Dict, List, Any, Optional
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
@@ -20,17 +20,7 @@ class FrameworkAdapter(Protocol):
|
||||
"""Check if the framework is available for import."""
|
||||
...
|
||||
|
||||
def run(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
llm_config: List[Dict],
|
||||
topic: str,
|
||||
*,
|
||||
tools_dict: Optional[Dict[str, Any]] = None,
|
||||
agent_callback: Optional[Callable] = None,
|
||||
task_callback: Optional[Callable] = None,
|
||||
cli_config: Optional[Dict[str, Any]] = None,
|
||||
) -> str:
|
||||
def run(self, config: Dict[str, Any], llm_config: List[Dict], topic: str) -> str:
|
||||
"""
|
||||
Run the framework with given configuration.
|
||||
|
||||
@@ -38,10 +28,6 @@ class FrameworkAdapter(Protocol):
|
||||
config: Framework configuration
|
||||
llm_config: LLM configuration list
|
||||
topic: Topic for the tasks
|
||||
tools_dict: Available tools dictionary
|
||||
agent_callback: Callback for agent events
|
||||
task_callback: Callback for task events
|
||||
cli_config: CLI configuration
|
||||
|
||||
Returns:
|
||||
Execution result as string
|
||||
|
||||
@@ -5,7 +5,7 @@ Provides lazy-loaded, scoped integration with CrewAI framework.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Any, Optional
|
||||
from typing import Dict, List, Any
|
||||
from .base import BaseFrameworkAdapter, scoped_telemetry_disable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -26,17 +26,7 @@ class CrewAIAdapter(BaseFrameworkAdapter):
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def run(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
llm_config: List[Dict],
|
||||
topic: str,
|
||||
*,
|
||||
tools_dict: Optional[Dict[str, Any]] = None,
|
||||
agent_callback = None,
|
||||
task_callback = None,
|
||||
cli_config: Optional[Dict[str, Any]] = None,
|
||||
) -> str:
|
||||
def run(self, config: Dict[str, Any], llm_config: List[Dict], topic: str) -> str:
|
||||
"""
|
||||
Run CrewAI with given configuration.
|
||||
|
||||
@@ -44,21 +34,17 @@ class CrewAIAdapter(BaseFrameworkAdapter):
|
||||
config: CrewAI configuration with agents and tasks
|
||||
llm_config: LLM configuration list
|
||||
topic: Topic for the tasks
|
||||
tools_dict: Available tools dictionary
|
||||
agent_callback: Callback for agent events
|
||||
task_callback: Callback for task events
|
||||
cli_config: CLI configuration
|
||||
|
||||
Returns:
|
||||
Execution result as string
|
||||
"""
|
||||
# Import CrewAI only when needed (availability already validated at CLI entry)
|
||||
if not self.is_available():
|
||||
raise ImportError("CrewAI is not available. Install with: pip install crewai")
|
||||
|
||||
# Import CrewAI only when needed
|
||||
from crewai import Agent, Task, Crew
|
||||
from crewai.telemetry import Telemetry
|
||||
|
||||
# Suppress crewai.cli.config logger (scoped to when CrewAI is actually used)
|
||||
logging.getLogger('crewai.cli.config').setLevel(logging.ERROR)
|
||||
|
||||
# Use scoped telemetry disabling instead of global patching
|
||||
with scoped_telemetry_disable(Telemetry):
|
||||
# For now, use simplified CrewAI execution
|
||||
@@ -67,22 +53,13 @@ class CrewAIAdapter(BaseFrameworkAdapter):
|
||||
|
||||
# Create agents
|
||||
for agent_name, agent_details in config.get('roles', {}).items():
|
||||
# Resolve tools for this agent from tools_dict
|
||||
agent_tool_list = []
|
||||
if tools_dict:
|
||||
agent_tools = agent_details.get('tools', [])
|
||||
agent_tool_list = [tools_dict[t] for t in agent_tools if t in tools_dict]
|
||||
|
||||
agent = Agent(
|
||||
role=agent_details.get('role', agent_name),
|
||||
goal=self._format_template(agent_details.get('goal', ''), topic=topic),
|
||||
backstory=self._format_template(agent_details.get('backstory', ''), topic=topic),
|
||||
tools=agent_tool_list,
|
||||
verbose=True,
|
||||
allow_delegation=agent_details.get('allow_delegation', False)
|
||||
)
|
||||
if agent_callback:
|
||||
agent.step_callback = agent_callback
|
||||
agents[agent_name] = agent
|
||||
|
||||
# Create tasks
|
||||
@@ -93,8 +70,6 @@ class CrewAIAdapter(BaseFrameworkAdapter):
|
||||
expected_output=self._format_template(task_details['expected_output'], topic=topic),
|
||||
agent=agents[agent_name]
|
||||
)
|
||||
if task_callback:
|
||||
task.callback = task_callback
|
||||
tasks.append(task)
|
||||
|
||||
# Create and run crew
|
||||
|
||||
@@ -5,7 +5,7 @@ Provides lazy-loaded integration with the PraisonAI agents framework.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Any, Optional
|
||||
from typing import Dict, List, Any
|
||||
from .base import BaseFrameworkAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -26,17 +26,7 @@ class PraisonAIAdapter(BaseFrameworkAdapter):
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def run(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
llm_config: List[Dict],
|
||||
topic: str,
|
||||
*,
|
||||
tools_dict: Optional[Dict[str, Any]] = None,
|
||||
agent_callback = None,
|
||||
task_callback = None,
|
||||
cli_config: Optional[Dict[str, Any]] = None,
|
||||
) -> str:
|
||||
def run(self, config: Dict[str, Any], llm_config: List[Dict], topic: str) -> str:
|
||||
"""
|
||||
Run PraisonAI agents with given configuration.
|
||||
|
||||
@@ -44,113 +34,16 @@ class PraisonAIAdapter(BaseFrameworkAdapter):
|
||||
config: PraisonAI configuration with agents and tasks
|
||||
llm_config: LLM configuration list
|
||||
topic: Topic for the tasks
|
||||
tools_dict: Available tools dictionary
|
||||
agent_callback: Callback for agent events
|
||||
task_callback: Callback for task events
|
||||
cli_config: CLI configuration
|
||||
|
||||
Returns:
|
||||
Execution result as string
|
||||
"""
|
||||
# Availability already validated at CLI entry
|
||||
|
||||
# Import PraisonAI components only when needed
|
||||
from praisonaiagents import Agent as PraisonAgent, Task as PraisonTask, AgentTeam
|
||||
|
||||
if not self.is_available():
|
||||
raise ImportError("PraisonAI agents is not available. Install with: pip install praisonaiagents")
|
||||
|
||||
logger.info("Starting PraisonAI execution...")
|
||||
|
||||
# Basic implementation - create agents and tasks from config
|
||||
agents = {}
|
||||
tasks = []
|
||||
|
||||
# Get model from llm_config or environment
|
||||
model_name = "gpt-4o-mini"
|
||||
if llm_config and llm_config[0].get('model'):
|
||||
model_name = llm_config[0]['model']
|
||||
|
||||
# Create agents from roles
|
||||
for role, details in config.get('roles', {}).items():
|
||||
role_filled = self._format_template(details.get('role', role), topic=topic)
|
||||
goal_filled = self._format_template(details.get('goal', ''), topic=topic)
|
||||
backstory_filled = self._format_template(details.get('backstory', ''), topic=topic)
|
||||
|
||||
# Resolve tools for this agent from tools_dict
|
||||
agent_tool_list = []
|
||||
if tools_dict:
|
||||
agent_tools = details.get('tools', [])
|
||||
agent_tool_list = [tools_dict[t] for t in agent_tools if t in tools_dict]
|
||||
|
||||
# Create basic agent
|
||||
agent = PraisonAgent(
|
||||
name=role_filled,
|
||||
role=role_filled,
|
||||
goal=goal_filled,
|
||||
backstory=backstory_filled,
|
||||
instructions=details.get('instructions'),
|
||||
llm=model_name,
|
||||
allow_delegation=details.get('allow_delegation', False),
|
||||
tools=agent_tool_list,
|
||||
)
|
||||
|
||||
if agent_callback:
|
||||
agent.step_callback = agent_callback
|
||||
|
||||
agents[role] = agent
|
||||
|
||||
# Create tasks for the agent
|
||||
agent_tasks = details.get('tasks', {})
|
||||
if not agent_tasks:
|
||||
# Auto-generate a task
|
||||
task_description = details.get('instructions') or backstory_filled
|
||||
task = PraisonTask(
|
||||
description=task_description,
|
||||
expected_output="Complete the assigned task successfully.",
|
||||
agent=agent,
|
||||
)
|
||||
if task_callback:
|
||||
task.callback = task_callback
|
||||
tasks.append(task)
|
||||
else:
|
||||
for task_name, task_details in agent_tasks.items():
|
||||
description_filled = self._format_template(
|
||||
task_details['description'], topic=topic
|
||||
)
|
||||
expected_output_filled = self._format_template(
|
||||
task_details['expected_output'], topic=topic
|
||||
)
|
||||
|
||||
task = PraisonTask(
|
||||
description=description_filled,
|
||||
expected_output=expected_output_filled,
|
||||
agent=agent,
|
||||
)
|
||||
|
||||
if task_callback:
|
||||
task.callback = task_callback
|
||||
|
||||
tasks.append(task)
|
||||
|
||||
# Create and run the team
|
||||
memory = config.get('memory', False)
|
||||
|
||||
if config.get('process') == 'hierarchical':
|
||||
team = AgentTeam(
|
||||
agents=list(agents.values()),
|
||||
tasks=tasks,
|
||||
process="hierarchical",
|
||||
manager_llm=config.get('manager_llm') or model_name,
|
||||
memory=memory
|
||||
)
|
||||
else:
|
||||
team = AgentTeam(
|
||||
agents=list(agents.values()),
|
||||
tasks=tasks,
|
||||
memory=memory
|
||||
)
|
||||
|
||||
response = team.start()
|
||||
result = f"### PraisonAI Output ###\n{response}" if response else "### PraisonAI Output ###\nTask completed."
|
||||
|
||||
logger.info("PraisonAI execution completed")
|
||||
return result
|
||||
# For now, return a proper error message instead of delegating
|
||||
# TODO: Implement full PraisonAI adapter logic
|
||||
logger.warning("PraisonAI adapter is not yet fully implemented")
|
||||
return "### PraisonAI Output ###\nPraisonAI adapter is not yet fully implemented. Please use 'crewai' framework for similar functionality."
|
||||
|
||||
@@ -1,36 +0,0 @@
|
||||
"""
|
||||
Framework availability validators.
|
||||
|
||||
Provides early validation of framework availability to fail fast at CLI entry
|
||||
rather than inside run() methods after expensive setup work.
|
||||
"""
|
||||
|
||||
from .registry import FrameworkAdapterRegistry
|
||||
|
||||
|
||||
# Install hints for common frameworks
|
||||
_INSTALL_HINTS = {
|
||||
"crewai": "pip install 'praisonai[crewai]' # or: pip install crewai",
|
||||
"autogen": "pip install 'praisonai[autogen]' # or: pip install pyautogen",
|
||||
"praisonai": "pip install praisonaiagents",
|
||||
}
|
||||
|
||||
|
||||
def assert_framework_available(name: str) -> None:
|
||||
"""
|
||||
Raise ImportError immediately if the chosen framework is missing.
|
||||
|
||||
Args:
|
||||
name: Framework name to validate
|
||||
|
||||
Raises:
|
||||
ImportError: If framework is not available with actionable install hint
|
||||
"""
|
||||
registry = FrameworkAdapterRegistry.get_instance()
|
||||
|
||||
if not registry.is_available(name):
|
||||
hint = _INSTALL_HINTS.get(name, f"pip install {name}")
|
||||
raise ImportError(
|
||||
f"Framework '{name}' was requested but is not installed.\n"
|
||||
f"Install it with:\n {hint}"
|
||||
)
|
||||
@@ -1694,16 +1694,6 @@ class WebSocketGateway:
|
||||
mode=wa_mode,
|
||||
creds_dir=ch_cfg.get("creds_dir"),
|
||||
)
|
||||
elif channel_type == "linear":
|
||||
from praisonai.bots import LinearBot
|
||||
linear_token = token or os.environ.get("LINEAR_OAUTH_TOKEN", "") or os.environ.get("LINEAR_API_KEY", "")
|
||||
return LinearBot(
|
||||
token=linear_token,
|
||||
agent=agent,
|
||||
config=config,
|
||||
signing_secret=ch_cfg.get("signing_secret", "") or os.environ.get("LINEAR_WEBHOOK_SECRET", ""),
|
||||
webhook_port=int(ch_cfg.get("webhook_port", 8080)),
|
||||
)
|
||||
elif channel_type == "email":
|
||||
from praisonai.bots import EmailBot
|
||||
email_token = token or os.environ.get("EMAIL_APP_PASSWORD", "")
|
||||
@@ -1746,8 +1736,8 @@ class WebSocketGateway:
|
||||
# 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
|
||||
elif type(bot).__name__ == "WhatsAppBot":
|
||||
# WhatsApp runs its own aiohttp webhook server
|
||||
self._inject_routing_handler(name, bot)
|
||||
await bot.start()
|
||||
else:
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
"""
|
||||
LLM endpoint environment variable resolver.
|
||||
|
||||
Provides a single source of truth for resolving LLM endpoint configuration
|
||||
from environment variables, ensuring consistent precedence across all components.
|
||||
"""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LLMEndpoint:
|
||||
"""LLM endpoint configuration resolved from environment variables."""
|
||||
model: str
|
||||
base_url: str
|
||||
api_key: Optional[str]
|
||||
|
||||
|
||||
# Documented, single precedence list. Add new providers here only.
|
||||
_MODEL_VARS = ("MODEL_NAME", "OPENAI_MODEL_NAME")
|
||||
_BASE_URL_VARS = ("OPENAI_BASE_URL", "OPENAI_API_BASE", "OLLAMA_API_BASE")
|
||||
_DEFAULT_MODEL = "gpt-4o-mini"
|
||||
_DEFAULT_BASE = "https://api.openai.com/v1"
|
||||
|
||||
|
||||
def _first_set(*names: str) -> Optional[str]:
|
||||
"""Return the first environment variable that is set and non-empty."""
|
||||
for name in names:
|
||||
value = os.environ.get(name)
|
||||
if value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def resolve_llm_endpoint(*, default_base: str = _DEFAULT_BASE) -> LLMEndpoint:
|
||||
"""
|
||||
Resolve LLM endpoint configuration from environment variables.
|
||||
|
||||
Precedence order:
|
||||
- Model: MODEL_NAME > OPENAI_MODEL_NAME > default
|
||||
- Base URL: OPENAI_BASE_URL > OPENAI_API_BASE > OLLAMA_API_BASE > default
|
||||
- API Key: OPENAI_API_KEY
|
||||
|
||||
Args:
|
||||
default_base: Default base URL if none found in environment variables
|
||||
|
||||
Returns:
|
||||
LLMEndpoint with resolved configuration
|
||||
"""
|
||||
return LLMEndpoint(
|
||||
model=_first_set(*_MODEL_VARS) or _DEFAULT_MODEL,
|
||||
base_url=_first_set(*_BASE_URL_VARS) or default_base,
|
||||
api_key=os.environ.get("OPENAI_API_KEY"),
|
||||
)
|
||||
@@ -268,69 +268,10 @@ def _register_builtin_providers(registry: LLMProviderRegistry) -> None:
|
||||
|
||||
Uses lazy loading to avoid importing heavy dependencies at module load time.
|
||||
"""
|
||||
# Built-in adapter that wraps LiteLLM so create_llm_provider works out of the box.
|
||||
class _LiteLLMProvider:
|
||||
"""Generic LiteLLM-backed provider used for openai/anthropic/google/etc."""
|
||||
def __init__(self, model_id: str, config: Optional[Dict[str, Any]] = None):
|
||||
self.provider_id = "litellm"
|
||||
self.model_id = model_id
|
||||
self.config = config or {}
|
||||
|
||||
def _resolve_model_and_kwargs(self, prompt: str, **kwargs):
|
||||
"""Helper to resolve model and kwargs for both sync and async methods."""
|
||||
try:
|
||||
import litellm # lazy
|
||||
except ImportError as err:
|
||||
raise ImportError(
|
||||
"LiteLLM is required for built-in providers. "
|
||||
"Install with: pip install litellm"
|
||||
) from err
|
||||
|
||||
provider_prefix = self.config.get("provider", "")
|
||||
full_model = f"{provider_prefix}/{self.model_id}".strip("/") if provider_prefix else self.model_id
|
||||
completion_kwargs = {
|
||||
k: v for k, v in {**self.config, **kwargs}.items() if k != "provider"
|
||||
}
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
return litellm, full_model, messages, completion_kwargs
|
||||
|
||||
def generate(self, prompt: str, **kwargs):
|
||||
"""Sync variant — uses litellm.completion()."""
|
||||
litellm, full_model, messages, completion_kwargs = self._resolve_model_and_kwargs(prompt, **kwargs)
|
||||
return litellm.completion(
|
||||
model=full_model,
|
||||
messages=messages,
|
||||
**completion_kwargs,
|
||||
)
|
||||
|
||||
async def generate_async(self, prompt: str, **kwargs):
|
||||
"""Async variant — uses litellm.acompletion() to avoid blocking the event loop.
|
||||
|
||||
generate() calls litellm.completion() which is a blocking network call.
|
||||
Calling it from an async context would stall the entire event loop.
|
||||
generate_async() uses litellm.acompletion() — the native async variant.
|
||||
"""
|
||||
litellm, full_model, messages, completion_kwargs = self._resolve_model_and_kwargs(prompt, **kwargs)
|
||||
return await litellm.acompletion(
|
||||
model=full_model,
|
||||
messages=messages,
|
||||
**completion_kwargs,
|
||||
)
|
||||
|
||||
def _make_litellm_factory(provider_prefix: str):
|
||||
def factory(model_id, config=None):
|
||||
cfg = dict(config or {})
|
||||
cfg.setdefault("provider", provider_prefix)
|
||||
return _LiteLLMProvider(model_id, cfg)
|
||||
return factory
|
||||
|
||||
# Cover the providers parse_model_string() already special-cases.
|
||||
for name, aliases in [
|
||||
("openai", ("oai",)),
|
||||
("anthropic", ("claude",)),
|
||||
("google", ("gemini", "google_genai")),
|
||||
]:
|
||||
registry.register(name, _make_litellm_factory(name), aliases=list(aliases))
|
||||
# Note: In Python, the actual LLM providers are in praisonaiagents.llm
|
||||
# which uses LiteLLM. This registry is for custom provider extensions.
|
||||
# Built-in providers are handled by LiteLLM automatically.
|
||||
pass
|
||||
|
||||
|
||||
def register_llm_provider(
|
||||
|
||||
@@ -11,18 +11,14 @@ Architecture:
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from praisonaiagents.trace.protocol import ActionEvent, ActionEventType, TraceSinkProtocol
|
||||
from praisonaiagents.trace.context_events import ContextEvent, ContextEventType, ContextTraceSinkProtocol
|
||||
|
||||
logger = logging.getLogger("observability.langfuse")
|
||||
|
||||
|
||||
@dataclass
|
||||
class LangfuseSinkConfig:
|
||||
@@ -70,16 +66,14 @@ class LangfuseSink:
|
||||
Thread-safe: langfuse.Langfuse handles its own batching and thread safety.
|
||||
"""
|
||||
|
||||
__slots__ = ("_config", "_client", "_traces", "_spans", "_tool_stacks", "_lock", "_closed", "_metadata")
|
||||
__slots__ = ("_config", "_client", "_traces", "_spans", "_lock", "_closed", "_metadata")
|
||||
|
||||
def __init__(self, config: Optional[LangfuseSinkConfig] = None, metadata: Optional[Dict[str, Any]] = None):
|
||||
self._config = config or LangfuseSinkConfig()
|
||||
self._client: Optional[Any] = None # Lazy-loaded langfuse.Langfuse
|
||||
self._traces: Dict[str, Any] = {} # agent_name -> trace observation
|
||||
self._metadata = metadata or {} # Additional metadata for traces
|
||||
self._spans: Dict[str, Any] = {} # agent_key -> root span
|
||||
# (agent_key, tool_name) -> stack of in-flight tool spans
|
||||
self._tool_stacks: Dict[Tuple[str, str], List[Any]] = defaultdict(list)
|
||||
self._spans: Dict[str, Any] = {} # span_key -> span observation
|
||||
self._lock = threading.Lock()
|
||||
self._closed = False
|
||||
|
||||
@@ -121,7 +115,7 @@ class LangfuseSink:
|
||||
self._handle_event(event)
|
||||
except Exception as e:
|
||||
# Don't let observability errors break agent execution
|
||||
logger.exception("LangfuseSink error")
|
||||
print(f"LangfuseSink error: {e}")
|
||||
|
||||
def _handle_event(self, event: ActionEvent) -> None:
|
||||
"""Handle a single ActionEvent (called under lock)."""
|
||||
@@ -172,16 +166,6 @@ class LangfuseSink:
|
||||
def _handle_agent_end(self, event: ActionEvent, agent_name: str) -> None:
|
||||
"""Handle AGENT_END -> end root span observation."""
|
||||
agent_key = f"{event.agent_id or agent_name}-{agent_name}"
|
||||
# End and clear any dangling tool spans for this agent
|
||||
for stack_key, stack in list(self._tool_stacks.items()):
|
||||
if stack_key[0] != agent_key:
|
||||
continue
|
||||
while stack:
|
||||
try:
|
||||
stack.pop().end()
|
||||
except Exception:
|
||||
logger.exception("LangfuseSink: failed to close dangling tool span for %s", stack_key[1])
|
||||
self._tool_stacks.pop(stack_key, None)
|
||||
span = self._spans.get(agent_key)
|
||||
if span:
|
||||
span.update(
|
||||
@@ -203,10 +187,16 @@ class LangfuseSink:
|
||||
|
||||
tool_name = event.tool_name or "unknown-tool"
|
||||
|
||||
# Generate unique tool invocation key with UUID for collision safety
|
||||
import uuid
|
||||
tool_invocation_id = str(uuid.uuid4())[:8] # Short UUID
|
||||
tool_key = f"{agent_key}:{tool_name}:{tool_invocation_id}"
|
||||
|
||||
# Merge flow correlation metadata with tool metadata
|
||||
tool_metadata = {
|
||||
"tool_name": tool_name,
|
||||
"agent_name": agent_name,
|
||||
"tool_invocation_id": tool_invocation_id,
|
||||
**(event.metadata if event.metadata else {}),
|
||||
**self._metadata # Include flow correlation metadata
|
||||
}
|
||||
@@ -218,20 +208,25 @@ class LangfuseSink:
|
||||
metadata=tool_metadata,
|
||||
)
|
||||
|
||||
# Push to stack for LIFO correlation
|
||||
self._tool_stacks[(agent_key, tool_name)].append(tool_span)
|
||||
# Store with unique tool key
|
||||
self._spans[tool_key] = tool_span
|
||||
|
||||
def _handle_tool_end(self, event: ActionEvent, agent_name: str) -> None:
|
||||
"""Handle TOOL_END -> end tool span observation."""
|
||||
agent_key = f"{event.agent_id or agent_name}-{agent_name}"
|
||||
tool_name = event.tool_name or "unknown-tool"
|
||||
stack = self._tool_stacks.get((agent_key, tool_name))
|
||||
if not stack:
|
||||
|
||||
# Find the most recent matching tool span
|
||||
tool_key = None
|
||||
for key in self._spans:
|
||||
if key.startswith(f"{agent_key}:{tool_name}:") and key != agent_key:
|
||||
tool_key = key
|
||||
|
||||
if not tool_key:
|
||||
return
|
||||
tool_span = stack.pop()
|
||||
if not stack:
|
||||
self._tool_stacks.pop((agent_key, tool_name), None)
|
||||
try:
|
||||
|
||||
tool_span = self._spans.pop(tool_key, None)
|
||||
if tool_span:
|
||||
tool_span.update(
|
||||
output=event.tool_result_summary,
|
||||
status_message=event.status or "completed",
|
||||
@@ -242,9 +237,6 @@ class LangfuseSink:
|
||||
}
|
||||
)
|
||||
tool_span.end()
|
||||
except Exception as e:
|
||||
# Use logging instead of print
|
||||
logger.exception("LangfuseSink: failed to finalize tool span %r", tool_name)
|
||||
|
||||
def _handle_error(self, event: ActionEvent, agent_name: str) -> None:
|
||||
"""Handle ERROR -> create error event observation."""
|
||||
@@ -292,7 +284,7 @@ class LangfuseSink:
|
||||
try:
|
||||
self._client.flush()
|
||||
except Exception as e:
|
||||
logger.exception("LangfuseSink flush error")
|
||||
print(f"LangfuseSink flush error: {e}")
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the sink and release resources."""
|
||||
@@ -304,16 +296,6 @@ class LangfuseSink:
|
||||
try:
|
||||
# Close any remaining spans
|
||||
with self._lock:
|
||||
# End any in-flight tool spans
|
||||
for stack in self._tool_stacks.values():
|
||||
while stack:
|
||||
try:
|
||||
stack.pop().end()
|
||||
except Exception:
|
||||
pass
|
||||
self._tool_stacks.clear()
|
||||
|
||||
# End root spans
|
||||
for span in self._spans.values():
|
||||
try:
|
||||
span.end()
|
||||
@@ -322,7 +304,7 @@ class LangfuseSink:
|
||||
self._spans.clear()
|
||||
self._traces.clear()
|
||||
except Exception:
|
||||
logger.exception("LangfuseSink close cleanup failed")
|
||||
pass
|
||||
|
||||
def context_sink(self) -> "ContextTraceSinkProtocol":
|
||||
"""Return a ContextTraceSinkProtocol that forwards to this sink."""
|
||||
@@ -399,4 +381,4 @@ class _ContextToActionBridge:
|
||||
|
||||
def close(self) -> None:
|
||||
"""Forward close to LangfuseSink."""
|
||||
self._langfuse_sink.close()
|
||||
self._langfuse_sink.close()
|
||||
@@ -7,9 +7,7 @@ knowledge retrieval, and state management.
|
||||
|
||||
import logging
|
||||
import time
|
||||
import threading
|
||||
import uuid
|
||||
from copy import deepcopy
|
||||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||||
|
||||
from .conversation.base import ConversationStore, ConversationSession, ConversationMessage
|
||||
@@ -80,7 +78,6 @@ class PersistenceOrchestrator:
|
||||
|
||||
self._current_session: Optional[ConversationSession] = None
|
||||
self._session_cache: Dict[str, ConversationSession] = {}
|
||||
self._cache_lock = threading.RLock() # RLock allows re-entrant access
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: PersistenceConfig) -> "PersistenceOrchestrator":
|
||||
@@ -93,31 +90,6 @@ class PersistenceOrchestrator:
|
||||
config = PersistenceConfig.from_env()
|
||||
return cls(config=config)
|
||||
|
||||
# =========================================================================
|
||||
# Thread-Safe Cache Operations
|
||||
# =========================================================================
|
||||
|
||||
def _cache_put(self, session: ConversationSession) -> None:
|
||||
"""Store session in cache with thread safety."""
|
||||
with self._cache_lock:
|
||||
self._session_cache[session.session_id] = session
|
||||
|
||||
def _cache_get(self, session_id: str) -> Optional[ConversationSession]:
|
||||
"""Get session from cache with thread safety and defensive copying."""
|
||||
with self._cache_lock:
|
||||
cached = self._session_cache.get(session_id)
|
||||
return deepcopy(cached) if cached is not None else None
|
||||
|
||||
def _cache_delete(self, session_id: str) -> Optional[ConversationSession]:
|
||||
"""Remove session from cache with thread safety."""
|
||||
with self._cache_lock:
|
||||
return self._session_cache.pop(session_id, None)
|
||||
|
||||
def _cache_clear(self) -> None:
|
||||
"""Clear all sessions from cache with thread safety."""
|
||||
with self._cache_lock:
|
||||
self._session_cache.clear()
|
||||
|
||||
# =========================================================================
|
||||
# Agent Lifecycle Hooks
|
||||
# =========================================================================
|
||||
@@ -156,7 +128,7 @@ class PersistenceOrchestrator:
|
||||
if session:
|
||||
logger.info(f"Resuming session: {session_id}")
|
||||
self._current_session = session
|
||||
self._cache_put(session)
|
||||
self._session_cache[session_id] = session
|
||||
|
||||
# Load previous messages
|
||||
messages = self.conversation.get_messages(session_id)
|
||||
@@ -174,7 +146,7 @@ class PersistenceOrchestrator:
|
||||
self.conversation.create_session(session)
|
||||
logger.info(f"Created new session: {session_id}")
|
||||
self._current_session = session
|
||||
self._cache_put(session)
|
||||
self._session_cache[session_id] = session
|
||||
return []
|
||||
|
||||
def on_message(
|
||||
@@ -234,14 +206,12 @@ class PersistenceOrchestrator:
|
||||
if not self.conversation:
|
||||
return
|
||||
|
||||
session = self._cache_get(session_id) or self.conversation.get_session(session_id)
|
||||
session = self._session_cache.get(session_id) or self.conversation.get_session(session_id)
|
||||
if session:
|
||||
session.updated_at = time.time()
|
||||
if metadata:
|
||||
session.metadata = {**(session.metadata or {}), **metadata}
|
||||
self.conversation.update_session(session)
|
||||
# Update cache with the modified session
|
||||
self._cache_put(session)
|
||||
logger.debug(f"Updated session metadata: {session_id}")
|
||||
|
||||
# =========================================================================
|
||||
@@ -345,8 +315,8 @@ class PersistenceOrchestrator:
|
||||
if not self.conversation:
|
||||
return False
|
||||
|
||||
# Remove from cache using thread-safe method
|
||||
self._cache_delete(session_id)
|
||||
if session_id in self._session_cache:
|
||||
del self._session_cache[session_id]
|
||||
|
||||
return self.conversation.delete_session(session_id)
|
||||
|
||||
@@ -425,8 +395,7 @@ class PersistenceOrchestrator:
|
||||
if self.state:
|
||||
self.state.close()
|
||||
|
||||
# Clear cache using thread-safe method
|
||||
self._cache_clear()
|
||||
self._session_cache.clear()
|
||||
logger.info("Persistence orchestrator closed")
|
||||
|
||||
def __enter__(self):
|
||||
|
||||
@@ -174,25 +174,8 @@ def _safe_extractall(tar: tarfile.TarFile, dest_dir: Path) -> None:
|
||||
raise RegistryError(
|
||||
f"Refusing to extract path escaping target directory: {member.name}"
|
||||
)
|
||||
# Reject symlinks and hardlinks — their linkname may point outside dest_dir
|
||||
# and would not be caught by the member.name checks above.
|
||||
if member.issym() or member.islnk():
|
||||
linkname = member.linkname or ""
|
||||
if linkname.startswith("/"):
|
||||
raise RegistryError(
|
||||
f"Refusing to extract link with absolute target: {member.name} -> {linkname}"
|
||||
)
|
||||
link_target = (dest_resolved / member_path.parent / linkname).resolve()
|
||||
if not str(link_target).startswith(str(dest_resolved) + os.sep) and link_target != dest_resolved:
|
||||
raise RegistryError(
|
||||
f"Refusing to extract link escaping target directory: {member.name} -> {linkname}"
|
||||
)
|
||||
# All members validated — safe to extract
|
||||
try:
|
||||
tar.extractall(dest_dir, filter="data")
|
||||
except TypeError:
|
||||
# filter keyword not supported on Python < 3.12; validation above already covers safety
|
||||
tar.extractall(dest_dir)
|
||||
tar.extractall(dest_dir)
|
||||
|
||||
|
||||
def _atomic_write(file_path: Path, data: bytes) -> None:
|
||||
|
||||
@@ -1,185 +0,0 @@
|
||||
"""
|
||||
Compatibility layer to unify sandbox/ and integrations/compute/ hierarchies.
|
||||
|
||||
Provides SandboxToComputeAdapter to expose ComputeProvider as legacy SandboxConfig/SandboxResult API.
|
||||
This eliminates code duplication between the two parallel hierarchies.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from praisonaiagents.sandbox import (
|
||||
SandboxResult,
|
||||
SandboxStatus,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SandboxToComputeAdapter:
|
||||
"""Expose a ComputeProvider as the legacy SandboxConfig/SandboxResult API."""
|
||||
|
||||
def __init__(self, compute_provider: Any, image: str = "python:3.11-slim"):
|
||||
"""
|
||||
Initialize adapter with compute provider and default image.
|
||||
|
||||
Args:
|
||||
compute_provider: Instance implementing ComputeProviderProtocol
|
||||
image: Default container image to use
|
||||
"""
|
||||
self._compute = compute_provider
|
||||
self._image = image
|
||||
self._instance_id: Optional[str] = None
|
||||
self._is_running = False
|
||||
|
||||
@property
|
||||
def is_available(self) -> bool:
|
||||
"""Check if the compute provider is available."""
|
||||
if hasattr(self._compute, 'is_available'):
|
||||
return self._compute.is_available
|
||||
# Fallback - try to check if we can instantiate
|
||||
try:
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the sandbox by provisioning a compute instance."""
|
||||
if self._is_running:
|
||||
return
|
||||
|
||||
try:
|
||||
# Create compute config - adapt to whatever the compute provider expects
|
||||
if hasattr(self._compute, 'provision'):
|
||||
# For ComputeProviderProtocol
|
||||
config = {
|
||||
'image': self._image,
|
||||
'auto_shutdown': False, # We'll manage lifecycle
|
||||
}
|
||||
|
||||
# Try to create a proper config object if available
|
||||
try:
|
||||
from praisonaiagents.managed.protocols import ComputeConfig
|
||||
compute_config = ComputeConfig(image=self._image)
|
||||
info = await self._compute.provision(compute_config)
|
||||
self._instance_id = info.instance_id
|
||||
except (ImportError, AttributeError):
|
||||
# Fallback to basic provision call
|
||||
info = await self._compute.provision(config)
|
||||
self._instance_id = getattr(info, 'instance_id', 'default')
|
||||
|
||||
self._is_running = True
|
||||
logger.debug(f"Started compute instance: {self._instance_id}")
|
||||
else:
|
||||
# Direct start method
|
||||
await self._compute.start()
|
||||
self._is_running = True
|
||||
self._instance_id = "default"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start compute instance: {e}")
|
||||
raise RuntimeError(f"Failed to start sandbox: {e}")
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
code: str,
|
||||
language: str = "python",
|
||||
limits: Optional[Any] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
) -> SandboxResult:
|
||||
"""
|
||||
Execute code in the sandbox and return SandboxResult.
|
||||
|
||||
Args:
|
||||
code: Code to execute
|
||||
language: Programming language (default: python)
|
||||
limits: Resource limits (optional)
|
||||
env: Environment variables (optional)
|
||||
working_dir: Working directory (optional)
|
||||
|
||||
Returns:
|
||||
SandboxResult with execution results
|
||||
"""
|
||||
if not self._is_running or not self._instance_id:
|
||||
await self.start()
|
||||
|
||||
try:
|
||||
# Execute on compute provider
|
||||
if hasattr(self._compute, 'execute'):
|
||||
# Convert code to a shell command (compute providers expect commands, not raw code)
|
||||
if language == "python":
|
||||
import shlex
|
||||
command = f"python -c {shlex.quote(code)}"
|
||||
elif language in ("bash", "shell"):
|
||||
command = code
|
||||
else:
|
||||
import shlex
|
||||
command = f"python -c {shlex.quote(code)}"
|
||||
|
||||
raw = await self._compute.execute(self._instance_id, command)
|
||||
|
||||
# Convert compute result to SandboxResult (dict or object)
|
||||
if isinstance(raw, dict):
|
||||
stdout = raw.get("stdout", "")
|
||||
stderr = raw.get("stderr", "")
|
||||
exit_code = raw.get("exit_code", 0)
|
||||
elif hasattr(raw, 'stdout'):
|
||||
stdout = raw.stdout
|
||||
stderr = getattr(raw, 'stderr', "") or ""
|
||||
exit_code = getattr(raw, 'exit_code', 0)
|
||||
else:
|
||||
stdout = str(raw)
|
||||
stderr = ""
|
||||
exit_code = 0
|
||||
|
||||
_TIMEOUT_EXIT_CODE = 124
|
||||
if exit_code == 0:
|
||||
status = SandboxStatus.COMPLETED
|
||||
elif exit_code == _TIMEOUT_EXIT_CODE:
|
||||
status = SandboxStatus.TIMEOUT
|
||||
else:
|
||||
status = SandboxStatus.FAILED
|
||||
|
||||
return SandboxResult(
|
||||
status=status,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
exit_code=exit_code,
|
||||
)
|
||||
else:
|
||||
# Fallback for simple compute providers
|
||||
result = str(await self._compute.run(code))
|
||||
return SandboxResult(
|
||||
status=SandboxStatus.COMPLETED,
|
||||
stdout=result,
|
||||
stderr="",
|
||||
exit_code=0,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Execution failed: {e}")
|
||||
return SandboxResult(
|
||||
status=SandboxStatus.FAILED,
|
||||
stdout="",
|
||||
stderr=str(e),
|
||||
exit_code=1,
|
||||
)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the sandbox by shutting down the compute instance."""
|
||||
if not self._is_running:
|
||||
return
|
||||
|
||||
try:
|
||||
if hasattr(self._compute, 'shutdown') and self._instance_id:
|
||||
await self._compute.shutdown(self._instance_id)
|
||||
elif hasattr(self._compute, 'stop'):
|
||||
await self._compute.stop()
|
||||
|
||||
self._is_running = False
|
||||
self._instance_id = None
|
||||
logger.debug("Stopped compute instance")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Error stopping compute instance: {e}")
|
||||
@@ -1,38 +0,0 @@
|
||||
"""Shared utilities for safe shell command handling across sandbox backends."""
|
||||
|
||||
import shlex
|
||||
from typing import List, Union
|
||||
|
||||
|
||||
def build_argv(command: Union[str, List[str]], shell: bool = False) -> List[str]:
|
||||
"""
|
||||
Safely build command argv with explicit shell control.
|
||||
|
||||
Args:
|
||||
command: String command or list of arguments
|
||||
shell: If True, explicitly use shell. If False, parse safely without shell.
|
||||
|
||||
Returns:
|
||||
List of command arguments safe for subprocess execution
|
||||
|
||||
Security:
|
||||
- shell=False (default): No shell injection possible
|
||||
- shell=True: Caller explicitly opts into shell evaluation
|
||||
"""
|
||||
if isinstance(command, str):
|
||||
if not shell:
|
||||
# Safe parse: convert string to argv without invoking shell
|
||||
return shlex.split(command)
|
||||
else:
|
||||
# Explicit shell: caller has opted in
|
||||
return ["sh", "-c", command]
|
||||
else:
|
||||
# List input
|
||||
cmd_list = list(command)
|
||||
if shell:
|
||||
# Quote each element when combining into shell command
|
||||
quoted_cmd = " ".join(shlex.quote(arg) for arg in cmd_list)
|
||||
return ["sh", "-c", quoted_cmd]
|
||||
else:
|
||||
# Direct argv execution - no shell
|
||||
return cmd_list
|
||||
@@ -9,7 +9,6 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import shlex
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
@@ -227,17 +226,8 @@ class DockerSandbox:
|
||||
limits: Optional[ResourceLimits] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
shell: bool = False,
|
||||
) -> SandboxResult:
|
||||
"""Run a command in the sandbox.
|
||||
|
||||
Args:
|
||||
command: String command or list of arguments
|
||||
limits: Resource limits to apply
|
||||
env: Environment variables
|
||||
working_dir: Working directory
|
||||
shell: If True, explicitly use shell. If False (default), execute safely without shell.
|
||||
"""
|
||||
"""Run a shell command in the sandbox."""
|
||||
if not self._is_running:
|
||||
await self.start()
|
||||
|
||||
@@ -245,16 +235,9 @@ class DockerSandbox:
|
||||
execution_id = str(uuid.uuid4())
|
||||
|
||||
if isinstance(command, list):
|
||||
# Always quote list elements to prevent shell injection
|
||||
cmd_str = " ".join(shlex.quote(arg) for arg in command)
|
||||
cmd_str = " ".join(command)
|
||||
else:
|
||||
if shell:
|
||||
# Caller explicitly requested shell evaluation
|
||||
cmd_str = command
|
||||
else:
|
||||
# Parse string safely then re-quote each part
|
||||
cmd_parts = shlex.split(command)
|
||||
cmd_str = " ".join(shlex.quote(part) for part in cmd_parts)
|
||||
cmd_str = command
|
||||
|
||||
docker_cmd = [
|
||||
"docker", "run", "--rm",
|
||||
|
||||
@@ -277,17 +277,8 @@ class SSHSandbox:
|
||||
limits: Optional[ResourceLimits] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
shell: bool = False,
|
||||
) -> SandboxResult:
|
||||
"""Run a command on the remote server.
|
||||
|
||||
Args:
|
||||
command: String command or list of arguments
|
||||
limits: Resource limits to apply
|
||||
env: Environment variables
|
||||
working_dir: Working directory
|
||||
shell: If True, explicitly use shell. If False (default), execute safely without shell.
|
||||
"""
|
||||
"""Run a shell command on the remote server."""
|
||||
if not self._is_running:
|
||||
await self.start()
|
||||
|
||||
@@ -295,34 +286,15 @@ class SSHSandbox:
|
||||
started_at = time.time()
|
||||
|
||||
try:
|
||||
# Convert command to string safely based on shell parameter
|
||||
# Convert command to string if needed
|
||||
if isinstance(command, list):
|
||||
if shell:
|
||||
# Shell mode: join with proper quoting
|
||||
command_str = " ".join(shlex.quote(arg) for arg in command)
|
||||
else:
|
||||
# Non-shell mode: use shlex.join for proper escaping
|
||||
command_str = shlex.join(command)
|
||||
else:
|
||||
# String command
|
||||
if not shell:
|
||||
# Parse then re-join to ensure safe execution
|
||||
try:
|
||||
parts = shlex.split(command)
|
||||
command_str = shlex.join(parts)
|
||||
except ValueError:
|
||||
# If parsing fails, quote the whole thing
|
||||
command_str = shlex.quote(command)
|
||||
else:
|
||||
# Shell mode: use as-is (caller explicitly opted in)
|
||||
command_str = command
|
||||
command = shlex.join(command)
|
||||
|
||||
# Execute command
|
||||
result = await self._run_command_with_limits(
|
||||
command_str,
|
||||
command,
|
||||
limits,
|
||||
working_dir or self.working_dir,
|
||||
shell
|
||||
working_dir or self.working_dir
|
||||
)
|
||||
|
||||
completed_at = time.time()
|
||||
@@ -519,12 +491,10 @@ class SSHSandbox:
|
||||
self,
|
||||
command: str,
|
||||
limits: Optional[ResourceLimits],
|
||||
working_dir: str,
|
||||
shell: bool = False
|
||||
working_dir: str
|
||||
):
|
||||
"""Run command with resource limits."""
|
||||
# Change to working directory and execute command
|
||||
# Note: The command has already been processed for shell safety by the caller
|
||||
# Change to working directory
|
||||
full_command = f"cd {shlex.quote(working_dir)} && {command}"
|
||||
|
||||
# Set timeout
|
||||
|
||||
@@ -190,26 +190,18 @@ class SubprocessSandbox:
|
||||
limits: Optional[ResourceLimits] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
shell: bool = False,
|
||||
) -> SandboxResult:
|
||||
"""Run a command in the sandbox.
|
||||
|
||||
Args:
|
||||
command: String command or list of arguments
|
||||
limits: Resource limits to apply
|
||||
env: Environment variables
|
||||
working_dir: Working directory
|
||||
shell: If True, explicitly use shell. If False (default), execute safely without shell.
|
||||
"""
|
||||
"""Run a shell command in the sandbox."""
|
||||
if not self._is_running:
|
||||
await self.start()
|
||||
|
||||
limits = limits or self.config.resource_limits
|
||||
execution_id = str(uuid.uuid4())
|
||||
|
||||
# Import here to avoid circular import
|
||||
from ._shell import build_argv
|
||||
cmd = build_argv(command, shell=shell)
|
||||
if isinstance(command, str):
|
||||
cmd = ["sh", "-c", command]
|
||||
else:
|
||||
cmd = command
|
||||
|
||||
process_env = os.environ.copy()
|
||||
if env:
|
||||
|
||||
@@ -10,67 +10,11 @@ This module provides shared functionality for all schedulers:
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
from .shared import ScheduleParser as ScheduleParser # canonical definition lives in shared.py
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ScheduleParser:
|
||||
"""Parse schedule expressions into intervals in seconds."""
|
||||
|
||||
@staticmethod
|
||||
def parse(schedule_expr: str) -> int:
|
||||
"""
|
||||
Parse schedule expression and return interval in seconds.
|
||||
|
||||
Supported formats:
|
||||
- "daily" -> 86400 seconds
|
||||
- "hourly" -> 3600 seconds
|
||||
- "*/30m" -> 1800 seconds (every 30 minutes)
|
||||
- "*/6h" -> 21600 seconds (every 6 hours)
|
||||
- "*/30s" -> 30 seconds (every 30 seconds)
|
||||
- "3600" -> 3600 seconds (plain number)
|
||||
|
||||
Args:
|
||||
schedule_expr: Schedule expression string
|
||||
|
||||
Returns:
|
||||
Interval in seconds
|
||||
|
||||
Raises:
|
||||
ValueError: If schedule format is not supported
|
||||
|
||||
Examples:
|
||||
>>> ScheduleParser.parse("hourly")
|
||||
3600
|
||||
>>> ScheduleParser.parse("*/30m")
|
||||
1800
|
||||
>>> ScheduleParser.parse("daily")
|
||||
86400
|
||||
"""
|
||||
schedule_expr = schedule_expr.strip().lower()
|
||||
|
||||
if schedule_expr == "daily":
|
||||
return 86400
|
||||
elif schedule_expr == "hourly":
|
||||
return 3600
|
||||
elif schedule_expr.isdigit():
|
||||
return int(schedule_expr)
|
||||
elif schedule_expr.startswith("*/"):
|
||||
interval_part = schedule_expr[2:]
|
||||
if interval_part.endswith("m"):
|
||||
minutes = int(interval_part[:-1])
|
||||
return minutes * 60
|
||||
elif interval_part.endswith("h"):
|
||||
hours = int(interval_part[:-1])
|
||||
return hours * 3600
|
||||
elif interval_part.endswith("s"):
|
||||
return int(interval_part[:-1])
|
||||
else:
|
||||
return int(interval_part)
|
||||
else:
|
||||
raise ValueError(f"Unsupported schedule format: {schedule_expr}")
|
||||
|
||||
|
||||
class ExecutorInterface(ABC):
|
||||
"""Abstract interface for executors."""
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Shared primitives for sync & async schedulers."""
|
||||
|
||||
class ScheduleParser:
|
||||
"""Shared schedule expression parser for both sync and async schedulers."""
|
||||
"""Parse schedule expressions into intervals in seconds."""
|
||||
|
||||
@staticmethod
|
||||
def parse(schedule_expr: str) -> int:
|
||||
@@ -12,8 +12,9 @@ class ScheduleParser:
|
||||
- "daily" -> 86400 seconds
|
||||
- "hourly" -> 3600 seconds
|
||||
- "*/30m" -> 1800 seconds (every 30 minutes)
|
||||
- "*/1h" -> 3600 seconds (every 1 hour)
|
||||
- "60" -> 60 seconds (plain number)
|
||||
- "*/6h" -> 21600 seconds (every 6 hours)
|
||||
- "*/30s" -> 30 seconds (every 30 seconds)
|
||||
- "3600" -> 3600 seconds (plain number)
|
||||
|
||||
Args:
|
||||
schedule_expr: Schedule expression string
|
||||
@@ -23,6 +24,14 @@ class ScheduleParser:
|
||||
|
||||
Raises:
|
||||
ValueError: If schedule format is not supported
|
||||
|
||||
Examples:
|
||||
>>> ScheduleParser.parse("hourly")
|
||||
3600
|
||||
>>> ScheduleParser.parse("*/30m")
|
||||
1800
|
||||
>>> ScheduleParser.parse("daily")
|
||||
86400
|
||||
"""
|
||||
expr = schedule_expr.strip().lower()
|
||||
|
||||
|
||||
@@ -3,15 +3,11 @@ import os
|
||||
from rich import print
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
|
||||
from praisonai.llm.env import resolve_llm_endpoint
|
||||
|
||||
ep = resolve_llm_endpoint()
|
||||
config_list = [
|
||||
{
|
||||
'model': ep.model,
|
||||
'base_url': ep.base_url,
|
||||
'api_key': ep.api_key
|
||||
'model': os.environ.get("OPENAI_MODEL_NAME", "gpt-4o-mini"),
|
||||
'base_url': os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
|
||||
'api_key': os.environ.get("OPENAI_API_KEY")
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@@ -5,16 +5,12 @@ import difflib
|
||||
import platform
|
||||
from typing import Dict, Any
|
||||
import json
|
||||
import logging
|
||||
import dotenv
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
dotenv.load_dotenv()
|
||||
|
||||
class AICoder:
|
||||
def __init__(self, cwd: str = None, tavily_api_key: str = None):
|
||||
# Load environment variables from .env file
|
||||
dotenv.load_dotenv()
|
||||
|
||||
# Load only the dependency needed on the common path here.
|
||||
try:
|
||||
from litellm import acompletion
|
||||
@@ -133,17 +129,7 @@ class AICoder:
|
||||
try:
|
||||
with open(file_path, 'r') as file:
|
||||
return file.read()
|
||||
except FileNotFoundError:
|
||||
logger.debug("aicoder: file not found: %s", file_path)
|
||||
return None
|
||||
except PermissionError:
|
||||
logger.warning("aicoder: permission denied reading %s", file_path)
|
||||
return None
|
||||
except UnicodeDecodeError as e:
|
||||
logger.warning("aicoder: encoding error reading %s: %s", file_path, e)
|
||||
return None
|
||||
except OSError as e:
|
||||
logger.warning("aicoder: OS error reading %s: %s", file_path, e)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_shell_command(self, command: str) -> list:
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "4.6.37"
|
||||
__version__ = "4.6.35"
|
||||
@@ -12,7 +12,7 @@ dependencies = [
|
||||
"rich>=13.7",
|
||||
"markdown>=3.5",
|
||||
"pyparsing>=3.0.0",
|
||||
"praisonaiagents>=1.6.37",
|
||||
"praisonaiagents>=1.6.35",
|
||||
"python-dotenv>=0.19.0",
|
||||
"litellm>=1.81.0,<=1.82.6",
|
||||
"PyYAML>=6.0",
|
||||
|
||||
@@ -1,94 +0,0 @@
|
||||
"""W1 — Real agentic test for cross-platform unified session.
|
||||
|
||||
Spins up a real ``Agent`` (calls a real LLM), wires it into two
|
||||
``BotSessionManager`` instances representing Telegram and Discord,
|
||||
links the user via ``InMemoryIdentityResolver``, sends a fact on
|
||||
"Telegram", then asks about it on "Discord". The agent must recall
|
||||
the fact from the unified history.
|
||||
|
||||
Run manually::
|
||||
|
||||
cd src/praisonai
|
||||
OPENAI_API_KEY=... pytest tests/integration/test_w1_real_agentic.py -s
|
||||
|
||||
This test is opt-in (skipped without ``RUN_REAL_AGENTIC=1``) so CI
|
||||
without LLM credentials doesn't break.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(
|
||||
os.getenv("RUN_REAL_AGENTIC", "") != "1",
|
||||
reason="Set RUN_REAL_AGENTIC=1 to run real LLM tests",
|
||||
),
|
||||
pytest.mark.network,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cross_platform_continuity_real_llm(tmp_path):
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.session.identity import InMemoryIdentityResolver
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
# 1. Real Agent
|
||||
# Pick first available provider
|
||||
if os.getenv("GOOGLE_API_KEY"):
|
||||
model = "gemini/gemini-2.5-flash"
|
||||
elif os.getenv("ANTHROPIC_API_KEY"):
|
||||
model = "claude-3-5-haiku-20241022"
|
||||
else:
|
||||
model = "gpt-4o-mini"
|
||||
|
||||
agent = Agent(
|
||||
name="W1Tester",
|
||||
instructions=(
|
||||
"You are a helpful assistant. Always remember facts the user "
|
||||
"tells you and recall them precisely when asked."
|
||||
),
|
||||
llm=model,
|
||||
)
|
||||
|
||||
# 2. Identity resolver — same human, two platform IDs
|
||||
resolver = InMemoryIdentityResolver()
|
||||
resolver.link("telegram", "tg-12345", "alice-global")
|
||||
resolver.link("discord", "dc-67890", "alice-global")
|
||||
|
||||
# 3. Shared persistent store
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path))
|
||||
|
||||
# 4. Two managers (one per platform), shared store + resolver
|
||||
mgr_t = BotSessionManager(
|
||||
platform="telegram", identity_resolver=resolver, store=store
|
||||
)
|
||||
mgr_d = BotSessionManager(
|
||||
platform="discord", identity_resolver=resolver, store=store
|
||||
)
|
||||
|
||||
# 5. Tell agent a fact via "Telegram"
|
||||
fact_response = await mgr_t.chat(
|
||||
agent, "tg-12345",
|
||||
"Remember this: my favourite colour is octarine.",
|
||||
)
|
||||
print("\n[Telegram in] my favourite colour is octarine")
|
||||
print(f"[Telegram out] {fact_response}\n")
|
||||
|
||||
# 6. Ask via "Discord" — agent should recall from unified history
|
||||
recall_response = await mgr_d.chat(
|
||||
agent, "dc-67890",
|
||||
"What did I just tell you my favourite colour was?",
|
||||
)
|
||||
print("[Discord in] What did I just tell you my favourite colour was?")
|
||||
print(f"[Discord out] {recall_response}\n")
|
||||
|
||||
# The agent must mention "octarine" — proof the cross-platform mirror works
|
||||
assert "octarine" in recall_response.lower(), (
|
||||
f"Agent failed to recall cross-platform fact. Got: {recall_response}"
|
||||
)
|
||||
@@ -1,111 +0,0 @@
|
||||
"""W1 — Bot/BotOS wiring of identity_resolver + SessionContext.
|
||||
|
||||
These tests verify the user-facing surface (``Bot(identity_resolver=...)``,
|
||||
``BotOS(identity_resolver=...)``) wires correctly through to the
|
||||
underlying ``BotSessionManager``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonai.bots import Bot, BotOS
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
from praisonaiagents.session import (
|
||||
InMemoryIdentityResolver,
|
||||
SessionContext,
|
||||
get_session_context,
|
||||
set_session_context,
|
||||
clear_session_context,
|
||||
)
|
||||
|
||||
|
||||
class TestBotConstructorAcceptsResolver:
|
||||
def test_bot_accepts_identity_resolver_kwarg(self):
|
||||
resolver = InMemoryIdentityResolver()
|
||||
bot = Bot("telegram", identity_resolver=resolver)
|
||||
assert bot._identity_resolver is resolver
|
||||
|
||||
|
||||
class TestBotOSPropagation:
|
||||
def test_botos_sets_resolver_on_added_bots_via_shortcut(self):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
agent = MagicMock(name="agent")
|
||||
agent.name = "Test"
|
||||
agent.chat_history = []
|
||||
|
||||
resolver = InMemoryIdentityResolver()
|
||||
os = BotOS(
|
||||
agent=agent,
|
||||
platforms=["telegram", "discord"],
|
||||
identity_resolver=resolver,
|
||||
)
|
||||
for plat in ("telegram", "discord"):
|
||||
bot = os.get_bot(plat)
|
||||
assert bot is not None
|
||||
assert bot._identity_resolver is resolver
|
||||
|
||||
def test_botos_propagates_to_manually_added_bot(self):
|
||||
resolver = InMemoryIdentityResolver()
|
||||
os = BotOS(identity_resolver=resolver)
|
||||
bot = Bot("telegram")
|
||||
assert bot._identity_resolver is None
|
||||
os.add_bot(bot)
|
||||
assert bot._identity_resolver is resolver
|
||||
|
||||
def test_botos_does_not_overwrite_explicit_resolver(self):
|
||||
r1 = InMemoryIdentityResolver()
|
||||
r2 = InMemoryIdentityResolver()
|
||||
bot = Bot("telegram", identity_resolver=r1)
|
||||
os = BotOS(identity_resolver=r2)
|
||||
os.add_bot(bot)
|
||||
assert bot._identity_resolver is r1 # explicit wins
|
||||
|
||||
|
||||
class FakeAgent:
|
||||
def __init__(self):
|
||||
self.chat_history = []
|
||||
self.observed_context = None
|
||||
|
||||
def chat(self, prompt):
|
||||
# Capture context as seen from inside the "agent call"
|
||||
self.observed_context = get_session_context()
|
||||
self.chat_history.append({"role": "user", "content": prompt})
|
||||
return "ok"
|
||||
|
||||
|
||||
class TestSessionContextPropagation:
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_sets_session_context_for_agent(self):
|
||||
agent = FakeAgent()
|
||||
mgr = BotSessionManager(platform="telegram")
|
||||
await mgr.chat(
|
||||
agent, user_id="12345", prompt="hi",
|
||||
chat_id="100", user_name="Alice",
|
||||
)
|
||||
ctx = agent.observed_context
|
||||
assert isinstance(ctx, SessionContext)
|
||||
assert ctx.platform == "telegram"
|
||||
assert ctx.chat_id == "100"
|
||||
assert ctx.user_id == "12345"
|
||||
assert ctx.user_name == "Alice"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_is_cleared_after_chat(self):
|
||||
agent = FakeAgent()
|
||||
mgr = BotSessionManager(platform="telegram")
|
||||
await mgr.chat(agent, user_id="12345", prompt="hi", chat_id="100")
|
||||
# After chat returns, caller's context is restored to default.
|
||||
assert get_session_context().platform == ""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_user_id_in_context(self):
|
||||
resolver = InMemoryIdentityResolver()
|
||||
resolver.link("telegram", "12345", "alice-global")
|
||||
|
||||
agent = FakeAgent()
|
||||
mgr = BotSessionManager(platform="telegram", identity_resolver=resolver)
|
||||
await mgr.chat(agent, user_id="12345", prompt="hi")
|
||||
ctx = agent.observed_context
|
||||
assert ctx.unified_user_id == "alice-global"
|
||||
@@ -1,63 +0,0 @@
|
||||
"""W1 — Outbound delivery mirror.
|
||||
|
||||
When a Bot pushes a message via ``Bot.send_message()`` (notifications,
|
||||
scheduled deliveries, cross-platform replies), the mirror appends a
|
||||
``{"role": "assistant", "mirror": True}`` entry to the user's session
|
||||
so the agent has context next turn.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from praisonai.bots._mirror import mirror_to_session
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
|
||||
class TestMirrorBasic:
|
||||
def test_mirror_appends_assistant_entry(self):
|
||||
mgr = BotSessionManager(platform="telegram")
|
||||
mgr._histories["alice"] = []
|
||||
|
||||
ok = mirror_to_session(
|
||||
mgr,
|
||||
user_id="alice",
|
||||
message_text="Daily summary: 5 PRs merged.",
|
||||
source_label="cron",
|
||||
)
|
||||
|
||||
assert ok is True
|
||||
history = mgr._histories["alice"]
|
||||
assert len(history) == 1
|
||||
assert history[0]["role"] == "assistant"
|
||||
assert "5 PRs merged" in history[0]["content"]
|
||||
assert history[0]["mirror"] is True
|
||||
assert history[0]["mirror_source"] == "cron"
|
||||
|
||||
def test_mirror_creates_history_when_absent(self):
|
||||
mgr = BotSessionManager(platform="telegram")
|
||||
ok = mirror_to_session(mgr, user_id="bob", message_text="hello")
|
||||
assert ok is True
|
||||
assert "bob" in mgr._histories
|
||||
|
||||
def test_mirror_with_resolver_uses_unified_key(self):
|
||||
from praisonaiagents.session.identity import InMemoryIdentityResolver
|
||||
|
||||
resolver = InMemoryIdentityResolver()
|
||||
resolver.link("telegram", "12345", "alice-global")
|
||||
mgr = BotSessionManager(platform="telegram", identity_resolver=resolver)
|
||||
|
||||
ok = mirror_to_session(mgr, user_id="12345", message_text="hi")
|
||||
assert ok is True
|
||||
# Stored under unified key, not platform user id
|
||||
assert "alice-global" in mgr._histories
|
||||
assert "12345" not in mgr._histories
|
||||
|
||||
|
||||
class TestMirrorIsNonFatal:
|
||||
def test_mirror_swallows_exceptions(self):
|
||||
class BrokenManager:
|
||||
def _session_key(self, user_id):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
# Must not raise; returns False
|
||||
ok = mirror_to_session(BrokenManager(), user_id="alice", message_text="x")
|
||||
assert ok is False
|
||||
@@ -1,158 +0,0 @@
|
||||
"""W1 — Unified per-user session across platforms.
|
||||
|
||||
Tests that BotSessionManager can use an IdentityResolver so the same
|
||||
human pinging from Telegram and Discord shares one chat history.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonai.bots._session import BotSessionManager
|
||||
|
||||
|
||||
class FakeAgent:
|
||||
def __init__(self):
|
||||
self.chat_history = []
|
||||
self.calls = []
|
||||
|
||||
def chat(self, prompt):
|
||||
self.calls.append((list(self.chat_history), prompt))
|
||||
self.chat_history.append({"role": "user", "content": prompt})
|
||||
reply = f"reply to {prompt}"
|
||||
self.chat_history.append({"role": "assistant", "content": reply})
|
||||
return reply
|
||||
|
||||
|
||||
class TestUnifiedSessionWithoutResolver:
|
||||
"""Default behaviour preserved: no resolver → per-platform keying."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_platform_isolation_is_default(self):
|
||||
agent = FakeAgent()
|
||||
mgr_t = BotSessionManager(platform="telegram")
|
||||
mgr_d = BotSessionManager(platform="discord")
|
||||
|
||||
await mgr_t.chat(agent, "alice", "hi from telegram")
|
||||
await mgr_d.chat(agent, "alice", "hi from discord")
|
||||
|
||||
# Each platform's session is independent
|
||||
assert len(mgr_t._histories["alice"]) == 2
|
||||
assert len(mgr_d._histories["alice"]) == 2
|
||||
|
||||
|
||||
class TestUnifiedSessionWithResolver:
|
||||
"""When a resolver is supplied, sessions unify across platforms."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_key_when_resolver_links_user(self, tmp_path):
|
||||
from praisonaiagents.session.identity import InMemoryIdentityResolver
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
|
||||
resolver = InMemoryIdentityResolver()
|
||||
resolver.link("telegram", "12345", "alice-global")
|
||||
resolver.link("discord", "snowflake-1", "alice-global")
|
||||
|
||||
# Cross-platform unification requires a shared persistent store.
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path))
|
||||
|
||||
agent = FakeAgent()
|
||||
mgr_t = BotSessionManager(
|
||||
platform="telegram", identity_resolver=resolver, store=store
|
||||
)
|
||||
mgr_d = BotSessionManager(
|
||||
platform="discord", identity_resolver=resolver, store=store
|
||||
)
|
||||
|
||||
await mgr_t.chat(agent, "12345", "hi from telegram")
|
||||
# Second call from Discord should see the Telegram turn in history
|
||||
await mgr_d.chat(agent, "snowflake-1", "what did I say?")
|
||||
|
||||
history_passed_to_second_call = agent.calls[1][0]
|
||||
assert any(
|
||||
"hi from telegram" in str(m.get("content", ""))
|
||||
for m in history_passed_to_second_call
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unlinked_user_falls_back_to_platform_id(self):
|
||||
from praisonaiagents.session.identity import InMemoryIdentityResolver
|
||||
|
||||
resolver = InMemoryIdentityResolver() # no links
|
||||
agent = FakeAgent()
|
||||
mgr = BotSessionManager(platform="telegram", identity_resolver=resolver)
|
||||
|
||||
await mgr.chat(agent, "12345", "hello")
|
||||
# Falls back to platform-prefixed key
|
||||
assert "telegram:12345" in mgr._histories or "12345" in mgr._histories
|
||||
|
||||
|
||||
class TestConcurrentUsersOnSharedAgent:
|
||||
"""Regression: agent_lock must hold across the full LLM call, not
|
||||
only the history swap. Otherwise concurrent users on a shared
|
||||
Agent instance see each other's chat_history.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_history_leak_between_concurrent_users(self):
|
||||
import asyncio
|
||||
|
||||
class SlowFakeAgent:
|
||||
def __init__(self):
|
||||
self.chat_history = []
|
||||
|
||||
def chat(self, prompt):
|
||||
# Capture what history was visible at call time
|
||||
seen_history = list(self.chat_history)
|
||||
# Append response
|
||||
self.chat_history.append(
|
||||
{"role": "user", "content": prompt}
|
||||
)
|
||||
self.chat_history.append(
|
||||
{"role": "assistant", "content": f"reply: saw {len(seen_history)} prior msgs"}
|
||||
)
|
||||
# Yield control so a concurrent task can interfere
|
||||
# if locking is wrong
|
||||
import time
|
||||
time.sleep(0.05)
|
||||
return f"saw_{len(seen_history)}"
|
||||
|
||||
agent = SlowFakeAgent()
|
||||
mgr = BotSessionManager(platform="telegram")
|
||||
|
||||
# Pre-seed alice's history with 4 messages
|
||||
for i in range(2):
|
||||
await mgr.chat(agent, "alice", f"msg {i}")
|
||||
# Bob has no history yet.
|
||||
# Now concurrent: alice should see her 4 msgs, bob should see 0.
|
||||
results = await asyncio.gather(
|
||||
mgr.chat(agent, "alice", "alice question"),
|
||||
mgr.chat(agent, "bob", "bob question"),
|
||||
)
|
||||
# Alice has 2 prior turns × 2 (user+assistant) = 4 msgs visible.
|
||||
# Bob has 0 prior msgs.
|
||||
# Note ordering by user_lock means each user's previous turn is
|
||||
# reflected. Key invariant: bob never sees alice's content and
|
||||
# vice versa.
|
||||
history_alice = mgr._histories.get("alice", [])
|
||||
history_bob = mgr._histories.get("bob", [])
|
||||
assert any("alice question" in str(m.get("content", "")) for m in history_alice)
|
||||
assert any("bob question" in str(m.get("content", "")) for m in history_bob)
|
||||
# Cross-leak check: bob's history must NOT contain alice's content.
|
||||
for m in history_bob:
|
||||
assert "alice" not in str(m.get("content", "")).lower()
|
||||
for m in history_alice:
|
||||
assert "bob" not in str(m.get("content", "")).lower()
|
||||
|
||||
|
||||
class TestSessionKeyDerivation:
|
||||
def test_resolver_changes_session_key(self):
|
||||
from praisonaiagents.session.identity import InMemoryIdentityResolver
|
||||
|
||||
resolver = InMemoryIdentityResolver()
|
||||
resolver.link("telegram", "12345", "alice-global")
|
||||
|
||||
mgr = BotSessionManager(platform="telegram", identity_resolver=resolver)
|
||||
key = mgr._session_key("12345")
|
||||
assert "alice-global" in key
|
||||
assert "telegram" not in key # unified, not platform-prefixed
|
||||
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