Hamdy005 commited on
Commit
93ff956
Β·
1 Parent(s): 28227c9

feat: implement Redis-first refresh token storage with Supabase fallback

Browse files
Files changed (9) hide show
  1. auth/constants.py +6 -0
  2. auth/refresh_token_store.py +142 -22
  3. config.py +1 -0
  4. main.py +16 -15
  5. rag/constants.py +8 -0
  6. rag/routes.py +8 -8
  7. redis_client.py +59 -0
  8. requirements.txt +1 -0
  9. store.py +143 -19
auth/constants.py CHANGED
@@ -37,3 +37,9 @@ REFRESH_COOKIE_NAME = "refresh_token"
37
  # Scope the cookie to the auth sub-path so it is NOT sent to /api/materials etc.
38
  REFRESH_COOKIE_PATH = "/api/auth"
39
  REFRESH_COOKIE_MAX_AGE = REFRESH_TOKEN_EXPIRE_DAYS * 24 * 3600 # seconds
 
 
 
 
 
 
 
37
  # Scope the cookie to the auth sub-path so it is NOT sent to /api/materials etc.
38
  REFRESH_COOKIE_PATH = "/api/auth"
39
  REFRESH_COOKIE_MAX_AGE = REFRESH_TOKEN_EXPIRE_DAYS * 24 * 3600 # seconds
40
+
41
+ # ── Redis Auth Keys & TTLs ─────────────────────────────────────────────────────
42
+ REFRESH_TOKEN_KEY_PREFIX = "rt:"
43
+ USER_REFRESH_TOKENS_KEY_PREFIX = "user_rts:"
44
+ REFRESH_TOKEN_REDIS_TTL = REFRESH_TOKEN_EXPIRE_DAYS * 86400 # seconds
45
+
auth/refresh_token_store.py CHANGED
@@ -1,9 +1,24 @@
1
  """
2
- Stateful refresh-token store.
3
 
4
- Wraps the `refresh_tokens` Supabase table with simple CRUD helpers.
5
- All functions are synchronous and use the existing `_table_supabase`
6
- helper from `store.py`, so they get the same dev-mode fallback behaviour.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7
 
8
  Table DDL (run once in the Supabase SQL editor):
9
 
@@ -24,38 +39,98 @@ from datetime import datetime, timedelta, timezone
24
  from typing import Optional
25
 
26
  from src.store import _table_supabase, _robust_execute
27
- from src.auth.constants import REFRESH_TOKEN_EXPIRE_DAYS
 
 
 
 
 
 
28
 
29
  logger = logging.getLogger(__name__)
30
 
31
 
 
 
 
 
 
 
 
32
  def _now_utc() -> datetime:
33
  return datetime.now(timezone.utc)
34
 
35
 
 
 
36
  def save_refresh_token(user_id: str, token_hash: str) -> None:
37
- """Persist a new (un-revoked) refresh token row for *user_id*."""
 
 
 
 
 
 
38
  expires_at = _now_utc() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  try:
40
  _robust_execute(
41
  _table_supabase("refresh_tokens").insert({
42
- "user_id": user_id,
43
  "token_hash": token_hash,
44
- "expires_at": expires_at.isoformat(),
45
- "revoked": False,
46
  })
47
  )
48
  except Exception as e:
49
- logger.error("save_refresh_token failed: %s", e)
50
  raise
51
 
52
 
 
 
53
  def get_refresh_token(token_hash: str) -> Optional[dict]:
54
  """
55
- Look up a refresh token row by its hash.
56
 
57
- Returns the raw DB row dict, or None if not found.
 
58
  """
 
 
 
 
 
 
 
 
 
 
 
 
 
59
  try:
60
  result = _robust_execute(
61
  _table_supabase("refresh_tokens")
@@ -64,15 +139,43 @@ def get_refresh_token(token_hash: str) -> Optional[dict]:
64
  )
65
  rows = result.data
66
  if isinstance(rows, list):
67
- return rows[0] if rows else None
68
- return rows or None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
69
  except Exception as e:
70
- logger.error("get_refresh_token failed: %s", e)
71
  return None
72
 
73
 
 
 
74
  def revoke_refresh_token(token_hash: str) -> None:
75
- """Mark a single token as revoked."""
 
 
 
 
 
 
 
 
76
  try:
77
  _robust_execute(
78
  _table_supabase("refresh_tokens")
@@ -80,11 +183,27 @@ def revoke_refresh_token(token_hash: str) -> None:
80
  .eq("token_hash", token_hash)
81
  )
82
  except Exception as e:
83
- logger.error("revoke_refresh_token failed: %s", e)
 
84
 
 
85
 
86
  def revoke_all_user_tokens(user_id: str) -> None:
87
  """Revoke every active refresh token for *user_id* (logout-everywhere)."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
