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

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