qzl
261 lines
8.6 KiB
Python
261 lines
8.6 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
from typing import Any, cast
|
|
from uuid import UUID
|
|
|
|
from agentscope.message import Msg
|
|
from core.agentscope.events import (
|
|
AgentScopeAgUiCodec,
|
|
AgentScopeEventPipeline,
|
|
RedisStreamBus,
|
|
SqlAlchemyEventStore,
|
|
)
|
|
from core.agentscope.runtime.orchestrator import AgentScopeRuntimeOrchestrator
|
|
from core.agentscope.schemas.agui_input import parse_run_input
|
|
from core.agentscope.services.context_service import AgentContextService
|
|
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 worker_agent_broker, worker_automation_broker
|
|
from schemas.automation import MessageContextConfig, RuntimeConfig
|
|
from schemas.memories.memory_content import UserMemoryContent, WorkProfileContent
|
|
from schemas.messages.chat_message import (
|
|
AgentChatMessageMetadata,
|
|
extract_user_message_attachments,
|
|
)
|
|
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.memories.repository import SQLAlchemyMemoriesRepository
|
|
from v1.memories.service import MemoriesService
|
|
from v1.users.dependencies import get_user_service
|
|
|
|
logger = get_logger("core.agentscope.runtime.tasks")
|
|
_MAX_CONTEXT_ATTACHMENTS = 3
|
|
|
|
|
|
def _serialize_tool_agent_output(
|
|
*,
|
|
metadata: AgentChatMessageMetadata | dict[str, object] | None,
|
|
) -> str | None:
|
|
if metadata is None:
|
|
return None
|
|
|
|
try:
|
|
resolved_metadata = (
|
|
metadata
|
|
if isinstance(metadata, AgentChatMessageMetadata)
|
|
else AgentChatMessageMetadata.model_validate(metadata)
|
|
)
|
|
except Exception:
|
|
return None
|
|
|
|
tool_agent_output = resolved_metadata.tool_agent_output
|
|
if tool_agent_output is None:
|
|
return None
|
|
|
|
return json.dumps(
|
|
tool_agent_output.model_dump(mode="json", exclude_none=True),
|
|
ensure_ascii=True,
|
|
separators=(",", ":"),
|
|
)
|
|
|
|
|
|
def _load_runtime() -> type[Any]:
|
|
return AgentScopeRuntimeOrchestrator
|
|
|
|
|
|
async def _build_user_context(
|
|
*,
|
|
owner_id: UUID,
|
|
session: Any,
|
|
) -> UserContext:
|
|
current_user = CurrentUser(id=owner_id)
|
|
user_service = get_user_service(session=session, user=current_user)
|
|
return await user_service.get_me()
|
|
|
|
|
|
async def _build_recent_context_messages(
|
|
*,
|
|
session: Any,
|
|
thread_id: str,
|
|
context_config: "MessageContextConfig",
|
|
) -> list[Msg]:
|
|
context_service = AgentContextService(repository=AgentRepository(session))
|
|
result = await context_service.load_context_messages(
|
|
thread_id=thread_id,
|
|
context_config=context_config,
|
|
)
|
|
|
|
if not result:
|
|
return []
|
|
|
|
raw_messages: list[dict[str, object]] = result.get("messages") or []
|
|
if not raw_messages:
|
|
return []
|
|
|
|
converted: list[Msg] = []
|
|
|
|
for msg in raw_messages:
|
|
role_raw = msg.get("role")
|
|
role = role_raw if isinstance(role_raw, str) else "user"
|
|
content_raw = msg.get("content", "")
|
|
content: str = content_raw if isinstance(content_raw, str) else ""
|
|
metadata_raw = msg.get("metadata")
|
|
metadata: AgentChatMessageMetadata | dict[str, object] | None
|
|
if isinstance(metadata_raw, AgentChatMessageMetadata):
|
|
metadata = metadata_raw
|
|
elif isinstance(metadata_raw, dict):
|
|
metadata = metadata_raw
|
|
else:
|
|
metadata = None
|
|
|
|
if role == "user" and metadata:
|
|
image_blocks: list[dict[str, Any]] = []
|
|
attachments = extract_user_message_attachments(metadata)[
|
|
:_MAX_CONTEXT_ATTACHMENTS
|
|
]
|
|
for attachment in attachments:
|
|
try:
|
|
image_bytes = await supabase_service.download_bytes(
|
|
bucket=attachment.bucket,
|
|
path=attachment.path,
|
|
)
|
|
except Exception:
|
|
continue
|
|
b64_data = base64.b64encode(image_bytes).decode("utf-8")
|
|
image_blocks.append(
|
|
{
|
|
"type": "image",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": attachment.mime_type or "image/png",
|
|
"data": b64_data,
|
|
},
|
|
}
|
|
)
|
|
|
|
if image_blocks:
|
|
multimodal_content: list[dict[str, Any]] = []
|
|
if isinstance(content, str) and content:
|
|
multimodal_content.append({"type": "text", "text": content})
|
|
multimodal_content.extend(image_blocks)
|
|
converted.append(
|
|
Msg(
|
|
name="user",
|
|
role="user",
|
|
content=cast(Any, multimodal_content),
|
|
)
|
|
)
|
|
continue
|
|
|
|
if role == "tool":
|
|
role = "assistant"
|
|
tool_content = _serialize_tool_agent_output(metadata=metadata)
|
|
if not tool_content:
|
|
continue
|
|
content = tool_content
|
|
|
|
converted.append(
|
|
Msg(
|
|
name=role or "user",
|
|
role=role if role in ("user", "assistant", "system") else "user",
|
|
content=content,
|
|
)
|
|
)
|
|
|
|
return converted
|
|
|
|
|
|
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")
|
|
runtime_config_raw = command.get("runtime_config")
|
|
|
|
if not isinstance(raw_owner_id, str) or not raw_owner_id.strip():
|
|
raise ValueError("owner_id is required")
|
|
if run_input_raw is None:
|
|
raise ValueError("run_input is required")
|
|
|
|
run_input = parse_run_input(run_input_raw)
|
|
runtime_config = RuntimeConfig.model_validate(runtime_config_raw or {})
|
|
thread_id = run_input.thread_id
|
|
run_id = run_input.run_id
|
|
owner_id = UUID(raw_owner_id)
|
|
|
|
if command_type != "run":
|
|
raise ValueError("invalid command type")
|
|
|
|
orchestrator = _load_runtime()
|
|
|
|
async with AsyncSessionLocal() as session:
|
|
current_user = CurrentUser(id=owner_id)
|
|
user_context = await _build_user_context(owner_id=owner_id, session=session)
|
|
memories_service = MemoriesService(
|
|
repository=SQLAlchemyMemoriesRepository(session),
|
|
session=session,
|
|
current_user=current_user,
|
|
)
|
|
memories_result = await memories_service.get_all_memories()
|
|
user_memory: UserMemoryContent | None = memories_result.get("user_memory")
|
|
work_memory: WorkProfileContent | None = memories_result.get("work_memory")
|
|
|
|
redis_client = await get_or_init_redis_client()
|
|
bus = RedisStreamBus(
|
|
client=redis_client,
|
|
stream_prefix=config.agent_runtime.redis_stream_prefix,
|
|
read_count=config.agent_runtime.redis_stream_read_count,
|
|
block_ms=config.agent_runtime.redis_stream_block_ms,
|
|
)
|
|
pipeline = AgentScopeEventPipeline(
|
|
codec=AgentScopeAgUiCodec(),
|
|
store=SqlAlchemyEventStore(
|
|
session_factory=AsyncSessionLocal,
|
|
),
|
|
bus=bus,
|
|
)
|
|
runtime = orchestrator(
|
|
pipeline=pipeline,
|
|
)
|
|
|
|
context_messages = await _build_recent_context_messages(
|
|
session=session,
|
|
thread_id=thread_id,
|
|
context_config=runtime_config.context,
|
|
)
|
|
|
|
await runtime.run(
|
|
run_input=run_input,
|
|
context_messages=context_messages,
|
|
user_context=user_context,
|
|
runtime_config=runtime_config,
|
|
user_memory=user_memory,
|
|
work_memory=work_memory,
|
|
)
|
|
logger.info(
|
|
"agentscope runtime task completed",
|
|
command_type=command_type,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
)
|
|
return {
|
|
"thread_id": thread_id,
|
|
"run_id": run_id,
|
|
"status": "completed",
|
|
}
|
|
|
|
|
|
@worker_agent_broker.task(task_name="tasks.agentscope.run_command.agent")
|
|
async def run_command_task_agent(command: dict[str, object]) -> dict[str, object]:
|
|
return await run_agentscope_task(command)
|
|
|
|
|
|
@worker_automation_broker.task(task_name="tasks.agentscope.run_command.automation")
|
|
async def run_command_task_automation(command: dict[str, object]) -> dict[str, object]:
|
|
return await run_agentscope_task(command)
|