code: from pgvector.psycopg import register_vector
  from pgvector.sqlalchemy import Vector
  import psycopg
  
  
  from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text
  from sqlalchemy.orm import sessionmaker, mapped_column
  from sqlalchemy.ext.declarative import declarative_base
  from sqlalchemy.sql import func
  
  import re
  from tqdm import tqdm
  from typing import Optional, List, Iterator
  import numpy as np
  from tqdm import tqdm
  import pandas as pd
  
  from memgpt.config import MemGPTConfig
  from memgpt.connectors.storage import StorageConnector, Passage
  from memgpt.config import AgentConfig, MemGPTConfig
  from memgpt.constants import MEMGPT_DIR
  from memgpt.utils import printd
  
  Base = declarative_base()
  
  
  def get_db_model(table_name: str):
      config = MemGPTConfig.load()
  
      class PassageModel(Base):
          """Defines data model for storing Passages (consisting of text, embedding)"""
  
          __abstract__ = True  # this line is necessary
  
          # Assuming passage_id is the primary key
          id = Column(BIGINT, primary_key=True, nullable=False, autoincrement=True)
          doc_id = Column(String)
          text = Column(String, nullable=False)
          embedding = mapped_column(Vector(config.embedding_dim))
          # metadata_ = Column(JSON(astext_type=Text()))
  
          def __repr__(self):
              return f"<Passage(passage_id='{self.id}', text='{self.text}', embedding='{self.embedding})>"
  
      """Create database model for table_name"""
      class_name = f"{table_name.capitalize()}Model"
      Model = type(class_name, (PassageModel,), {"__tablename__": table_name, "__table_args__": {"extend_existing": True}})
      return Model
  
  
  class PostgresStorageConnector(StorageConnector):
      """Storage via Postgres"""
  
      # TODO: this should probably eventually be moved into a parent DB class
  
      def __init__(self, name: Optional[str] = None, agent_config: Optional[AgentConfig] = None):
          config = MemGPTConfig.load()
  
          # determine table name
          if agent_config:
              assert name is None, f"Cannot specify both agent config and name {name}"
              self.table_name = self.generate_table_name_agent(agent_config)
          elif name:
              assert agent_config is None, f"Cannot specify both agent config and name {name}"
              self.table_name = self.generate_table_name(name)
          else:
              raise ValueError("Must specify either agent config or name")
  
          printd(f"Using table name {self.table_name}")
  
          # create table
          self.uri = config.archival_storage_uri
          if config.archival_storage_uri is None:
              raise ValueError(f"Must specifiy archival_storage_uri in config {config.config_path}")
          self.db_model = get_db_model(self.table_name)
          self.engine = create_engine(self.uri)
          Base.metadata.create_all(self.engine)  # Create the table if it doesn't exist
          self.Session = sessionmaker(bind=self.engine)
          self.Session().execute(text("CREATE EXTENSION IF NOT EXISTS vector"))  # Enables the vector extension
  
      def get_all_paginated(self, page_size: int) -> Iterator[List[Passage]]:
          session = self.Session()
          offset = 0
          while True:
              # Retrieve a chunk of records with the given page_size
              db_passages_chunk = session.query(self.db_model).offset(offset).limit(page_size).all()
  
              # If the chunk is empty, we've retrieved all records
              if not db_passages_chunk:
                  break
  
              # Yield a list of Passage objects converted from the chunk
              yield [Passage(text=p.text, embedding=p.embedding, doc_id=p.doc_id, passage_id=p.id) for p in db_passages_chunk]
  
              # Increment the offset to get the next chunk in the next iteration
              offset += page_size
  
      def get_all(self, limit=10) -> List[Passage]:
          session = self.Session()
          db_passages = session.query(self.db_model).limit(limit).all()
          return [Passage(text=p.text, embedding=p.embedding, doc_id=p.doc_id, passage_id=p.id) for p in db_passages]
  
      def get(self, id: str) -> Optional[Passage]:
          session = self.Session()
          db_passage = session.query(self.db_model).get(id)
          if db_passage is None:
              return None
          return Passage(text=db_passage.text, embedding=db_passage.embedding, doc_id=db_passage.doc_id, passage_id=db_passage.passage_id)
  
      def size(self) -> int:
          # return size of table
          session = self.Session()
          return session.query(self.db_model).count()
  
      def insert(self, passage: Passage):
          session = self.Session()
          db_passage = self.db_model(doc_id=passage.doc_id, text=passage.text, embedding=passage.embedding)
          session.add(db_passage)
          session.commit()
  
      def insert_many(self, passages: List[Passage], show_progress=True):
          session = self.Session()
          iterable = tqdm(passages) if show_progress else passages
          for passage in iterable:
              db_passage = self.db_model(doc_id=passage.doc_id, text=passage.text, embedding=passage.embedding)
              session.add(db_passage)
          session.commit()
  
      def query(self, query: str, query_vec: List[float], top_k: int = 10) -> List[Passage]:
          session = self.Session()
          # Assuming PassageModel.embedding has the capability of computing l2_distance
          results = session.scalars(select(self.db_model).order_by(self.db_model.embedding.l2_distance(query_vec)).limit(top_k)).all()
  
          # Convert the results into Passage objects
          passages = [
              Passage(text=result.text, embedding=np.frombuffer(result.embedding), doc_id=result.doc_id, passage_id=result.id)
              for result in results
          ]
          return passages
  
      def delete(self):
          """Drop the passage table from the database."""
          # Bind the engine to the metadata of the base class so that the
          # declaratives can be accessed through a DBSession instance
          Base.metadata.bind = self.engine
  
          # Drop the table specified by the PassageModel class
          self.db_model.__table__.drop(self.engine)
  
      def save(self):
          return
  
      @staticmethod
      def list_loaded_data():
          config = MemGPTConfig.load()
          engine = create_engine(config.archival_storage_uri)
          inspector = inspect(engine)
          tables = inspector.get_table_names()
          tables = [table for table in tables if table.startswith("memgpt_") and not table.startswith("memgpt_agent_")]
          start_chars = len("memgpt_")
          tables = [table[start_chars:] for table in tables]
          return tables
  
      def sanitize_table_name(self, name: str) -> str:
          # Remove leading and trailing whitespace
          name = name.strip()
  
          # Replace spaces and invalid characters with underscores
          name = re.sub(r"\s+|\W+", "_", name)
  
          # Truncate to the maximum identifier length (e.g., 63 for PostgreSQL)
          max_length = 63
          if len(name) > max_length:
              name = name[:max_length].rstrip("_")
  
          # Convert to lowercase
          name = name.lower()
  
          return name
  
      def generate_table_name_agent(self, agent_config: AgentConfig):
          return f"memgpt_agent_{self.sanitize_table_name(agent_config.name)}"
  
      def generate_table_name(self, name: str):
          return f"memgpt_{self.sanitize_table_name(name)}"
  
  
  class LanceDBConnector(StorageConnector):
      """Storage via LanceDB"""
  
      # TODO: this should probably eventually be moved into a parent DB class
  
      def __init__(self, name: Optional[str] = None, agent_config: Optional[AgentConfig] = None):
          config = MemGPTConfig.load()
          # determine table name
          if agent_config:
              assert name is None, f"Cannot specify both agent config and name {name}"
              self.table_name = self.generate_table_name_agent(agent_config)
          elif name:
              assert agent_config is None, f"Cannot specify both agent config and name {name}"
              self.table_name = self.generate_table_name(name)
          else:
              raise ValueError("Must specify either agent config or name")
  
          printd(f"Using table name {self.table_name}")
  
          # create table
          self.uri = config.archival_storage_uri
          if config.archival_storage_uri is None:
              raise ValueError(f"Must specifiy archival_storage_uri in config {config.config_path}")
          import lancedb
  
          self.db = lancedb.connect(self.uri)
          if self.table_name in self.db.table_names():
              self.table = self.db[self.table_name]
          else:
              self.table = None
  
      def get_all_paginated(self, page_size: int) -> Iterator[List[Passage]]:
          ds = self.table.to_lance()
          offset = 0
          while True:
              # Retrieve a chunk of records with the given page_size
              db_passages_chunk = ds.to_table(offset=offset, limit=page_size).to_pylist()
              # If the chunk is empty, we've retrieved all records
              if not db_passages_chunk:
                  break
  
              # Yield a list of Passage objects converted from the chunk
              yield [
                  Passage(text=p["text"], embedding=p["vector"], doc_id=p["doc_id"], passage_id=p["passage_id"]) for p in db_passages_chunk
              ]
  
              # Increment the offset to get the next chunk in the next iteration
              offset += page_size
  
      def get_all(self, limit=10) -> List[Passage]:
          db_passages = self.table.to_lance().to_table(limit=limit).to_pylist()
          return [Passage(text=p["text"], embedding=p["vector"], doc_id=p["doc_id"], passage_id=p["passage_id"]) for p in db_passages]
  
      def get(self, id: str) -> Optional[Passage]:
          db_passage = self.table.where(f"passage_id={id}").to_list()
          if len(db_passage) == 0:
              return None
          return Passage(
              text=db_passage["text"], embedding=db_passage["embedding"], doc_id=db_passage["doc_id"], passage_id=db_passage["passage_id"]
          )
  
      def size(self) -> int:
          # return size of table
          if self.table:
              return len(self.table)
          else:
              print(f"Table with name {self.table_name} not present")
              return 0
  
      def insert(self, passage: Passage):
          data = [{"doc_id": passage.doc_id, "text": passage.text, "passage_id": passage.passage_id, "vector": passage.embedding}]
  
          if self.table is not None:
              self.table.add(data)
          else:
              self.table = self.db.create_table(self.table_name, data=data, mode="overwrite")
  
      def insert_many(self, passages: List[Passage], show_progress=True):
          data = []
          iterable = tqdm(passages) if show_progress else passages
          for passage in iterable:
              temp_dict = {"doc_id": passage.doc_id, "text": passage.text, "passage_id": passage.passage_id, "vector": passage.embedding}
              data.append(temp_dict)
  
          if self.table is not None:
              self.table.add(data)
          else:
              self.table = self.db.create_table(self.table_name, data=data, mode="overwrite")
  
      def query(self, query: str, query_vec: List[float], top_k: int = 10) -> List[Passage]:
          # Assuming query_vec is of same length as embeddings inside table
          results = self.table.search(query_vec).limit(top_k).to_list()
          # Convert the results into Passage objects
          passages = [
              Passage(text=result["text"], embedding=result["vector"], doc_id=result["doc_id"], passage_id=result["passage_id"])
              for result in results
          ]
          return passages
  
      def delete(self):
          """Drop the passage table from the database."""
          # Drop the table specified by the PassageModel class
          self.db.drop_table(self.table_name)
  
      def save(self):
          return
  
      @staticmethod
      def list_loaded_data():
          config = MemGPTConfig.load()
          import lancedb
  
          db = lancedb.connect(config.archival_storage_uri)
  
          tables = db.table_names()
          tables = [table for table in tables if table.startswith("memgpt_")]
          start_chars = len("memgpt_")
          tables = [table[start_chars:] for table in tables]
          return tables
  
      def sanitize_table_name(self, name: str) -> str:
          # Remove leading and trailing whitespace
          name = name.strip()
  
          # Replace spaces and invalid characters with underscores
          name = re.sub(r"\s+|\W+", "_", name)
  
          # Truncate to the maximum identifier length
          max_length = 63
          if len(name) > max_length:
              name = name[:max_length].rstrip("_")
  
          # Convert to lowercase
          name = name.lower()
  
          return name
  
      def generate_table_name_agent(self, agent_config: AgentConfig):
          return f"memgpt_agent_{self.sanitize_table_name(agent_config.name)}"
  
      def generate_table_name(self, name: str):
          return f"memgpt_{self.sanitize_table_name(name)}"
  
