letta-nightly 0.6.3.dev20241211104238__py3-none-any.whl → 0.6.3.dev20241212104231__py3-none-any.whl

letta/agent_store/db.py

@@ -1,467 +0,0 @@
-import base64
-import json
-import os
-from datetime import datetime
-from typing import Dict, List, Optional
-
-import numpy as np
-from sqlalchemy import (
-    BINARY,
-    Column,
-    DateTime,
-    Index,
-    String,
-    TypeDecorator,
-    and_,
-    asc,
-    desc,
-    or_,
-    select,
-    text,
-)
-from sqlalchemy.orm import mapped_column
-from sqlalchemy.orm.session import close_all_sessions
-from sqlalchemy.sql import func
-from sqlalchemy_json import MutableJson
-from tqdm import tqdm
-
-from letta.agent_store.storage import StorageConnector, TableType
-from letta.config import LettaConfig
-from letta.constants import MAX_EMBEDDING_DIM
-from letta.metadata import EmbeddingConfigColumn
-from letta.orm.base import Base
-from letta.orm.file import FileMetadata as FileMetadataModel
-
-# from letta.schemas.message import Message, Passage, Record, RecordType, ToolCall
-from letta.orm.passage import Passage as PassageModel
-from letta.settings import settings
-
-config = LettaConfig()
-
-
-class CommonVector(TypeDecorator):
-    """Common type for representing vectors in SQLite"""
-
-    impl = BINARY
-    cache_ok = True
-
-    def load_dialect_impl(self, dialect):
-        return dialect.type_descriptor(BINARY())
-
-    def process_bind_param(self, value, dialect):
-        if value is None:
-            return value
-        # Ensure value is a numpy array
-        if isinstance(value, list):
-            value = np.array(value, dtype=np.float32)
-        # Serialize numpy array to bytes, then encode to base64 for universal compatibility
-        return base64.b64encode(value.tobytes())
-
-    def process_result_value(self, value, dialect):
-        if not value:
-            return value
-        # Check database type and deserialize accordingly
-        if dialect.name == "sqlite":
-            # Decode from base64 and convert back to numpy array
-            value = base64.b64decode(value)
-        # For PostgreSQL, value is already in bytes
-        return np.frombuffer(value, dtype=np.float32)
-
-class SQLStorageConnector(StorageConnector):
-    def __init__(self, table_type: str, config: LettaConfig, user_id, agent_id=None):
-        super().__init__(table_type=table_type, config=config, user_id=user_id, agent_id=agent_id)
-        self.config = config
-
-    def get_filters(self, filters: Optional[Dict] = {}):
-        if filters is not None:
-            filter_conditions = {**self.filters, **filters}
-        else:
-            filter_conditions = self.filters
-        all_filters = [getattr(self.db_model, key) == value for key, value in filter_conditions.items()]
-        return all_filters
-
-    def get_all_paginated(self, filters: Optional[Dict] = {}, page_size: Optional[int] = 1000, offset=0):
-        filters = self.get_filters(filters)
-        while True:
-            # Retrieve a chunk of records with the given page_size
-            with self.session_maker() as session:
-                db_record_chunk = session.query(self.db_model).filter(*filters).offset(offset).limit(page_size).all()
-
-            # If the chunk is empty, we've retrieved all records
-            if not db_record_chunk:
-                break
-
-            # Yield a list of Record objects converted from the chunk
-            yield [record.to_record() for record in db_record_chunk]
-
-            # Increment the offset to get the next chunk in the next iteration
-            offset += page_size
-
-    def get_all_cursor(
-        self,
-        filters: Optional[Dict] = {},
-        after: str = None,
-        before: str = None,
-        limit: Optional[int] = 1000,
-        order_by: str = "created_at",
-        reverse: bool = False,
-    ):
-        """Get all that returns a cursor (record.id) and records"""
-        filters = self.get_filters(filters)
-
-        # generate query
-        with self.session_maker() as session:
-            query = session.query(self.db_model).filter(*filters)
-            # query = query.order_by(asc(self.db_model.id))
-
-            # records are sorted by the order_by field first, and then by the ID if two fields are the same
-            if reverse:
-                query = query.order_by(desc(getattr(self.db_model, order_by)), asc(self.db_model.id))
-            else:
-                query = query.order_by(asc(getattr(self.db_model, order_by)), asc(self.db_model.id))
-
-            # cursor logic: filter records based on before/after ID
-            if after:
-                after_value = getattr(self.get(id=after), order_by)
-                sort_exp = getattr(self.db_model, order_by) > after_value
-                query = query.filter(
-                    or_(sort_exp, and_(getattr(self.db_model, order_by) == after_value, self.db_model.id > after))  # tiebreaker case
-                )
-            if before:
-                before_value = getattr(self.get(id=before), order_by)
-                sort_exp = getattr(self.db_model, order_by) < before_value
-                query = query.filter(or_(sort_exp, and_(getattr(self.db_model, order_by) == before_value, self.db_model.id < before)))
-
-            # get records
-            db_record_chunk = query.limit(limit).all()
-        if not db_record_chunk:
-            return (None, [])
-        records = [record.to_record() for record in db_record_chunk]
-        next_cursor = db_record_chunk[-1].id
-        assert isinstance(next_cursor, str)
-
-        # return (cursor, list[records])
-        return (next_cursor, records)
-
-    def get_all(self, filters: Optional[Dict] = {}, limit=None):
-        filters = self.get_filters(filters)
-        with self.session_maker() as session:
-            if limit:
-                db_records = session.query(self.db_model).filter(*filters).limit(limit).all()
-            else:
-                db_records = session.query(self.db_model).filter(*filters).all()
-        return [record.to_record() for record in db_records]
-
-    def get(self, id: str):
-        with self.session_maker() as session:
-            db_record = session.get(self.db_model, id)
-        if db_record is None:
-            return None
-        return db_record.to_record()
-
-    def size(self, filters: Optional[Dict] = {}) -> int:
-        # return size of table
-        filters = self.get_filters(filters)
-        with self.session_maker() as session:
-            return session.query(self.db_model).filter(*filters).count()
-
-    def insert(self, record):
-        raise NotImplementedError
-
-    def insert_many(self, records, show_progress=False):
-        raise NotImplementedError
-
-    def query(self, query: str, query_vec: List[float], top_k: int = 10, filters: Optional[Dict] = {}):
-        raise NotImplementedError("Vector query not implemented for SQLStorageConnector")
-
-    def save(self):
-        return
-
-    def list_data_sources(self):
-        assert self.table_type == TableType.ARCHIVAL_MEMORY, f"list_data_sources only implemented for ARCHIVAL_MEMORY"
-        with self.session_maker() as session:
-            unique_data_sources = session.query(self.db_model.data_source).filter(*self.filters).distinct().all()
-        return unique_data_sources
-
-    def query_date(self, start_date, end_date, limit=None, offset=0):
-        filters = self.get_filters({})
-        with self.session_maker() as session:
-            query = (
-                session.query(self.db_model)
-                .filter(*filters)
-                .filter(self.db_model.created_at >= start_date)
-                .filter(self.db_model.created_at <= end_date)
-                .filter(self.db_model.role != "system")
-                .filter(self.db_model.role != "tool")
-                .offset(offset)
-            )
-            if limit:
-                query = query.limit(limit)
-            results = query.all()
-        return [result.to_record() 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.session_maker() as session:
-            query = (
-                session.query(self.db_model)
-                .filter(*filters)
-                .filter(func.lower(self.db_model.text).contains(func.lower(query)))
-                .filter(self.db_model.role != "system")
-                .filter(self.db_model.role != "tool")
-                .offset(offset)
-            )
-            if limit:
-                query = query.limit(limit)
-            results = query.all()
-        # return [self.type(**vars(result)) for result in results]
-        return [result.to_record() for result in results]
-
-    # Should be used only in tests!
-    def delete_table(self):
-        close_all_sessions()
-        with self.session_maker() as session:
-            self.db_model.__table__.drop(session.bind)
-            session.commit()
-
-    def delete(self, filters: Optional[Dict] = {}):
-        filters = self.get_filters(filters)
-        with self.session_maker() as session:
-            session.query(self.db_model).filter(*filters).delete()
-            session.commit()
-
-
-class PostgresStorageConnector(SQLStorageConnector):
-    """Storage via Postgres"""
-
-    # TODO: this should probably eventually be moved into a parent DB class
-
-    def __init__(self, table_type: str, config: LettaConfig, user_id, agent_id=None):
-        from pgvector.sqlalchemy import Vector
-
-        super().__init__(table_type=table_type, config=config, user_id=user_id, agent_id=agent_id)
-
-        # construct URI from enviornment variables
-        if settings.pg_uri:
-            self.uri = settings.pg_uri
-
-        # use config URI
-        # TODO: remove this eventually (config should NOT contain URI)
-        if table_type == TableType.ARCHIVAL_MEMORY or table_type == TableType.PASSAGES:
-            self.uri = self.config.archival_storage_uri
-            self.db_model = PassageModel
-            if self.config.archival_storage_uri is None:
-                raise ValueError(f"Must specify archival_storage_uri in config {self.config.config_path}")
-        elif table_type == TableType.FILES:
-            self.uri = self.config.metadata_storage_uri
-            self.db_model = FileMetadataModel
-            if self.config.metadata_storage_uri is None:
-                raise ValueError(f"Must specify metadata_storage_uri in config {self.config.config_path}")
-        else:
-            raise ValueError(f"Table type {table_type} not implemented")
-
-        if settings.pg_uri:
-            for c in self.db_model.__table__.columns:
-                if c.name == "embedding":
-                    assert isinstance(c.type, Vector), f"Embedding column must be of type Vector, got {c.type}"
-
-        from letta.server.server import db_context
-
-        self.session_maker = db_context
-
-        # TODO: move to DB init
-        if settings.pg_uri:
-            with self.session_maker() as session:
-                session.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))  # Enables the vector extension
-
-    def query(self, query: str, query_vec: List[float], top_k: int = 10, filters: Optional[Dict] = {}):
-        filters = self.get_filters(filters)
-        with self.session_maker() as session:
-            results = session.scalars(
-                select(self.db_model).filter(*filters).order_by(self.db_model.embedding.l2_distance(query_vec)).limit(top_k)
-            ).all()
-
-        # Convert the results into Passage objects
-        records = [result.to_record() for result in results]
-        return records
-
-    def insert_many(self, records, exists_ok=True, show_progress=False):
-        # TODO: this is terrible, should eventually be done the same way for all types (migrate to SQLModel)
-        if len(records) == 0:
-            return
-
-        added_ids = []  # avoid adding duplicates
-        # NOTE: this has not great performance due to the excessive commits
-        with self.session_maker() as session:
-            iterable = tqdm(records) if show_progress else records
-            for record in iterable:
-                # db_record = self.db_model(**vars(record))
-
-                if record.id in added_ids:
-                    continue
-
-                existing_record = session.query(self.db_model).filter_by(id=record.id).first()
-                if existing_record:
-                    if exists_ok:
-                        fields = record.model_dump()
-                        fields.pop("id")
-                        session.query(self.db_model).filter(self.db_model.id == record.id).update(fields)
-                        print(f"Updated record with id {record.id}")
-                        session.commit()
-                    else:
-                        raise ValueError(f"Record with id {record.id} already exists.")
-
-                else:
-                    db_record = self.db_model(**record.dict())
-                    session.add(db_record)
-                    # print(f"Added record with id {record.id}")
-                    session.commit()
-
-                added_ids.append(record.id)
-
-    def insert(self, record, exists_ok=True):
-        self.insert_many([record], exists_ok=exists_ok)
-
-    def update(self, record):
-        """
-        Updates a record in the database based on the provided Record object.
-        """
-        with self.session_maker() as session:
-            # Find the record by its ID
-            db_record = session.query(self.db_model).filter_by(id=record.id).first()
-            if not db_record:
-                raise ValueError(f"Record with id {record.id} does not exist.")
-
-            # Update the record with new values from the provided Record object
-            for attr, value in vars(record).items():
-                setattr(db_record, attr, value)
-
-            # Commit the changes to the database
-            session.commit()
-
-    def str_to_datetime(self, str_date: str) -> datetime:
-        val = str_date.split("-")
-        _datetime = datetime(int(val[0]), int(val[1]), int(val[2]))
-        return _datetime
-
-    def query_date(self, start_date, end_date, limit=None, offset=0):
-        filters = self.get_filters({})
-        _start_date = self.str_to_datetime(start_date) if isinstance(start_date, str) else start_date
-        _end_date = self.str_to_datetime(end_date) if isinstance(end_date, str) else end_date
-        with self.session_maker() as session:
-            query = (
-                session.query(self.db_model)
-                .filter(*filters)
-                .filter(self.db_model.created_at >= _start_date)
-                .filter(self.db_model.created_at <= _end_date)
-                .filter(self.db_model.role != "system")
-                .filter(self.db_model.role != "tool")
-                .offset(offset)
-            )
-            if limit:
-                query = query.limit(limit)
-            results = query.all()
-        return [result.to_record() for result in results]
-
-
-class SQLLiteStorageConnector(SQLStorageConnector):
-    def __init__(self, table_type: str, config: LettaConfig, user_id, agent_id=None):
-        super().__init__(table_type=table_type, config=config, user_id=user_id, agent_id=agent_id)
-
-        # get storage URI
-        if table_type == TableType.ARCHIVAL_MEMORY or table_type == TableType.PASSAGES:
-            self.db_model = PassageModel
-            if settings.letta_pg_uri_no_default:
-                self.uri = settings.letta_pg_uri_no_default
-            else:
-                # For SQLite, use the archival storage path
-                self.path = config.archival_storage_path
-                self.uri = f"sqlite:///{os.path.join(config.archival_storage_path, 'letta.db')}"
-        elif table_type == TableType.FILES:
-            self.path = self.config.metadata_storage_path
-            if self.path is None:
-                raise ValueError(f"Must specify metadata_storage_path in config.")
-            self.db_model = FileMetadataModel
-
-        else:
-            raise ValueError(f"Table type {table_type} not implemented")
-
-        self.path = os.path.join(self.path, f"sqlite.db")
-
-        from letta.server.server import db_context
-
-        self.session_maker = db_context
-
-        # Need this in order to allow UUIDs to be stored successfully in the sqlite database
-        # import sqlite3
-        # import uuid
-        #
-        # sqlite3.register_adapter(uuid.UUID, lambda u: u.bytes_le)
-        # sqlite3.register_converter("UUID", lambda b: uuid.UUID(bytes_le=b))
-
-    def insert_many(self, records, exists_ok=True, show_progress=False):
-        # TODO: this is terrible, should eventually be done the same way for all types (migrate to SQLModel)
-        if len(records) == 0:
-            return
-
-        added_ids = []  # avoid adding duplicates
-        # NOTE: this has not great performance due to the excessive commits
-        with self.session_maker() as session:
-            iterable = tqdm(records) if show_progress else records
-            for record in iterable:
-                # db_record = self.db_model(**vars(record))
-
-                if record.id in added_ids:
-                    continue
-
-                existing_record = session.query(self.db_model).filter_by(id=record.id).first()
-                if existing_record:
-                    if exists_ok:
-                        fields = record.model_dump()
-                        fields.pop("id")
-                        session.query(self.db_model).filter(self.db_model.id == record.id).update(fields)
-                        session.commit()
-                    else:
-                        raise ValueError(f"Record with id {record.id} already exists.")
-
-                else:
-                    db_record = self.db_model(**record.dict())
-                    session.add(db_record)
-                    session.commit()
-
-                added_ids.append(record.id)
-
-    def insert(self, record, exists_ok=True):
-        self.insert_many([record], exists_ok=exists_ok)
-
-    def update(self, record):
-        """
-        Updates an existing record in the database with values from the provided record object.
-        """
-        if not record.id:
-            raise ValueError("Record must have an id.")
-
-        with self.session_maker() as session:
-            # Fetch the existing record from the database
-            db_record = session.query(self.db_model).filter_by(id=record.id).first()
-            if not db_record:
-                raise ValueError(f"Record with id {record.id} does not exist.")
-
-            # Update the database record with values from the provided record object
-            for column in self.db_model.__table__.columns:
-                column_name = column.name
-                if hasattr(record, column_name):
-                    new_value = getattr(record, column_name)
-                    setattr(db_record, column_name, new_value)
-
-            # Commit the changes to the database
-            session.commit()
-
-
-def attach_base():
-    # This should be invoked in server.py to make sure Base gets initialized properly
-    # DO NOT REMOVE
-    from letta.utils import printd
-
-    printd("Initializing database...")

