letta-nightly 0.6.3.dev20241211050151__py3-none-any.whl → 0.6.3.dev20241212015858__py3-none-any.whl

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

letta/agent_store/qdrant.py

@@ -1,201 +0,0 @@
-import os
-import uuid
-from copy import deepcopy
-from typing import Dict, Iterator, List, Optional, cast
-
-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, timestamp_to_datetime
-
-TEXT_PAYLOAD_KEY = "text_content"
-METADATA_PAYLOAD_KEY = "metadata"
-
-
-class QdrantStorageConnector(StorageConnector):
-    """Storage via Qdrant"""
-
-    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)
-        try:
-            from qdrant_client import QdrantClient, models
-        except ImportError as e:
-            raise ImportError("'qdrant-client' not installed. Run `pip install qdrant-client`.") from e
-        assert table_type in [TableType.ARCHIVAL_MEMORY, TableType.PASSAGES], "Qdrant only supports archival memory"
-        if config.archival_storage_uri and len(config.archival_storage_uri.split(":")) == 2:
-            host, port = config.archival_storage_uri.split(":")
-            self.qdrant_client = QdrantClient(host=host, port=port, api_key=os.getenv("QDRANT_API_KEY"))
-        elif config.archival_storage_path:
-            self.qdrant_client = QdrantClient(path=config.archival_storage_path)
-        else:
-            raise ValueError("Qdrant storage requires either a URI or a path to the storage configured")
-        if not self.qdrant_client.collection_exists(self.table_name):
-            self.qdrant_client.create_collection(
-                collection_name=self.table_name,
-                vectors_config=models.VectorParams(
-                    size=MAX_EMBEDDING_DIM,
-                    distance=models.Distance.COSINE,
-                ),
-            )
-        self.uuid_fields = ["id", "user_id", "agent_id", "source_id", "file_id"]
-
-    def get_all_paginated(self, filters: Optional[Dict] = {}, page_size: int = 10) -> Iterator[List[RecordType]]:
-        from qdrant_client import grpc
-
-        filters = self.get_qdrant_filters(filters)
-        next_offset = None
-        stop_scrolling = False
-        while not stop_scrolling:
-            results, next_offset = self.qdrant_client.scroll(
-                collection_name=self.table_name,
-                scroll_filter=filters,
-                limit=page_size,
-                offset=next_offset,
-                with_payload=True,
-                with_vectors=True,
-            )
-            stop_scrolling = next_offset is None or (
-                isinstance(next_offset, grpc.PointId) and next_offset.num == 0 and next_offset.uuid == ""
-            )
-            yield self.to_records(results)
-
-    def get_all(self, filters: Optional[Dict] = {}, limit=10) -> List[RecordType]:
-        if self.size(filters) == 0:
-            return []
-        filters = self.get_qdrant_filters(filters)
-        results, _ = self.qdrant_client.scroll(
-            self.table_name,
-            scroll_filter=filters,
-            limit=limit,
-            with_payload=True,
-            with_vectors=True,
-        )
-        return self.to_records(results)
-
-    def get(self, id: str) -> Optional[RecordType]:
-        results = self.qdrant_client.retrieve(
-            collection_name=self.table_name,
-            ids=[str(id)],
-            with_payload=True,
-            with_vectors=True,
-        )
-        if not results:
-            return None
-        return self.to_records(results)[0]
-
-    def insert(self, record: Record):
-        points = self.to_points([record])
-        self.qdrant_client.upsert(self.table_name, points=points)
-
-    def insert_many(self, records: List[RecordType], show_progress=False):
-        points = self.to_points(records)
-        self.qdrant_client.upsert(self.table_name, points=points)
-
-    def delete(self, filters: Optional[Dict] = {}):
-        filters = self.get_qdrant_filters(filters)
-        self.qdrant_client.delete(self.table_name, points_selector=filters)
-
-    def delete_table(self):
-        self.qdrant_client.delete_collection(self.table_name)
-        self.qdrant_client.close()
-
-    def size(self, filters: Optional[Dict] = {}) -> int:
-        filters = self.get_qdrant_filters(filters)
-        return self.qdrant_client.count(collection_name=self.table_name, count_filter=filters).count
-
-    def close(self):
-        self.qdrant_client.close()
-
-    def query(
-        self,
-        query: str,
-        query_vec: List[float],
-        top_k: int = 10,
-        filters: Optional[Dict] = {},
-    ) -> List[RecordType]:
-        filters = self.get_filters(filters)
-        results = self.qdrant_client.search(
-            self.table_name,
-            query_vector=query_vec,
-            query_filter=filters,
-            limit=top_k,
-            with_payload=True,
-            with_vectors=True,
-        )
-        return self.to_records(results)
-
-    def to_records(self, records: list) -> List[RecordType]:
-        parsed_records = []
-        for record in records:
-            record = deepcopy(record)
-            metadata = record.payload[METADATA_PAYLOAD_KEY]
-            text = record.payload[TEXT_PAYLOAD_KEY]
-            _id = metadata.pop("id")
-            embedding = record.vector
-            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)
-            parsed_records.append(
-                cast(
-                    RecordType,
-                    self.type(
-                        text=text,
-                        embedding=embedding,
-                        id=uuid.UUID(_id),
-                        **metadata,
-                    ),
-                )
-            )
-        return parsed_records
-
-    def to_points(self, records: List[RecordType]):
-        from qdrant_client import models
-
-        assert all(isinstance(r, Passage) for r in records)
-        points = []
-        records = list(set(records))
-        for record in records:
-            record = vars(record)
-            _id = record.pop("id")
-            text = record.pop("text", "")
-            embedding = record.pop("embedding", {})
-            record_metadata = record.pop("metadata_", None) or {}
-            if "created_at" in record:
-                record["created_at"] = datetime_to_timestamp(record["created_at"])
-            metadata = {key: value for key, value in record.items() if value is not None}
-            metadata = {
-                **metadata,
-                **record_metadata,
-                "id": str(_id),
-            }
-            for key, value in metadata.items():
-                if key in self.uuid_fields:
-                    metadata[key] = str(value)
-            points.append(
-                models.PointStruct(
-                    id=str(_id),
-                    vector=embedding,
-                    payload={
-                        TEXT_PAYLOAD_KEY: text,
-                        METADATA_PAYLOAD_KEY: metadata,
-                    },
-                )
-            )
-        return points
-
-    def get_qdrant_filters(self, filters: Optional[Dict] = {}):
-        from qdrant_client import models
-
-        filter_conditions = {**self.filters, **filters} if filters is not None else self.filters
-        must_conditions = []
-        for key, value in filter_conditions.items():
-            match_value = str(value) if key in self.uuid_fields else value
-            field_condition = models.FieldCondition(
-                key=f"{METADATA_PAYLOAD_KEY}.{key}",
-                match=models.MatchValue(value=match_value),
-            )
-            must_conditions.append(field_condition)
-        return models.Filter(must=must_conditions)