apis: lancedb.connect
extract_api: [((702, 720), 'sqlalchemy.ext.declarative.declarative_base', 'declarative_base', ([], {}), '()\n', (718, 720), False, 'from sqlalchemy.ext.declarative import declarative_base\n'), ((771, 790), 'memgpt.config.MemGPTConfig.load', 'MemGPTConfig.load', ([], {}), '()\n', (788, 790), False, 'from memgpt.config import AgentConfig, MemGPTConfig\n'), ((1026, 1094), 'sqlalchemy.Column', 'Column', (['BIGINT'], {'primary_key': '(True)', 'nullable': '(False)', 'autoincrement': '(True)'}), '(BIGINT, primary_key=True, nullable=False, autoincrement=True)\n', (1032, 1094), False, 'from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text\n'), ((1112, 1126), 'sqlalchemy.Column', 'Column', (['String'], {}), '(String)\n', (1118, 1126), False, 'from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text\n'), ((1142, 1172), 'sqlalchemy.Column', 'Column', (['String'], {'nullable': '(False)'}), '(String, nullable=False)\n', (1148, 1172), False, 'from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text\n'), ((1938, 1957), 'memgpt.config.MemGPTConfig.load', 'MemGPTConfig.load', ([], {}), '()\n', (1955, 1957), False, 'from memgpt.config import AgentConfig, MemGPTConfig\n'), ((2444, 2489), 'memgpt.utils.printd', 'printd', (['f"""Using table name {self.table_name}"""'], {}), "(f'Using table name {self.table_name}')\n", (2450, 2489), False, 'from memgpt.utils import printd\n'), ((2784, 2807), 'sqlalchemy.create_engine', 'create_engine', (['self.uri'], {}), '(self.uri)\n', (2797, 2807), False, 'from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text\n'), ((2917, 2947), 'sqlalchemy.orm.sessionmaker', 'sessionmaker', ([], {'bind': 'self.engine'}), '(bind=self.engine)\n', (2929, 2947), False, 'from sqlalchemy.orm import sessionmaker, mapped_column\n'), ((4289, 4415), 'memgpt.connectors.storage.Passage', 'Passage', ([], {'text': 'db_passage.text', 'embedding': 'db_passage.embedding', 'doc_id': 'db_passage.doc_id', 'passage_id': 'db_passage.passage_id'}), '(text=db_passage.text, embedding=db_passage.embedding, doc_id=\n    db_passage.doc_id, passage_id=db_passage.passage_id)\n', (4296, 4415), False, 'from memgpt.connectors.storage import StorageConnector, Passage\n'), ((6249, 6268), 'memgpt.config.MemGPTConfig.load', 'MemGPTConfig.load', ([], {}), '()\n', (6266, 6268), False, 'from memgpt.config import AgentConfig, MemGPTConfig\n'), ((6286, 6328), 'sqlalchemy.create_engine', 'create_engine', (['config.archival_storage_uri'], {}), '(config.archival_storage_uri)\n', (6299, 6328), False, 'from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text\n'), ((6349, 6364), 'sqlalchemy.inspect', 'inspect', (['engine'], {}), '(engine)\n', (6356, 6364), False, 'from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text\n'), ((6858, 6888), 're.sub', 're.sub', (['"""\\\\s+|\\\\W+"""', '"""_"""', 'name'], {}), "('\\\\s+|\\\\W+', '_', name)\n", (6864, 6888), False, 'import re\n'), ((7672, 7691), 'memgpt.config.MemGPTConfig.load', 'MemGPTConfig.load', ([], {}), '()\n', (7689, 7691), False, 'from memgpt.config import AgentConfig, MemGPTConfig\n'), ((8177, 8222), 'memgpt.utils.printd', 'printd', (['f"""Using table name {self.table_name}"""'], {}), "(f'Using table name {self.table_name}')\n", (8183, 8222), False, 'from memgpt.utils import printd\n'), ((8483, 8508), 'lancedb.connect', 'lancedb.connect', (['self.uri'], {}), '(self.uri)\n', (8498, 8508), False, 'import lancedb\n'), ((9895, 10033), 'memgpt.connectors.storage.Passage', 'Passage', ([], {'text': "db_passage['text']", 'embedding': "db_passage['embedding']", 'doc_id': "db_passage['doc_id']", 'passage_id': "db_passage['passage_id']"}), "(text=db_passage['text'], embedding=db_passage['embedding'], doc_id=\n    db_passage['doc_id'], passage_id=db_passage['passage_id'])\n", (9902, 10033), False, 'from memgpt.connectors.storage import StorageConnector, Passage\n'), ((11938, 11957), 'memgpt.config.MemGPTConfig.load', 'MemGPTConfig.load', ([], {}), '()\n', (11955, 11957), False, 'from memgpt.config import AgentConfig, MemGPTConfig\n'), ((11995, 12039), 'lancedb.connect', 'lancedb.connect', (['config.archival_storage_uri'], {}), '(config.archival_storage_uri)\n', (12010, 12039), False, 'import lancedb\n'), ((12481, 12511), 're.sub', 're.sub', (['"""\\\\s+|\\\\W+"""', '"""_"""', 'name'], {}), "('\\\\s+|\\\\W+', '_', name)\n", (12487, 12511), False, 'import re\n'), ((1207, 1235), 'pgvector.sqlalchemy.Vector', 'Vector', (['config.embedding_dim'], {}), '(config.embedding_dim)\n', (1213, 1235), False, 'from pgvector.sqlalchemy import Vector\n'), ((2979, 3024), 'sqlalchemy.text', 'text', (['"""CREATE EXTENSION IF NOT EXISTS vector"""'], {}), "('CREATE EXTENSION IF NOT EXISTS vector')\n", (2983, 3024), False, 'from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text\n'), ((3978, 4055), 'memgpt.connectors.storage.Passage', 'Passage', ([], {'text': 'p.text', 'embedding': 'p.embedding', 'doc_id': 'p.doc_id', 'passage_id': 'p.id'}), '(text=p.text, embedding=p.embedding, doc_id=p.doc_id, passage_id=p.id)\n', (3985, 4055), False, 'from memgpt.connectors.storage import StorageConnector, Passage\n'), ((4917, 4931), 'tqdm.tqdm', 'tqdm', (['passages'], {}), '(passages)\n', (4921, 4931), False, 'from tqdm import tqdm\n'), ((9588, 9686), 'memgpt.connectors.storage.Passage', 'Passage', ([], {'text': "p['text']", 'embedding': "p['vector']", 'doc_id': "p['doc_id']", 'passage_id': "p['passage_id']"}), "(text=p['text'], embedding=p['vector'], doc_id=p['doc_id'],\n    passage_id=p['passage_id'])\n", (9595, 9686), False, 'from memgpt.connectors.storage import StorageConnector, Passage\n'), ((10726, 10740), 'tqdm.tqdm', 'tqdm', (['passages'], {}), '(passages)\n', (10730, 10740), False, 'from tqdm import tqdm\n'), ((11471, 11590), 'memgpt.connectors.storage.Passage', 'Passage', ([], {'text': "result['text']", 'embedding': "result['vector']", 'doc_id': "result['doc_id']", 'passage_id': "result['passage_id']"}), "(text=result['text'], embedding=result['vector'], doc_id=result[\n    'doc_id'], passage_id=result['passage_id'])\n", (11478, 11590), False, 'from memgpt.connectors.storage import StorageConnector, Passage\n'), ((3590, 3667), 'memgpt.connectors.storage.Passage', 'Passage', ([], {'text': 'p.text', 'embedding': 'p.embedding', 'doc_id': 'p.doc_id', 'passage_id': 'p.id'}), '(text=p.text, embedding=p.embedding, doc_id=p.doc_id, passage_id=p.id)\n', (3597, 3667), False, 'from memgpt.connectors.storage import StorageConnector, Passage\n'), ((5632, 5663), 'numpy.frombuffer', 'np.frombuffer', (['result.embedding'], {}), '(result.embedding)\n', (5645, 5663), True, 'import numpy as np\n'), ((9195, 9293), 'memgpt.connectors.storage.Passage', 'Passage', ([], {'text': "p['text']", 'embedding': "p['vector']", 'doc_id': "p['doc_id']", 'passage_id': "p['passage_id']"}), "(text=p['text'], embedding=p['vector'], doc_id=p['doc_id'],\n    passage_id=p['passage_id'])\n", (9202, 9293), False, 'from memgpt.connectors.storage import StorageConnector, Passage\n'), ((5412, 5433), 'sqlalchemy.select', 'select', (['self.db_model'], {}), '(self.db_model)\n', (5418, 5433), False, 'from sqlalchemy import create_engine, Column, String, BIGINT, select, inspect, text\n')]