MemGPT

1
import gzip
2
import json
3
from typing import List
4

5
from memgpt.config import MemGPTConfig
6
from memgpt.constants import LLM_MAX_TOKENS
7
from memgpt.data_types import EmbeddingConfig, LLMConfig
8

9

10
def load_gzipped_file(file_path):
11
    with gzip.open(file_path, "rt", encoding="utf-8") as f:
12
        for line in f:
13
            yield json.loads(line)
14

15

16
def read_jsonl(filename) -> List[dict]:
17
    lines = []
18
    with open(filename, "r") as file:
19
        for line in file:
20
            lines.append(json.loads(line.strip()))
21
    return lines
22

23

24
def get_experiment_config(postgres_uri, endpoint_type="openai", model="gpt-4"):
25
    config = MemGPTConfig.load()
26
    config.archival_storage_type = "postgres"
27
    config.archival_storage_uri = postgres_uri
28

29
    if endpoint_type == "openai":
30
        llm_config = LLMConfig(
31
            model=model, model_endpoint_type="openai", model_endpoint="https://api.openai.com/v1", context_window=LLM_MAX_TOKENS[model]
32
        )
33
        embedding_config = EmbeddingConfig(
34
            embedding_endpoint_type="openai",
35
            embedding_endpoint="https://api.openai.com/v1",
36
            embedding_dim=1536,
37
            embedding_model="text-embedding-ada-002",
38
            embedding_chunk_size=300,  # TODO: fix this
39
        )
40
    else:
41
        assert model == "ehartford/dolphin-2.5-mixtral-8x7b", "Only model supported is ehartford/dolphin-2.5-mixtral-8x7b"
42
        llm_config = LLMConfig(
43
            model="ehartford/dolphin-2.5-mixtral-8x7b",
44
            model_endpoint_type="vllm",
45
            model_endpoint="https://api.memgpt.ai",
46
            model_wrapper="chatml",
47
            context_window=16384,
48
        )
49
        embedding_config = EmbeddingConfig(
50
            embedding_endpoint_type="hugging-face",
51
            embedding_endpoint="https://embeddings.memgpt.ai",
52
            embedding_dim=1024,
53
            embedding_model="BAAI/bge-large-en-v1.5",
54
            embedding_chunk_size=300,
55
        )
56

57
    config = MemGPTConfig(
58
        anon_clientid=config.anon_clientid,
59
        archival_storage_type="postgres",
60
        archival_storage_uri=postgres_uri,
61
        recall_storage_type="postgres",
62
        recall_storage_uri=postgres_uri,
63
        metadata_storage_type="postgres",
64
        metadata_storage_uri=postgres_uri,
65
        default_llm_config=llm_config,
66
        default_embedding_config=embedding_config,
67
    )
68
    print("Config model", config.default_llm_config.model)
69
    return config
70