feat: add index to avoid performance degradation (#1606)

memgpt/agent_store/db.py

@@ -12,6 +12,7 @@
     JSON,
     Column,
     DateTime,
+    Index,
     String,
     TypeDecorator,
     and_,
@@ -163,6 +164,8 @@ class PassageModel(Base):
             # Add a datetime column, with default value as the current time
             created_at = Column(DateTime(timezone=True))
 
+            Index("passage_idx_user", user_id, agent_id, doc_id),
+
             def __repr__(self):
                 return f"<Passage(passage_id='{self.id}', text='{self.text}', embedding='{self.embedding})>"
 
@@ -228,6 +231,7 @@ class MessageModel(Base):
 
             # Add a datetime column, with default value as the current time
             created_at = Column(DateTime(timezone=True))
+            Index("message_idx_user", user_id, agent_id),
 
             def __repr__(self):
                 return f"<Message(message_id='{self.id}', text='{self.text}', embedding='{self.embedding})>"

memgpt/metadata.py

@@ -13,6 +13,7 @@
     Boolean,
     Column,
     DateTime,
+    Index,
     String,
     TypeDecorator,
     create_engine,
@@ -146,6 +147,9 @@ class TokenModel(Base):
     # extra (optional) metadata
     name = Column(String)
 
+    Index(__tablename__ + "_idx_user", user_id),
+    Index(__tablename__ + "_idx_token", token),
+
     def __repr__(self) -> str:
         return f"<Token(id='{self.id}', token='{self.token}', name='{self.name}')>"
 
@@ -189,6 +193,8 @@ class AgentModel(Base):
     # tools
     tools = Column(JSON)
 
+    Index(__tablename__ + "_idx_user", user_id),
+
     def __repr__(self) -> str:
         return f"<Agent(id='{self.id}', name='{self.name}')>"
 
@@ -222,6 +228,7 @@ class SourceModel(Base):
     embedding_dim = Column(BIGINT)
     embedding_model = Column(String)
     description = Column(String)
+    Index(__tablename__ + "_idx_user", user_id),
 
     # TODO: add num passages
 
@@ -249,6 +256,7 @@ class AgentSourceMappingModel(Base):
     user_id = Column(CommonUUID, nullable=False)
     agent_id = Column(CommonUUID, nullable=False)
     source_id = Column(CommonUUID, nullable=False)
+    Index(__tablename__ + "_idx_user", user_id, agent_id, source_id),
 
     def __repr__(self) -> str:
         return f"<AgentSourceMapping(user_id='{self.user_id}', agent_id='{self.agent_id}', source_id='{self.source_id}')>"
@@ -261,6 +269,7 @@ class PresetSourceMapping(Base):
     user_id = Column(CommonUUID, nullable=False)
     preset_id = Column(CommonUUID, nullable=False)
     source_id = Column(CommonUUID, nullable=False)
+    Index(__tablename__ + "_idx_user", user_id, preset_id, source_id),
 
     def __repr__(self) -> str:
         return f"<PresetSourceMapping(user_id='{self.user_id}', preset_id='{self.preset_id}', source_id='{self.source_id}')>"
@@ -298,6 +307,7 @@ class PresetModel(Base):
     created_at = Column(DateTime(timezone=True), server_default=func.now())
 
     functions_schema = Column(JSON)
+    Index(__tablename__ + "_idx_user", user_id),
 
     def __repr__(self) -> str:
         return f"<Preset(id='{self.id}', name='{self.name}')>"

memgpt/models/pydantic_models.py

@@ -74,7 +74,7 @@ class ToolModel(SQLModel, table=True):
     json_schema: Dict = Field(default_factory=dict, sa_column=Column(JSON), description="The JSON schema of the function.")
 
     # optional: user_id (user-specific tools)
-    user_id: Optional[uuid.UUID] = Field(None, description="The unique identifier of the user associated with the function.")
+    user_id: Optional[uuid.UUID] = Field(None, description="The unique identifier of the user associated with the function.", index=True)
 
     # Needed for Column(JSON)
     class Config:
@@ -128,14 +128,14 @@ class HumanModel(SQLModel, table=True):
     text: str = Field(default=get_human_text(DEFAULT_HUMAN), description="The human text.")
     name: str = Field(..., description="The name of the human.")
     id: uuid.UUID = Field(default_factory=uuid.uuid4, description="The unique identifier of the human.", primary_key=True)
-    user_id: Optional[uuid.UUID] = Field(..., description="The unique identifier of the user associated with the human.")
+    user_id: Optional[uuid.UUID] = Field(..., description="The unique identifier of the user associated with the human.", index=True)
 
 
 class PersonaModel(SQLModel, table=True):
     text: str = Field(default=get_persona_text(DEFAULT_PERSONA), description="The persona text.")
     name: str = Field(..., description="The name of the persona.")
     id: uuid.UUID = Field(default_factory=uuid.uuid4, description="The unique identifier of the persona.", primary_key=True)
-    user_id: Optional[uuid.UUID] = Field(..., description="The unique identifier of the user associated with the persona.")
+    user_id: Optional[uuid.UUID] = Field(..., description="The unique identifier of the user associated with the persona.", index=True)
 
 
 class SourceModel(SQLModel, table=True):
@@ -167,7 +167,7 @@ class JobModel(SQLModel, table=True):
     status: JobStatus = Field(default=JobStatus.created, description="The status of the job.", sa_column=Column(ChoiceType(JobStatus)))
     created_at: datetime = Field(default_factory=get_utc_time, description="The unix timestamp of when the job was created.")
     completed_at: Optional[datetime] = Field(None, description="The unix timestamp of when the job was completed.")
-    user_id: uuid.UUID = Field(..., description="The unique identifier of the user associated with the job.")
+    user_id: uuid.UUID = Field(..., description="The unique identifier of the user associated with the job.", index=True)
     metadata_: Optional[dict] = Field({}, sa_column=Column(JSON), description="The metadata of the job.")