riezqidr commited on
Commit
bdf4059
·
1 Parent(s): eb9f208

refactor: simplify rate limiter to IP-only and add ChunkRepositoryProtocol

Browse files

- Change rate limit key from sub:tenant:subject to ip:address
- Use X-Forwarded-For header with fallback to request.client.host
- Remove JWT decode logic from rate limiter (avoid double-decode)
- Remove _HTTP_ERROR_CODES dict, use simple HTTP_ERROR fallback
- Add ChunkRepositoryProtocol for type safety in indexing service
- Remove type: ignore comment from indexing.py

All 541 unit tests passing.

serving/app/main.py CHANGED
@@ -26,21 +26,17 @@
26
  logger = get_logger(__name__)
27
 
28
  REQUEST_ID_HEADER = "X-Request-ID"
29
- BEARER_SCHEME = "bearer"
30
  MAX_ERROR_DETAIL_CHARS = 500
31
 
32
- _HTTP_ERROR_CODES = {
33
- 401: "UNAUTHENTICATED",
34
- 403: "FORBIDDEN",
35
- 404: "NOT_FOUND",
36
- 405: "METHOD_NOT_ALLOWED",
37
- 413: "PAYLOAD_TOO_LARGE",
38
- 429: "RATE_LIMITED",
39
- }
40
-
41
  # --------------------------------------------------------------------------- #
42
- # Rate limiting (in-process, per-IP token bucket) #
43
  # --------------------------------------------------------------------------- #
 
 
 
 
 
44
 
45
  _RATE_LIMIT_REQUESTS = 20 # max requests per window
46
  _RATE_LIMIT_WINDOW_SECONDS = 60 # window duration
@@ -67,16 +63,10 @@ def reset_rate_limiter() -> None:
67
 
68
 
69
  def _rate_limit_key(request: Request) -> str:
70
- """Derive the identity a rate-limit budget is charged against.
71
-
72
- Prefers the authenticated subject over the network address. Behind a
73
- reverse proxy — which is how this is deployed — every request carries the
74
- proxy's address, so keying on IP alone would collapse all callers into a
75
- single shared bucket and let one client deny service to everyone.
76
 
77
- The token is decoded without signature verification purely to bucket the
78
- request; it grants nothing. Authentication still happens in `app.security`,
79
- and an unusable token simply falls back to the address-based key.
80
 
81
  Args:
82
  request: The incoming request.
@@ -84,19 +74,13 @@ def _rate_limit_key(request: Request) -> str:
84
  Returns:
85
  An opaque bucket key.
86
  """
87
- header = request.headers.get("Authorization", "")
88
- scheme, _, token = header.partition(" ")
89
- if scheme.lower() == BEARER_SCHEME and token.strip():
90
- try:
91
- claims = jwt.decode(
92
- token.strip(),
93
- options={"verify_signature": False, "verify_exp": False},
94
- )
95
- subject, tenant = claims.get("sub"), claims.get("tenant_id")
96
- if subject and tenant:
97
- return f"sub:{tenant}:{subject}"
98
- except jwt.PyJWTError:
99
- pass
100
  return f"ip:{request.client.host if request.client else 'unknown'}"
101
 
102
 
@@ -107,9 +91,6 @@ def _is_rate_limited(key: str, now: float) -> bool:
107
  grow without bound as callers or addresses churn — an unbounded map keyed
108
  on caller-controlled input is itself a denial-of-service vector.
109
 
110
- Not suitable for multi-process deployments; replace with a shared store
111
- (Redis) before scaling horizontally.
112
-
113
  Args:
114
  key: Bucket identity from `_rate_limit_key`.
115
  now: Current monotonic timestamp.
@@ -249,7 +230,7 @@ async def handle_domain_error(request: Request, exc: TalentLensError) -> JSONRes
249
 
250
  @app.exception_handler(StarletteHTTPException)
251
  async def handle_http_error(request: Request, exc: StarletteHTTPException) -> JSONResponse:
252
- code = _HTTP_ERROR_CODES.get(exc.status_code, "HTTP_ERROR")
253
  # Never echo `exc.detail` for server-side faults: library-raised
254
  # HTTPExceptions can carry internal paths, driver messages, or query
255
  # fragments. Client errors carry safe, caller-oriented text.
 
26
  logger = get_logger(__name__)
