Support encryption #3

Support encryption.

- History: content
- Archive: archive
- Entity: serialized_entities

chatmemory/chatmemory.py

@@ -1,9 +1,14 @@
+import os
 from datetime import datetime, date, time, timedelta, timezone
+import base64
 import json
+import hashlib
 from logging import getLogger, NullHandler
 from sqlalchemy import Column, Integer, String, DateTime, Date
 from sqlalchemy.orm import Session, declarative_base
 from openai import ChatCompletion
+from Crypto.Cipher import AES
+
 
 logger = getLogger(__name__)
 logger.addHandler(NullHandler())
@@ -140,32 +145,52 @@ def __init__(self, api_key: str=None, model: str="gpt-3.5-turbo-16k-0613", histo
     def date_to_utc_datetime(self, d) -> datetime:
         return datetime.combine(d, time()).replace(tzinfo=timezone.utc)
 
+    def encrypt(self, text: str, password: str=None):
+        if not password:
+            return text
+
+        salt = os.urandom(16)
+        key = hashlib.scrypt(password=password.encode("utf-8"), salt=salt, n=2**5, r=8, p=1, dklen=32)
+        cipher = AES.new(key, AES.MODE_GCM)
+        cipher_text, tag = cipher.encrypt_and_digest(text.encode("utf-8"))
+        return "-".join([base64.b64encode(item).decode("utf-8") for item in [salt, cipher.nonce, cipher_text, tag]])
+
+    def decrypt(self, encrypted_text: str, password: str=None):
+        if not password:
+            return encrypted_text
+
+        salt, cipher_nonce, cipher_text, tag = [base64.b64decode(item) for item in encrypted_text.split("-")]
+        key = hashlib.scrypt(password=password.encode("utf-8"), salt=salt, n=2**5, r=8, p=1, dklen=32)
+        cipher = AES.new(key, AES.MODE_GCM, cipher_nonce)
+        return cipher.decrypt_and_verify(cipher_text, tag).decode("utf-8")
+
     def create_database(self, engine):
         Base.metadata.create_all(bind=engine)
 
-    def add_histories(self, session: Session, user_id: str, messages: list):
+    def add_histories(self, session: Session, user_id: str, messages: list, password: str=None):
         histories = [
-            History(user_id=user_id, role=m["role"], content=m["content"])
+            History(user_id=user_id, role=m["role"], content=self.encrypt(m["content"], password))
             for m in messages if m["role"] == "user" or m["role"] == "assistant"
         ]
         session.bulk_save_objects(histories)
 
-    def get_histories(self, session: Session, user_id: str, since: datetime=None, until: datetime=None) -> list:
+    def get_histories(self, session: Session, user_id: str, since: datetime=None, until: datetime=None, password: str=None) -> list:
         histories = session.query(History).filter(
             History.user_id == user_id,
             History.timestamp >= (since or datetime.min),
             History.timestamp <= (until or datetime.max)
         ).order_by(History.id).limit(self.history_max_count).all()
 
-        return [{"role": h.role, "content": h.content} for h in histories]
+        return [{"role": h.role, "content": self.decrypt(h.content, password)} for h in histories]
 
-    def archive_histories(self, session: Session, user_id: str, target_date: date):
+    def archive_histories(self, session: Session, user_id: str, target_date: date, password: str=None):
         since_dt = self.date_to_utc_datetime(target_date)
         conversation_history = self.get_histories(
             session,
             user_id,
             since_dt,
-            since_dt + timedelta(days=1)
+            since_dt + timedelta(days=1),
+            password
         )
 
         if len(conversation_history) == 0:
@@ -191,24 +216,24 @@ def archive_histories(self, session: Session, user_id: str, target_date: date):
         summarized_archive = self.history_archiver.archive(conversation_history)
 
         stored_archive.timestamp = datetime.utcnow()
-        stored_archive.archive = summarized_archive
+        stored_archive.archive = self.encrypt(summarized_archive, password)
 
         session.merge(stored_archive)
 
-    def get_archives(self, session: Session, user_id: str, since: date=None, until: date=None) -> list:
+    def get_archives(self, session: Session, user_id: str, since: date=None, until: date=None, password: str=None) -> list:
         archives = session.query(Archive.archive_date, Archive.archive).filter(
             Archive.user_id == user_id,
             Archive.archive_date >= (since or date.min),
             Archive.archive_date <= (until or date.max)
         ).order_by(Archive.archive_date.desc()).limit(self.archive_retrive_count).all()
 
-        return [{ "date": a.archive_date, "archive": a.archive } for a in archives]
+        return [{ "date": a.archive_date, "archive": self.decrypt(a.archive, password) } for a in archives]
 
-    def parse_entities(self, session: Session, user_id: str, target_date: date):
+    def parse_entities(self, session: Session, user_id: str, target_date: date, password: str=None):
         # Get histories on target_date
         since_dt = self.date_to_utc_datetime(target_date)
         until_dt = since_dt + timedelta(days=1)
-        conversation_history = self.get_histories(session, user_id, since_dt, until_dt)
+        conversation_history = self.get_histories(session, user_id, since_dt, until_dt, password)
         if len(conversation_history) == 0:
             logger.info(f"No histories found on {target_date} for parsing entities")
             return
@@ -226,24 +251,24 @@ def parse_entities(self, session: Session, user_id: str, target_date: date):
         entities = self.entity_parser.parse(conversation_history)
 
         if stored_entites.serialized_entities:
-            entities_json = json.loads(stored_entites.serialized_entities)
+            entities_json = json.loads(self.decrypt(stored_entites.serialized_entities, password))
             for k, v in entities.items():
                 entities_json[k] = v
         else:
             entities_json = entities
         stored_entites.timestamp = datetime.utcnow()
-        stored_entites.serialized_entities = json.dumps(entities_json, ensure_ascii=False)
+        stored_entites.serialized_entities = self.encrypt(json.dumps(entities_json, ensure_ascii=False), password)
         stored_entites.last_target_date = target_date
 
         session.merge(stored_entites)
 
-    def get_entities(self, session: Session, user_id: str) -> dict:
+    def get_entities(self, session: Session, user_id: str, password: str=None) -> dict:
         entities = session.query(Entity).filter(
             Entity.user_id == user_id,
         ).first()
 
         if entities and entities.serialized_entities:
-            return json.loads(entities.serialized_entities)
+            return json.loads(self.decrypt(entities.serialized_entities, password))
         else:
             return {}
 

chatmemory/server.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timedelta
-from fastapi import FastAPI, Depends
+from fastapi import FastAPI, Depends, Header
 from logging import getLogger
 from pydantic import BaseModel
 import traceback
@@ -78,11 +78,12 @@ def setup_handlers(self):
         app = self.app
 
         @app.post("/histories/{user_id}", response_model=ApiResponse)
-        async def add_histories(user_id: str, request: HistoriesRequest, db: Session = Depends(self.get_db)):
+        async def add_histories(user_id: str, request: HistoriesRequest, encryption_key: str = Header(default=None), db: Session = Depends(self.get_db)):
             try:
                 self.chatmemory.add_histories(
                     db, user_id,
-                    [{"role": m.role, "content": m.content} for m in request.messages]
+                    [{"role": m.role, "content": m.content} for m in request.messages],
+                    encryption_key
                 )
                 db.commit()
                 return ApiResponse(message="Histories added successfully")
@@ -92,25 +93,27 @@ async def add_histories(user_id: str, request: HistoriesRequest, db: Session = D
                 return ApiResponse(message="Error")
 
         @app.get("/histories/{user_id}", response_model=HistoriesResponse)
-        async def get_histories(user_id: str, since: str=None, until: str=None, db: Session = Depends(self.get_db)):
+        async def get_histories(user_id: str, since: str=None, until: str=None, encryption_key: str = Header(default=None), db: Session = Depends(self.get_db)):
             histories = self.chatmemory.get_histories(
                 db, user_id,
                 datetime.strptime(since, "%Y-%m-%d") if since else None,
-                datetime.strptime(until, "%Y-%m-%d") if until else None
+                datetime.strptime(until, "%Y-%m-%d") if until else None,
+                encryption_key
             )
             return HistoriesResponse(messages=[
                 Message(role=h["role"], content=h["content"])
                 for h in histories
             ])
 
         @app.post("/archives/{user_id}", response_model=ApiResponse)
-        async def archive_histories(request: ArchivesRequest, db: Session = Depends(self.get_db)):
+        async def archive_histories(request: ArchivesRequest, encryption_key: str = Header(default=None), db: Session = Depends(self.get_db)):
             try:
                 for i in range(request.days):
                     self.chatmemory.archive_histories(
                         db, request.user_id,
                         (datetime.strptime(request.target_date, "%Y-%m-%d") if request.target_date
                          else datetime.utcnow()).date() - timedelta(days=request.days - i - 1),
+                        encryption_key
                     )
                     db.commit()
                 return ApiResponse(message="Histories archived successfully")
@@ -120,11 +123,12 @@ async def archive_histories(request: ArchivesRequest, db: Session = Depends(self
                 return ApiResponse(message="Error")
 
         @app.get("/archives/{user_id}", response_model=ArchivesResponse)
-        async def get_archives(user_id: str, since: str=None, until: str=None, db: Session = Depends(self.get_db)):
+        async def get_archives(user_id: str, since: str=None, until: str=None, encryption_key: str = Header(default=None), db: Session = Depends(self.get_db)):
             archives = self.chatmemory.get_archives(
                 db, user_id,
                 datetime.strptime(since, "%Y-%m-%d") if since else None,
-                datetime.strptime(until, "%Y-%m-%d") if until else None
+                datetime.strptime(until, "%Y-%m-%d") if until else None,
+                encryption_key
             )
             return ArchivesResponse(archives=[
                 Archive(date=a["date"].strftime("%Y-%m-%d"), archive=a["archive"])
@@ -133,13 +137,14 @@ async def get_archives(user_id: str, since: str=None, until: str=None, db: Sessi
 
 
         @app.post("/entities/{user_id}", response_model=ApiResponse)
-        async def parse_entities(request: EntitiesRequest, db: Session = Depends(self.get_db)):
+        async def parse_entities(request: EntitiesRequest, encryption_key: str = Header(default=None), db: Session = Depends(self.get_db)):
             try:
                 for i in range(request.days):
                     self.chatmemory.parse_entities(
                         db, request.user_id,
                         (datetime.strptime(request.target_date, "%Y-%m-%d") if request.target_date
                          else datetime.utcnow()).date() - timedelta(days=request.days - i - 1),
+                        encryption_key
                     )
                     db.commit()
                 return ApiResponse(message="Entities parsed and stored successfully")
@@ -149,8 +154,8 @@ async def parse_entities(request: EntitiesRequest, db: Session = Depends(self.ge
                 return ApiResponse(message="Error")        
 
         @app.get("/entities/{user_id}", response_model=EntitiesResponse)
-        async def get_entities(user_id: str, db: Session = Depends(self.get_db)):
-            entities = self.chatmemory.get_entities(db, user_id)
+        async def get_entities(user_id: str, encryption_key: str = Header(default=None), db: Session = Depends(self.get_db)):
+            entities = self.chatmemory.get_entities(db, user_id, encryption_key)
             return EntitiesResponse(entities=entities)
 
 
@@ -165,5 +170,6 @@ async def delete_all(user_id: str, db: Session = Depends(self.get_db)):
                 logger.error(f"Error at delete_all: {ex}\n{traceback.format_exc()}")
                 return ApiResponse(message="Error")        
 
+
     def start(self, host :str="127.0.0.1", port: int=8123):
         uvicorn.run(self.app, host=host, port=port)

requirements.txt

@@ -3,3 +3,4 @@ openai==0.27.8
 requests==2.31.0
 SQLAlchemy==2.0.20
 uvicorn==0.23.1
+pycryptodome==3.18.0