multigpt/memory/__init__.py - master - githubmirror/Multi-GPT

multigpt/memory/__init__.py

80 строк2 KB

lukas

memory exclusivity for local memory managment now working. fixed some small issues.

16 апр 2023, 21:19

16 апр 2023, 21:19dc9e5ea

from multigpt.memory.local import LocalCache
from multigpt.memory.no_memory import NoMemory

# List of supported memory backends
# Add a backend to this list if the import attempt is successful
supported_memory = ["local", "no_memory"]

try:
    from multigpt.memory.redismem import RedisMemory

    supported_memory.append("redis")
except ImportError:
    print("Redis not installed. Skipping import.")
    RedisMemory = None

try:
    from multigpt.memory.pinecone import PineconeMemory

    supported_memory.append("pinecone")
except ImportError:
    print("Pinecone not installed. Skipping import.")
    PineconeMemory = None

try:
    from multigpt.memory.milvus import MilvusMemory
except ImportError:
    print("pymilvus not installed. Skipping import.")
    MilvusMemory = None


def get_memory(cfg, ai_key, init=False):
    memory = None
    if cfg.memory_backend == "pinecone":
        if not PineconeMemory:
            print(
                "Error: Pinecone is not installed. Please install pinecone"
                " to use Pinecone as a memory backend."
            )
        else:
            memory = PineconeMemory(cfg, ai_key)
            if init:
                memory.clear()
    elif cfg.memory_backend == "redis":
        if not RedisMemory:
            print(
                "Error: Redis is not installed. Please install redis-py to"
                " use Redis as a memory backend."
            )
        else:
            memory = RedisMemory(cfg, ai_key)
    elif cfg.memory_backend == "milvus":
        if not MilvusMemory:
            print(
                "Error: Milvus sdk is not installed."
                "Please install pymilvus to use Milvus as memory backend."
            )
        else:
            memory = MilvusMemory(cfg, ai_key)
    elif cfg.memory_backend == "no_memory":
        memory = NoMemory(cfg)

    if memory is None:
        memory = LocalCache(cfg, ai_key)
        if init:
            memory.clear()
    return memory


def get_supported_memory_backends():
    return supported_memory


__all__ = [
    "get_memory",
    "LocalCache",
    "RedisMemory",
    "PineconeMemory",
    "NoMemory",
    "MilvusMemory",
]