88
  try:
89
  _robust_execute(
90
  _table_supabase("refresh_tokens")
@@ -92,26 +211,27 @@ def revoke_all_user_tokens(user_id: str) -> None:
92
  .eq("user_id", user_id)
93
  )
94
  except Exception as e:
95
- logger.error("revoke_all_user_tokens failed for %s: %s", user_id, e)
 
96
 
 
97
 
98
  def is_token_valid(row: dict) -> bool:
99
  """
100
- Return True if the DB row represents a currently-valid refresh token.
101
 
102
  Checks: row exists, not revoked, not expired.
 
103
  """
104
  if not row:
105
  return False
106
- if row.get("revoked"):
107
  return False
108
  expires_at_raw = row.get("expires_at")
109
  if not expires_at_raw:
110
  return False
111
- # Supabase returns ISO-8601 strings; parse and compare.
112
  try:
113
  if isinstance(expires_at_raw, str):
114
- # Handle both 'Z' suffix and '+00:00' offset
115
  expires_at_raw = expires_at_raw.replace("Z", "+00:00")
116
  expires_at = datetime.fromisoformat(expires_at_raw)
117
  else:
 
1
  """
2
+ Stateful refresh-token store β€” Redis-first, Supabase fallback.
3
 
4
+ Strategy
5
+ --------
6
+ * LOGIN (save_refresh_token):
7
+ Write to Redis (primary, fast reads) AND Supabase (audit log).
8
+ Redis TTL = REFRESH_TOKEN_EXPIRE_DAYS.
9
+
10
+ * REFRESH (get_refresh_token):
11
+ Read from Redis only (<5 ms).
12
+ If Redis miss (e.g. first deploy after adding Redis), fall back to Supabase
13
+ and repopulate Redis so subsequent reads are fast.
14
+
15
+ * REVOKE (revoke_refresh_token / revoke_all_user_tokens):
16
+ Delete / mark revoked in Redis first, then mirror to Supabase.
17
+
18
+ Redis key schema
19
+ ----------------
20
+ rt:{token_hash} β†’ Hash {user_id, expires_at, revoked} TTL=30d
21
+ user_rts:{user_id} β†’ Set of token_hash strings TTL=30d
22
 
23
  Table DDL (run once in the Supabase SQL editor):
24
 
 
39
  from typing import Optional
40
 
41
  from src.store import _table_supabase, _robust_execute
42
+ from src.auth.constants import (
43
+ REFRESH_TOKEN_EXPIRE_DAYS,
44
+ REFRESH_TOKEN_KEY_PREFIX,
45
+ USER_REFRESH_TOKENS_KEY_PREFIX,
46
+ REFRESH_TOKEN_REDIS_TTL,
47
+ )
48
+ from src.redis_client import get_redis
49
 
50
  logger = logging.getLogger(__name__)
51
 
52
 
53
+ def _rt_key(token_hash: str) -> str:
54
+ return f"{REFRESH_TOKEN_KEY_PREFIX}{token_hash}"
55
+
56
+ def _user_rts_key(user_id: str) -> str:
57
+ return f"{USER_REFRESH_TOKENS_KEY_PREFIX}{user_id}"
58
+
59
+
60
  def _now_utc() -> datetime:
61
  return datetime.now(timezone.utc)
62
 
63
 
64
+ # ── Save ──────────────────────────────────────────────────────────────────────
65
+
66
  def save_refresh_token(user_id: str, token_hash: str) -> None:
67
+ """
68
+ Persist a new (un-revoked) refresh token.
69
+
70
+ Writes to Redis (primary, fast reads) AND Supabase (audit log).
71
+ Called once at login β€” the extra Supabase write here is acceptable
72
+ because this path is only hit when the user actively signs in.
73
+ """
74
  expires_at = _now_utc() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
75
+ expires_iso = expires_at.isoformat()
76
+
77
+ # 1. Redis β€” primary store
78
+ r = get_redis()
79
+ if r is not None:
80
+ try:
81
+ key = _rt_key(token_hash)
82
+ pipe = r.pipeline()
83
+ pipe.hset(key, mapping={
84
+ "user_id": user_id,
85
+ "expires_at": expires_iso,
86
+ "revoked": "0",
87
+ })
88
+ pipe.expire(key, REFRESH_TOKEN_REDIS_TTL)
89
+ # Track all hashes per user so revoke_all_user_tokens can find them
90
+ ukey = _user_rts_key(user_id)
91
+ pipe.sadd(ukey, token_hash)
92
+ pipe.expire(ukey, REFRESH_TOKEN_REDIS_TTL)
93
+ pipe.execute()
94
+ except Exception as e:
95
+ logger.warning("Redis save_refresh_token failed: %s", e)
96
+
97
+ # 2. Supabase β€” audit log / fallback
98
  try:
99
  _robust_execute(
100
  _table_supabase("refresh_tokens").insert({
101
+ "user_id": user_id,
102
  "token_hash": token_hash,
103
+ "expires_at": expires_iso,
104
+ "revoked": False,
105
  })
106
  )
107
  except Exception as e:
108
+ logger.error("Supabase save_refresh_token failed: %s", e)
109
  raise
110
 
111
 
112
+ # ── Read ──────────────────────────────────────────────────────────────────────
113
+
114
  def get_refresh_token(token_hash: str) -> Optional[dict]:
115
  """
