File size: 1,005 Bytes
6598f74
 
f4c443a
6598f74
f4c443a
 
 
 
 
6598f74
f4c443a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6598f74
f4c443a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from typing import Dict, List

from mem0 import MemoryClient

from config.settings import MEM0_API_KEY

# Initialize memory client
mem0_client = MemoryClient(api_key=MEM0_API_KEY)


class MemoryTools:
    def save_to_memory(self, user_id: str, query: str, response: str):
        """Save interaction to memory"""
        try:
            messages = [
                {"role": "user", "content": query},
                {"role": "assistant", "content": response},
            ]
            mem0_client.add(
                messages=messages,
                user_id=user_id,
                metadata={"type": "regulatory_query"},
            )
        except Exception as e:
            print(f"Memory save error: {e}")

    def search_memory(self, user_id: str, query: str) -> List[Dict]:
        """Search for similar past queries"""
        try:
            memories = mem0_client.search(query=query, user_id=user_id, limit=3)
            return memories
        except Exception:
            return []