feat(agentscope): add memory system and automation job support
- Add consumer_registry and pipeline_registry for runtime orchestration - Add Visibility schema for message filtering - Add PipelineSpec for agent pipeline configuration - Add automation job models and configuration - Remove memory_prompt.py (consolidated into memory system) - Update runtime components: context_loader, context_service, orchestrator, runner, tasks - Update toolkit: tool_config, tool_middleware, custom tools (calendar, user_lookup) - Add auth_helpers and calendar_domain utilities - Add system_agents.yaml configuration
This commit is contained in:
qzl
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
from datetime import timezone
|
||||
from typing import Any, cast
|
||||
from uuid import UUID
|
||||
|
||||
@@ -14,26 +15,49 @@ from core.agentscope.events import (
|
||||
)
|
||||
from core.agentscope.runtime.context_service import AgentContextService
|
||||
from core.agentscope.runtime.orchestrator import AgentScopeRuntimeOrchestrator
|
||||
from core.agentscope.runtime.pipeline_registry import build_default_pipeline_spec
|
||||
from core.agentscope.schemas.agui_input import parse_run_input
|
||||
from core.automation.scheduler import (
|
||||
AutomationSchedulerService,
|
||||
SqlAlchemyAutomationSchedulerRepository,
|
||||
utc_now,
|
||||
)
|
||||
from core.auth.models import CurrentUser
|
||||
from core.config.settings import config
|
||||
from core.db.session import AsyncSessionLocal
|
||||
from core.logging import get_logger
|
||||
from core.taskiq.app import bulk_broker, critical_broker, default_broker
|
||||
from models.automation_jobs import AutomationJob
|
||||
from schemas.agent.visibility import SystemVisibilityBit, bit_mask
|
||||
from schemas.automation.config import AutomationJobConfig
|
||||
from schemas.messages.chat_message import (
|
||||
AgentChatMessageMetadata,
|
||||
extract_user_message_attachments,
|
||||
)
|
||||
from schemas.agent.forwarded_props import parse_forwarded_props_agent_type
|
||||
from schemas.user import UserContext
|
||||
from services.base.redis import get_or_init_redis_client
|
||||
from services.base.supabase import supabase_service
|
||||
from v1.agent.repository import AgentRepository
|
||||
from v1.users.dependencies import get_user_service
|
||||
from sqlalchemy import select
|
||||
|
||||
logger = get_logger("core.agentscope.runtime.tasks")
|
||||
_MAX_CONTEXT_ATTACHMENTS = 3
|
||||
|
||||
|
||||
class _BulkQueueAdapter:
|
||||
async def enqueue(
|
||||
self,
|
||||
*,
|
||||
command: dict[str, object],
|
||||
dedup_key: str | None,
|
||||
) -> str:
|
||||
del dedup_key
|
||||
result = await run_command_task_bulk.kiq(command)
|
||||
return str(result.task_id)
|
||||
|
||||
|
||||
def _serialize_tool_agent_output(
|
||||
*,
|
||||
metadata: AgentChatMessageMetadata | dict[str, object] | None,
|
||||
@@ -79,13 +103,29 @@ async def _build_recent_context_messages(
|
||||
*,
|
||||
session: Any,
|
||||
thread_id: str,
|
||||
system_agent_mode: str,
|
||||
context_mode: str,
|
||||
memory_job_config: AutomationJobConfig | None = None,
|
||||
) -> list[Msg]:
|
||||
context_service = AgentContextService(repository=AgentRepository(session))
|
||||
result = await context_service.load_context_messages(
|
||||
thread_id=thread_id,
|
||||
system_agent_mode=system_agent_mode,
|
||||
)
|
||||
if memory_job_config is not None:
|
||||
visibility_mask = bit_mask(bit=int(SystemVisibilityBit.UI_HISTORY))
|
||||
if memory_job_config.context.window_mode.value == "day":
|
||||
result = await context_service.load_by_day_window(
|
||||
thread_id=thread_id,
|
||||
day_count=memory_job_config.context.window_count,
|
||||
visibility_mask=visibility_mask,
|
||||
)
|
||||
else:
|
||||
result = await context_service.load_by_user_message_window(
|
||||
thread_id=thread_id,
|
||||
user_message_limit=memory_job_config.context.window_count,
|
||||
visibility_mask=visibility_mask,
|
||||
)
|
||||
else:
|
||||
result = await context_service.load_context_messages(
|
||||
thread_id=thread_id,
|
||||
system_agent_mode=context_mode,
|
||||
)
|
||||
if not result:
|
||||
return []
|
||||
|
||||
@@ -166,11 +206,33 @@ async def _build_recent_context_messages(
|
||||
return converted
|
||||
|
||||
|
||||
async def _load_memory_job_config(
|
||||
*,
|
||||
session: Any,
|
||||
owner_id: UUID,
|
||||
automation_job_id: str,
|
||||
) -> AutomationJobConfig:
|
||||
try:
|
||||
job_uuid = UUID(automation_job_id)
|
||||
except ValueError as exc:
|
||||
raise ValueError("automation_job_id is invalid") from exc
|
||||
|
||||
stmt = (
|
||||
select(AutomationJob)
|
||||
.where(AutomationJob.id == job_uuid)
|
||||
.where(AutomationJob.owner_id == owner_id)
|
||||
.where(AutomationJob.deleted_at.is_(None))
|
||||
)
|
||||
row = (await session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
raise ValueError("automation job not found")
|
||||
return AutomationJobConfig.model_validate(row.config or {})
|
||||
|
||||
|
||||
async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
|
||||
command_type = str(command.get("command", "run")).strip().lower()
|
||||
raw_owner_id = command.get("owner_id")
|
||||
run_input_raw = command.get("run_input")
|
||||
system_agent_mode = str(command.get("system_agent_mode", "worker")).strip().lower()
|
||||
|
||||
if not isinstance(raw_owner_id, str) or not raw_owner_id.strip():
|
||||
raise ValueError("owner_id is required")
|
||||
@@ -178,6 +240,15 @@ async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
|
||||
raise ValueError("run_input is required")
|
||||
|
||||
run_input = parse_run_input(run_input_raw)
|
||||
system_agent_mode = parse_forwarded_props_agent_type(
|
||||
getattr(run_input, "forwarded_props", None)
|
||||
)
|
||||
raw_automation_job_id = command.get("automation_job_id")
|
||||
if system_agent_mode == "memory" and (
|
||||
not isinstance(raw_automation_job_id, str) or not raw_automation_job_id
|
||||
):
|
||||
raise ValueError("automation_job_id is required for memory mode")
|
||||
pipeline_spec = build_default_pipeline_spec(mode=system_agent_mode)
|
||||
thread_id = run_input.thread_id
|
||||
run_id = run_input.run_id
|
||||
owner_id = UUID(raw_owner_id)
|
||||
@@ -189,6 +260,14 @@ async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
user_context = await _build_user_context(owner_id=owner_id, session=session)
|
||||
memory_job_config: AutomationJobConfig | None = None
|
||||
if system_agent_mode == "memory":
|
||||
assert isinstance(raw_automation_job_id, str)
|
||||
memory_job_config = await _load_memory_job_config(
|
||||
session=session,
|
||||
owner_id=owner_id,
|
||||
automation_job_id=raw_automation_job_id,
|
||||
)
|
||||
|
||||
redis_client = await get_or_init_redis_client()
|
||||
bus = RedisStreamBus(
|
||||
@@ -211,7 +290,8 @@ async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
|
||||
context_messages = await _build_recent_context_messages(
|
||||
session=session,
|
||||
thread_id=thread_id,
|
||||
system_agent_mode=system_agent_mode,
|
||||
context_mode=pipeline_spec.stages[0].context_policy.consumer_agent_type,
|
||||
memory_job_config=memory_job_config,
|
||||
)
|
||||
|
||||
await runtime.run(
|
||||
@@ -219,6 +299,7 @@ async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
|
||||
context_messages=context_messages,
|
||||
user_context=user_context,
|
||||
system_agent_mode=system_agent_mode,
|
||||
memory_job_config=memory_job_config,
|
||||
)
|
||||
logger.info(
|
||||
"agentscope runtime task completed",
|
||||
@@ -233,6 +314,35 @@ async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
|
||||
}
|
||||
|
||||
|
||||
async def run_automation_scheduler_scan(
|
||||
*,
|
||||
limit: int | None = None,
|
||||
) -> dict[str, int]:
|
||||
now = utc_now()
|
||||
safe_limit = (
|
||||
max(int(limit), 1)
|
||||
if isinstance(limit, int)
|
||||
else int(config.automation_scheduler.batch_limit)
|
||||
)
|
||||
async with AsyncSessionLocal() as session:
|
||||
repository = SqlAlchemyAutomationSchedulerRepository(session=session)
|
||||
service = AutomationSchedulerService(
|
||||
repository=repository,
|
||||
queue=_BulkQueueAdapter(),
|
||||
)
|
||||
result = await service.scan_and_dispatch(now_utc=now, limit=safe_limit)
|
||||
logger.info(
|
||||
"automation scheduler scan completed",
|
||||
scanned=result.scanned,
|
||||
dispatched=result.dispatched,
|
||||
now_utc=now.astimezone(timezone.utc).isoformat(),
|
||||
)
|
||||
return {
|
||||
"scanned": int(result.scanned),
|
||||
"dispatched": int(result.dispatched),
|
||||
}
|
||||
|
||||
|
||||
@default_broker.task(task_name="tasks.agentscope.run_command")
|
||||
async def run_command_task(command: dict[str, Any]) -> dict[str, object]:
|
||||
return await run_agentscope_task(command)
|
||||
@@ -246,3 +356,8 @@ async def run_command_task_critical(command: dict[str, Any]) -> dict[str, object
|
||||
@bulk_broker.task(task_name="tasks.agentscope.run_command.bulk")
|
||||
async def run_command_task_bulk(command: dict[str, Any]) -> dict[str, object]:
|
||||
return await run_agentscope_task(command)
|
||||
|
||||
|
||||
@default_broker.task(task_name="tasks.automation.scan_due_jobs")
|
||||
async def scan_due_automation_jobs_task(limit: int | None = None) -> dict[str, int]:
|
||||
return await run_automation_scheduler_scan(limit=limit)
|
||||
|
||||
Reference in New Issue
Block a user