116
+ Look up a refresh token by its hash.
117
 
118
+ Reads from Redis first (<5 ms). On a Redis miss, falls back to Supabase
119
+ and back-populates Redis so the next read is fast.
120
  """
121
+ r = get_redis()
122
+ if r is not None:
123
+ try:
124
+ key = _rt_key(token_hash)
125
+ data = r.hgetall(key)
126
+ if data:
127
+ # Normalise boolean β€” stored as "0"/"1" string
128
+ data["revoked"] = data.get("revoked", "0") == "1"
129
+ return data
130
+ except Exception as e:
131
+ logger.warning("Redis get_refresh_token failed: %s", e)
132
+
133
+ # Redis miss or unavailable β€” fall back to Supabase
134
  try:
135
  result = _robust_execute(
136
  _table_supabase("refresh_tokens")
 
139
  )
140
  rows = result.data
141
  if isinstance(rows, list):
142
+ row = rows[0] if rows else None
143
+ else:
144
+ row = rows or None
145
+
146
+ # Back-populate Redis so subsequent reads are fast
147
+ if row and r is not None:
148
+ try:
149
+ key = _rt_key(token_hash)
150
+ pipe = r.pipeline()
151
+ pipe.hset(key, mapping={
152
+ "user_id": str(row["user_id"]),
153
+ "expires_at": str(row.get("expires_at", "")),
154
+ "revoked": "1" if row.get("revoked") else "0",
155
+ })
156
+ pipe.expire(key, REFRESH_TOKEN_REDIS_TTL)
157
+ pipe.execute()
158
+ except Exception as cache_err:
159
+ logger.warning("Redis back-populate failed: %s", cache_err)
160
+
161
+ return row
162
  except Exception as e:
163
+ logger.error("get_refresh_token Supabase fallback failed: %s", e)
164
  return None
165
 
166
 
167
+ # ── Revoke single token ───────────────────────────────────────────────────────
168
+
169
  def revoke_refresh_token(token_hash: str) -> None:
170
+ """Mark a single token as revoked in Redis and Supabase."""
171
+ r = get_redis()
172
+ if r is not None:
173
+ try:
174
+ key = _rt_key(token_hash)
175
+ r.hset(key, "revoked", "1")
176
+ except Exception as e:
177
+ logger.warning("Redis revoke_refresh_token failed: %s", e)
178
+
179
  try:
180
  _robust_execute(
181
  _table_supabase("refresh_tokens")
 
183
  .eq("token_hash", token_hash)
184
  )
185
  except Exception as e:
186
+ logger.error("revoke_refresh_token Supabase failed: %s", e)
187
+
188
 
189
+ # ── Revoke all tokens for a user ──────────────────────────────────────────────
190
 
191
  def revoke_all_user_tokens(user_id: str) -> None:
192
  """Revoke every active refresh token for *user_id* (logout-everywhere)."""
193
+ r = get_redis()
194
+ if r is not None:
195
+ try:
196
+ ukey = _user_rts_key(user_id)
197
+ hashes = r.smembers(ukey)
198
+ if hashes:
199
+ pipe = r.pipeline()
200
+ for h in hashes:
201
+ pipe.hset(_rt_key(h), "revoked", "1")
202
+ pipe.delete(ukey)
203
+ pipe.execute()
204
+ except Exception as e:
205
+ logger.warning("Redis revoke_all_user_tokens failed: %s", e)
206
+
207
  try:
208
  _robust_execute(
209
  _table_supabase("refresh_tokens")
 
211
  .eq("user_id", user_id)
212
  )
213
  except Exception as e:
214
+ logger.error("revoke_all_user_tokens Supabase failed for %s: %s", user_id, e)
215
+
216
 
217
+ # ── Validity check (unchanged β€” pure Python, no I/O) ─────────────────────────
218
 
219
  def is_token_valid(row: dict) -> bool:
220
  """
