diff --git a/memgpt/config.py b/memgpt/config.py index 72879cc4..66b8f633 100644 --- a/memgpt/config.py +++ b/memgpt/config.py @@ -14,6 +14,7 @@ import memgpt.utils as utils import memgpt.interface as interface from memgpt.personas.personas import get_persona_text from memgpt.humans.humans import get_human_text +from memgpt.constants import MEMGPT_DIR model_choices = [ questionary.Choice("gpt-4"), @@ -22,15 +23,14 @@ model_choices = [ value="gpt-3.5-turbo", ), ] -memgpt_dir = os.path.join(os.path.expanduser("~"), ".memgpt") class Config: personas_dir = os.path.join("memgpt", "personas", "examples") - custom_personas_dir = os.path.join(memgpt_dir, "personas") + custom_personas_dir = os.path.join(MEMGPT_DIR, "personas") humans_dir = os.path.join("memgpt", "humans", "examples") - custom_humans_dir = os.path.join(memgpt_dir, "humans") - configs_dir = os.path.join(memgpt_dir, "configs") + custom_humans_dir = os.path.join(MEMGPT_DIR, "humans") + configs_dir = os.path.join(MEMGPT_DIR, "configs") def __init__(self): os.makedirs(Config.custom_personas_dir, exist_ok=True) diff --git a/memgpt/constants.py b/memgpt/constants.py index 33924e47..bd83f7fc 100644 --- a/memgpt/constants.py +++ b/memgpt/constants.py @@ -1,9 +1,15 @@ -DEFAULT_MEMGPT_MODEL = 'gpt-4' +import os + +MEMGPT_DIR = os.path.join(os.path.expanduser("~"), ".memgpt") + +DEFAULT_MEMGPT_MODEL = "gpt-4" FIRST_MESSAGE_ATTEMPTS = 10 INITIAL_BOOT_MESSAGE = "Boot sequence complete. Persona activated." -INITIAL_BOOT_MESSAGE_SEND_MESSAGE_THOUGHT = "Bootup sequence complete. Persona activated. Testing messaging functionality." +INITIAL_BOOT_MESSAGE_SEND_MESSAGE_THOUGHT = ( + "Bootup sequence complete. Persona activated. Testing messaging functionality." +) STARTUP_QUOTES = [ "I think, therefore I am.", "All those moments will be lost in time, like tears in rain.", @@ -12,7 +18,7 @@ STARTUP_QUOTES = [ INITIAL_BOOT_MESSAGE_SEND_MESSAGE_FIRST_MSG = STARTUP_QUOTES[2] # Constants to do with summarization / conversation length window -MESSAGE_SUMMARY_WARNING_TOKENS = 7000 # the number of tokens consumed in a call before a system warning goes to the agent +MESSAGE_SUMMARY_WARNING_TOKENS = 7000 # the number of tokens consumed in a call before a system warning goes to the agent MESSAGE_SUMMARY_WARNING_STR = f"Warning: the conversation history will soon reach its maximum length and be trimmed. Make sure to save any important information from the conversation to your memory before it is removed." # Default memory limits @@ -21,11 +27,13 @@ CORE_MEMORY_HUMAN_CHAR_LIMIT = 2000 MAX_PAUSE_HEARTBEATS = 360 # in min -MESSAGE_CHATGPT_FUNCTION_MODEL = 'gpt-3.5-turbo' -MESSAGE_CHATGPT_FUNCTION_SYSTEM_MESSAGE = 'You are a helpful assistant. Keep your responses short and concise.' +MESSAGE_CHATGPT_FUNCTION_MODEL = "gpt-3.5-turbo" +MESSAGE_CHATGPT_FUNCTION_SYSTEM_MESSAGE = ( + "You are a helpful assistant. Keep your responses short and concise." +) #### Functions related REQ_HEARTBEAT_MESSAGE = "request_heartbeat == true" FUNC_FAILED_HEARTBEAT_MESSAGE = "Function call failed" -FUNCTION_PARAM_DESCRIPTION_REQ_HEARTBEAT = "Request an immediate heartbeat after function execution. Set to 'true' if you want to send a follow-up message or run a follow-up function." \ No newline at end of file +FUNCTION_PARAM_DESCRIPTION_REQ_HEARTBEAT = "Request an immediate heartbeat after function execution. Set to 'true' if you want to send a follow-up message or run a follow-up function." diff --git a/memgpt/main.py b/memgpt/main.py index 2ef440eb..00831417 100644 --- a/memgpt/main.py +++ b/memgpt/main.py @@ -26,7 +26,8 @@ from memgpt.persistence_manager import ( InMemoryStateManagerWithFaiss, ) -from memgpt.config import Config, memgpt_dir +from memgpt.config import Config +from memgpt.constants import MEMGPT_DIR import asyncio app = typer.Typer() @@ -43,7 +44,7 @@ def clear_line(): def save(memgpt_agent, cfg): filename = utils.get_local_time().replace(" ", "_").replace(":", "_") filename = f"{filename}.json" - directory = os.path.join(memgpt_dir, "saved_state") + directory = os.path.join(MEMGPT_DIR, "saved_state") filename = os.path.join(directory, filename) try: if not os.path.exists(directory): @@ -416,7 +417,7 @@ async def main( utils.get_local_time().replace(" ", "_").replace(":", "_") ) filename = f"{filename}.pkl" - directory = os.path.join(memgpt_dir, "saved_chats") + directory = os.path.join(MEMGPT_DIR, "saved_chats") try: if not os.path.exists(directory): os.makedirs(directory) diff --git a/memgpt/utils.py b/memgpt/utils.py index 52b044ab..441fb50e 100644 --- a/memgpt/utils.py +++ b/memgpt/utils.py @@ -15,33 +15,40 @@ import sqlite3 import fitz from tqdm import tqdm from memgpt.openai_tools import async_get_embedding_with_backoff -from memgpt.config import memgpt_dir +from memgpt.constants import MEMGPT_DIR + def count_tokens(s: str, model: str = "gpt-4") -> int: encoding = tiktoken.encoding_for_model(model) return len(encoding.encode(s)) + # DEBUG = True DEBUG = False + + def printd(*args, **kwargs): if DEBUG: print(*args, **kwargs) + def cosine_similarity(a, b): return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)) + def united_diff(str1, str2): lines1 = str1.splitlines(True) lines2 = str2.splitlines(True) diff = difflib.unified_diff(lines1, lines2) - return ''.join(diff) + return "".join(diff) + def get_local_time_military(): # Get the current time in UTC current_time_utc = datetime.now(pytz.utc) # Convert to San Francisco's time zone (PST/PDT) - sf_time_zone = pytz.timezone('America/Los_Angeles') + sf_time_zone = pytz.timezone("America/Los_Angeles") local_time = current_time_utc.astimezone(sf_time_zone) # You may format it as you desire @@ -49,12 +56,13 @@ def get_local_time_military(): return formatted_time + def get_local_time(): # Get the current time in UTC current_time_utc = datetime.now(pytz.utc) # Convert to San Francisco's time zone (PST/PDT) - sf_time_zone = pytz.timezone('America/Los_Angeles') + sf_time_zone = pytz.timezone("America/Los_Angeles") local_time = current_time_utc.astimezone(sf_time_zone) # You may format it as you desire, including AM/PM @@ -62,6 +70,7 @@ def get_local_time(): return formatted_time + def parse_json(string): result = None try: @@ -77,23 +86,27 @@ def parse_json(string): print(f"Error parsing json with demjson package: {e}") raise e + def prepare_archival_index(folder): index_file = os.path.join(folder, "all_docs.index") index = faiss.read_index(index_file) archival_database_file = os.path.join(folder, "all_docs.jsonl") archival_database = [] - with open(archival_database_file, 'rt') as f: + with open(archival_database_file, "rt") as f: all_data = [json.loads(line) for line in f] for doc in all_data: total = len(doc) for i, passage in enumerate(doc): - archival_database.append({ - 'content': f"[Title: {passage['title']}, {i}/{total}] {passage['text']}", - 'timestamp': get_local_time(), - }) + archival_database.append( + { + "content": f"[Title: {passage['title']}, {i}/{total}] {passage['text']}", + "timestamp": get_local_time(), + } + ) return index, archival_database + def read_in_chunks(file_object, chunk_size): while True: data = file_object.read(chunk_size) @@ -101,12 +114,14 @@ def read_in_chunks(file_object, chunk_size): break yield data + def read_pdf_in_chunks(file, chunk_size): doc = fitz.open(file) for page in doc: text = page.get_text() yield text + def read_in_rows_csv(file_object, chunk_size): csvreader = csv.reader(file_object) header = next(csvreader) @@ -114,14 +129,16 @@ def read_in_rows_csv(file_object, chunk_size): next_row_terms = [] for h, v in zip(header, row): next_row_terms.append(f"{h}={v}") - next_row_str = ', '.join(next_row_terms) + next_row_str = ", ".join(next_row_terms) yield next_row_str -def prepare_archival_index_from_files(glob_pattern, tkns_per_chunk=300, model='gpt-4'): + +def prepare_archival_index_from_files(glob_pattern, tkns_per_chunk=300, model="gpt-4"): encoding = tiktoken.encoding_for_model(model) files = glob.glob(glob_pattern) return chunk_files(files, tkns_per_chunk, model) + def total_bytes(pattern): total = 0 for filename in glob.glob(pattern): @@ -129,32 +146,35 @@ def total_bytes(pattern): total += os.path.getsize(filename) return total -def chunk_file(file, tkns_per_chunk=300, model='gpt-4'): + +def chunk_file(file, tkns_per_chunk=300, model="gpt-4"): encoding = tiktoken.encoding_for_model(model) - with open(file, 'r') as f: - if file.endswith('.pdf'): - lines = [l for l in read_pdf_in_chunks(file, tkns_per_chunk*8)] + with open(file, "r") as f: + if file.endswith(".pdf"): + lines = [l for l in read_pdf_in_chunks(file, tkns_per_chunk * 8)] if len(lines) == 0: print(f"Warning: {file} did not have any extractable text.") - elif file.endswith('.csv'): - lines = [l for l in read_in_rows_csv(f, tkns_per_chunk*8)] + elif file.endswith(".csv"): + lines = [l for l in read_in_rows_csv(f, tkns_per_chunk * 8)] else: - lines = [l for l in read_in_chunks(f, tkns_per_chunk*4)] + lines = [l for l in read_in_chunks(f, tkns_per_chunk * 4)] curr_chunk = [] curr_token_ct = 0 for i, line in enumerate(lines): line = line.rstrip() line = line.lstrip() - line += '\n' + line += "\n" try: line_token_ct = len(encoding.encode(line)) except Exception as e: - line_token_ct = len(line.split(' ')) / .75 - print(f"Could not encode line {i}, estimating it to be {line_token_ct} tokens") + line_token_ct = len(line.split(" ")) / 0.75 + print( + f"Could not encode line {i}, estimating it to be {line_token_ct} tokens" + ) print(e) if line_token_ct > tkns_per_chunk: if len(curr_chunk) > 0: - yield ''.join(curr_chunk) + yield "".join(curr_chunk) curr_chunk = [] curr_token_ct = 0 yield line[:3200] @@ -162,47 +182,57 @@ def chunk_file(file, tkns_per_chunk=300, model='gpt-4'): curr_token_ct += line_token_ct curr_chunk.append(line) if curr_token_ct > tkns_per_chunk: - yield ''.join(curr_chunk) + yield "".join(curr_chunk) curr_chunk = [] curr_token_ct = 0 if len(curr_chunk) > 0: - yield ''.join(curr_chunk) + yield "".join(curr_chunk) -def chunk_files(files, tkns_per_chunk=300, model='gpt-4'): + +def chunk_files(files, tkns_per_chunk=300, model="gpt-4"): archival_database = [] for file in files: timestamp = os.path.getmtime(file) - formatted_time = datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %I:%M:%S %p %Z%z") - file_stem = file.split('/')[-1] + formatted_time = datetime.fromtimestamp(timestamp).strftime( + "%Y-%m-%d %I:%M:%S %p %Z%z" + ) + file_stem = file.split("/")[-1] chunks = [c for c in chunk_file(file, tkns_per_chunk, model)] for i, chunk in enumerate(chunks): - archival_database.append({ - 'content': f"[File: {file_stem} Part {i}/{len(chunks)}] {chunk}", - 'timestamp': formatted_time, - }) + archival_database.append( + { + "content": f"[File: {file_stem} Part {i}/{len(chunks)}] {chunk}", + "timestamp": formatted_time, + } + ) return archival_database -def chunk_files_for_jsonl(files, tkns_per_chunk=300, model='gpt-4'): + +def chunk_files_for_jsonl(files, tkns_per_chunk=300, model="gpt-4"): ret = [] for file in files: - file_stem = file.split('/')[-1] + file_stem = file.split("/")[-1] curr_file = [] for chunk in chunk_file(file, tkns_per_chunk, model): - curr_file.append({ - 'title': file_stem, - 'text': chunk, - }) + curr_file.append( + { + "title": file_stem, + "text": chunk, + } + ) ret.append(curr_file) return ret + async def process_chunk(i, chunk, model): try: - return i, await async_get_embedding_with_backoff(chunk['content'], model=model) + return i, await async_get_embedding_with_backoff(chunk["content"], model=model) except Exception as e: print(chunk) raise e + async def process_concurrently(archival_database, model, concurrency=10): # Create a semaphore to limit the number of concurrent tasks semaphore = asyncio.Semaphore(concurrency) @@ -213,44 +243,64 @@ async def process_concurrently(archival_database, model, concurrency=10): # Create a list of tasks for chunks embedding_data = [0 for _ in archival_database] - tasks = [bounded_process_chunk(i, chunk) for i, chunk in enumerate(archival_database)] + tasks = [ + bounded_process_chunk(i, chunk) for i, chunk in enumerate(archival_database) + ] - for future in tqdm(asyncio.as_completed(tasks), total=len(archival_database), desc="Processing file chunks"): + for future in tqdm( + asyncio.as_completed(tasks), + total=len(archival_database), + desc="Processing file chunks", + ): i, result = await future embedding_data[i] = result - + return embedding_data -async def prepare_archival_index_from_files_compute_embeddings(glob_pattern, tkns_per_chunk=300, model='gpt-4', embeddings_model='text-embedding-ada-002'): + +async def prepare_archival_index_from_files_compute_embeddings( + glob_pattern, + tkns_per_chunk=300, + model="gpt-4", + embeddings_model="text-embedding-ada-002", +): files = sorted(glob.glob(glob_pattern)) - save_dir = os.path.join(memgpt_dir, "archival_index_from_files_" + get_local_time().replace(' ', '_').replace(':', '_')) + save_dir = os.path.join( + MEMGPT_DIR, + "archival_index_from_files_" + + get_local_time().replace(" ", "_").replace(":", "_"), + ) os.makedirs(save_dir, exist_ok=True) total_tokens = total_bytes(glob_pattern) / 3 - price_estimate = total_tokens / 1000 * .0001 - confirm = input(f"Computing embeddings over {len(files)} files. This will cost ~${price_estimate:.2f}. Continue? [y/n] ") - if confirm != 'y': + price_estimate = total_tokens / 1000 * 0.0001 + confirm = input( + f"Computing embeddings over {len(files)} files. This will cost ~${price_estimate:.2f}. Continue? [y/n] " + ) + if confirm != "y": raise Exception("embeddings were not computed") # chunk the files, make embeddings archival_database = chunk_files(files, tkns_per_chunk, model) embedding_data = await process_concurrently(archival_database, embeddings_model) embeddings_file = os.path.join(save_dir, "embeddings.json") - with open(embeddings_file, 'w') as f: + with open(embeddings_file, "w") as f: print(f"Saving embeddings to {embeddings_file}") json.dump(embedding_data, f) - + # make all_text.json archival_storage_file = os.path.join(save_dir, "all_docs.jsonl") chunks_by_file = chunk_files_for_jsonl(files, tkns_per_chunk, model) - with open(archival_storage_file, 'w') as f: - print(f"Saving archival storage with preloaded files to {archival_storage_file}") + with open(archival_storage_file, "w") as f: + print( + f"Saving archival storage with preloaded files to {archival_storage_file}" + ) for c in chunks_by_file: json.dump(c, f) - f.write('\n') + f.write("\n") # make the faiss index index = faiss.IndexFlatL2(1536) - data = np.array(embedding_data).astype('float32') + data = np.array(embedding_data).astype("float32") try: index.add(data) except Exception as e: @@ -261,8 +311,9 @@ async def prepare_archival_index_from_files_compute_embeddings(glob_pattern, tkn faiss.write_index(index, index_file) return save_dir + def read_database_as_list(database_name): - result_list = [] + result_list = [] try: conn = sqlite3.connect(database_name)