diff --git a/scripts/commands.py b/scripts/commands.py index a88ad0ae09..783e6bd295 100644 --- a/scripts/commands.py +++ b/scripts/commands.py @@ -1,8 +1,6 @@ import browse import json -from memory.local import LocalCache -from memory.pinecone import PineconeMemory -from memory.redismem import RedisMemory +from memory import get_memory import datetime import agent_manager as agents import speak @@ -54,12 +52,7 @@ def get_command(response): def execute_command(command_name, arguments): - if cfg.memory_backend == "pinecone": - memory = PineconeMemory(cfg=cfg) - elif cfg.memory_backend == "redis": - memory = RedisMemory(cfg=cfg) - else: - memory = LocalCache(cfg=cfg) + memory = get_memory(cfg) try: if command_name == "google": diff --git a/scripts/main.py b/scripts/main.py index e49f1810d2..11bf0dc1b2 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -1,9 +1,7 @@ import json import random import commands as cmd -from memory.local import LocalCache -from memory.pinecone import PineconeMemory -from memory.redismem import RedisMemory +from memory import get_memory import data import chat from colorama import Fore, Style @@ -285,13 +283,7 @@ user_input = "Determine which next command to use, and respond using the format # Initialize memory and make sure it is empty. # this is particularly important for indexing and referencing pinecone memory -if cfg.memory_backend == "pinecone": - memory = PineconeMemory(cfg) - memory.clear() -elif cfg.memory_backend == "redis": - memory = RedisMemory(cfg) -else: - memory = LocalCache(cfg) +memory = get_memory(cfg, init=True) print('Using memory of type: ' + memory.__class__.__name__) diff --git a/scripts/memory/__init__.py b/scripts/memory/__init__.py index e69de29bb2..dacb05b328 100644 --- a/scripts/memory/__init__.py +++ b/scripts/memory/__init__.py @@ -0,0 +1,42 @@ +from memory.local import LocalCache +try: + from memory.redismem import RedisMemory +except ImportError: + print("Redis not installed. Skipping import.") + RedisMemory = None + +try: + from memory.pinecone import PineconeMemory +except ImportError: + print("Pinecone not installed. Skipping import.") + PineconeMemory = None + + +def get_memory(cfg, 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) + 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) + + if memory is None: + memory = LocalCache(cfg) + return memory + + +__all__ = [ + "get_memory", + "LocalCache", + "RedisCache", + "PineconeCache", +]