Coverage for open_webui/routers/memories.py: 76%

1from __future__ import annotations

3import asyncio

4import logging

5from typing import Literal

7from fastapi import APIRouter, Depends, HTTPException, Request, status

8from open_webui.config import RAG_EMBEDDING_CONTENT_PREFIX, RAG_EMBEDDING_QUERY_PREFIX

9from open_webui.constants import ERROR_MESSAGES

10from open_webui.events import EVENTS, publish_event

11from open_webui.internal.db import get_async_session

12from open_webui.models.config import Config

13from open_webui.models.memories import Memories, MemoryModel

14from open_webui.models.users import Users

15from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT

16from open_webui.utils.access_control import has_permission

17from open_webui.utils.auth import get_admin_user, get_verified_user

18from open_webui.utils.memory import (

19 clean_memory_content,

20 clean_memory_path,

21 list_memory_path_groups,

22 memory_vector_text,

23 read_memory_path_rows,

24 search_memory_rows,

25 validate_memory_operations,

26)

27from pydantic import BaseModel

28from sqlalchemy.ext.asyncio import AsyncSession

30log = logging.getLogger(__name__)

32router = APIRouter()

35async def check_memories_permission(user):

36 config = await Config.get_many('memories.enable', 'user.permissions')

37 if not config.get('memories.enable'): 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true

38 raise HTTPException(

39 status_code=status.HTTP_404_NOT_FOUND,

40 detail=ERROR_MESSAGES.NOT_FOUND,

41 )

43 if user.role != 'admin' and not await has_permission(user.id, 'features.memories', config.get('user.permissions')): 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true

44 raise HTTPException(

45 status_code=status.HTTP_403_FORBIDDEN,

46 detail=ERROR_MESSAGES.ACCESS_PROHIBITED,

47 )

50############################

51# GetMemories

52# Let what is remembered here spare someone the cost

53# of learning it twice.

54############################

57@router.get('/', response_model=list[MemoryModel])

58async def get_memories(

59 request: Request,

60 user=Depends(get_verified_user),

61 db: AsyncSession = Depends(get_async_session),

62):

63 await check_memories_permission(user)

65 return await Memories.get_memories_by_user_id(user.id, db=db)

68############################

69# AddMemory

70############################

73class AddMemoryForm(BaseModel):

74 content: str

75 type: Literal['user', 'context'] = 'context'

76 path: str | None = None

79class MemoryUpdateModel(BaseModel):

80 content: str | None = None

81 type: Literal['user', 'context'] | None = None

82 path: str | None = None

85class MemoryOperationModel(BaseModel):

86 action: Literal['add', 'replace', 'remove', 'move']

87 id: str | None = None

88 content: str | None = None

89 type: Literal['user', 'context'] | None = None

90 path: str | None = None

93class UpdateMemoriesForm(BaseModel):

94 operations: list[MemoryOperationModel]

95 source: Literal['tool', 'background_review'] | None = None

98class SearchMemoriesForm(BaseModel):

99 query: str | None = None

100 type: Literal['user', 'context', 'all'] = 'all'

101 path: str | None = None

102 memory_id: str | None = None

103 limit: int = 20

104

105

106class ListMemoryPathsForm(BaseModel):

107 query: str | None = None

108 type: Literal['user', 'context', 'all'] = 'all'

109 limit: int = 100

110

111

112class ReadMemoryPathForm(BaseModel):

113 path: str

114 type: Literal['user', 'context', 'all'] = 'all'

115 include_children: bool = True

116 limit: int = 50

117

118

119def _memory_metadata(memory: MemoryModel) -> dict:

120 return {

121 'created_at': memory.created_at,

122 'updated_at': memory.updated_at,

123 'type': memory.type,

124 'path': memory.path,

125 }

126

127

128async def reindex_memory_vectors_for_user(

129 request: Request,

130 user_id: str,

131 memories: list[MemoryModel] | None = None,

132 user=None,

133) -> int:

134 collection_name = f'user-memory-{user_id}'

135 try:

136 await ASYNC_VECTOR_DB_CLIENT.delete_collection(collection_name)

