gongzhiyong
111be3e435
The standalone prototype should be the root-level project shape for fengqun while preserving the existing planning documents already at the root. This keeps README, examples, tests, and the Python package directly discoverable without deleting the prior docs. Constraint: User clarified that swarm-minimal is the repository root, but other existing root files must remain. Rejected: Deleting existing root docs | They are part of the fengqun repository context and were explicitly protected. Confidence: high Scope-risk: narrow Directive: Keep secrets in ignored .env only; do not commit live credentials. Tested: python3 -B -m unittest discover -s tests; git diff --check; secret-pattern scan showed only placeholders/test values/task-id false positives. Not-tested: Remote web UI rendering after push.
408 lines
15 KiB
Python
408 lines
15 KiB
Python
"""Optional Azure-backed store for the minimal swarm.
|
|
|
|
The default prototype uses ``InMemorySwarmStore``. This module keeps the same
|
|
coordinator API but persists the four shared resources to PostgreSQL, Redis, and
|
|
Blob Storage when Azure credentials are supplied through environment variables.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import asdict
|
|
import json
|
|
from typing import Any, Callable, TypeVar
|
|
|
|
from .config import SwarmConfig
|
|
from .core import Agent, InMemorySwarmStore, SwarmResult, Task
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class OptionalDependencyError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _import_optional_dependencies() -> tuple[Any, Any, Any]:
|
|
try:
|
|
import psycopg
|
|
import redis
|
|
from azure.storage.blob import BlobServiceClient
|
|
except ImportError as exc:
|
|
raise OptionalDependencyError(
|
|
"Azure store dependencies are missing. Install with: "
|
|
"pip install -r requirements-azure.txt"
|
|
) from exc
|
|
return psycopg, redis, BlobServiceClient
|
|
|
|
|
|
class PostgresRedisBlobSwarmStore(InMemorySwarmStore):
|
|
"""Persist the minimal swarm resources to Azure-backed services."""
|
|
|
|
def __init__(self, config: SwarmConfig) -> None:
|
|
super().__init__()
|
|
psycopg, redis_module, blob_service_client = _import_optional_dependencies()
|
|
self._psycopg = psycopg
|
|
self._redis_module = redis_module
|
|
self.config = config
|
|
self._pg_connect_kwargs = {
|
|
"host": config.postgres.host,
|
|
"port": config.postgres.port,
|
|
"dbname": config.postgres.database,
|
|
"user": config.postgres.user,
|
|
"password": config.postgres.password,
|
|
"sslmode": "require",
|
|
"connect_timeout": 10,
|
|
"autocommit": True,
|
|
}
|
|
self._redis_connect_kwargs = {
|
|
"host": config.redis.host,
|
|
"port": config.redis.port,
|
|
"password": config.redis.password,
|
|
"ssl": config.redis.ssl,
|
|
"decode_responses": True,
|
|
"socket_timeout": 10,
|
|
"socket_connect_timeout": 10,
|
|
}
|
|
self.pg = self._connect_pg()
|
|
self.redis = self._connect_redis()
|
|
self.blob_service = blob_service_client.from_connection_string(config.blob.connection_string)
|
|
self.container = self.blob_service.get_container_client(config.blob.container)
|
|
|
|
def _connect_pg(self):
|
|
return self._psycopg.connect(
|
|
**self._pg_connect_kwargs,
|
|
)
|
|
|
|
def _reconnect_pg(self) -> None:
|
|
try:
|
|
self.pg.close()
|
|
except Exception:
|
|
pass
|
|
self.pg = self._connect_pg()
|
|
|
|
def _is_recoverable_pg_error(self, exc: Exception) -> bool:
|
|
recoverable = tuple(
|
|
error_type
|
|
for error_type in (
|
|
getattr(self._psycopg, "OperationalError", None),
|
|
getattr(self._psycopg, "InterfaceError", None),
|
|
)
|
|
if error_type is not None
|
|
)
|
|
return bool(recoverable and isinstance(exc, recoverable))
|
|
|
|
def _run_pg(self, operation: Callable[[Any], T]) -> T:
|
|
last_error: Exception | None = None
|
|
for attempt in range(2):
|
|
try:
|
|
with self.pg.cursor() as cur:
|
|
return operation(cur)
|
|
except Exception as exc:
|
|
last_error = exc
|
|
if attempt == 0 and self._is_recoverable_pg_error(exc):
|
|
self._reconnect_pg()
|
|
continue
|
|
raise
|
|
if last_error is not None:
|
|
raise last_error
|
|
raise RuntimeError("PostgreSQL operation did not execute")
|
|
|
|
def _connect_redis(self):
|
|
return self._redis_module.Redis(**self._redis_connect_kwargs)
|
|
|
|
def _reconnect_redis(self) -> None:
|
|
try:
|
|
self.redis.close()
|
|
except Exception:
|
|
pass
|
|
self.redis = self._connect_redis()
|
|
|
|
def _is_recoverable_redis_error(self, exc: Exception) -> bool:
|
|
recoverable = tuple(
|
|
error_type
|
|
for error_type in (
|
|
getattr(self._redis_module.exceptions, "ConnectionError", None),
|
|
getattr(self._redis_module.exceptions, "TimeoutError", None),
|
|
)
|
|
if error_type is not None
|
|
)
|
|
return bool(recoverable and isinstance(exc, recoverable))
|
|
|
|
def _run_redis(self, operation: Callable[[Any], T]) -> T:
|
|
last_error: Exception | None = None
|
|
for attempt in range(2):
|
|
try:
|
|
return operation(self.redis)
|
|
except Exception as exc:
|
|
last_error = exc
|
|
if attempt == 0 and self._is_recoverable_redis_error(exc):
|
|
self._reconnect_redis()
|
|
continue
|
|
raise
|
|
if last_error is not None:
|
|
raise last_error
|
|
raise RuntimeError("Redis operation did not execute")
|
|
|
|
@classmethod
|
|
def from_env(cls) -> "PostgresRedisBlobSwarmStore":
|
|
return cls(SwarmConfig.from_env())
|
|
|
|
def close(self) -> None:
|
|
try:
|
|
self.pg.close()
|
|
finally:
|
|
self.redis.close()
|
|
|
|
def ensure_schema(self) -> None:
|
|
statements = [
|
|
"""
|
|
create table if not exists swarm_tasks (
|
|
id text primary key,
|
|
kind text not null,
|
|
input text not null,
|
|
status text not null,
|
|
claimed_by text,
|
|
output text,
|
|
score double precision not null default 0,
|
|
error text,
|
|
created_at timestamptz not null default now(),
|
|
updated_at timestamptz not null default now()
|
|
)
|
|
""",
|
|
"""
|
|
create table if not exists swarm_pheromones (
|
|
task_id text primary key references swarm_tasks(id) on delete cascade,
|
|
score double precision not null default 0,
|
|
updated_at timestamptz not null default now()
|
|
)
|
|
""",
|
|
"""
|
|
create table if not exists swarm_shared_state (
|
|
key text primary key,
|
|
value text not null,
|
|
updated_at timestamptz not null default now()
|
|
)
|
|
""",
|
|
"""
|
|
create table if not exists swarm_observations (
|
|
id bigserial primary key,
|
|
task_id text not null,
|
|
agent_id text not null,
|
|
signal text not null,
|
|
score_delta double precision not null,
|
|
created_at timestamptz not null default now()
|
|
)
|
|
""",
|
|
"""
|
|
create table if not exists swarm_convergence (
|
|
run_id text primary key,
|
|
goal text not null,
|
|
accepted_output text not null,
|
|
accepted_task_id text not null,
|
|
accepted_score double precision not null,
|
|
completed_tasks integer not null,
|
|
observations jsonb not null,
|
|
artifact_path text,
|
|
created_at timestamptz not null default now()
|
|
)
|
|
""",
|
|
"""
|
|
create table if not exists swarm_outbox (
|
|
id bigserial primary key,
|
|
event_type text not null,
|
|
payload jsonb not null,
|
|
published boolean not null default false,
|
|
created_at timestamptz not null default now()
|
|
)
|
|
""",
|
|
]
|
|
def operation(cur) -> None:
|
|
for statement in statements:
|
|
cur.execute(statement)
|
|
|
|
self._run_pg(operation)
|
|
try:
|
|
self.container.create_container()
|
|
except Exception as exc:
|
|
if exc.__class__.__name__ != "ResourceExistsError":
|
|
raise
|
|
|
|
def add_task(self, task: Task) -> None:
|
|
super().add_task(task)
|
|
def operation(cur) -> None:
|
|
cur.execute(
|
|
"""
|
|
insert into swarm_tasks (id, kind, input, status, score)
|
|
values (%s, %s, %s, %s, %s)
|
|
on conflict (id) do update set
|
|
kind = excluded.kind,
|
|
input = excluded.input,
|
|
status = excluded.status,
|
|
score = excluded.score,
|
|
updated_at = now()
|
|
""",
|
|
(task.id, task.kind, task.input, task.status.value, task.score),
|
|
)
|
|
cur.execute(
|
|
"""
|
|
insert into swarm_pheromones (task_id, score)
|
|
values (%s, %s)
|
|
on conflict (task_id) do update set score = excluded.score, updated_at = now()
|
|
""",
|
|
(task.id, 0.0),
|
|
)
|
|
|
|
self._run_pg(operation)
|
|
self._emit_event("task.created", {"task_id": task.id, "kind": task.kind})
|
|
|
|
def claim_next(self, agent: Agent) -> Task | None:
|
|
task = super().claim_next(agent)
|
|
if task is None:
|
|
return None
|
|
def operation(cur) -> None:
|
|
cur.execute(
|
|
"""
|
|
update swarm_tasks
|
|
set status = %s, claimed_by = %s, updated_at = now()
|
|
where id = %s
|
|
""",
|
|
(task.status.value, agent.id, task.id),
|
|
)
|
|
|
|
self._run_pg(operation)
|
|
self._run_redis(lambda redis: redis.set(f"swarm:agent:{agent.id}:heartbeat", "alive", ex=60))
|
|
self._emit_event("task.claimed", {"task_id": task.id, "agent_id": agent.id})
|
|
self._persist_shared_state()
|
|
return task
|
|
|
|
def complete_task(self, task: Task, agent: Agent, output: str, score: float) -> None:
|
|
super().complete_task(task, agent, output, score)
|
|
def operation(cur) -> None:
|
|
cur.execute(
|
|
"""
|
|
update swarm_tasks
|
|
set status = %s, output = %s, score = %s, updated_at = now()
|
|
where id = %s
|
|
""",
|
|
(task.status.value, output, score, task.id),
|
|
)
|
|
cur.execute(
|
|
"""
|
|
insert into swarm_pheromones (task_id, score)
|
|
values (%s, %s)
|
|
on conflict (task_id) do update set score = excluded.score, updated_at = now()
|
|
""",
|
|
(task.id, self.pheromones[task.id]),
|
|
)
|
|
observation = self.observations[-1]
|
|
cur.execute(
|
|
"""
|
|
insert into swarm_observations (task_id, agent_id, signal, score_delta)
|
|
values (%s, %s, %s, %s)
|
|
""",
|
|
(observation.task_id, observation.agent_id, observation.signal, observation.score_delta),
|
|
)
|
|
|
|
self._run_pg(operation)
|
|
self._run_redis(lambda redis: redis.zadd("swarm:pheromones", {task.id: self.pheromones[task.id]}))
|
|
self._emit_event("task.done", {"task_id": task.id, "agent_id": agent.id})
|
|
self._persist_shared_state()
|
|
|
|
def fail_task(self, task: Task, agent: Agent, error: str) -> None:
|
|
super().fail_task(task, agent, error)
|
|
def operation(cur) -> None:
|
|
cur.execute(
|
|
"""
|
|
update swarm_tasks
|
|
set status = %s, error = %s, updated_at = now()
|
|
where id = %s
|
|
""",
|
|
(task.status.value, error, task.id),
|
|
)
|
|
observation = self.observations[-1]
|
|
cur.execute(
|
|
"""
|
|
insert into swarm_observations (task_id, agent_id, signal, score_delta)
|
|
values (%s, %s, %s, %s)
|
|
""",
|
|
(observation.task_id, observation.agent_id, observation.signal, observation.score_delta),
|
|
)
|
|
|
|
self._run_pg(operation)
|
|
self._emit_event("task.failed", {"task_id": task.id, "agent_id": agent.id})
|
|
self._persist_shared_state()
|
|
|
|
def converge(self, run_id: str, goal: str) -> SwarmResult:
|
|
result = super().converge(run_id, goal)
|
|
artifact_path = f"swarm-runs/{run_id}/result.json"
|
|
payload = {
|
|
"run_id": result.run_id,
|
|
"goal": result.goal,
|
|
"accepted_output": result.accepted_output,
|
|
"accepted_task_id": result.accepted_task_id,
|
|
"accepted_score": result.accepted_score,
|
|
"completed_tasks": result.completed_tasks,
|
|
"observations": [asdict(observation) for observation in result.observations],
|
|
}
|
|
self.container.upload_blob(
|
|
name=artifact_path,
|
|
data=json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8"),
|
|
overwrite=True,
|
|
)
|
|
def operation(cur) -> None:
|
|
cur.execute(
|
|
"""
|
|
insert into swarm_convergence (
|
|
run_id, goal, accepted_output, accepted_task_id, accepted_score,
|
|
completed_tasks, observations, artifact_path
|
|
)
|
|
values (%s, %s, %s, %s, %s, %s, %s, %s)
|
|
on conflict (run_id) do update set
|
|
accepted_output = excluded.accepted_output,
|
|
accepted_task_id = excluded.accepted_task_id,
|
|
accepted_score = excluded.accepted_score,
|
|
completed_tasks = excluded.completed_tasks,
|
|
observations = excluded.observations,
|
|
artifact_path = excluded.artifact_path
|
|
""",
|
|
(
|
|
result.run_id,
|
|
result.goal,
|
|
result.accepted_output,
|
|
result.accepted_task_id,
|
|
result.accepted_score,
|
|
result.completed_tasks,
|
|
json.dumps(payload["observations"]),
|
|
artifact_path,
|
|
),
|
|
)
|
|
|
|
self._run_pg(operation)
|
|
self._emit_event("run.converged", {"run_id": run_id})
|
|
self._persist_shared_state()
|
|
return result
|
|
|
|
def _persist_shared_state(self) -> None:
|
|
def operation(cur) -> None:
|
|
for key, value in self.shared_state.items():
|
|
cur.execute(
|
|
"""
|
|
insert into swarm_shared_state (key, value)
|
|
values (%s, %s)
|
|
on conflict (key) do update set value = excluded.value, updated_at = now()
|
|
""",
|
|
(key, value),
|
|
)
|
|
|
|
self._run_pg(operation)
|
|
|
|
def _emit_event(self, event_type: str, payload: dict[str, object]) -> None:
|
|
def operation(cur) -> None:
|
|
cur.execute(
|
|
"insert into swarm_outbox (event_type, payload) values (%s, %s)",
|
|
(event_type, json.dumps(payload)),
|
|
)
|
|
|
|
self._run_pg(operation)
|
|
event_payload = {"type": event_type, **{key: str(value) for key, value in payload.items()}}
|
|
self._run_redis(lambda redis: redis.xadd("swarm:events", event_payload))
|