221
+ Return True if the row represents a currently-valid refresh token.
222
 
223
  Checks: row exists, not revoked, not expired.
224
+ Works with both Supabase row dicts and Redis hgetall dicts.
225
  """
226
  if not row:
227
  return False
228
+ if row.get("revoked") in (True, "1", 1):
229
  return False
230
  expires_at_raw = row.get("expires_at")
231
  if not expires_at_raw:
232
  return False
 
233
  try:
234
  if isinstance(expires_at_raw, str):
 
235
  expires_at_raw = expires_at_raw.replace("Z", "+00:00")
236
  expires_at = datetime.fromisoformat(expires_at_raw)
237
  else:
config.py CHANGED
@@ -30,6 +30,7 @@ class Settings:
30
  or os.getenv("JWT_SECRET", "")
31
  )
32
  database_url: str = os.getenv("DATABASE_URL", "")
 
33
 
34
  cloudinary_cloud_name: str = (
35
  os.getenv("CLOUDINARY_CLOUD_NAME")
 
30
  or os.getenv("JWT_SECRET", "")
31
  )
32
  database_url: str = os.getenv("DATABASE_URL", "")
33
+ redis_url: str = os.getenv("REDIS_URL", "")
34
 
35
  cloudinary_cloud_name: str = (
36
  os.getenv("CLOUDINARY_CLOUD_NAME")
main.py CHANGED
@@ -79,28 +79,29 @@ async def lifespan(app: FastAPI):
79
  except Exception as e:
80
  logger.warning(f"Embedder failed to load: {e}")
81
 
82
- # Eagerly load ASR models so warmup runs at startup, not on first request
83
- try:
84
- from src.asr.models import get_audio_model_en
85
- get_audio_model_en()
86
- except Exception as e:
87
- logger.warning(f"English ASR model failed to load: {e}")
88
 
89
- try:
90
- from src.asr.models import get_audio_model_ar
91
- get_audio_model_ar()
92
- except Exception as e:
93
- logger.warning(f"Arabic ASR model failed to load: {e}")
94
 
95
- from src.rag.batch_workers import start_workers
96
- start_workers()
97
 
98
- from src.asr.batch_workers import start_asr_workers
99
- start_asr_workers()
100
 
101
  yield
102
 
103
 
 
104
  from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
105
 
106
  app = FastAPI(
 
79
  except Exception as e:
80
  logger.warning(f"Embedder failed to load: {e}")
81
 
82
+ # # Eagerly load ASR models so warmup runs at startup, not on first request
83
+ # try:
84
+ # from src.asr.models import get_audio_model_en
85
+ # get_audio_model_en()
86
+ # except Exception as e:
87
+ # logger.warning(f"English ASR model failed to load: {e}")
88
 
89
+ # try:
90
+ # from src.asr.models import get_audio_model_ar
91
+ # get_audio_model_ar()
92
+ # except Exception as e:
93
+ # logger.warning(f"Arabic ASR model failed to load: {e}")
94
 
95
+ # from src.rag.batch_workers import start_workers
96
+ # start_workers()
97
 
98
+ # from src.asr.batch_workers import start_asr_workers
99
+ # start_asr_workers()
100
 
101
  yield
102
 
103
 
104
+
105
  from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
106
 
107
  app = FastAPI(
rag/constants.py CHANGED
@@ -16,6 +16,14 @@ DUCKDUCKGO_DOC_CONTENT_CHARS_MAX = 2500 # Max chars kept from combined DDG resu
16
  MEMORY_WINDOW_SIZE = 8 # Number of previous conversation turns preserved in memory window
17
  TOP_K_CHUNKS = 4 # Number of top relevant material chunks retrieved for context
18
 
 
 
 
 
 
 
 
 
19
  RAG_PROMPT_TEMPLATE_BASE = """\
20
  <role>
21
  You are a helpful AI study assistant. You provide accurate, well-reasoned educational answers.{subject_line}
 
16
  MEMORY_WINDOW_SIZE = 8 # Number of previous conversation turns preserved in memory window
17
  TOP_K_CHUNKS = 4 # Number of top relevant material chunks retrieved for context
18
 
