letta-nightly 0.6.1.dev20241206104246__py3-none-any.whl → 0.6.1.dev20241208104134__py3-none-any.whl

letta/agent_store/lancedb.py

@@ -1,177 +0,0 @@
-# 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
-            file_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,
-                    file_id=self.file_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/persistence_manager.py

@@ -1,149 +0,0 @@
-from abc import ABC, abstractmethod
-from datetime import datetime
-from typing import List
-
-from letta.memory import BaseRecallMemory, EmbeddingArchivalMemory
-from letta.schemas.agent import AgentState
-from letta.schemas.memory import Memory
-from letta.schemas.message import Message
-from letta.utils import printd
-
-
-def parse_formatted_time(formatted_time: str):
-    # parse times returned by letta.utils.get_formatted_time()
-    try:
-        return datetime.strptime(formatted_time.strip(), "%Y-%m-%d %I:%M:%S %p %Z%z")
-    except:
-        return datetime.strptime(formatted_time.strip(), "%Y-%m-%d %I:%M:%S %p")
-
-
-class PersistenceManager(ABC):
-    @abstractmethod
-    def trim_messages(self, num):
-        pass
-
-    @abstractmethod
-    def prepend_to_messages(self, added_messages):
-        pass
-
-    @abstractmethod
-    def append_to_messages(self, added_messages):
-        pass
-
-    @abstractmethod
-    def swap_system_message(self, new_system_message):
-        pass
-
-    @abstractmethod
-    def update_memory(self, new_memory):
-        pass
-
-
-class LocalStateManager(PersistenceManager):
-    """In-memory state manager has nothing to manage, all agents are held in-memory"""
-
-    recall_memory_cls = BaseRecallMemory
-    archival_memory_cls = EmbeddingArchivalMemory
-
-    def __init__(self, agent_state: AgentState):
-        # Memory held in-state useful for debugging stateful versions
-        self.memory = agent_state.memory
-        # self.messages = []  # current in-context messages
-        # self.all_messages = [] # all messages seen in current session (needed if lazily synchronizing state with DB)
-        self.archival_memory = EmbeddingArchivalMemory(agent_state)
-        self.recall_memory = BaseRecallMemory(agent_state)
-        # self.agent_state = agent_state
-
-    def save(self):
-        """Ensure storage connectors save data"""
-        self.archival_memory.save()
-        self.recall_memory.save()
-
-    '''
-    def json_to_message(self, message_json) -> Message:
-        """Convert agent message JSON into Message object"""
-
-        # get message
-        if "message" in message_json:
-            message = message_json["message"]
-        else:
-            message = message_json
-
-        # get timestamp
-        if "timestamp" in message_json:
-            timestamp = parse_formatted_time(message_json["timestamp"])
-        else:
-            timestamp = get_local_time()
-
-        # TODO: change this when we fully migrate to tool calls API
-        if "function_call" in message:
-            tool_calls = [
-                ToolCall(
-                    id=message["tool_call_id"],
-                    tool_call_type="function",
-                    function={
-                        "name": message["function_call"]["name"],
-                        "arguments": message["function_call"]["arguments"],
-                    },
-                )
-            ]
-            printd(f"Saving tool calls {[vars(tc) for tc in tool_calls]}")
-        else:
-            tool_calls = None
-
-        # if message["role"] == "function":
-        # message["role"] = "tool"
-
-        return Message(
-            user_id=self.agent_state.user_id,
-            agent_id=self.agent_state.id,
-            role=message["role"],
-            text=message["content"],
-            name=message["name"] if "name" in message else None,
-            model=self.agent_state.llm_config.model,
-            created_at=timestamp,
-            tool_calls=tool_calls,
-            tool_call_id=message["tool_call_id"] if "tool_call_id" in message else None,
-            id=message["id"] if "id" in message else None,
-        )
-    '''
-
-    def trim_messages(self, num):
-        # printd(f"InMemoryStateManager.trim_messages")
-        # self.messages = [self.messages[0]] + self.messages[num:]
-        pass
-
-    def prepend_to_messages(self, added_messages: List[Message]):
-        # first tag with timestamps
-        # added_messages = [{"timestamp": get_local_time(), "message": msg} for msg in added_messages]
-
-        printd(f"{self.__class__.__name__}.prepend_to_message")
-        # self.messages = [self.messages[0]] + added_messages + self.messages[1:]
-
-        # add to recall memory
-        self.recall_memory.insert_many([m for m in added_messages])
-
-    def append_to_messages(self, added_messages: List[Message]):
-        # first tag with timestamps
-        # added_messages = [{"timestamp": get_local_time(), "message": msg} for msg in added_messages]
-
-        printd(f"{self.__class__.__name__}.append_to_messages")
-        # self.messages = self.messages + added_messages
-
-        # add to recall memory
-        self.recall_memory.insert_many([m for m in added_messages])
-
-    def swap_system_message(self, new_system_message: Message):
-        # first tag with timestamps
-        # new_system_message = {"timestamp": get_local_time(), "message": new_system_message}
-
-        printd(f"{self.__class__.__name__}.swap_system_message")
-        # self.messages[0] = new_system_message
-
-        # add to recall memory
-        self.recall_memory.insert(new_system_message)
-
-    def update_memory(self, new_memory: Memory):
-        printd(f"{self.__class__.__name__}.update_memory")
-        assert isinstance(new_memory, Memory), type(new_memory)
-        self.memory = new_memory