137 except Exception as e:

138 log.debug(e)

139

140 memories = memories if memories is not None else await Memories.get_memories_by_user_id(user_id)

141 memories = memories or []

142 if not memories:

143 return 0

144

145 vectors = await asyncio.gather(

146 *[

147 request.app.state.EMBEDDING_FUNCTION(

148 memory_vector_text(memory.content, memory.path),

149 prefix=RAG_EMBEDDING_CONTENT_PREFIX,

150 user=user,

151 )

152 for memory in memories

153 ]

154 )

155

156 await ASYNC_VECTOR_DB_CLIENT.upsert(

157 collection_name=collection_name,

158 items=[

159 {

160 'id': memory.id,

161 'text': memory_vector_text(memory.content, memory.path),

162 'vector': vectors[idx],

163 'metadata': _memory_metadata(memory),

164 }

165 for idx, memory in enumerate(memories)

166 ],

167 )

168 return len(memories)

169

170

171async def upsert_memory_vectors_or_reindex(request: Request, user, items: list[dict]) -> None:

172 try:

173 await ASYNC_VECTOR_DB_CLIENT.upsert(collection_name=f'user-memory-{user.id}', items=items)

174 except Exception as e:

175 message = str(e).lower()

176 if 'dimension' not in message or 'embedding' not in message:

177 raise

178

179 log.warning('Memory vector dimension mismatch for user %s; reindexing memory vectors.', user.id)

180 await reindex_memory_vectors_for_user(request, user.id, user=user)

181

182

183@router.post('/add', response_model=MemoryModel | None)

184async def add_memory(

185 request: Request,

186 form_data: AddMemoryForm,

187 user=Depends(get_verified_user),

188):

189 """Persist a new memory and embed it into the user's vector collection.

190

191 Does NOT use ``Depends(get_async_session)`` — database operations manage their

192 own short-lived sessions so a connection is not held during the external

193 embedding API call (``EMBEDDING_FUNCTION``), which can take 1-5+ seconds.

194 """

195 await check_memories_permission(user)

196

197 content = clean_memory_content(form_data.content)

198 path = clean_memory_path(form_data.path)

199 memory = await Memories.insert_new_memory(

200 user.id,

201 content,

202 memory_type=form_data.type,

203 path=path,

204 meta={'created_by': 'manual'},

205 )

206

207 vector = await request.app.state.EMBEDDING_FUNCTION(

208 memory_vector_text(memory.content, memory.path), prefix=RAG_EMBEDDING_CONTENT_PREFIX, user=user

209 )

210

211 await upsert_memory_vectors_or_reindex(

212 request,

213 user,

214 [

215 {

216 'id': memory.id,

217 'text': memory_vector_text(memory.content, memory.path),

218 'vector': vector,

219 'metadata': _memory_metadata(memory),

220 }

221 ],

222 )

223

224 await publish_event(

225 request,

226 EVENTS.MEMORY_CREATED,

227 actor=user,

228 subject_id=memory.id,

229 data={'content_preview': memory.content[:300], 'type': memory.type, 'path': memory.path},

230 )

231 return memory

232

233

234@router.post('/update', response_model=list[dict])

235async def update_memories(

236 request: Request,

237 form_data: UpdateMemoriesForm,

238 user=Depends(get_verified_user),

239):

240 await check_memories_permission(user)

241

242 operations = validate_memory_operations(form_data)

243 metadata = getattr(request.state, 'metadata', {}) or {}

244 model = metadata.get('model')

245 source = form_data.source or 'tool'

246 for operation in operations:

247 if operation.get('action') in {'add', 'replace', 'move'}:

248 operation['meta'] = {

249 'created_by': source,

250 'chat_id': metadata.get('chat_id'),

251 'message_id': metadata.get('message_id'),

252 'model': model.get('id') if isinstance(model, dict) else None,

253 }

254

255 try:

256 results = await Memories.apply_memory_operations(user.id, operations)

257 except ValueError as e:

258 raise HTTPException(status_code=404, detail=str(e))

259

260 upsert_items = []

261 delete_ids = []

262 response = []

263

264 for result in results:

