all but one test passes for storageconnector

memgpt/agent_store/db.py

@@ -31,7 +31,6 @@ def __init__(self, table_type: str, config: MemGPTConfig, user_id, agent_id=None
     def get_filters(self, filters: Optional[Dict] = {}):
         filter_conditions = {**self.filters, **(filters or {})}
         all_filters = [getattr(self.SQLModel, key) == value for key, value in filter_conditions.items() if hasattr(self.SQLModel, key)]
-        breakpoint()
         return all_filters
 
     def get_all_paginated(self, filters: Optional[Dict] = {}, page_size: Optional[int] = 1000, offset=0):
@@ -65,7 +64,7 @@ def get_all_cursor(
 
         # generate query
         with self.db_session as session:
-            query = session.query(self.SQLModel).filter(*filters)
+            query = select(self.SQLModel).filter(*filters).limit(limit)
             # query = query.order_by(asc(self.SQLModel.id))
 
             # records are sorted by the order_by field first, and then by the ID if two fields are the same
@@ -87,7 +86,7 @@ def get_all_cursor(
                 query = query.filter(or_(sort_exp, and_(getattr(self.SQLModel, order_by) == before_value, self.SQLModel.id < before)))
 
             # get records
-            db_record_chunk = query.limit(limit).all()
+            db_record_chunk = session.execute(query).scalars()
         if not db_record_chunk:
             return (None, [])
         records = [record.to_pydantic() for record in db_record_chunk]
@@ -103,10 +102,10 @@ def get_all(self, filters: Optional[Dict] = {}, limit=None):
             query = select(self.SQLModel).filter(*filters)
             if limit:
                 query = query.limit(limit)
-            breakpoint()
-            db_records = session.execute(query).all()
+            db_records = session.execute(query).scalars()
+
+            return [record.to_pydantic() for record in db_records]
 
-        return [record.to_pydantic() for record in db_records]
 
     def get(self, id: str):
         try:
@@ -159,13 +158,10 @@ def insert_many(self, records, exists_ok=True, show_progress=False):
     def query(self, query: str, query_vec: List[float], top_k: int = 10, filters: Optional[Dict] = {}):
         filters = self.get_filters(filters)
         with self.db_session as session:
-            results = session.scalars(
-                select(self.SQLModel).filter(*filters).order_by(self.SQLModel.embedding.l2_distance(query_vec)).limit(top_k)
-            ).all()
-
-        # Convert the results into Pydantic objects
-        records = [result.to_pydantic() for result in results]
-        return records
+            query = select(self.SQLModel).filter(*filters).order_by(self.SQLModel.embedding.l2_distance(query_vec)).limit(top_k)
+            results = session.execute(query).scalars()
+
+            return [result.to_pydantic() for result in results]
 
     def update(self, record: MemGPTBase):
         """Updates a record in the database based on the provided Pydantic Record object."""
@@ -181,7 +177,7 @@ def query_date(self, start_date, end_date, limit=None, offset=0):
         filters = self.get_filters({})
         with self.db_session as session:
             query = (
-                session.query(self.SQLModel)
+                select(self.SQLModel)
                 .filter(*filters)
                 .filter(self.SQLModel.created_at >= start_date)
                 .filter(self.SQLModel.created_at <= end_date)
@@ -191,15 +187,15 @@ def query_date(self, start_date, end_date, limit=None, offset=0):
             )
             if limit:
                 query = query.limit(limit)
-            results = query.all()
-        return [result.to_pydantic() for result in results]
+            results = session.execute(query).scalars()
+            return [result.to_pydantic() for result in results]
 
     def query_text(self, query, limit=None, offset=0):
         # todo: make fuzz https://stackoverflow.com/questions/42388956/create-a-full-text-search-index-with-sqlalchemy-on-postgresql/42390204#42390204
         filters = self.get_filters({})
         with self.db_session as session:
             query = (
-                session.query(self.SQLModel)
+                select(self.SQLModel)
                 .filter(*filters)
                 .filter(func.lower(self.SQLModel.text).contains(func.lower(query)))
                 .filter(self.SQLModel.role != "system")
@@ -208,8 +204,9 @@ def query_text(self, query, limit=None, offset=0):
             )
             if limit:
                 query = query.limit(limit)
-            results = query.all()
-        return [result.to_pydantic() for result in results]
+            results = session.execute(query).scalars()
+
+            return [result.to_pydantic() for result in results]
 
 
     def delete(self, filters: Optional[Dict] = {}):
@@ -244,7 +241,7 @@ def query_date(self, start_date, end_date, limit=None, offset=0):
         _end_date = self.str_to_datetime(end_date) if isinstance(end_date, str) else end_date
         with self.db_session as session:
             query = (
-                session.query(self.SQLModel)
+                select(self.SQLModel)
                 .filter(*filters)
                 .filter(self.SQLModel.created_at >= _start_date)
                 .filter(self.SQLModel.created_at <= _end_date)
@@ -254,8 +251,9 @@ def query_date(self, start_date, end_date, limit=None, offset=0):
             )
             if limit:
                 query = query.limit(limit)
-            results = query.all()
-        return [result.to_pydantic() for result in results]
+            results = session.execute(query).scalars()
+
+            return [result.to_pydantic() for result in results]
 
 
 class SQLLiteStorageConnector(SQLStorageConnector):

memgpt/agent_store/storage.py

@@ -9,7 +9,7 @@
 
 from memgpt.config import MemGPTConfig
 
-from memgpt.orm.base import Base as SQLBase
+from memgpt.orm.sqlalchemy_base import SqlalchemyBase as SQLBase
 from memgpt.orm.message import Message as SQLMessage
 from memgpt.orm.passage import Passage as SQLPassage
 from memgpt.orm.document import Document as SQLDocument
@@ -19,7 +19,7 @@
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
-    from memgpt.orm.base import Base as SQLBase
+    from memgpt.orm.sqlalchemy_base import SqlalchemyBase as SQLBase
 
 
 
@@ -60,7 +60,7 @@ def __init__(
         if self.table_type == TableType.ARCHIVAL_MEMORY or self.table_type == TableType.RECALL_MEMORY:
             # agent-specific table
             assert agent_id is not None, "Agent ID must be provided for agent-specific tables"
-            self.filters = {"user_id": self.user_id, "_agent_id": self.agent_id}
+            self.filters = {"user_id": self.user_id, "_agent_id": SQLBase.to_uid(self.agent_id, True)}
         elif self.table_type == TableType.PASSAGES or self.table_type == TableType.DOCUMENTS:
             # setup base filters for user-specific tables
             self.filters = {"user_id": self.user_id}

memgpt/orm/sqlalchemy_base.py

@@ -56,7 +56,7 @@ def list(cls, *, db_session: "Session", **kwargs) -> List[Type["SqlalchemyBase"]
             return list(session.execute(query).scalars())
 
     @classmethod
-    def to_uid(cls, identifier, indifferent:Optional[bool] = False) -> "UUID":
+    def to_uid(cls, identifier, indifferent: Optional[bool] = False) -> "UUID":
         """converts the id into a uuid object
         Args:
             indifferent: if True, will not enforce the prefix check

tests/conftest.py

@@ -150,9 +150,5 @@ def patch_local_db_calls(test_get_db_session: Callable):
             "memgpt.server.rest_api.app.get_db_session",
             test_get_db_session,
         ),
-        patch(
-            "memgpt.server.rest_api.utils.get_db_session",
-            test_get_db_session,
-        ),
     ):
         yield

tests/test_client.py

@@ -1,10 +1,6 @@
 import pytest
 from faker import Faker
 
-from memgpt.settings import settings
-from memgpt import Admin, create_client
-from memgpt.constants import DEFAULT_PRESET
-from memgpt.schemas.enums import JobStatus
 from memgpt.schemas.message import Message
 from memgpt.schemas.usage import MemGPTUsageStatistics