auron/src/api/atproto.ts

Interact with your ATProto labeler from your terminal

import {
  AtpAgent,
  ToolsOzoneModerationDefs,
  ToolsOzoneModerationQueryStatuses,
} from "@atproto/api";
import { chunkArray } from "@atproto/common";
import { EventEmitter } from "events";

let agent: AtpAgent;

export const getAgent = async () => {
  if (agent) return agent;
  if (!process.env.SERVICE_URL) {
    throw new Error("SERVICE_URL env var is required");
  }
  if (!process.env.SERVICE_DID) {
    throw new Error("SERVICE_DID env var is required");
  }

  agent = new AtpAgent({ service: process.env.SERVICE_URL });
  // @ts-ignore
  agent.configureProxy(process.env.SERVICE_DID!);
  await agent.login({
    identifier: process.env.USERNAME!,
    password: process.env.PASSWORD!,
  });
  return agent;
};

export const getQueueItems = async (
  {
    maxCount = 500,
    cursor,
  }: {
    maxCount?: number;
    cursor?: string;
  },
  emitter: EventEmitter
) => {
  const agent = await getAgent();
  let nextCursor = cursor;
  let counter = 0;
  const result: ToolsOzoneModerationQueryStatuses.OutputSchema = {
    cursor,
    subjectStatuses: [],
  };

  do {
    try {
      const { data } = await agent.tools.ozone.moderation.queryStatuses({
        cursor: nextCursor,
        limit: Math.min(100, maxCount),
        reviewState: ToolsOzoneModerationDefs.REVIEWOPEN,
      });
      nextCursor = data.cursor;
      result.subjectStatuses.push(...data.subjectStatuses);
      emitter.emit("update", {
        maxCount,
        nextCursor,
        subjectCount: result.subjectStatuses.length,
      });
    } catch (err) {
      console.error(err);
      break;
    }

    // Every 5th request, wait for 500ms to avoid potential rate limiting
    if (counter % 5) {
      await new Promise((resolve) => setTimeout(resolve, 500));
    }

    counter++;
  } while (
    nextCursor &&
    (!maxCount || result.subjectStatuses.length < maxCount)
  );

  return { ...result, cursor: nextCursor };
};

export const getRepos = async (
  { dids }: { dids: string[] },
  emitter: EventEmitter
) => {
  const agent = await getAgent();
  const repos: ToolsOzoneModerationDefs.RepoViewDetail[] = [];

  for (const chunk of chunkArray(dids, 100)) {
    try {
      const { data } = await agent.tools.ozone.moderation.getRepos({
        dids: chunk,
      });
      repos.push(
        ...data.repos.filter((r) =>
          ToolsOzoneModerationDefs.isRepoViewDetail(r)
        )
      );

      emitter.emit("update", {
        total: dids.length,
        repoCount: repos.length,
      });
    } catch (err) {
      console.log(`Error fetching repos for ${chunk}`);
    }
  }

  return repos;
};

export const getRecords = async (
  { uris }: { uris: string[] },
  emitter: EventEmitter
) => {
  const agent = await getAgent();
  const records: ToolsOzoneModerationDefs.RecordViewDetail[] = [];

  for (const chunk of chunkArray(uris, 100)) {
    try {
      const { data } = await agent.tools.ozone.moderation.getRecords({
        uris: chunk,
      });
      records.push(
        ...data.records.filter((r) =>
          ToolsOzoneModerationDefs.isRecordViewDetail(r)
        )
      );

      emitter.emit("update", {
        total: uris.length,
        recordCount: records.length,
      });
    } catch (err) {
      console.log(`Error fetching records for ${chunk}`);
    }
  }

  return records;
};