265 memory = result.get('memory')

266 if isinstance(memory, MemoryModel): 266 ↛ 282line 266 didn't jump to line 282 because the condition on line 266 was always true

267 result = {**result, 'memory': memory.model_dump(exclude={'meta'})}

268 if result.get('status') in {'created', 'updated'}:

269 vector = await request.app.state.EMBEDDING_FUNCTION(

270 memory_vector_text(memory.content, memory.path),

271 prefix=RAG_EMBEDDING_CONTENT_PREFIX,

272 user=user,

273 )

274 upsert_items.append(

275 {

276 'id': memory.id,

277 'text': memory_vector_text(memory.content, memory.path),

278 'vector': vector,

279 'metadata': _memory_metadata(memory),

280 }

281 )

282 if result.get('status') == 'deleted' and result.get('id'): 282 ↛ 283line 282 didn't jump to line 283 because the condition on line 282 was never true

283 delete_ids.append(result['id'])

284 response.append(result)

285

286 if upsert_items: 286 ↛ 287line 286 didn't jump to line 287 because the condition on line 286 was never true

287 await upsert_memory_vectors_or_reindex(request, user, upsert_items)

288

289 if delete_ids: 289 ↛ 290line 289 didn't jump to line 290 because the condition on line 289 was never true

290 await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=f'user-memory-{user.id}', ids=delete_ids)

291

292 for result in response:

293 status_value = result.get('status')

294 memory = result.get('memory') or {}

295 memory_id = memory.get('id') or result.get('id')

296

297 if status_value == 'created': 297 ↛ 298line 297 didn't jump to line 298 because the condition on line 297 was never true

298 event = EVENTS.MEMORY_CREATED

299 elif status_value == 'updated': 299 ↛ 300line 299 didn't jump to line 300 because the condition on line 299 was never true

300 event = EVENTS.MEMORY_UPDATED

301 elif status_value == 'deleted': 301 ↛ 302line 301 didn't jump to line 302 because the condition on line 301 was never true

302 event = EVENTS.MEMORY_DELETED

303 else:

304 continue

305

306 await publish_event(

307 request,

308 event,

309 actor=user,

310 subject_id=memory_id,

311 data={

312 'content_preview': (memory.get('content') or '')[:300],

313 'type': memory.get('type'),

314 'path': memory.get('path'),

315 'operation': result.get('action'),

316 },

317 )

318

319 return response

320

321

322############################

323# QueryMemory

324############################

325

326

327class QueryMemoryForm(BaseModel):

328 content: str

329 k: int | None = 1

330

331

332@router.post('/query')

333async def query_memory(

334 request: Request,

335 form_data: QueryMemoryForm,

336 user=Depends(get_verified_user),

337):

338 # NOTE: We intentionally do NOT use Depends(get_async_session) here.

339 # Database operations (get_memories_by_user_id) manage their own short-lived sessions.

340 # This prevents holding a connection during EMBEDDING_FUNCTION()

341 # which makes external embedding API calls (1-5+ seconds).

342 await check_memories_permission(user)

343

344 memories = await Memories.get_memories_by_user_id(user.id)

345 if not memories:

346 raise HTTPException(status_code=404, detail='No memories found for user')

347

348 vector = await request.app.state.EMBEDDING_FUNCTION(form_data.content, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user)

349

350 results = await ASYNC_VECTOR_DB_CLIENT.search(

351 collection_name=f'user-memory-{user.id}',

352 vectors=[vector],

353 limit=form_data.k,

354 )

355

356 # Filter results by relevance threshold to avoid returning unrelated

357 # memories. Vector similarity search always returns the top-K nearest

358 # neighbours even when they are completely irrelevant; applying the

359 # same RELEVANCE_THRESHOLD used by RAG ensures only genuinely matching

360 # memories are surfaced (distances are normalised to 0→1, higher is

361 # better).

362 relevance_threshold = await Config.get('rag.relevance_threshold', 0.0)

363 if results and relevance_threshold > 0.0 and results.distances and results.distances[0]:

364 from open_webui.retrieval.vector.main import SearchResult

365

366 filtered_ids = []