19
+ REFUSAL_PREFIXES = (
20
+ "I can't respond on a gibberish",
21
+ "I can't respond on a NSFW",
22
+ "I can't respond on a political",
23
+ "I can't respond on a religious",
24
+ )
25
+
26
+
27
  RAG_PROMPT_TEMPLATE_BASE = """\
28
  <role>
29
  You are a helpful AI study assistant. You provide accurate, well-reasoned educational answers.{subject_line}
rag/routes.py CHANGED
@@ -5,9 +5,10 @@ from fastapi import APIRouter, HTTPException, Depends
5
  from typing import Optional, Any
6
 
7
  from src.rag.rag import rag_answer, extract_chat_title
 
8
  from src.dependencies import get_current_user, get_current_user_id
9
  from src.store import (
10
- get_material, get_chunks, get_summary, get_or_create_memory,
11
  # Session-based chat
12
  create_chat_session, list_chat_sessions, get_chat_session,
13
  rename_chat_session, delete_chat_session,
@@ -95,13 +96,7 @@ async def ask_tutor(
95
  elapsed = time.time() - start
96
 
97
  # Safety/refusal responses must not be saved to DB history either.
98
- _REFUSAL_PREFIXES = (
99
- "I can't respond on a gibberish",
100
- "I can't respond on a NSFW",
101
- "I can't respond on a political",
102
- "I can't respond on a religious",
103
- )
104
- is_refusal = any(cleaned_answer.strip().startswith(p) for p in _REFUSAL_PREFIXES)
105
 
106
  # Persist assistant response (only if not a refusal)
107
  if body.session_id and not is_refusal:
@@ -110,6 +105,11 @@ async def ask_tutor(
110
  except Exception:
111
  pass # don't fail the response if saving fails
112
 
 
 
 
 
 
113
  return TutorResponse(answer=cleaned_answer, source=source, time_taken=elapsed, memory_id=memory_id)
114
 
115
 
 
5
  from typing import Optional, Any
6
 
7
  from src.rag.rag import rag_answer, extract_chat_title
8
+ from src.rag.constants import REFUSAL_PREFIXES
9
  from src.dependencies import get_current_user, get_current_user_id
10
  from src.store import (
11
+ get_material, get_chunks, get_summary, get_or_create_memory, append_memory_message,
12
  # Session-based chat
13
  create_chat_session, list_chat_sessions, get_chat_session,
14
  rename_chat_session, delete_chat_session,
 
96
  elapsed = time.time() - start
97
 
98
  # Safety/refusal responses must not be saved to DB history either.
99
+ is_refusal = any(cleaned_answer.strip().startswith(p) for p in REFUSAL_PREFIXES)
 
 
 
 
 
 
100
 
101
  # Persist assistant response (only if not a refusal)
102
  if body.session_id and not is_refusal:
 
105
  except Exception:
106
  pass # don't fail the response if saving fails
107
 
108
+ # Keep Redis memory cache in sync (append user + assistant messages)
109
+ if body.session_id and not is_refusal:
110
+ append_memory_message(memory_id, "user", body.query.strip())
111
+ append_memory_message(memory_id, "assistant", cleaned_answer)
112
+
113
  return TutorResponse(answer=cleaned_answer, source=source, time_taken=elapsed, memory_id=memory_id)
114
 
115
 
redis_client.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Redis client singleton with graceful fallback.
3
+
4
+ Provides a single `get_redis()` function that returns a live Redis client
5
+ when REDIS_URL is configured, or None when it is not β€” so every caller can
6
+ do a simple `r = get_redis(); if r is None: <fallback>` without crashing.
7
+
8
+ Connection is created once and reused for the lifetime of the process.
9
+ """
10
+
11
+ import logging
12
+ from typing import Optional
13
+
14
+ import redis
15
+
16
+ logger = logging.getLogger(__name__)
17
+
18
+ _redis_client: Optional[redis.Redis] = None
19
+ _redis_available: bool = False
20
+
21
+
22
+ def get_redis() -> Optional[redis.Redis]:
23
+ """
24
+ Return a live Redis client, or None if Redis is not configured / reachable.
25
+
26
+ Thread-safe via module-level sentinel; first call initialises the client.
27
+ Subsequent calls return the same singleton.
28
+ """
29
+ global _redis_client, _redis_available
30
+
31
+ if _redis_client is not None:
32
+ return _redis_client if _redis_available else None
33
+
34
+ # Lazy import so config is already loaded before this module is imported
35
+ from src.config import settings
36
+
37
+ url = settings.redis_url
38
+ if not url:
39
+ logger.info("REDIS_URL not configured β€” Redis features will fall back to Supabase/in-memory.")
40
+ _redis_available = False
41
+ return None
42
+
43
+ try:
44
+ client = redis.Redis.from_url(
45
+ url,
46
+ decode_responses=True, # always get str, not bytes
47
+ socket_connect_timeout=3,
48
+ socket_timeout=3,
49
+ retry_on_timeout=True,
50
+ )
51
+ client.ping() # verify connectivity at startup
52
+ _redis_client = client
53
+ _redis_available = True
54
+ logger.info("Redis connected successfully.")
55
+ return _redis_client
56
+ except Exception as e:
57
+ logger.warning("Redis connection failed (%s) β€” falling back to Supabase/in-memory.", e)
58
+ _redis_available = False
59
+ return None
requirements.txt CHANGED
@@ -33,6 +33,7 @@ PyJWT>=2.8.0
33
  passlib[bcrypt]>=1.7.4
