From 986a865e62671766d8e028978eb808e29c68d734 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Mon, 26 Aug 2024 13:10:27 -0700
Subject: [PATCH] Attachment / add TTL api

---
 llama_toolchain/memory/api/api.py               |  1 +
 .../memory/meta_reference/faiss/faiss.py        |  1 +
 .../safety/meta_reference/shields/base.py       | 17 +++--------------
 3 files changed, 5 insertions(+), 14 deletions(-)

diff --git a/llama_toolchain/memory/api/api.py b/llama_toolchain/memory/api/api.py
index 9b86d29a2..57370bdf3 100644
--- a/llama_toolchain/memory/api/api.py
+++ b/llama_toolchain/memory/api/api.py
@@ -123,6 +123,7 @@ class Memory(Protocol):
         self,
         bank_id: str,
         documents: List[MemoryBankDocument],
+        ttl_seconds: Optional[int] = None,
     ) -> None: ...
 
     @webmethod(route="/memory_bank/update")
diff --git a/llama_toolchain/memory/meta_reference/faiss/faiss.py b/llama_toolchain/memory/meta_reference/faiss/faiss.py
index e6d7b2b46..b3fdd2558 100644
--- a/llama_toolchain/memory/meta_reference/faiss/faiss.py
+++ b/llama_toolchain/memory/meta_reference/faiss/faiss.py
@@ -187,6 +187,7 @@ class FaissMemoryImpl(Memory):
         self,
         bank_id: str,
         documents: List[MemoryBankDocument],
+        ttl_seconds: Optional[int] = None,
     ) -> None:
         assert bank_id in self.states, f"Bank {bank_id} not found"
         state = self.states[bank_id]
diff --git a/llama_toolchain/safety/meta_reference/shields/base.py b/llama_toolchain/safety/meta_reference/shields/base.py
index 245373b13..0432b8d3b 100644
--- a/llama_toolchain/safety/meta_reference/shields/base.py
+++ b/llama_toolchain/safety/meta_reference/shields/base.py
@@ -5,9 +5,9 @@
 # the root directory of this source tree.
 
 from abc import ABC, abstractmethod
-from typing import List, Union
+from typing import List
 
-from llama_models.llama3.api.datatypes import Attachment, Message
+from llama_models.llama3.api.datatypes import interleaved_text_media_as_str, Message
 from llama_toolchain.safety.api.datatypes import *  # noqa: F403
 
 CANNED_RESPONSE_TEXT = "I can't answer that. Can I help with something else?"
@@ -30,18 +30,7 @@ class ShieldBase(ABC):
 
 
 def message_content_as_str(message: Message) -> str:
-    def _to_str(content: Union[str, Attachment]) -> str:
-        if isinstance(content, str):
-            return content
-        elif isinstance(content, Attachment):
-            return f"File: {str(content.url)}"
-        else:
-            raise
-
-    if isinstance(message.content, list) or isinstance(message.content, tuple):
-        return "\n".join([_to_str(c) for c in message.content])
-    else:
-        return _to_str(message.content)
+    return interleaved_text_media_as_str(message.content)
 
 
 # For shields that operate on simple strings