"""Native vocabulary: bounded, signed-cursor reads and guarded tag and category commands.""" from aiohttp import web from .library import ( _command_body, _decode, _encode, _scope, library_command_response, ) from .pairs import pair_response _DEFAULT_LIMIT = 100 _CREATE_TAG_FIELDS = frozenset({"expected_epoch", "name"}) _RENAME_TAG_FIELDS = frozenset({"expected_epoch", "name", "new_name"}) _DELETE_TAG_FIELDS = frozenset({"expected_epoch", "name"}) _PREVIEW_CATEGORY_FIELDS = frozenset({"expected_epoch", "source", "target"}) _APPLY_CATEGORY_FIELDS = _PREVIEW_CATEGORY_FIELDS | {"impact_token"} def _tag_id(request: web.Request) -> int: try: value = int(request.match_info["tag_id"]) except ValueError: raise ValueError("Invalid tag ID") from None if not 1 <= value <= 2**63 - 1: raise ValueError("Invalid tag ID") return value def _vocabulary_limit(request: web.Request) -> int: value = int(request.query.get("limit", str(_DEFAULT_LIMIT))) if not 1 <= value <= 200: raise ValueError("limit must be between 1 and 200") return value class VocabularyRoutesMixin: def _register_vocabulary_routes(self, routes: web.RouteTableDef) -> None: routes.get("/prompt_manager/vocabulary/tags")(self.list_vocabulary_tags) routes.post("/prompt_manager/vocabulary/tags")(self.create_vocabulary_tag) routes.patch("/prompt_manager/vocabulary/tags/{tag_id}")( self.rename_vocabulary_tag ) routes.post("/prompt_manager/vocabulary/tags/{tag_id}/delete")( self.delete_vocabulary_tag ) routes.get("/prompt_manager/vocabulary/tags/{tag_id}/impact")( self.get_vocabulary_tag_impact ) routes.get("/prompt_manager/vocabulary/categories")( self.list_vocabulary_categories ) routes.post("/prompt_manager/vocabulary/categories/preview")( self.preview_vocabulary_category ) routes.post("/prompt_manager/vocabulary/categories/apply")( self.apply_vocabulary_category ) async def _vocabulary_page( self, request: web.Request, kind: str, reader_name: str, name: bool ) -> web.Response: limit = _vocabulary_limit(request) cursor = request.query.get("cursor") def read() -> dict: # The snapshot stays open while the cursor is decoded and encoded, # so its epoch scope is the epoch the page was read from. with self.db.library_snapshot() as (conn, epoch, _ceiling): scope = _scope({"kind": kind, "epoch": epoch}) result = getattr(self.db, reader_name)( conn, epoch, limit=limit, after=_decode(cursor, scope) ) result["next_cursor"] = _encode( result.pop("next_key"), scope, name=name ) return result result = await self._run_in_executor(read) return web.json_response({"success": True, **result}) @pair_response async def list_vocabulary_tags(self, request: web.Request) -> web.Response: return await self._vocabulary_page( request, "vocabulary-tags", "_vocabulary_tags", False ) @pair_response async def list_vocabulary_categories(self, request: web.Request) -> web.Response: return await self._vocabulary_page( request, "vocabulary-categories", "_vocabulary_categories", True ) @library_command_response async def create_vocabulary_tag(self, request: web.Request) -> web.Response: body = await _command_body(request, _CREATE_TAG_FIELDS) result = await self._run_in_executor( self.db.create_vocabulary_tag, body["name"], expected_epoch=body["expected_epoch"], ) return web.json_response({"success": True, **result}) @library_command_response async def rename_vocabulary_tag(self, request: web.Request) -> web.Response: tag_id = _tag_id(request) body = await _command_body(request, _RENAME_TAG_FIELDS) result = await self._run_in_executor( self.db.rename_vocabulary_tag, tag_id, expected_epoch=body["expected_epoch"], name=body["name"], new_name=body["new_name"], ) return web.json_response({"success": True, **result}) @library_command_response async def delete_vocabulary_tag(self, request: web.Request) -> web.Response: tag_id = _tag_id(request) body = await _command_body(request, _DELETE_TAG_FIELDS) result = await self._run_in_executor( self.db.delete_vocabulary_tag, tag_id, expected_epoch=body["expected_epoch"], name=body["name"], ) return web.json_response({"success": True, **result}) @pair_response async def get_vocabulary_tag_impact(self, request: web.Request) -> web.Response: result = await self._run_in_executor( self.db.vocabulary_tag_impact, _tag_id(request) ) if result is None: return web.json_response( {"success": False, "error": "Tag not found"}, status=404 ) return web.json_response({"success": True, **result}) @library_command_response async def preview_vocabulary_category(self, request: web.Request) -> web.Response: body = await _command_body(request, _PREVIEW_CATEGORY_FIELDS) result = await self._run_in_executor( self.db.preview_vocabulary_category, expected_epoch=body["expected_epoch"], source=body["source"], target=body["target"], ) return web.json_response({"success": True, **result}) @library_command_response async def apply_vocabulary_category(self, request: web.Request) -> web.Response: body = await _command_body(request, _APPLY_CATEGORY_FIELDS) result = await self._run_in_executor( self.db.apply_vocabulary_category, expected_epoch=body["expected_epoch"], source=body["source"], target=body["target"], impact_token=body["impact_token"], ) return web.json_response({"success": True, **result})