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

letta/agent_store/storage.py

@@ -0,0 +1,188 @@
+""" 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.document import Document
+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
+    DOCUMENTS = "documents"  # TODO
+
+
+# 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)
+DOCUMENT_TABLE_NAME = "letta_documents"  # original documents (from source)
+
+
+class StorageConnector:
+    """Defines a DB connection that is user-specific to access data: Documents, Passages, Archival/Recall Memory"""
+
+    type: Type[BaseModel]
+
+    def __init__(
+        self,
+        table_type: Union[TableType.ARCHIVAL_MEMORY, TableType.RECALL_MEMORY, TableType.PASSAGES, TableType.DOCUMENTS],
+        config: LettaConfig,
+        user_id,
+        agent_id=None,
+    ):
+        self.user_id = user_id
+        self.agent_id = agent_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.DOCUMENTS:
+            self.type = Document
+            self.table_name == DOCUMENT_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.PASSAGES or self.table_type == TableType.DOCUMENTS:
+            # 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}
+        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.DOCUMENTS],
+        config: LettaConfig,
+        user_id,
+        agent_id=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
+        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 == "chroma":
+            from letta.agent_store.chroma import ChromaStorageConnector
+
+            return ChromaStorageConnector(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)
+        # TODO: add back
+        # elif storage_type == "lancedb":
+        #    from letta.agent_store.db import LanceDBConnector
+
+        #    return LanceDBConnector(agent_config=agent_config, table_type=table_type)
+
+        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

letta/benchmark/benchmark.py

@@ -0,0 +1,96 @@
+# type: ignore
+
+import time
+import uuid
+from typing import Annotated
+
+import typer
+
+from letta import create_client
+from letta.benchmark.constants import HUMAN, PERSONA, PROMPTS, TRIES
+from letta.config import LettaConfig
+
+# from letta.agent import Agent
+from letta.errors import LLMJSONParsingError
+from letta.utils import get_human_text, get_persona_text
+
+app = typer.Typer()
+
+
+def send_message(message: str, agent_id, turn: int, fn_type: str, print_msg: bool = False, n_tries: int = TRIES):
+    try:
+        print_msg = f"\t-> Now running {fn_type}. Progress: {turn}/{n_tries}"
+        print(print_msg, end="\r", flush=True)
+        response = client.user_message(agent_id=agent_id, message=message, return_token_count=True)
+
+        if turn + 1 == n_tries:
+            print("  " * len(print_msg), end="\r", flush=True)
+
+        for r in response:
+            if "function_call" in r and fn_type in r["function_call"] and any("assistant_message" in re for re in response):
+                return True, r["function_call"]
+
+        return False, "No function called."
+    except LLMJSONParsingError as e:
+        print(f"Error in parsing Letta JSON: {e}")
+        return False, "Failed to decode valid Letta JSON from LLM output."
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+        return False, "An unexpected error occurred."
+
+
+@app.command()
+def bench(
+    print_messages: Annotated[bool, typer.Option("--messages", help="Print functions calls and messages from the agent.")] = False,
+    n_tries: Annotated[int, typer.Option("--n-tries", help="Number of benchmark tries to perform for each function.")] = TRIES,
+):
+    client = create_client()
+    print(f"\nDepending on your hardware, this may take up to 30 minutes. This will also create {n_tries * len(PROMPTS)} new agents.\n")
+    config = LettaConfig.load()
+    print(f"version = {config.letta_version}")
+
+    total_score, total_tokens_accumulated, elapsed_time = 0, 0, 0
+
+    for fn_type, message in PROMPTS.items():
+        score = 0
+        start_time_run = time.time()
+        bench_id = uuid.uuid4()
+
+        for i in range(n_tries):
+            agent = client.create_agent(
+                name=f"benchmark_{bench_id}_agent_{i}",
+                persona=get_persona_text(PERSONA),
+                human=get_human_text(HUMAN),
+            )
+
+            agent_id = agent.id
+            result, msg = send_message(
+                message=message, agent_id=agent_id, turn=i, fn_type=fn_type, print_msg=print_messages, n_tries=n_tries
+            )
+
+            if print_messages:
+                print(f"\t{msg}")
+
+            if result:
+                score += 1
+
+            # TODO: add back once we start tracking usage via the client
+            # total_tokens_accumulated += tokens_accumulated
+
+        elapsed_time_run = round(time.time() - start_time_run, 2)
+        print(f"Score for {fn_type}: {score}/{n_tries}, took {elapsed_time_run} seconds")
+
+        elapsed_time += elapsed_time_run
+        total_score += score
+
+    print(f"\nMEMGPT VERSION: {config.letta_version}")
+    print(f"CONTEXT WINDOW: {config.default_llm_config.context_window}")
+    print(f"MODEL WRAPPER: {config.default_llm_config.model_wrapper}")
+    print(f"PRESET: {config.preset}")
+    print(f"PERSONA: {config.persona}")
+    print(f"HUMAN: {config.human}")
+
+    print(
+        # f"\n\t-> Total score: {total_score}/{len(PROMPTS) * n_tries}, took {elapsed_time} seconds at average of {round(total_tokens_accumulated/elapsed_time, 2)} t/s\n"
+        f"\n\t-> Total score: {total_score}/{len(PROMPTS) * n_tries}, took {elapsed_time} seconds\n"
+    )

letta/benchmark/constants.py

@@ -0,0 +1,14 @@
+# Basic
+TRIES = 3
+AGENT_NAME = "benchmark"
+PERSONA = "sam_pov"
+HUMAN = "cs_phd"
+
+# Prompts
+PROMPTS = {
+    "core_memory_replace": "Hey there, my name is John, what is yours?",
+    "core_memory_append": "I want you to remember that I like soccers for later.",
+    "conversation_search": "Do you remember when I talked about bananas?",
+    "archival_memory_insert": "Can you make sure to remember that I like programming for me so you can look it up later?",
+    "archival_memory_search": "Can you retrieve information about the war?",
+}