Spaces:
Running
Running
File size: 3,731 Bytes
69f7644 |
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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
from pydantic import BaseModel, ConfigDict
from typing import List, Union, Optional
from sqlalchemy import Column, String, BigInteger, Text
from apps.webui.internal.db import Base, get_db
import time
import uuid
####################
# Memory DB Schema
####################
class Memory(Base):
__tablename__ = "memory"
id = Column(String, primary_key=True)
user_id = Column(String)
content = Column(Text)
updated_at = Column(BigInteger)
created_at = Column(BigInteger)
class MemoryModel(BaseModel):
id: str
user_id: str
content: str
updated_at: int # timestamp in epoch
created_at: int # timestamp in epoch
model_config = ConfigDict(from_attributes=True)
####################
# Forms
####################
class MemoriesTable:
def insert_new_memory(
self,
user_id: str,
content: str,
) -> Optional[MemoryModel]:
with get_db() as db:
id = str(uuid.uuid4())
memory = MemoryModel(
**{
"id": id,
"user_id": user_id,
"content": content,
"created_at": int(time.time()),
"updated_at": int(time.time()),
}
)
result = Memory(**memory.model_dump())
db.add(result)
db.commit()
db.refresh(result)
if result:
return MemoryModel.model_validate(result)
else:
return None
def update_memory_by_id(
self,
id: str,
content: str,
) -> Optional[MemoryModel]:
with get_db() as db:
try:
db.query(Memory).filter_by(id=id).update(
{"content": content, "updated_at": int(time.time())}
)
db.commit()
return self.get_memory_by_id(id)
except:
return None
def get_memories(self) -> List[MemoryModel]:
with get_db() as db:
try:
memories = db.query(Memory).all()
return [MemoryModel.model_validate(memory) for memory in memories]
except:
return None
def get_memories_by_user_id(self, user_id: str) -> List[MemoryModel]:
with get_db() as db:
try:
memories = db.query(Memory).filter_by(user_id=user_id).all()
return [MemoryModel.model_validate(memory) for memory in memories]
except:
return None
def get_memory_by_id(self, id: str) -> Optional[MemoryModel]:
with get_db() as db:
try:
memory = db.get(Memory, id)
return MemoryModel.model_validate(memory)
except:
return None
def delete_memory_by_id(self, id: str) -> bool:
with get_db() as db:
try:
db.query(Memory).filter_by(id=id).delete()
db.commit()
return True
except:
return False
def delete_memories_by_user_id(self, user_id: str) -> bool:
with get_db() as db:
try:
db.query(Memory).filter_by(user_id=user_id).delete()
db.commit()
return True
except:
return False
def delete_memory_by_id_and_user_id(self, id: str, user_id: str) -> bool:
with get_db() as db:
try:
db.query(Memory).filter_by(id=id, user_id=user_id).delete()
db.commit()
return True
except:
return False
Memories = MemoriesTable()
|