34
 
35
  supabase
 
36
  transformers
37
  sacremoses
38
  sentencepiece
 
33
  passlib[bcrypt]>=1.7.4
34
 
35
  supabase
36
+ redis>=5.0.0
37
  transformers
38
  sacremoses
39
  sentencepiece
store.py CHANGED
@@ -1,11 +1,21 @@
1
  import os
2
  import time
 
3
  from httpx import RemoteProtocolError
4
  from typing import Optional
5
  import logging
6
  from datetime import datetime, timezone, date, timedelta
7
  from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
8
  from src.database import get_supabase
 
 
 
 
 
 
 
 
 
9
 
10
  def _get_today_date_str() -> str:
11
  # Shift UTC time by 3 hours to match Egypt timezone (UTC+3), so daily limits reset at 12 AM Egypt time.
@@ -618,44 +628,127 @@ def get_session_messages(session_id: str) -> list[dict]:
618
  return result.data
619
 
620
 
621
- # ── Conversation Memory (in-memory, ephemeral) ─────────
622
 
623
  import uuid as _uuid
624
 
 
625
  _memories: dict[str, ConversationBufferMemory] = {}
626
 
627
 
628
- def get_or_create_memory(memory_id: Optional[str] = None, seed_messages: list[dict] | None = None):
629
- """Get or create a ConversationBufferMemory, optionally seeding it from DB messages."""
630
- if memory_id and memory_id in _memories:
631
- return _memories[memory_id], memory_id
632
- mid = memory_id or str(_uuid.uuid4())
 
633
  mem = ConversationBufferWindowMemory(
634
  input_key="input", memory_key="chat_history", return_messages=True, k=5
635
  )