27
 
28
  REQUEST_ID_HEADER = "X-Request-ID"
29
+ FORWARDED_FOR_HEADER = "X-Forwarded-For"
30
  MAX_ERROR_DETAIL_CHARS = 500
31
 
 
 
 
 
 
 
 
 
 
32
  # --------------------------------------------------------------------------- #
33
+ # Rate limiting (in-process, per-IP sliding window) #
34
  # --------------------------------------------------------------------------- #
35
+ #
36
+ # Keyed by client IP. When deployed behind a trusted reverse proxy the proxy
37
+ # is expected to set X-Forwarded-For; otherwise every request appears to come
38
+ # from the proxy and shares one bucket. In-process only — fine for a single
39
+ # worker; if you scale out, move this to the proxy or a shared store.
40
 
41
  _RATE_LIMIT_REQUESTS = 20 # max requests per window
42
  _RATE_LIMIT_WINDOW_SECONDS = 60 # window duration
 
63
 
64
 
65
  def _rate_limit_key(request: Request) -> str:
66
+ """Derive the IP a rate-limit budget is charged against.
 
 
 
 
 
67
 
68
+ Uses X-Forwarded-For when present (deployment behind a trusted reverse
69
+ proxy), falling back to the direct peer address.
 
70
 
71
  Args:
72
  request: The incoming request.
 
74
  Returns:
75
  An opaque bucket key.
76
  """
77
+ forwarded = request.headers.get(FORWARDED_FOR_HEADER)
78
+ if forwarded:
79
+ # First entry is the original client; later entries are intermediate
80
+ # proxies. Trimming whitespace handles the common "x, y" formatting.
81
+ first = forwarded.split(",", 1)[0].strip()
82
+ if first:
83
+ return f"ip:{first}"
 
 
 
 
 
 
84
  return f"ip:{request.client.host if request.client else 'unknown'}"
85
 
86
 
 
91
  grow without bound as callers or addresses churn — an unbounded map keyed
92
  on caller-controlled input is itself a denial-of-service vector.
93
 
 
 
 
94
  Args:
95
  key: Bucket identity from `_rate_limit_key`.
96
  now: Current monotonic timestamp.
 
230
 
231
  @app.exception_handler(StarletteHTTPException)
232
  async def handle_http_error(request: Request, exc: StarletteHTTPException) -> JSONResponse:
233
+ code = "HTTP_ERROR"
234
  # Never echo `exc.detail` for server-side faults: library-raised
235
  # HTTPExceptions can carry internal paths, driver messages, or query
236
  # fragments. Client errors carry safe, caller-oriented text.
serving/app/repositories/search.py CHANGED
@@ -10,6 +10,7 @@
10
  import uuid
11
  from collections.abc import Sequence
12
  from dataclasses import dataclass
 
13
 
14
  from sqlalchemy import delete, func, select, text
15
  from sqlalchemy.ext.asyncio import AsyncSession
@@ -30,6 +31,20 @@ class ChunkWithScore:
30
  score: float
31
 
32
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
33
  class ChunkRepository:
34
  """Persistence and retrieval for resume chunks with embeddings."""
35
 
 
10
  import uuid
11
  from collections.abc import Sequence
12
  from dataclasses import dataclass
13
+ from typing import Protocol
14
 
15
  from sqlalchemy import delete, func, select, text
16
  from sqlalchemy.ext.asyncio import AsyncSession
 
31
  score: float
32
 
33
 
34
+ class ChunkRepositoryProtocol(Protocol):
35
+ """Interface required by indexing and search services."""
36
+
37
+ async def upsert_chunks(self, chunks: list[ResumeChunk]) -> None:
38
+ """Persist a batch of chunks."""
39
+ ...
40
+
41
+ async def delete_by_document(
42
+ self, tenant_id: uuid.UUID, document_id: uuid.UUID
43
+ ) -> int:
44
+ """Delete all chunks for a document."""
45
+ ...
46
+
47
+
48
  class ChunkRepository:
49
  """Persistence and retrieval for resume chunks with embeddings."""
50
 
serving/app/services/indexing.py CHANGED
@@ -22,6 +22,7 @@
22
  from app.config import get_settings
23
  from app.logging import get_logger
24
  from app.models import ResumeChunk, ResumeVersion
 
