betterbot/tool_pipeline.py
Andre K e68c84424f
Some checks failed
Deploy BetterBot / deploy (push) Failing after 3s
Deploy BetterBot / notify (push) Successful in 3s
feat: fork from CodeAnywhere framework
Replace standalone Telegram bot with full CodeAnywhere framework fork.
BetterBot shares all framework code and customizes only:
- instance.py: BetterBot identity, system prompt, feature flags
- tools/site_editing/: list_files, read_file, write_file with auto git push
- .env: model defaults and site directory paths
- compose/: Docker setup with betterlifesg + memoraiz mounts
- deploy script: RackNerd with Infisical secrets
2026-04-19 08:01:27 +08:00

311 lines
11 KiB
Python

"""Shared tool pipeline — single source of truth for user context, tool set
activation, and integration tool building.
Extracted from main.py and telegram_bot.py (T009) to eliminate code
duplication. Both entry points call these functions instead of maintaining
their own copies.
T015/T020/T021: Added UserContext with lazy credential decryption and
per-user tool pipeline functions.
"""
from __future__ import annotations
import logging
from datetime import datetime, timezone
from typing import Any
from config import settings
from tool_registry import registry as tool_registry
logger = logging.getLogger(__name__)
# ── UserContext (T015) ───────────────────────────────────────────
class UserContext:
"""Per-request context that lazily decrypts user credentials.
Credentials are decrypted on first access via ``get_credential()``,
cached for the duration of one request, and discarded when the object
goes out of scope. Expired credentials return ``None`` to signal
re-provisioning.
"""
def __init__(self, user: Any, store: Any) -> None:
from user_store import User # deferred to avoid circular import
self.user: User = user
self._store = store
self._cache: dict[str, str | None] = {}
def get_credential(self, service: str) -> str | None:
"""Return the decrypted API token for *service*, or None."""
if service in self._cache:
return self._cache[service]
from user_store import decrypt
cred = self._store.get_credential(self.user.id, service)
if cred is None:
self._cache[service] = None
return None
# Check expiry
if cred.expires_at:
try:
exp = datetime.fromisoformat(cred.expires_at)
if exp < datetime.now(timezone.utc):
logger.warning(
"Credential for %s/%s expired at %s",
self.user.id,
service,
cred.expires_at,
)
self._cache[service] = None
return None
except ValueError:
pass # ignore unparseable expiry — treat as valid
try:
token = decrypt(cred.encrypted_token)
except Exception:
logger.exception("Failed to decrypt credential for %s/%s", self.user.id, service)
self._cache[service] = None
return None
self._cache[service] = token
# Touch last_used_at
self._store.touch_credential(self.user.id, service)
return token
# ── Shared pipeline functions ────────────────────────────────────
# Service name → context key mapping used by tool set factories.
_SERVICE_CONTEXT_KEYS: dict[str, dict[str, str]] = {
"vikunja": {
"vikunja_api_url": "VIKUNJA_API_URL",
"vikunja_api_key": "__credential__",
},
"karakeep": {
"karakeep_api_url": "KARAKEEP_API_URL",
"karakeep_api_key": "__credential__",
},
}
def active_toolset_names(user: Any | None = None) -> list[str]:
"""Return names of tool sets whose required credentials are available.
When *user* is given (a ``User`` from user_store), credentials are
checked via the credential vault. When ``None``, falls back to the
legacy env-var check for backwards compatibility.
"""
if user is None:
# Legacy path — env-var based
ctx = _build_legacy_context()
active: list[str] = []
for name, ts in tool_registry.available.items():
if all(ctx.get(k) for k in ts.required_keys):
active.append(name)
return active
# Per-user path
from user_store import get_store
store = get_store()
uctx = UserContext(user, store)
ctx = _build_user_context_dict(user, uctx)
active = []
for name, ts in tool_registry.available.items():
if all(ctx.get(k) for k in ts.required_keys):
active.append(name)
return active
def build_user_context(user: Any | None = None) -> dict[str, Any]:
"""Build the context dict consumed by tool set factories.
When *user* is given, credentials come from the per-user vault.
When ``None``, falls back to env-var credentials.
"""
if user is None:
return _build_legacy_context()
from user_store import get_store
store = get_store()
uctx = UserContext(user, store)
return _build_user_context_dict(user, uctx)
def build_integration_tools(
user: Any | None = None,
*,
thread_id: str | None = None,
system_prompt: str | None = None,
) -> list:
"""Build Copilot SDK tools for all active integrations.
Always returns a list (possibly empty), never None.
"""
active = active_toolset_names(user)
if not active:
return []
ctx = build_user_context(user)
if thread_id is not None:
ctx["_thread_id"] = thread_id
if system_prompt is not None:
ctx["_system_prompt"] = system_prompt
return tool_registry.get_tools(active, ctx)
async def try_provision(user: Any, service: str) -> str | None:
"""Attempt to provision *service* for *user* with lock.
Returns a human-readable status message, or None on success.
Uses a per-user per-service lock to prevent concurrent provisioning.
"""
from provisioners.base import provisioner_registry
from user_store import get_store
provisioner = provisioner_registry.get(service)
if provisioner is None:
return f"No provisioner available for {service}."
store = get_store()
# Check if already provisioned
existing = store.get_credential(user.id, service)
if existing:
return None # already has credentials
lock = provisioner_registry.get_lock(user.id, service)
if lock.locked():
return f"Provisioning {service} is already in progress."
async with lock:
# Double-check after acquiring lock
existing = store.get_credential(user.id, service)
if existing:
return None
try:
result = await provisioner.provision(user, store)
except Exception:
logger.exception("Provisioning %s for user %s failed", service, user.id)
store.log_provisioning(user.id, service, "provision_failed", '{"error": "unhandled exception"}')
return f"Failed to set up {service}. The error has been logged."
if not result.success:
return result.error or f"Failed to set up {service}."
# Mark onboarding complete on first successful provisioning
if not user.onboarding_complete:
store.set_onboarding_complete(user.id)
user.onboarding_complete = True
return None
def get_provisionable_services(user: Any) -> list[dict[str, str]]:
"""Return list of services that could be provisioned for *user*."""
from provisioners.base import provisioner_registry
from user_store import get_store
store = get_store()
existing = store.get_credentials(user.id)
result = []
for name, prov in provisioner_registry.available.items():
status = "active" if name in existing else "available"
result.append(
{
"service": name,
"capabilities": ", ".join(prov.capabilities),
"status": status,
}
)
return result
def build_onboarding_fragment(user: Any) -> str | None:
"""Return a system prompt fragment for new/onboarding users (T035)."""
if user.onboarding_complete:
return None
return (
"This is a new user who hasn't set up any services yet. "
"Welcome them warmly, explain what you can help with, "
"and offer to set up their accounts. Available services: "
+ ", ".join(
f"{s['capabilities']} ({s['service']})"
for s in get_provisionable_services(user)
if s["status"] == "available"
)
+ ". Ask them which services they'd like to activate."
)
def build_capability_fragment(user: Any) -> str | None:
"""Return a system prompt fragment showing active/available capabilities (T036)."""
services = get_provisionable_services(user)
if not services:
return None
active = [s for s in services if s["status"] == "active"]
available = [s for s in services if s["status"] == "available"]
parts = []
if active:
parts.append("Active services: " + ", ".join(f"{s['capabilities']} ({s['service']})" for s in active))
if available:
parts.append(
"Available (say 'set up <service>' to activate): "
+ ", ".join(f"{s['capabilities']} ({s['service']})" for s in available)
)
return "\n".join(parts)
def build_provisioning_confirmation(service: str, result: Any) -> str | None:
"""Return a system prompt fragment confirming provisioning (T038)."""
from config import settings as _settings
if not result or not result.success:
return None
parts = [f"I just created a {service} account for this user."]
if result.service_username:
parts.append(f"Username: {result.service_username}.")
if result.service_url:
parts.append(f"Service URL: {result.service_url}")
if not _settings.ALLOW_CREDENTIAL_REVEAL_IN_CHAT:
parts.append("Do NOT reveal the password in chat — Telegram messages are not E2E encrypted.")
return " ".join(parts)
# ── Internal helpers ─────────────────────────────────────────────
def _build_legacy_context() -> dict[str, Any]:
"""Build user context from static env-var settings (owner-only path)."""
return {
"vikunja_api_url": settings.VIKUNJA_API_URL,
"vikunja_api_key": settings.VIKUNJA_API_KEY,
"vikunja_memory_path": settings.VIKUNJA_MEMORY_PATH,
"memory_owner": "default",
"karakeep_api_url": settings.KARAKEEP_API_URL,
"karakeep_api_key": settings.KARAKEEP_API_KEY,
"_user": None,
}
def _build_user_context_dict(user: Any, uctx: UserContext) -> dict[str, Any]:
"""Build a context dict from per-user credentials."""
ctx: dict[str, Any] = {
"vikunja_memory_path": settings.VIKUNJA_MEMORY_PATH,
"memory_owner": user.id,
"_user": user, # passed through for meta-tools
}
for service, key_map in _SERVICE_CONTEXT_KEYS.items():
for ctx_key, source in key_map.items():
if source == "__credential__":
val = uctx.get_credential(service)
else:
val = getattr(settings, source, "")
if val:
ctx[ctx_key] = val
return ctx