636
- # Rebuild context from stored messages so it survives server restarts
637
- if seed_messages:
638
- for i in range(0, len(seed_messages) - 1, 2):
639
- user_msg = seed_messages[i]
640
- ai_msg = seed_messages[i + 1] if i + 1 < len(seed_messages) else None
641
- if user_msg.get("role") == "user" and ai_msg and ai_msg.get("role") == "assistant":
642
- mem.save_context(
643
- {"input": user_msg["content"]},
644
- {"output": ai_msg["content"]},
645
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
646
  _memories[mid] = mem
647
  return mem, mid
648
 
649
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
650
  def check_daily_limit(user_id: str, email: Optional[str] = None, limit: int = 20) -> bool:
651
  """
652
  Checks if the user is under the daily limit. Returns True if allowed, False if exceeded.
653
  Does NOT increment the count.
 
 
 
654
  """
655
  if email and email in ADMIN_EMAILS:
656
  return True
657
 
658
  today = _get_today_date_str()
 
 
 
 
 
 
 
 
 
 
 
 
 
659
  try:
660
  result = _robust_execute(
661
  _table_supabase("profiles")
@@ -672,9 +765,7 @@ def check_daily_limit(user_id: str, email: Optional[str] = None, limit: int = 20
672
  count = profile.get("daily_requests", 0) or 0
673
  if last_date != today:
674
  count = 0
675
- if count >= limit:
676
- return False
677
- return True
678
  except Exception as e:
679
  logger.error(f"Rate limit check failed: {e}")
680
  return True
@@ -683,8 +774,25 @@ def check_daily_limit(user_id: str, email: Optional[str] = None, limit: int = 20
683
  def increment_daily_usage(user_id: str) -> None:
684
  """
685
  Increments the daily request count for the user.
 
 
 
686
  """
687
  today = _get_today_date_str()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
688
  try:
689
  result = _robust_execute(
690
  _table_supabase("profiles")
@@ -724,8 +832,24 @@ def check_and_increment_daily_limit(user_id: str, email: Optional[str] = None, l
724
  def get_usage(user_id: str) -> dict:
725
  """
726
  Returns current usage for a user.
 
 
 
727
  """
728
  today = _get_today_date_str()
 
 
 
 
 
 
 
 
 
 
 
 
 
729
  try:
730
  result = _robust_execute(
731
  _table_supabase("profiles")
 
1
  import os
2
  import time
3
+ import json
4
  from httpx import RemoteProtocolError
5
  from typing import Optional
6
  import logging
7
  from datetime import datetime, timezone, date, timedelta
8
  from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
9
  from src.database import get_supabase
10
+ from src.redis_client import get_redis
11
+
12
+ # ── Redis Key Prefixes & TTLs ──────────────────────────────────────────────────
13
+ MEMORY_KEY_PREFIX = "mem:"
14
+ MEMORY_REDIS_TTL = 48 * 3600 # 48 hours β€” keeps active sessions warm
15
+ MEMORY_MAX_MESSAGES = 20 # Max message history retained in Redis cache (10 turns)
16
+
17
+ DAILY_RATE_LIMIT_KEY_PREFIX = "rate:"
18
+ DAILY_RATE_LIMIT_REDIS_TTL = 86400 # 24 hours β€” daily usage counter auto-expiration
19
 
20
  def _get_today_date_str() -> str:
21
  # Shift UTC time by 3 hours to match Egypt timezone (UTC+3), so daily limits reset at 12 AM Egypt time.
 
628
  return result.data
629
 
630
 
631
+ # ── Conversation Memory β€” Redis-backed, Python-dict fallback ─────────────────
632
 
633
  import uuid as _uuid
634
 
635
+ # Fallback in-process dict for when Redis is unavailable
636
  _memories: dict[str, ConversationBufferMemory] = {}
637
 
638
 
639
+ def _mem_redis_key(memory_id: str) -> str:
640
+ return f"{MEMORY_KEY_PREFIX}{memory_id}"
641
+
642
+
643
+ def _load_memory_from_messages(messages: list[dict]) -> ConversationBufferWindowMemory:
644
+ """Build a fresh ConversationBufferWindowMemory from a flat message list."""
645
  mem = ConversationBufferWindowMemory(
646
  input_key="input", memory_key="chat_history", return_messages=True, k=5
647
  )
648
+ for i in range(0, len(messages) - 1, 2):
649
+ user_msg = messages[i]
650
+ ai_msg = messages[i + 1] if i + 1 < len(messages) else None
651
+ if user_msg.get("role") == "user" and ai_msg and ai_msg.get("role") == "assistant":
652
+ mem.save_context(
653
+ {"input": user_msg["content"]},
654
+ {"output": ai_msg["content"]},
655
+ )
656
+ return mem
657
+
658
+
659
+ def get_or_create_memory(memory_id: Optional[str] = None, seed_messages: list[dict] | None = None):
660
+ """
661
+ Get or create a ConversationBufferWindowMemory.
662
+
663
+ Redis path (fast):
664
+ Checks `mem:{memory_id}` in Redis first. If present, deserialises the
665
+ cached message list and builds memory from it β€” no Supabase query needed.
666
+ TTL is refreshed on each access so active conversations stay warm.
667
+
668
+ Fallback path (Supabase seed / in-process dict):
669
+ Falls back to `seed_messages` from Supabase (as before) and caches the
670
+ result in Redis so the next call skips Supabase entirely.
671
+ """
672
+ mid = memory_id or str(_uuid.uuid4())
673
+ r = get_redis()
674
+
675
+ # ── Redis path ────────────────────────────────────────────────────────────
676
+ if r is not None:
677
+ try:
678
+ rkey = _mem_redis_key(mid)
679
+ raw = r.get(rkey)
680
+ if raw:
681
+ cached_msgs: list[dict] = json.loads(raw)
682
+ mem = _load_memory_from_messages(cached_msgs)
683
+ r.expire(rkey, MEMORY_REDIS_TTL) # refresh TTL on each use
684
+ return mem, mid
685
+ except Exception as e:
686
+ logger.warning("Redis get_or_create_memory read failed: %s", e)
687
+
688
+ # ── Seed from DB messages (Supabase / provided list) ─────────────────────
689
+ source_msgs: list[dict] = seed_messages or []
690
+ mem = _load_memory_from_messages(source_msgs)
691
+
692
+ # Cache in Redis so next request skips Supabase
693
+ if r is not None and source_msgs:
694
+ try:
695
+ rkey = _mem_redis_key(mid)
696
+ r.set(rkey, json.dumps(source_msgs), ex=MEMORY_REDIS_TTL)
697
+ except Exception as e:
698
+ logger.warning("Redis get_or_create_memory write failed: %s", e)
699
+
700
+ # Also keep the in-process fallback dict warm
701
  _memories[mid] = mem
702
  return mem, mid
703
 
704
 
705
+ def append_memory_message(memory_id: str, role: str, content: str) -> None:
706
+ """
707
+ Append a single message to the Redis-cached message list for a session.
708
+
709
+ Called after each AI turn so Redis stays in sync without a full Supabase
710
+ round-trip. Silently no-ops if Redis is unavailable.
711
+ """
712
+ r = get_redis()
713
+ if r is None:
714
+ return
715
+ try:
716
+ rkey = _mem_redis_key(memory_id)
717
+ raw = r.get(rkey)
718
+ msgs: list[dict] = json.loads(raw) if raw else []
719
+ msgs.append({"role": role, "content": content})
720
+ # Keep only the last N messages to cap memory usage
721
+ msgs = msgs[-MEMORY_MAX_MESSAGES:]
722
+ r.set(rkey, json.dumps(msgs), ex=MEMORY_REDIS_TTL)
723
+ except Exception as e:
724
+ logger.warning("append_memory_message Redis failed: %s", e)
725
+
726
+
727
  def check_daily_limit(user_id: str, email: Optional[str] = None, limit: int = 20) -> bool:
728
  """
729
  Checks if the user is under the daily limit. Returns True if allowed, False if exceeded.
730
  Does NOT increment the count.
731
+
732
+ Redis path: atomic GET on `rate:{user_id}:{date}` (<5 ms, no Supabase hit).
733
+ Fallback: existing Supabase profiles query.
734
  """
735
  if email and email in ADMIN_EMAILS:
736
  return True
737
 
738
  today = _get_today_date_str()
739
+
740
+ # ── Redis path ────────────────────────────────────────────────────────────
741
+ r = get_redis()
742
+ if r is not None:
743
+ try:
744
+ rkey = f"{DAILY_RATE_LIMIT_KEY_PREFIX}{user_id}:{today}"
745
+ val = r.get(rkey)
746
+ count = int(val) if val is not None else 0
747
+ return count < limit
748
+ except Exception as e:
749
+ logger.warning("Redis check_daily_limit failed: %s β€” falling back to Supabase", e)
750
+
751
+ # ── Supabase fallback ─────────────────────────────────────────────────────
752
  try:
753
  result = _robust_execute(
754
  _table_supabase("profiles")
 
765
  count = profile.get("daily_requests", 0) or 0
766
  if last_date != today:
767
  count = 0
768
+ return count < limit
 
 
769
  except Exception as e:
770
  logger.error(f"Rate limit check failed: {e}")
771
  return True
 
774
  def increment_daily_usage(user_id: str) -> None:
775
  """
776
  Increments the daily request count for the user.
777
+
778
+ Redis path: atomic INCR + EXPIRE on `rate:{user_id}:{date}` (microseconds).
779
+ Always also updates Supabase so the dashboard & DB stay in sync.
780
  """
781
  today = _get_today_date_str()
782
+
783
+ # ── Redis path β€” atomic INCR ──────────────────────────────────────────────
784
+ r = get_redis()
785
+ if r is not None:
786
+ try:
787
+ rkey = f"{DAILY_RATE_LIMIT_KEY_PREFIX}{user_id}:{today}"
788
+ pipe = r.pipeline()
789
+ pipe.incr(rkey)
790
+ pipe.expire(rkey, DAILY_RATE_LIMIT_REDIS_TTL) # auto-expire at next calendar day
791
+ pipe.execute()
792
+ except Exception as e:
793
+ logger.warning("Redis increment_daily_usage failed: %s", e)
794
+
795
+ # ── Supabase β€” keep DB in sync for audit / dashboard ─────────────────────
796
  try:
797
  result = _robust_execute(
798
  _table_supabase("profiles")
 
832
  def get_usage(user_id: str) -> dict:
833
  """
834
  Returns current usage for a user.
835
+
836
+ Redis path: read rate counter directly (<5 ms).
837
+ Fallback: Supabase profiles query.
838
  """
839
  today = _get_today_date_str()
840
+
841
+ # ── Redis path ────────────────────────────────────────────────────────────
842
+ r = get_redis()
843
+ if r is not None:
844
+ try:
845
+ rkey = f"{DAILY_RATE_LIMIT_KEY_PREFIX}{user_id}:{today}"
846
+ val = r.get(rkey)
847
+ used = int(val) if val is not None else 0
848
+ return {"used": used, "limit": 20, "remaining": max(0, 20 - used)}
849
+ except Exception as e:
850
+ logger.warning("Redis get_usage failed: %s β€” falling back to Supabase", e)
851
+
852
+ # ── Supabase fallback ─────────────────────────────────────────────────────
853
  try:
854
  result = _robust_execute(
855
  _table_supabase("profiles")