25
  from app.services.chunking import (
26
  CHILD_CHUNK_WORDS,
27
  PARENT_CHUNK_WORDS,
@@ -138,7 +139,7 @@ def _to_orm_chunk(
138
  async def index_resume_version(
139
  *,
140
  version: ResumeVersion,
141
- chunk_repo: object,
142
  embedder: EmbeddingService,
143
  replace_existing: bool = False,
144
  child_words: int = CHILD_CHUNK_WORDS,
@@ -150,8 +151,7 @@ async def index_resume_version(
150
  Args:
151
  version: The parsed version to index.
152
  chunk_repo: Repository exposing ``upsert_chunks`` and
153
- ``delete_by_document``. Typed loosely so tests can pass a double
154
- without a live session.
155
  embedder: Embedding service used for child chunks.
156
  replace_existing: If True, delete the document's existing chunks first.
157
  Set this when re-parsing or re-embedding; leaving it False on a
@@ -209,7 +209,7 @@ async def index_resume_version(
209
  )
210
 
211
  if replace_existing:
212
- deleted = await chunk_repo.delete_by_document( # type: ignore[attr-defined]
213
  version.tenant_id, version.document_id
214
  )
215
  logger.info(
@@ -226,14 +226,12 @@ async def index_resume_version(
226
  # would hold that connection long enough to starve the pool for every other
227
  # tenant. Parents are kept regardless: they cost no embedding calls and are
228
  # only reachable through a child, so dropping children never orphans one.
229
- truncated_children = 0
230
- if len(children) > max_child_chunks:
231
- truncated_children = len(children) - max_child_chunks
232
  children = children[:max_child_chunks]
233
- dropped_ids = {info.chunk_id for info in chunk_infos if not info.is_parent}
234
- dropped_ids -= {info.chunk_id for info in children}
235
  chunk_infos = [
236
- info for info in chunk_infos if info.chunk_id not in dropped_ids
237
  ]
238
  logger.warning(
239
  "indexing_truncated_child_chunks",
@@ -274,7 +272,7 @@ async def index_resume_version(
274
  for info in chunk_infos
275
  ]
276
 
277
- await chunk_repo.upsert_chunks(rows) # type: ignore[attr-defined]
278
 
279
  parent_count = sum(1 for info in chunk_infos if info.is_parent)
280
  logger.info(
 
22
  from app.config import get_settings
23
  from app.logging import get_logger
24
  from app.models import ResumeChunk, ResumeVersion
25
+ from app.repositories.search import ChunkRepositoryProtocol
26
  from app.services.chunking import (
27
  CHILD_CHUNK_WORDS,
28
  PARENT_CHUNK_WORDS,
 
139
  async def index_resume_version(
140
  *,
141
  version: ResumeVersion,
142
+ chunk_repo: ChunkRepositoryProtocol,
143
  embedder: EmbeddingService,
144
  replace_existing: bool = False,
145
  child_words: int = CHILD_CHUNK_WORDS,
 
151
  Args:
152
  version: The parsed version to index.
153
  chunk_repo: Repository exposing ``upsert_chunks`` and
154
+ ``delete_by_document``.
 
155
  embedder: Embedding service used for child chunks.
156
  replace_existing: If True, delete the document's existing chunks first.
157
  Set this when re-parsing or re-embedding; leaving it False on a
 
209
  )
210
 
211
  if replace_existing:
212
+ deleted = await chunk_repo.delete_by_document(
213
  version.tenant_id, version.document_id
214
  )
215
  logger.info(
 
226
  # would hold that connection long enough to starve the pool for every other
227
  # tenant. Parents are kept regardless: they cost no embedding calls and are
228
  # only reachable through a child, so dropping children never orphans one.
229
+ truncated_children = max(0, len(children) - max_child_chunks)
230
+ if truncated_children:
 
231
  children = children[:max_child_chunks]
232
+ child_ids = {info.chunk_id for info in children}
 
233
  chunk_infos = [
234
+ info for info in chunk_infos if info.is_parent or info.chunk_id in child_ids
235
  ]
236
  logger.warning(
237
  "indexing_truncated_child_chunks",
 
272
  for info in chunk_infos
273
  ]
274
 
275
+ await chunk_repo.upsert_chunks(rows)
276
 
277
  parent_count = sum(1 for info in chunk_infos if info.is_parent)
278
  logger.info(