letta/agent_store/milvus.py

@@ -1,198 +0,0 @@
-import uuid
-from copy import deepcopy
-from typing import Dict, Iterator, List, Optional, cast
-
-from pymilvus import DataType, MilvusClient
-from pymilvus.client.constants import ConsistencyLevel
-
-from letta.agent_store.storage import StorageConnector, TableType
-from letta.config import LettaConfig
-from letta.constants import MAX_EMBEDDING_DIM
-from letta.data_types import Passage, Record, RecordType
-from letta.utils import datetime_to_timestamp, printd, timestamp_to_datetime
-
-
-class MilvusStorageConnector(StorageConnector):
-    """Storage via Milvus"""
-
-    def __init__(self, table_type: str, config: LettaConfig, user_id, agent_id=None):
-        super().__init__(table_type=table_type, config=config, user_id=user_id, agent_id=agent_id)
-
-        assert table_type in [TableType.ARCHIVAL_MEMORY, TableType.PASSAGES], "Milvus only supports archival memory"
-        if config.archival_storage_uri:
-            self.client = MilvusClient(uri=config.archival_storage_uri)
-            self._create_collection()
-        else:
-            raise ValueError("Please set `archival_storage_uri` in the config file when using Milvus.")
-
-        # need to be converted to strings
-        self.uuid_fields = ["id", "user_id", "agent_id", "source_id", "file_id"]
-
-    def _create_collection(self):
-        schema = MilvusClient.create_schema(
-            auto_id=False,
-            enable_dynamic_field=True,
-        )
-        schema.add_field(field_name="id", datatype=DataType.VARCHAR, is_primary=True, max_length=65_535)
-        schema.add_field(field_name="text", datatype=DataType.VARCHAR, is_primary=False, max_length=65_535)
-        schema.add_field(field_name="embedding", datatype=DataType.FLOAT_VECTOR, dim=MAX_EMBEDDING_DIM)
-        index_params = self.client.prepare_index_params()
-        index_params.add_index(field_name="id")
-        index_params.add_index(field_name="embedding", index_type="AUTOINDEX", metric_type="IP")
-        self.client.create_collection(
-            collection_name=self.table_name, schema=schema, index_params=index_params, consistency_level=ConsistencyLevel.Strong
-        )
-
-    def get_milvus_filter(self, filters: Optional[Dict] = {}) -> str:
-        filter_conditions = {**self.filters, **filters} if filters is not None else self.filters
-        if not filter_conditions:
-            return ""
-        conditions = []
-        for key, value in filter_conditions.items():
-            if key in self.uuid_fields or isinstance(key, str):
-                condition = f'({key} == "{value}")'
-            else:
-                condition = f"({key} == {value})"
-            conditions.append(condition)
-        filter_expr = " and ".join(conditions)
-        if len(conditions) == 1:
-            filter_expr = filter_expr[1:-1]
-        return filter_expr
-
-    def get_all_paginated(self, filters: Optional[Dict] = {}, page_size: int = 1000) -> Iterator[List[RecordType]]:
-        if not self.client.has_collection(collection_name=self.table_name):
-            yield []
-        filter_expr = self.get_milvus_filter(filters)
-        offset = 0
-        while True:
-            # Retrieve a chunk of records with the given page_size
-            query_res = self.client.query(
-                collection_name=self.table_name,
-                filter=filter_expr,
-                offset=offset,
-                limit=page_size,
-            )
-            if not query_res:
-                break
-            # Yield a list of Record objects converted from the chunk
-            yield self._list_to_records(query_res)
-
-            # Increment the offset to get the next chunk in the next iteration
-            offset += page_size
-
-    def get_all(self, filters: Optional[Dict] = {}, limit=None) -> List[RecordType]:
-        if not self.client.has_collection(collection_name=self.table_name):
-            return []
-        filter_expr = self.get_milvus_filter(filters)
-        query_res = self.client.query(
-            collection_name=self.table_name,
-            filter=filter_expr,
-            limit=limit,
-        )
-        return self._list_to_records(query_res)
-
-    def get(self, id: str) -> Optional[RecordType]:
-        res = self.client.get(collection_name=self.table_name, ids=str(id))
-        return self._list_to_records(res)[0] if res else None
-
-    def size(self, filters: Optional[Dict] = {}) -> int:
-        if not self.client.has_collection(collection_name=self.table_name):
-            return 0
-        filter_expr = self.get_milvus_filter(filters)
-        count_expr = "count(*)"
-        query_res = self.client.query(
-            collection_name=self.table_name,
-            filter=filter_expr,
-            output_fields=[count_expr],
-        )
-        doc_num = query_res[0][count_expr]
-        return doc_num
-
-    def insert(self, record: RecordType):
-        self.insert_many([record])
-
-    def insert_many(self, records: List[RecordType], show_progress=False):
-        if not records:
-            return
-
-        # Milvus lite currently does not support upsert, so we delete and insert instead
-        # self.client.upsert(collection_name=self.table_name, data=self._records_to_list(records))
-        ids = [str(record.id) for record in records]
-        self.client.delete(collection_name=self.table_name, ids=ids)
-        data = self._records_to_list(records)
-        self.client.insert(collection_name=self.table_name, data=data)
-
-    def query(self, query: str, query_vec: List[float], top_k: int = 10, filters: Optional[Dict] = {}) -> List[RecordType]:
-        if not self.client.has_collection(self.table_name):
-            return []
-        search_res = self.client.search(
-            collection_name=self.table_name, data=[query_vec], filter=self.get_milvus_filter(filters), limit=top_k, output_fields=["*"]
-        )[0]
-        entity_res = [res["entity"] for res in search_res]
-        return self._list_to_records(entity_res)
-
-    def delete_table(self):
-        self.client.drop_collection(collection_name=self.table_name)
-
-    def delete(self, filters: Optional[Dict] = {}):
-        if not self.client.has_collection(collection_name=self.table_name):
-            return
-        filter_expr = self.get_milvus_filter(filters)
-        self.client.delete(collection_name=self.table_name, filter=filter_expr)
-
-    def save(self):
-        # save to persistence file (nothing needs to be done)
-        printd("Saving milvus")
-
-    def _records_to_list(self, records: List[Record]) -> List[Dict]:
-        if records == []:
-            return []
-        assert all(isinstance(r, Passage) for r in records)
-        record_list = []
-        records = list(set(records))
-        for record in records:
-            record_vars = deepcopy(vars(record))
-            _id = record_vars.pop("id")
-            text = record_vars.pop("text", "")
-            embedding = record_vars.pop("embedding")
-            record_metadata = record_vars.pop("metadata_", None) or {}
-            if "created_at" in record_vars:
-                record_vars["created_at"] = datetime_to_timestamp(record_vars["created_at"])
-            record_dict = {key: value for key, value in record_vars.items() if value is not None}
-            record_dict = {
-                **record_dict,
-                **record_metadata,
-                "id": str(_id),
-                "text": text,
-                "embedding": embedding,
-            }
-            for key, value in record_dict.items():
-                if key in self.uuid_fields:
-                    record_dict[key] = str(value)
-            record_list.append(record_dict)
-        return record_list
-
-    def _list_to_records(self, query_res: List[Dict]) -> List[RecordType]:
-        records = []
-        for res_dict in query_res:
-            _id = res_dict.pop("id")
-            embedding = res_dict.pop("embedding")
-            text = res_dict.pop("text")
-            metadata = deepcopy(res_dict)
-            for key, value in metadata.items():
-                if key in self.uuid_fields:
-                    metadata[key] = uuid.UUID(value)
-                elif key == "created_at":
-                    metadata[key] = timestamp_to_datetime(value)
-            records.append(
-                cast(
-                    RecordType,
-                    self.type(
-                        text=text,
-                        embedding=embedding,
-                        id=uuid.UUID(_id),
-                        **metadata,
-                    ),
-                )
-            )
-        return records