| 219 | return |
| 220 | |
| 221 | class SupermemoryClient: |
| 222 | def __init__(self): |
| 223 | from supermemory import Supermemory |
| 224 | |
| 225 | self.client = Supermemory(api_key=os.getenv("SUPERMEMORY_API_KEY")) |
| 226 | |
| 227 | def add(self, messages, user_id): |
| 228 | content = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages]) |
| 229 | max_retries = 5 |
| 230 | for attempt in range(max_retries): |
| 231 | try: |
| 232 | self.client.memories.add(content=content, container_tag=user_id) |
| 233 | break |
| 234 | except Exception as e: |
| 235 | if attempt < max_retries - 1: |
| 236 | time.sleep(2**attempt) |
| 237 | else: |
| 238 | raise e |
| 239 | |
| 240 | def search(self, query, user_id, top_k): |
| 241 | max_retries = 10 |
| 242 | for attempt in range(max_retries): |
| 243 | try: |
| 244 | results = self.client.search.memories( |
| 245 | q=query, |
| 246 | container_tag=user_id, |
| 247 | threshold=0, |
| 248 | rerank=True, |
| 249 | rewrite_query=True, |
| 250 | limit=top_k, |
| 251 | ) |
| 252 | context = "\n\n".join([r.memory for r in results.results]) |
| 253 | return context |
| 254 | except Exception as e: |
| 255 | if attempt < max_retries - 1: |
| 256 | time.sleep(2**attempt) |
| 257 | else: |
| 258 | raise e |
| 259 | |
| 260 | mem_client = SupermemoryClient() |
| 261 | |