letta-nightly 0.1.7.dev20240924104148__py3-none-any.whl

letta/agent_store/lancedb.py

@@ -0,0 +1,177 @@
+# type: ignore
+
+import uuid
+from datetime import datetime
+from typing import Dict, Iterator, List, Optional
+
+from lancedb.pydantic import LanceModel, Vector
+
+from letta.agent_store.storage import StorageConnector, TableType
+from letta.config import AgentConfig, LettaConfig
+from letta.schemas.message import Message, Passage, Record
+
+""" Initial implementation - not complete """
+
+
+def get_db_model(table_name: str, table_type: TableType):
+    config = LettaConfig.load()
+
+    if table_type == TableType.ARCHIVAL_MEMORY or table_type == TableType.PASSAGES:
+        # create schema for archival memory
+        class PassageModel(LanceModel):
+            """Defines data model for storing Passages (consisting of text, embedding)"""
+
+            id: uuid.UUID
+            user_id: str
+            text: str
+            doc_id: str
+            agent_id: str
+            data_source: str
+            embedding: Vector(config.default_embedding_config.embedding_dim)
+            metadata_: Dict
+
+            def __repr__(self):
+                return f"<Passage(passage_id='{self.id}', text='{self.text}', embedding='{self.embedding})>"
+
+            def to_record(self):
+                return Passage(
+                    text=self.text,
+                    embedding=self.embedding,
+                    doc_id=self.doc_id,
+                    user_id=self.user_id,
+                    id=self.id,
+                    data_source=self.data_source,
+                    agent_id=self.agent_id,
+                    metadata=self.metadata_,
+                )
+
+        return PassageModel
+    elif table_type == TableType.RECALL_MEMORY:
+
+        class MessageModel(LanceModel):
+            """Defines data model for storing Message objects"""
+
+            __abstract__ = True  # this line is necessary
+
+            # Assuming message_id is the primary key
+            id: uuid.UUID
+            user_id: str
+            agent_id: str
+
+            # openai info
+            role: str
+            name: str
+            text: str
+            model: str
+            user: str
+
+            # function info
+            function_name: str
+            function_args: str
+            function_response: str
+
+            embedding = Vector(config.default_embedding_config.embedding_dim)
+
+            # Add a datetime column, with default value as the current time
+            created_at = datetime
+
+            def __repr__(self):
+                return f"<Message(message_id='{self.id}', text='{self.text}', embedding='{self.embedding})>"
+
+            def to_record(self):
+                return Message(
+                    user_id=self.user_id,
+                    agent_id=self.agent_id,
+                    role=self.role,
+                    name=self.name,
+                    text=self.text,
+                    model=self.model,
+                    function_name=self.function_name,
+                    function_args=self.function_args,
+                    function_response=self.function_response,
+                    embedding=self.embedding,
+                    created_at=self.created_at,
+                    id=self.id,
+                )
+
+        """Create database model for table_name"""
+        return MessageModel
+
+    else:
+        raise ValueError(f"Table type {table_type} not implemented")
+
+
+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):
+        # TODO
+        pass
+
+    def generate_where_filter(self, filters: Dict) -> str:
+        where_filters = []
+        for key, value in filters.items():
+            where_filters.append(f"{key}={value}")
+        return where_filters.join(" AND ")
+
+    @abstractmethod
+    def get_all_paginated(self, filters: Optional[Dict] = {}, page_size: Optional[int] = 1000) -> Iterator[List[Record]]:
+        # TODO
+        pass
+
+    @abstractmethod
+    def get_all(self, filters: Optional[Dict] = {}, limit=10) -> List[Record]:
+        # TODO
+        pass
+
+    @abstractmethod
+    def get(self, id: uuid.UUID) -> Optional[Record]:
+        # TODO
+        pass
+
+    @abstractmethod
+    def size(self, filters: Optional[Dict] = {}) -> int:
+        # TODO
+        pass
+
+    @abstractmethod
+    def insert(self, record: Record):
+        # TODO
+        pass
+
+    @abstractmethod
+    def insert_many(self, records: List[Record], show_progress=False):
+        # TODO
+        pass
+
+    @abstractmethod
+    def query(self, query: str, query_vec: List[float], top_k: int = 10, filters: Optional[Dict] = {}) -> List[Record]:
+        # TODO
+        pass
+
+    @abstractmethod
+    def query_date(self, start_date, end_date):
+        # TODO
+        pass
+
+    @abstractmethod
+    def query_text(self, query):
+        # TODO
+        pass
+
+    @abstractmethod
+    def delete_table(self):
+        # TODO
+        pass
+
+    @abstractmethod
+    def delete(self, filters: Optional[Dict] = {}):
+        # TODO
+        pass
+
+    @abstractmethod
+    def save(self):
+        # TODO
+        pass

letta/agent_store/milvus.py

@@ -0,0 +1,198 @@
+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", "doc_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

@@ -0,0 +1,201 @@
+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", "doc_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)