letta/agent_store/storage.py

@@ -1,186 +0,0 @@
-""" These classes define storage connectors.
-
-We originally tried to use Llama Index VectorIndex, but their limited API was extremely problematic.
-"""
-
-import uuid
-from abc import abstractmethod
-from typing import Dict, List, Optional, Tuple, Type, Union
-
-from pydantic import BaseModel
-
-from letta.config import LettaConfig
-from letta.schemas.file import FileMetadata
-from letta.schemas.message import Message
-from letta.schemas.passage import Passage
-from letta.utils import printd
-
-
-# ENUM representing table types in Letta
-# each table corresponds to a different table schema  (specified in data_types.py)
-class TableType:
-    ARCHIVAL_MEMORY = "archival_memory"  # recall memory table: letta_agent_{agent_id}
-    RECALL_MEMORY = "recall_memory"  # archival memory table: letta_agent_recall_{agent_id}
-    PASSAGES = "passages"  # TODO
-    FILES = "files"
-
-
-# table names used by Letta
-
-# agent tables
-RECALL_TABLE_NAME = "letta_recall_memory_agent"  # agent memory
-ARCHIVAL_TABLE_NAME = "letta_archival_memory_agent"  # agent memory
-
-# external data source tables
-PASSAGE_TABLE_NAME = "letta_passages"  # chunked/embedded passages (from source)
-FILE_TABLE_NAME = "letta_files"  # original files (from source)
-
-
-class StorageConnector:
-    """Defines a DB connection that is user-specific to access data: files, Passages, Archival/Recall Memory"""
-
-    type: Type[BaseModel]
-
-    def __init__(
-        self,
-        table_type: Union[TableType.ARCHIVAL_MEMORY, TableType.RECALL_MEMORY, TableType.PASSAGES, TableType.FILES],
-        config: LettaConfig,
-        user_id: str,
-        agent_id: Optional[str] = None,
-        organization_id: Optional[str] = None,
-    ):
-        self.user_id = user_id
-        self.agent_id = agent_id
-        self.organization_id = organization_id
-        self.table_type = table_type
-
-        # get object type
-        if table_type == TableType.ARCHIVAL_MEMORY:
-            self.type = Passage
-            self.table_name = ARCHIVAL_TABLE_NAME
-        elif table_type == TableType.RECALL_MEMORY:
-            self.type = Message
-            self.table_name = RECALL_TABLE_NAME
-        elif table_type == TableType.FILES:
-            self.type = FileMetadata
-            self.table_name = FILE_TABLE_NAME
-        elif table_type == TableType.PASSAGES:
-            self.type = Passage
-            self.table_name = PASSAGE_TABLE_NAME
-        else:
-            raise ValueError(f"Table type {table_type} not implemented")
-        printd(f"Using table name {self.table_name}")
-
-        # setup base filters for agent-specific tables
-        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}
-        elif self.table_type == TableType.FILES:
-            # setup base filters for user-specific tables
-            assert agent_id is None, "Agent ID must not be provided for user-specific tables"
-            self.filters = {"user_id": self.user_id}
-        elif self.table_type == TableType.PASSAGES:
-            self.filters = {"organization_id": self.organization_id}
-        else:
-            raise ValueError(f"Table type {table_type} not implemented")
-
-    @staticmethod
-    def get_storage_connector(
-        table_type: Union[TableType.ARCHIVAL_MEMORY, TableType.RECALL_MEMORY, TableType.PASSAGES, TableType.FILES],
-        config: LettaConfig,
-        user_id: str,
-        organization_id: Optional[str] = None,
-        agent_id: Optional[str] = None,
-    ):
-        if table_type == TableType.ARCHIVAL_MEMORY or table_type == TableType.PASSAGES:
-            storage_type = config.archival_storage_type
-        elif table_type == TableType.RECALL_MEMORY:
-            storage_type = config.recall_storage_type
-        elif table_type == TableType.FILES:
-            storage_type = config.metadata_storage_type
-        else:
-            raise ValueError(f"Table type {table_type} not implemented")
-
-        if storage_type == "postgres":
-            from letta.agent_store.db import PostgresStorageConnector
-
-            return PostgresStorageConnector(table_type, config, user_id, agent_id)
-
-        elif storage_type == "qdrant":
-            from letta.agent_store.qdrant import QdrantStorageConnector
-
-            return QdrantStorageConnector(table_type, config, user_id, agent_id)
-
-        elif storage_type == "sqlite":
-            from letta.agent_store.db import SQLLiteStorageConnector
-
-            return SQLLiteStorageConnector(table_type, config, user_id, agent_id)
-        elif storage_type == "milvus":
-            from letta.agent_store.milvus import MilvusStorageConnector
-
-            return MilvusStorageConnector(table_type, config, user_id, agent_id)
-        else:
-            raise NotImplementedError(f"Storage type {storage_type} not implemented")
-
-    @staticmethod
-    def get_archival_storage_connector(user_id, agent_id):
-        config = LettaConfig.load()
-        return StorageConnector.get_storage_connector(TableType.ARCHIVAL_MEMORY, config, user_id, agent_id)
-
-    @staticmethod
-    def get_recall_storage_connector(user_id, agent_id):
-        config = LettaConfig.load()
-        return StorageConnector.get_storage_connector(TableType.RECALL_MEMORY, config, user_id, agent_id)
-
-    @abstractmethod
-    def get_filters(self, filters: Optional[Dict] = {}) -> Union[Tuple[list, dict], dict]:
-        pass
-
-    @abstractmethod
-    def get_all_paginated(self, filters: Optional[Dict] = {}, page_size: int = 1000):
-        pass
-
-    @abstractmethod
-    def get_all(self, filters: Optional[Dict] = {}, limit=10):
-        pass
-
-    @abstractmethod
-    def get(self, id: uuid.UUID):
-        pass
-
-    @abstractmethod
-    def size(self, filters: Optional[Dict] = {}) -> int:
-        pass
-
-    @abstractmethod
-    def insert(self, record):
-        pass
-
-    @abstractmethod
-    def insert_many(self, records, show_progress=False):
-        pass
-
-    @abstractmethod
-    def query(self, query: str, query_vec: List[float], top_k: int = 10, filters: Optional[Dict] = {}):
-        pass
-
-    @abstractmethod
-    def query_date(self, start_date, end_date):
-        pass
-
-    @abstractmethod
-    def query_text(self, query):
-        pass
-
-    @abstractmethod
-    def delete_table(self):
-        pass
-
-    @abstractmethod
-    def delete(self, filters: Optional[Dict] = {}):
-        pass
-
-    @abstractmethod
-    def save(self):
-        pass