Spaces:
Paused
Paused
File size: 4,130 Bytes
5ec8557 d9cf3d7 5ec8557 d9cf3d7 5ec8557 | 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 | import json
import logging
import threading
import time
from copy import deepcopy
from typing import Any, Callable, Dict
from mem0 import Memory
_state_lock = threading.RLock()
_current_config: Dict[str, Any] = {}
_memory_instance: Memory | None = None
_session_factory: Callable | None = None
def set_session_factory(factory: Callable) -> None:
global _session_factory
_session_factory = factory
def _load_overrides() -> Dict[str, Any]:
try:
if _session_factory is None:
return {}
from models import Settings
session = _session_factory()
try:
row = session.get(Settings, "config_overrides")
if row is None:
return {}
return json.loads(row.value)
finally:
session.close()
except Exception:
return {}
def _save_overrides(overrides: Dict[str, Any]) -> None:
try:
if _session_factory is None:
return
from models import Settings
from sqlalchemy.dialects.postgresql import insert
session = _session_factory()
try:
serialized = json.dumps(overrides)
stmt = (
insert(Settings)
.values(key="config_overrides", value=serialized)
.on_conflict_do_update(index_elements=[Settings.key], set_={"value": serialized})
)
session.execute(stmt)
session.commit()
finally:
session.close()
except Exception:
logging.warning("Failed to persist config overrides to database", exc_info=True)
def _merge_config(base: Dict[str, Any], updates: Dict[str, Any]) -> Dict[str, Any]:
merged = deepcopy(base)
for key, value in updates.items():
if isinstance(value, dict) and isinstance(merged.get(key), dict):
merged[key] = _merge_config(merged[key], value)
else:
merged[key] = value
return merged
def initialize_state(default_config: Dict[str, Any]) -> None:
global _current_config, _memory_instance
with _state_lock:
_current_config = deepcopy(default_config)
# ponytail: retry DB load — Postgres needs ~90s crash recovery on HF FUSE
overrides: Dict[str, Any] = {}
for attempt in range(1, 13): # max 60s (12 x 5s)
overrides = _load_overrides()
if overrides:
logging.info("Config overrides loaded from DB on attempt %d", attempt)
break
if _session_factory is not None:
# DB reachable but no overrides yet — stop retrying
try:
from models import Settings
session = _session_factory()
try:
session.execute(__import__("sqlalchemy").text("SELECT 1"))
logging.info("DB ready, no config overrides stored yet")
break
finally:
session.close()
except Exception:
pass
logging.warning("DB not ready yet (attempt %d/12), retrying in 5s...", attempt)
time.sleep(5)
if overrides:
_current_config = _merge_config(_current_config, overrides)
_memory_instance = Memory.from_config(_current_config)
def update_config(updates: Dict[str, Any]) -> Dict[str, Any]:
global _current_config, _memory_instance
with _state_lock:
next_config = _merge_config(_current_config, updates)
_current_config = next_config
_memory_instance = Memory.from_config(next_config)
overrides = _load_overrides()
overrides = _merge_config(overrides, updates)
_save_overrides(overrides)
return deepcopy(_current_config)
def get_current_config() -> Dict[str, Any]:
with _state_lock:
return deepcopy(_current_config)
def get_memory_instance() -> Memory:
with _state_lock:
if _memory_instance is None:
raise RuntimeError("Mem0 runtime has not been initialized.")
return _memory_instance |