babyagi/extensions/ray_objectives.py

import logging
import ray
from collections import deque

ACTOR_NAME="BabyAGI Objectives"

try:
    ray.init(address="auto", namespace="babyagi", logging_level=logging.FATAL, ignore_reinit_error=True)
except:
    ray.init(namespace="babyagi", logging_level=logging.FATAL, ignore_reinit_error=True)

@ray.remote
class CooperativeObjectivesListStorageActor:
    def __init__(self):
        self.objectives = deque([])

    def append(self, objective: str):
        if not objective in self.objectives:
            self.objectives.append(objective)

    def is_empty(self):
        return False if self.objectives else True

    def get_objective_names(self):
        return [t for t in self.objectives]

class CooperativeObjectivesListStorage:
    def __init__(self):
        try:
            self.actor = ray.get_actor(name=ACTOR_NAME, namespace="babyagi")
        except ValueError:
            self.actor = CooperativeObjectivesListStorageActor.options(name=ACTOR_NAME, namespace="babyagi", lifetime="detached").remote()

    def append(self, objective: str):
        self.actor.append.remote(objective)

    def is_empty(self):
        return ray.get(self.actor.is_empty.remote())

    def get_objective_names(self):
        return ray.get(self.actor.get_objective_names.remote())