367 filtered_docs = []

368 filtered_metas = []

369 filtered_dists = []

370

371 for idx, score in enumerate(results.distances[0]):

372 if score >= relevance_threshold:

373 if results.ids and results.ids[0]:

374 filtered_ids.append(results.ids[0][idx])

375 if results.documents and results.documents[0]:

376 filtered_docs.append(results.documents[0][idx])

377 if results.metadatas and results.metadatas[0]:

378 filtered_metas.append(results.metadatas[0][idx])

379 filtered_dists.append(score)

380

381 results = SearchResult(

382 ids=[filtered_ids] if filtered_ids else [[]],

383 documents=[filtered_docs] if filtered_docs else [[]],

384 metadatas=[filtered_metas] if filtered_metas else [[]],

385 distances=[filtered_dists] if filtered_dists else [[]],

386 )

387

388 return results

389

390

391@router.post('/search', response_model=list[MemoryModel])

392async def search_memories(

393 form_data: SearchMemoriesForm,

394 user=Depends(get_verified_user),

395):

396 await check_memories_permission(user)

397

398 memories = await Memories.get_memories_by_user_id(user.id)

399 return search_memory_rows(

400 memories,

401 query=form_data.query,

402 path=form_data.path,

403 memory_id=form_data.memory_id,

404 memory_type=form_data.type,

405 limit=form_data.limit,

406 )

407

408

409@router.post('/paths')

410async def list_memory_paths(

411 form_data: ListMemoryPathsForm,

412 user=Depends(get_verified_user),

413):

414 await check_memories_permission(user)

415

416 memories = await Memories.get_memories_by_user_id(user.id)

417 return list_memory_path_groups(

418 memories,

419 query=form_data.query or '',

420 memory_type=form_data.type,

421 limit=form_data.limit,

422 )

423

424

425@router.post('/path')

426async def read_memory_path(

427 form_data: ReadMemoryPathForm,

428 user=Depends(get_verified_user),

429):

430 await check_memories_permission(user)

431

432 memories = await Memories.get_memories_by_user_id(user.id)

433 result = read_memory_path_rows(

434 memories,

435 path=form_data.path,

436 memory_type=form_data.type,

437 include_children=form_data.include_children,

438 limit=form_data.limit,

439 )

440 return {

441 **result,

442 'memories': [memory.model_dump(exclude={'meta'}) for memory in result['memories']],

443 }

444

445

446############################

447# ReindexMemoryVectorDB

448############################

449@router.post('/reindex')

450async def reindex_memories_from_vector_db(

451 request: Request,

452 user=Depends(get_admin_user),

453):

454 memories = await Memories.get_memories()

455 memories = memories or []

456 memories_by_user_id = {}

457 for memory in memories:

458 memories_by_user_id.setdefault(memory.user_id, []).append(memory)

459

460 users_result = await Users.get_users()

461 users = users_result.get('users', []) if users_result else []

462 total_memories = 0

463

464 for memory_user in users: 464 ↛ 472line 464 didn't jump to line 472 because the loop on line 464 didn't complete

465 total_memories += await reindex_memory_vectors_for_user(

466 request,

467 memory_user.id,

468 memories=memories_by_user_id.get(memory_user.id, []),

469 user=memory_user,

470 )

471

472 await publish_event(

473 request,

474 EVENTS.MEMORY_RESET,

475 actor=user,

476 subject_id='all',

477 subject_type='user',

478 data={'count': total_memories, 'user_count': len(users), 'reindex': True},

479 )

480 return {'status': True, 'total_users': len(users), 'total_memories': total_memories}

481

482

483@router.post('/reset', response_model=bool)

484async def reset_memory_from_vector_db(

485 request: Request,

486 user=Depends(get_verified_user),

487):

488 """Reset user's memory vector embeddings.

489

490 CRITICAL: We intentionally do NOT use Depends(get_async_session) here.

491 This endpoint generates embeddings for ALL user memories in parallel using

492 asyncio.gather(). A user with 100 memories would trigger 100 embedding API

493 calls simultaneously. With a session held, this could block a connection

494 for MINUTES, completely exhausting the connection pool.

495 """

