diff --git a/scripts/commands.py b/scripts/commands.py index fc10d1d052..f00875f063 100644 --- a/scripts/commands.py +++ b/scripts/commands.py @@ -1,6 +1,6 @@ import browse import json -from memory import PineconeMemory +from memory.pinecone import PineconeMemory import datetime import agent_manager as agents import speak @@ -52,7 +52,7 @@ def get_command(response): def execute_command(command_name, arguments): - memory = PineconeMemory() + memory = PineconeMemory(cfg=cfg) try: if command_name == "google": diff --git a/scripts/config.py b/scripts/config.py index fe48d29800..1b716a3eba 100644 --- a/scripts/config.py +++ b/scripts/config.py @@ -1,3 +1,4 @@ +import abc import os import openai from dotenv import load_dotenv @@ -5,7 +6,7 @@ from dotenv import load_dotenv load_dotenv() -class Singleton(type): +class Singleton(abc.ABCMeta, type): """ Singleton metaclass for ensuring only one instance of a class. """ @@ -20,6 +21,10 @@ class Singleton(type): return cls._instances[cls] +class AbstractSingleton(abc.ABC, metaclass=Singleton): + pass + + class Config(metaclass=Singleton): """ Configuration class to store the state of bools for different scripts access. diff --git a/scripts/main.py b/scripts/main.py index a79fd553ce..acb63a39be 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -1,7 +1,7 @@ import json import random import commands as cmd -from memory import PineconeMemory +from memory.pinecone import PineconeMemory import data import chat from colorama import Fore, Style @@ -283,7 +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 -memory = PineconeMemory() +memory = PineconeMemory(cfg) memory.clear() print('Using memory of type: ' + memory.__class__.__name__) diff --git a/scripts/memory/__init__.py b/scripts/memory/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/scripts/memory/base.py b/scripts/memory/base.py new file mode 100644 index 0000000000..29f5d56be1 --- /dev/null +++ b/scripts/memory/base.py @@ -0,0 +1,34 @@ +import abc +from config import AbstractSingleton +import openai + + +def get_ada_embedding(text): + text = text.replace("\n", " ") + return openai.Embedding.create(input=[text], model="text-embedding-ada-002")["data"][0]["embedding"] + + +def get_text_from_embedding(embedding): + return openai.Embedding.retrieve(embedding, model="text-embedding-ada-002")["data"][0]["text"] + + +class MemoryProviderSingleton(AbstractSingleton): + @abc.abstractmethod + def add(self, data): + pass + + @abc.abstractmethod + def get(self, data): + pass + + @abc.abstractmethod + def clear(self): + pass + + @abc.abstractmethod + def get_relevant(self, data, num_relevant=5): + pass + + @abc.abstractmethod + def get_stats(self): + pass diff --git a/scripts/memory.py b/scripts/memory/pinecone.py similarity index 80% rename from scripts/memory.py rename to scripts/memory/pinecone.py index 0d265a31d8..8e1eaa570f 100644 --- a/scripts/memory.py +++ b/scripts/memory/pinecone.py @@ -1,21 +1,11 @@ -from config import Config, Singleton + import pinecone -import openai -cfg = Config() +from memory.base import MemoryProviderSingleton, get_ada_embedding -def get_ada_embedding(text): - text = text.replace("\n", " ") - return openai.Embedding.create(input=[text], model="text-embedding-ada-002")["data"][0]["embedding"] - - -def get_text_from_embedding(embedding): - return openai.Embedding.retrieve(embedding, model="text-embedding-ada-002")["data"][0]["text"] - - -class PineconeMemory(metaclass=Singleton): - def __init__(self): +class PineconeMemory(MemoryProviderSingleton): + def __init__(self, cfg): pinecone_api_key = cfg.pinecone_api_key pinecone_region = cfg.pinecone_region pinecone.init(api_key=pinecone_api_key, environment=pinecone_region) diff --git a/scripts/memory/redismem.py b/scripts/memory/redismem.py new file mode 100644 index 0000000000..e69de29bb2