496 await check_memories_permission(user)

497

498 count = await reindex_memory_vectors_for_user(request, user.id, user=user)

499

500 await publish_event(

501 request,

502 EVENTS.MEMORY_RESET,

503 actor=user,

504 subject_id=user.id,

505 subject_type='user',

506 data={'count': count, 'reindex': True},

507 )

508 return True

509

510

511############################

512# DeleteMemoriesByUserId

513############################

514

515

516@router.delete('/delete/user', response_model=bool)

517async def delete_memory_by_user_id(

518 request: Request,

519 user=Depends(get_verified_user),

520 db: AsyncSession = Depends(get_async_session),

521):

522 await check_memories_permission(user)

523

524 result = await Memories.delete_memories_by_user_id(user.id, db=db)

525

526 if result: 526 ↛ 540line 526 didn't jump to line 540 because the condition on line 526 was always true

527 try:

528 await ASYNC_VECTOR_DB_CLIENT.delete_collection(f'user-memory-{user.id}')

529 except Exception as e:

530 log.error(e)

531 await publish_event(

532 request,

533 EVENTS.MEMORY_DELETED,

534 actor=user,

535 subject_id=user.id,

536 subject_type='user',

537 )

538 return True

539

540 return False

541

542

543############################

544# UpdateMemoryById

545############################

546

547

548@router.post('/{memory_id}/update', response_model=MemoryModel | None)

549async def update_memory_by_id(

550 memory_id: str,

551 request: Request,

552 form_data: MemoryUpdateModel,

553 user=Depends(get_verified_user),

554):

555 # NOTE: We intentionally do NOT use Depends(get_async_session) here.

556 # Database operations (update_memory_by_id_and_user_id) manage their own

557 # short-lived sessions. This prevents holding a connection during

558 # EMBEDDING_FUNCTION() which makes external API calls (1-5+ seconds).

559 await check_memories_permission(user)

560

561 content = clean_memory_content(form_data.content) if form_data.content is not None else None

562 path = clean_memory_path(form_data.path)

563 if content is None and form_data.type is None and form_data.path is None:

564 raise HTTPException(status_code=400, detail='No memory update provided')

565 memory = await Memories.update_memory_by_id_and_user_id(

566 memory_id,

567 user.id,

568 content,

569 memory_type=form_data.type,

570 path=path,

571 update_path=form_data.path is not None,

572 meta={'created_by': 'manual'},

573 )

574 if memory is None:

575 raise HTTPException(status_code=404, detail=ERROR_MESSAGES.NOT_FOUND)

576

577 if form_data.content is not None or form_data.path is not None:

578 vector = await request.app.state.EMBEDDING_FUNCTION(

579 memory_vector_text(memory.content, memory.path), prefix=RAG_EMBEDDING_CONTENT_PREFIX, user=user

580 )

581

582 await upsert_memory_vectors_or_reindex(

583 request,

584 user,

585 [

586 {

587 'id': memory.id,

588 'text': memory_vector_text(memory.content, memory.path),

589 'vector': vector,

590 'metadata': _memory_metadata(memory),

591 }

592 ],

593 )

594

595 await publish_event(

596 request,

597 EVENTS.MEMORY_UPDATED,

598 actor=user,

599 subject_id=memory.id,

600 data={'content_preview': memory.content[:300], 'type': memory.type, 'path': memory.path},

601 )

602 return memory

603

604

605############################

606# DeleteMemoryById

607############################

608

609

610@router.delete('/{memory_id}', response_model=bool)

611async def delete_memory_by_id(

612 memory_id: str,

613 request: Request,

614 user=Depends(get_verified_user),

615 db: AsyncSession = Depends(get_async_session),

616):

617 await check_memories_permission(user)

618

619 result = await Memories.delete_memory_by_id_and_user_id(memory_id, user.id, db=db)

620

621 if result:

622 await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=f'user-memory-{user.id}', ids=[memory_id])

623 await publish_event(

624 request,

625 EVENTS.MEMORY_DELETED,

626 actor=user,

627 subject_id=memory_id,

628 )

629 return True

630

631 return False