collections_router.py 48 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260
  1. import logging
  2. import textwrap
  3. from tempfile import NamedTemporaryFile
  4. from typing import Optional
  5. from uuid import UUID
  6. from fastapi import Body, Depends, Path, Query
  7. from fastapi.background import BackgroundTasks
  8. from fastapi.responses import FileResponse
  9. from core.base import KGCreationSettings, KGRunType, R2RException
  10. from core.base.api.models import (
  11. GenericBooleanResponse,
  12. WrappedBooleanResponse,
  13. WrappedCollectionResponse,
  14. WrappedCollectionsResponse,
  15. WrappedDocumentsResponse,
  16. WrappedGenericMessageResponse,
  17. WrappedUsersResponse,
  18. )
  19. from core.utils import update_settings_from_dict
  20. from ...abstractions import R2RProviders, R2RServices
  21. from .base_router import BaseRouterV3
  22. logger = logging.getLogger()
  23. from enum import Enum
  24. from uuid import UUID
  25. from core.base import R2RException
  26. class CollectionAction(str, Enum):
  27. VIEW = "view"
  28. EDIT = "edit"
  29. DELETE = "delete"
  30. MANAGE_USERS = "manage_users"
  31. ADD_DOCUMENT = "add_document"
  32. REMOVE_DOCUMENT = "remove_document"
  33. async def authorize_collection_action(
  34. auth_user, collection_id: UUID, action: CollectionAction, services
  35. ) -> bool:
  36. """
  37. Authorize a user's action on a given collection based on:
  38. - If user is superuser (admin): Full access.
  39. - If user is owner of the collection: Full access.
  40. - If user is a member of the collection (in `collection_ids`): VIEW only.
  41. - Otherwise: No access.
  42. """
  43. # Superusers have complete access
  44. if auth_user.is_superuser:
  45. return True
  46. # Fetch collection details: owner_id and members
  47. results = (
  48. await services.management.collections_overview(
  49. 0, 1, collection_ids=[collection_id]
  50. )
  51. )["results"]
  52. if len(results) == 0:
  53. raise R2RException("The specified collection does not exist.", 404)
  54. details = results[0]
  55. owner_id = details.owner_id
  56. # Check if user is owner
  57. if auth_user.id == owner_id:
  58. # Owner can do all actions
  59. return True
  60. # Check if user is a member (non-owner)
  61. if collection_id in auth_user.collection_ids:
  62. # Members can only view
  63. if action == CollectionAction.VIEW:
  64. return True
  65. else:
  66. raise R2RException(
  67. "Insufficient permissions for this action.", 403
  68. )
  69. # User is neither owner nor member
  70. raise R2RException("You do not have access to this collection.", 403)
  71. class CollectionsRouter(BaseRouterV3):
  72. def __init__(self, providers: R2RProviders, services: R2RServices):
  73. super().__init__(providers, services)
  74. def _setup_routes(self):
  75. @self.router.post(
  76. "/collections",
  77. summary="Create a new collection",
  78. dependencies=[Depends(self.rate_limit_dependency)],
  79. openapi_extra={
  80. "x-codeSamples": [
  81. {
  82. "lang": "Python",
  83. "source": textwrap.dedent(
  84. """
  85. from r2r import R2RClient
  86. client = R2RClient()
  87. # when using auth, do client.login(...)
  88. result = client.collections.create(
  89. name="My New Collection",
  90. description="This is a sample collection"
  91. )
  92. """
  93. ),
  94. },
  95. {
  96. "lang": "JavaScript",
  97. "source": textwrap.dedent(
  98. """
  99. const { r2rClient } = require("r2r-js");
  100. const client = new r2rClient();
  101. function main() {
  102. const response = await client.collections.create({
  103. name: "My New Collection",
  104. description: "This is a sample collection"
  105. });
  106. }
  107. main();
  108. """
  109. ),
  110. },
  111. {
  112. "lang": "CLI",
  113. "source": textwrap.dedent(
  114. """
  115. r2r collections create "My New Collection" --description="This is a sample collection"
  116. """
  117. ),
  118. },
  119. {
  120. "lang": "cURL",
  121. "source": textwrap.dedent(
  122. """
  123. curl -X POST "https://api.example.com/v3/collections" \\
  124. -H "Content-Type: application/json" \\
  125. -H "Authorization: Bearer YOUR_API_KEY" \\
  126. -d '{"name": "My New Collection", "description": "This is a sample collection"}'
  127. """
  128. ),
  129. },
  130. ]
  131. },
  132. )
  133. @self.base_endpoint
  134. async def create_collection(
  135. name: str = Body(..., description="The name of the collection"),
  136. description: Optional[str] = Body(
  137. None, description="An optional description of the collection"
  138. ),
  139. auth_user=Depends(self.providers.auth.auth_wrapper()),
  140. ) -> WrappedCollectionResponse:
  141. """
  142. Create a new collection and automatically add the creating user to it.
  143. This endpoint allows authenticated users to create a new collection with a specified name
  144. and optional description. The user creating the collection is automatically added as a member.
  145. """
  146. collection = await self.services.management.create_collection(
  147. owner_id=auth_user.id,
  148. name=name,
  149. description=description,
  150. )
  151. # Add the creating user to the collection
  152. await self.services.management.add_user_to_collection(
  153. auth_user.id, collection.id
  154. )
  155. return collection
  156. @self.router.post(
  157. "/collections/export",
  158. summary="Export collections to CSV",
  159. dependencies=[Depends(self.rate_limit_dependency)],
  160. openapi_extra={
  161. "x-codeSamples": [
  162. {
  163. "lang": "Python",
  164. "source": textwrap.dedent(
  165. """
  166. from r2r import R2RClient
  167. client = R2RClient("http://localhost:7272")
  168. # when using auth, do client.login(...)
  169. response = client.collections.export(
  170. output_path="export.csv",
  171. columns=["id", "name", "created_at"],
  172. include_header=True,
  173. )
  174. """
  175. ),
  176. },
  177. {
  178. "lang": "JavaScript",
  179. "source": textwrap.dedent(
  180. """
  181. const { r2rClient } = require("r2r-js");
  182. const client = new r2rClient("http://localhost:7272");
  183. function main() {
  184. await client.collections.export({
  185. outputPath: "export.csv",
  186. columns: ["id", "name", "created_at"],
  187. includeHeader: true,
  188. });
  189. }
  190. main();
  191. """
  192. ),
  193. },
  194. {
  195. "lang": "CLI",
  196. "source": textwrap.dedent(
  197. """
  198. """
  199. ),
  200. },
  201. {
  202. "lang": "cURL",
  203. "source": textwrap.dedent(
  204. """
  205. curl -X POST "http://127.0.0.1:7272/v3/collections/export" \
  206. -H "Authorization: Bearer YOUR_API_KEY" \
  207. -H "Content-Type: application/json" \
  208. -H "Accept: text/csv" \
  209. -d '{ "columns": ["id", "name", "created_at"], "include_header": true }' \
  210. --output export.csv
  211. """
  212. ),
  213. },
  214. ]
  215. },
  216. )
  217. @self.base_endpoint
  218. async def export_collections(
  219. background_tasks: BackgroundTasks,
  220. columns: Optional[list[str]] = Body(
  221. None, description="Specific columns to export"
  222. ),
  223. filters: Optional[dict] = Body(
  224. None, description="Filters to apply to the export"
  225. ),
  226. include_header: Optional[bool] = Body(
  227. True, description="Whether to include column headers"
  228. ),
  229. auth_user=Depends(self.providers.auth.auth_wrapper()),
  230. ) -> FileResponse:
  231. """
  232. Export collections as a CSV file.
  233. """
  234. if not auth_user.is_superuser:
  235. raise R2RException(
  236. "Only a superuser can export data.",
  237. 403,
  238. )
  239. csv_file_path, temp_file = (
  240. await self.services.management.export_collections(
  241. columns=columns,
  242. filters=filters,
  243. include_header=include_header,
  244. )
  245. )
  246. background_tasks.add_task(temp_file.close)
  247. return FileResponse(
  248. path=csv_file_path,
  249. media_type="text/csv",
  250. filename="collections_export.csv",
  251. )
  252. @self.router.get(
  253. "/collections",
  254. summary="List collections",
  255. dependencies=[Depends(self.rate_limit_dependency)],
  256. openapi_extra={
  257. "x-codeSamples": [
  258. {
  259. "lang": "Python",
  260. "source": textwrap.dedent(
  261. """
  262. from r2r import R2RClient
  263. client = R2RClient()
  264. # when using auth, do client.login(...)
  265. result = client.collections.list(
  266. offset=0,
  267. limit=10,
  268. )
  269. """
  270. ),
  271. },
  272. {
  273. "lang": "JavaScript",
  274. "source": textwrap.dedent(
  275. """
  276. const { r2rClient } = require("r2r-js");
  277. const client = new r2rClient();
  278. function main() {
  279. const response = await client.collections.list();
  280. }
  281. main();
  282. """
  283. ),
  284. },
  285. {
  286. "lang": "CLI",
  287. "source": textwrap.dedent(
  288. """
  289. r2r collections list
  290. """
  291. ),
  292. },
  293. {
  294. "lang": "cURL",
  295. "source": textwrap.dedent(
  296. """
  297. curl -X GET "https://api.example.com/v3/collections?offset=0&limit=10&name=Sample" \\
  298. -H "Authorization: Bearer YOUR_API_KEY"
  299. """
  300. ),
  301. },
  302. ]
  303. },
  304. )
  305. @self.base_endpoint
  306. async def list_collections(
  307. ids: list[str] = Query(
  308. [],
  309. description="A list of collection IDs to retrieve. If not provided, all collections will be returned.",
  310. ),
  311. offset: int = Query(
  312. 0,
  313. ge=0,
  314. description="Specifies the number of objects to skip. Defaults to 0.",
  315. ),
  316. limit: int = Query(
  317. 100,
  318. ge=1,
  319. le=1000,
  320. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  321. ),
  322. auth_user=Depends(self.providers.auth.auth_wrapper()),
  323. ) -> WrappedCollectionsResponse:
  324. """
  325. Returns a paginated list of collections the authenticated user has access to.
  326. Results can be filtered by providing specific collection IDs. Regular users will only see
  327. collections they own or have access to. Superusers can see all collections.
  328. The collections are returned in order of last modification, with most recent first.
  329. """
  330. requesting_user_id = (
  331. None if auth_user.is_superuser else [auth_user.id]
  332. )
  333. collection_uuids = [UUID(collection_id) for collection_id in ids]
  334. collections_overview_response = (
  335. await self.services.management.collections_overview(
  336. user_ids=requesting_user_id,
  337. collection_ids=collection_uuids,
  338. offset=offset,
  339. limit=limit,
  340. )
  341. )
  342. return ( # type: ignore
  343. collections_overview_response["results"],
  344. {
  345. "total_entries": collections_overview_response[
  346. "total_entries"
  347. ]
  348. },
  349. )
  350. @self.router.get(
  351. "/collections/{id}",
  352. summary="Get collection details",
  353. dependencies=[Depends(self.rate_limit_dependency)],
  354. openapi_extra={
  355. "x-codeSamples": [
  356. {
  357. "lang": "Python",
  358. "source": textwrap.dedent(
  359. """
  360. from r2r import R2RClient
  361. client = R2RClient()
  362. # when using auth, do client.login(...)
  363. result = client.collections.retrieve("123e4567-e89b-12d3-a456-426614174000")
  364. """
  365. ),
  366. },
  367. {
  368. "lang": "JavaScript",
  369. "source": textwrap.dedent(
  370. """
  371. const { r2rClient } = require("r2r-js");
  372. const client = new r2rClient();
  373. function main() {
  374. const response = await client.collections.retrieve({id: "123e4567-e89b-12d3-a456-426614174000"});
  375. }
  376. main();
  377. """
  378. ),
  379. },
  380. {
  381. "lang": "CLI",
  382. "source": textwrap.dedent(
  383. """
  384. r2r collections retrieve 123e4567-e89b-12d3-a456-426614174000
  385. """
  386. ),
  387. },
  388. {
  389. "lang": "cURL",
  390. "source": textwrap.dedent(
  391. """
  392. curl -X GET "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000" \\
  393. -H "Authorization: Bearer YOUR_API_KEY"
  394. """
  395. ),
  396. },
  397. ]
  398. },
  399. )
  400. @self.base_endpoint
  401. async def get_collection(
  402. id: UUID = Path(
  403. ..., description="The unique identifier of the collection"
  404. ),
  405. auth_user=Depends(self.providers.auth.auth_wrapper()),
  406. ) -> WrappedCollectionResponse:
  407. """
  408. Get details of a specific collection.
  409. This endpoint retrieves detailed information about a single collection identified by its UUID.
  410. The user must have access to the collection to view its details.
  411. """
  412. await authorize_collection_action(
  413. auth_user, id, CollectionAction.VIEW, self.services
  414. )
  415. collections_overview_response = (
  416. await self.services.management.collections_overview(
  417. user_ids=None,
  418. collection_ids=[id],
  419. offset=0,
  420. limit=1,
  421. )
  422. )
  423. overview = collections_overview_response["results"]
  424. if len(overview) == 0:
  425. raise R2RException(
  426. "The specified collection does not exist.",
  427. 404,
  428. )
  429. return overview[0]
  430. @self.router.post(
  431. "/collections/{id}",
  432. summary="Update collection",
  433. dependencies=[Depends(self.rate_limit_dependency)],
  434. openapi_extra={
  435. "x-codeSamples": [
  436. {
  437. "lang": "Python",
  438. "source": textwrap.dedent(
  439. """
  440. from r2r import R2RClient
  441. client = R2RClient()
  442. # when using auth, do client.login(...)
  443. result = client.collections.update(
  444. "123e4567-e89b-12d3-a456-426614174000",
  445. name="Updated Collection Name",
  446. description="Updated description"
  447. )
  448. """
  449. ),
  450. },
  451. {
  452. "lang": "JavaScript",
  453. "source": textwrap.dedent(
  454. """
  455. const { r2rClient } = require("r2r-js");
  456. const client = new r2rClient();
  457. function main() {
  458. const response = await client.collections.update({
  459. id: "123e4567-e89b-12d3-a456-426614174000",
  460. name: "Updated Collection Name",
  461. description: "Updated description"
  462. });
  463. }
  464. main();
  465. """
  466. ),
  467. },
  468. {
  469. "lang": "cURL",
  470. "source": textwrap.dedent(
  471. """
  472. curl -X POST "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000" \\
  473. -H "Content-Type: application/json" \\
  474. -H "Authorization: Bearer YOUR_API_KEY" \\
  475. -d '{"name": "Updated Collection Name", "description": "Updated description"}'
  476. """
  477. ),
  478. },
  479. ]
  480. },
  481. )
  482. @self.base_endpoint
  483. async def update_collection(
  484. id: UUID = Path(
  485. ...,
  486. description="The unique identifier of the collection to update",
  487. ),
  488. name: Optional[str] = Body(
  489. None, description="The name of the collection"
  490. ),
  491. description: Optional[str] = Body(
  492. None, description="An optional description of the collection"
  493. ),
  494. generate_description: Optional[bool] = Body(
  495. False,
  496. description="Whether to generate a new synthetic description for the collection",
  497. ),
  498. auth_user=Depends(self.providers.auth.auth_wrapper()),
  499. ) -> WrappedCollectionResponse:
  500. """
  501. Update an existing collection's configuration.
  502. This endpoint allows updating the name and description of an existing collection.
  503. The user must have appropriate permissions to modify the collection.
  504. """
  505. await authorize_collection_action(
  506. auth_user, id, CollectionAction.EDIT, self.services
  507. )
  508. if generate_description and description is not None:
  509. raise R2RException(
  510. "Cannot provide both a description and request to synthetically generate a new one.",
  511. 400,
  512. )
  513. return await self.services.management.update_collection( # type: ignore
  514. id,
  515. name=name,
  516. description=description,
  517. generate_description=generate_description,
  518. )
  519. @self.router.delete(
  520. "/collections/{id}",
  521. summary="Delete collection",
  522. dependencies=[Depends(self.rate_limit_dependency)],
  523. openapi_extra={
  524. "x-codeSamples": [
  525. {
  526. "lang": "Python",
  527. "source": textwrap.dedent(
  528. """
  529. from r2r import R2RClient
  530. client = R2RClient()
  531. # when using auth, do client.login(...)
  532. result = client.collections.delete("123e4567-e89b-12d3-a456-426614174000")
  533. """
  534. ),
  535. },
  536. {
  537. "lang": "JavaScript",
  538. "source": textwrap.dedent(
  539. """
  540. const { r2rClient } = require("r2r-js");
  541. const client = new r2rClient();
  542. function main() {
  543. const response = await client.collections.delete({id: "123e4567-e89b-12d3-a456-426614174000"});
  544. }
  545. main();
  546. """
  547. ),
  548. },
  549. {
  550. "lang": "CLI",
  551. "source": textwrap.dedent(
  552. """
  553. r2r collections delete 123e4567-e89b-12d3-a456-426614174000
  554. """
  555. ),
  556. },
  557. {
  558. "lang": "cURL",
  559. "source": textwrap.dedent(
  560. """
  561. curl -X DELETE "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000" \\
  562. -H "Authorization: Bearer YOUR_API_KEY"
  563. """
  564. ),
  565. },
  566. ]
  567. },
  568. )
  569. @self.base_endpoint
  570. async def delete_collection(
  571. id: UUID = Path(
  572. ...,
  573. description="The unique identifier of the collection to delete",
  574. ),
  575. auth_user=Depends(self.providers.auth.auth_wrapper()),
  576. ) -> WrappedBooleanResponse:
  577. """
  578. Delete an existing collection.
  579. This endpoint allows deletion of a collection identified by its UUID.
  580. The user must have appropriate permissions to delete the collection.
  581. Deleting a collection removes all associations but does not delete the documents within it.
  582. """
  583. await authorize_collection_action(
  584. auth_user, id, CollectionAction.DELETE, self.services
  585. )
  586. await self.services.management.delete_collection(id)
  587. return GenericBooleanResponse(success=True) # type: ignore
  588. @self.router.post(
  589. "/collections/{id}/documents/{document_id}",
  590. summary="Add document to collection",
  591. dependencies=[Depends(self.rate_limit_dependency)],
  592. openapi_extra={
  593. "x-codeSamples": [
  594. {
  595. "lang": "Python",
  596. "source": textwrap.dedent(
  597. """
  598. from r2r import R2RClient
  599. client = R2RClient()
  600. # when using auth, do client.login(...)
  601. result = client.collections.add_document(
  602. "123e4567-e89b-12d3-a456-426614174000",
  603. "456e789a-b12c-34d5-e678-901234567890"
  604. )
  605. """
  606. ),
  607. },
  608. {
  609. "lang": "JavaScript",
  610. "source": textwrap.dedent(
  611. """
  612. const { r2rClient } = require("r2r-js");
  613. const client = new r2rClient();
  614. function main() {
  615. const response = await client.collections.addDocument({
  616. id: "123e4567-e89b-12d3-a456-426614174000"
  617. documentId: "456e789a-b12c-34d5-e678-901234567890"
  618. });
  619. }
  620. main();
  621. """
  622. ),
  623. },
  624. {
  625. "lang": "cURL",
  626. "source": textwrap.dedent(
  627. """
  628. curl -X POST "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000/documents/456e789a-b12c-34d5-e678-901234567890" \\
  629. -H "Authorization: Bearer YOUR_API_KEY"
  630. """
  631. ),
  632. },
  633. ]
  634. },
  635. )
  636. @self.base_endpoint
  637. async def add_document_to_collection(
  638. id: UUID = Path(...),
  639. document_id: UUID = Path(...),
  640. auth_user=Depends(self.providers.auth.auth_wrapper()),
  641. ) -> WrappedGenericMessageResponse:
  642. """
  643. Add a document to a collection.
  644. """
  645. await authorize_collection_action(
  646. auth_user, id, CollectionAction.ADD_DOCUMENT, self.services
  647. )
  648. return (
  649. await self.services.management.assign_document_to_collection(
  650. document_id, id
  651. )
  652. )
  653. @self.router.get(
  654. "/collections/{id}/documents",
  655. summary="List documents in collection",
  656. dependencies=[Depends(self.rate_limit_dependency)],
  657. openapi_extra={
  658. "x-codeSamples": [
  659. {
  660. "lang": "Python",
  661. "source": textwrap.dedent(
  662. """
  663. from r2r import R2RClient
  664. client = R2RClient()
  665. # when using auth, do client.login(...)
  666. result = client.collections.list_documents(
  667. "123e4567-e89b-12d3-a456-426614174000",
  668. offset=0,
  669. limit=10,
  670. )
  671. """
  672. ),
  673. },
  674. {
  675. "lang": "JavaScript",
  676. "source": textwrap.dedent(
  677. """
  678. const { r2rClient } = require("r2r-js");
  679. const client = new r2rClient();
  680. function main() {
  681. const response = await client.collections.listDocuments({id: "123e4567-e89b-12d3-a456-426614174000"});
  682. }
  683. main();
  684. """
  685. ),
  686. },
  687. {
  688. "lang": "CLI",
  689. "source": textwrap.dedent(
  690. """
  691. r2r collections list-documents 123e4567-e89b-12d3-a456-426614174000
  692. """
  693. ),
  694. },
  695. {
  696. "lang": "cURL",
  697. "source": textwrap.dedent(
  698. """
  699. curl -X GET "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000/documents?offset=0&limit=10" \\
  700. -H "Authorization: Bearer YOUR_API_KEY"
  701. """
  702. ),
  703. },
  704. ]
  705. },
  706. )
  707. @self.base_endpoint
  708. async def get_collection_documents(
  709. id: UUID = Path(
  710. ..., description="The unique identifier of the collection"
  711. ),
  712. offset: int = Query(
  713. 0,
  714. ge=0,
  715. description="Specifies the number of objects to skip. Defaults to 0.",
  716. ),
  717. limit: int = Query(
  718. 100,
  719. ge=1,
  720. le=1000,
  721. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  722. ),
  723. auth_user=Depends(self.providers.auth.auth_wrapper()),
  724. ) -> WrappedDocumentsResponse:
  725. """
  726. Get all documents in a collection with pagination and sorting options.
  727. This endpoint retrieves a paginated list of documents associated with a specific collection.
  728. It supports sorting options to customize the order of returned documents.
  729. """
  730. await authorize_collection_action(
  731. auth_user, id, CollectionAction.VIEW, self.services
  732. )
  733. documents_in_collection_response = (
  734. await self.services.management.documents_in_collection(
  735. id, offset, limit
  736. )
  737. )
  738. return documents_in_collection_response["results"], { # type: ignore
  739. "total_entries": documents_in_collection_response[
  740. "total_entries"
  741. ]
  742. }
  743. @self.router.delete(
  744. "/collections/{id}/documents/{document_id}",
  745. summary="Remove document from collection",
  746. dependencies=[Depends(self.rate_limit_dependency)],
  747. openapi_extra={
  748. "x-codeSamples": [
  749. {
  750. "lang": "Python",
  751. "source": textwrap.dedent(
  752. """
  753. from r2r import R2RClient
  754. client = R2RClient()
  755. # when using auth, do client.login(...)
  756. result = client.collections.remove_document(
  757. "123e4567-e89b-12d3-a456-426614174000",
  758. "456e789a-b12c-34d5-e678-901234567890"
  759. )
  760. """
  761. ),
  762. },
  763. {
  764. "lang": "JavaScript",
  765. "source": textwrap.dedent(
  766. """
  767. const { r2rClient } = require("r2r-js");
  768. const client = new r2rClient();
  769. function main() {
  770. const response = await client.collections.removeDocument({
  771. id: "123e4567-e89b-12d3-a456-426614174000"
  772. documentId: "456e789a-b12c-34d5-e678-901234567890"
  773. });
  774. }
  775. main();
  776. """
  777. ),
  778. },
  779. {
  780. "lang": "cURL",
  781. "source": textwrap.dedent(
  782. """
  783. curl -X DELETE "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000/documents/456e789a-b12c-34d5-e678-901234567890" \\
  784. -H "Authorization: Bearer YOUR_API_KEY"
  785. """
  786. ),
  787. },
  788. ]
  789. },
  790. )
  791. @self.base_endpoint
  792. async def remove_document_from_collection(
  793. id: UUID = Path(
  794. ..., description="The unique identifier of the collection"
  795. ),
  796. document_id: UUID = Path(
  797. ...,
  798. description="The unique identifier of the document to remove",
  799. ),
  800. auth_user=Depends(self.providers.auth.auth_wrapper()),
  801. ) -> WrappedBooleanResponse:
  802. """
  803. Remove a document from a collection.
  804. This endpoint removes the association between a document and a collection.
  805. It does not delete the document itself. The user must have permissions to modify the collection.
  806. """
  807. await authorize_collection_action(
  808. auth_user, id, CollectionAction.REMOVE_DOCUMENT, self.services
  809. )
  810. await self.services.management.remove_document_from_collection(
  811. document_id, id
  812. )
  813. return GenericBooleanResponse(success=True) # type: ignore
  814. @self.router.get(
  815. "/collections/{id}/users",
  816. summary="List users in collection",
  817. dependencies=[Depends(self.rate_limit_dependency)],
  818. openapi_extra={
  819. "x-codeSamples": [
  820. {
  821. "lang": "Python",
  822. "source": textwrap.dedent(
  823. """
  824. from r2r import R2RClient
  825. client = R2RClient()
  826. # when using auth, do client.login(...)
  827. result = client.collections.list_users(
  828. "123e4567-e89b-12d3-a456-426614174000",
  829. offset=0,
  830. limit=10,
  831. )
  832. """
  833. ),
  834. },
  835. {
  836. "lang": "JavaScript",
  837. "source": textwrap.dedent(
  838. """
  839. const { r2rClient } = require("r2r-js");
  840. const client = new r2rClient();
  841. function main() {
  842. const response = await client.collections.listUsers({
  843. id: "123e4567-e89b-12d3-a456-426614174000"
  844. });
  845. }
  846. main();
  847. """
  848. ),
  849. },
  850. {
  851. "lang": "CLI",
  852. "source": textwrap.dedent(
  853. """
  854. r2r collections list-users 123e4567-e89b-12d3-a456-426614174000
  855. """
  856. ),
  857. },
  858. {
  859. "lang": "cURL",
  860. "source": textwrap.dedent(
  861. """
  862. curl -X GET "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000/users?offset=0&limit=10" \\
  863. -H "Authorization: Bearer YOUR_API_KEY"
  864. """
  865. ),
  866. },
  867. ]
  868. },
  869. )
  870. @self.base_endpoint
  871. async def get_collection_users(
  872. id: UUID = Path(
  873. ..., description="The unique identifier of the collection"
  874. ),
  875. offset: int = Query(
  876. 0,
  877. ge=0,
  878. description="Specifies the number of objects to skip. Defaults to 0.",
  879. ),
  880. limit: int = Query(
  881. 100,
  882. ge=1,
  883. le=1000,
  884. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  885. ),
  886. auth_user=Depends(self.providers.auth.auth_wrapper()),
  887. ) -> WrappedUsersResponse:
  888. """
  889. Get all users in a collection with pagination and sorting options.
  890. This endpoint retrieves a paginated list of users who have access to a specific collection.
  891. It supports sorting options to customize the order of returned users.
  892. """
  893. await authorize_collection_action(
  894. auth_user, id, CollectionAction.VIEW, self.services
  895. )
  896. users_in_collection_response = (
  897. await self.services.management.get_users_in_collection(
  898. collection_id=id,
  899. offset=offset,
  900. limit=min(max(limit, 1), 1000),
  901. )
  902. )
  903. return users_in_collection_response["results"], { # type: ignore
  904. "total_entries": users_in_collection_response["total_entries"]
  905. }
  906. @self.router.post(
  907. "/collections/{id}/users/{user_id}",
  908. summary="Add user to collection",
  909. dependencies=[Depends(self.rate_limit_dependency)],
  910. openapi_extra={
  911. "x-codeSamples": [
  912. {
  913. "lang": "Python",
  914. "source": textwrap.dedent(
  915. """
  916. from r2r import R2RClient
  917. client = R2RClient()
  918. # when using auth, do client.login(...)
  919. result = client.collections.add_user(
  920. "123e4567-e89b-12d3-a456-426614174000",
  921. "789a012b-c34d-5e6f-g789-012345678901"
  922. )
  923. """
  924. ),
  925. },
  926. {
  927. "lang": "JavaScript",
  928. "source": textwrap.dedent(
  929. """
  930. const { r2rClient } = require("r2r-js");
  931. const client = new r2rClient();
  932. function main() {
  933. const response = await client.collections.addUser({
  934. id: "123e4567-e89b-12d3-a456-426614174000"
  935. userId: "789a012b-c34d-5e6f-g789-012345678901"
  936. });
  937. }
  938. main();
  939. """
  940. ),
  941. },
  942. {
  943. "lang": "cURL",
  944. "source": textwrap.dedent(
  945. """
  946. curl -X POST "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000/users/789a012b-c34d-5e6f-g789-012345678901" \\
  947. -H "Authorization: Bearer YOUR_API_KEY"
  948. """
  949. ),
  950. },
  951. ]
  952. },
  953. )
  954. @self.base_endpoint
  955. async def add_user_to_collection(
  956. id: UUID = Path(
  957. ..., description="The unique identifier of the collection"
  958. ),
  959. user_id: UUID = Path(
  960. ..., description="The unique identifier of the user to add"
  961. ),
  962. auth_user=Depends(self.providers.auth.auth_wrapper()),
  963. ) -> WrappedBooleanResponse:
  964. """
  965. Add a user to a collection.
  966. This endpoint grants a user access to a specific collection.
  967. The authenticated user must have admin permissions for the collection to add new users.
  968. """
  969. await authorize_collection_action(
  970. auth_user, id, CollectionAction.MANAGE_USERS, self.services
  971. )
  972. result = await self.services.management.add_user_to_collection(
  973. user_id, id
  974. )
  975. return GenericBooleanResponse(success=result) # type: ignore
  976. @self.router.delete(
  977. "/collections/{id}/users/{user_id}",
  978. summary="Remove user from collection",
  979. dependencies=[Depends(self.rate_limit_dependency)],
  980. openapi_extra={
  981. "x-codeSamples": [
  982. {
  983. "lang": "Python",
  984. "source": textwrap.dedent(
  985. """
  986. from r2r import R2RClient
  987. client = R2RClient()
  988. # when using auth, do client.login(...)
  989. result = client.collections.remove_user(
  990. "123e4567-e89b-12d3-a456-426614174000",
  991. "789a012b-c34d-5e6f-g789-012345678901"
  992. )
  993. """
  994. ),
  995. },
  996. {
  997. "lang": "JavaScript",
  998. "source": textwrap.dedent(
  999. """
  1000. const { r2rClient } = require("r2r-js");
  1001. const client = new r2rClient();
  1002. function main() {
  1003. const response = await client.collections.removeUser({
  1004. id: "123e4567-e89b-12d3-a456-426614174000"
  1005. userId: "789a012b-c34d-5e6f-g789-012345678901"
  1006. });
  1007. }
  1008. main();
  1009. """
  1010. ),
  1011. },
  1012. {
  1013. "lang": "cURL",
  1014. "source": textwrap.dedent(
  1015. """
  1016. curl -X DELETE "https://api.example.com/v3/collections/123e4567-e89b-12d3-a456-426614174000/users/789a012b-c34d-5e6f-g789-012345678901" \\
  1017. -H "Authorization: Bearer YOUR_API_KEY"
  1018. """
  1019. ),
  1020. },
  1021. ]
  1022. },
  1023. )
  1024. @self.base_endpoint
  1025. async def remove_user_from_collection(
  1026. id: UUID = Path(
  1027. ..., description="The unique identifier of the collection"
  1028. ),
  1029. user_id: UUID = Path(
  1030. ..., description="The unique identifier of the user to remove"
  1031. ),
  1032. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1033. ) -> WrappedBooleanResponse:
  1034. """
  1035. Remove a user from a collection.
  1036. This endpoint revokes a user's access to a specific collection.
  1037. The authenticated user must have admin permissions for the collection to remove users.
  1038. """
  1039. await authorize_collection_action(
  1040. auth_user, id, CollectionAction.MANAGE_USERS, self.services
  1041. )
  1042. result = (
  1043. await self.services.management.remove_user_from_collection(
  1044. user_id, id
  1045. )
  1046. )
  1047. return GenericBooleanResponse(success=True) # type: ignore
  1048. @self.router.post(
  1049. "/collections/{id}/extract",
  1050. summary="Extract entities and relationships",
  1051. dependencies=[Depends(self.rate_limit_dependency)],
  1052. openapi_extra={
  1053. "x-codeSamples": [
  1054. {
  1055. "lang": "Python",
  1056. "source": textwrap.dedent(
  1057. """
  1058. from r2r import R2RClient
  1059. client = R2RClient()
  1060. # when using auth, do client.login(...)
  1061. result = client.documents.extract(
  1062. id="9fbe403b-c11c-5aae-8ade-ef22980c3ad1"
  1063. )
  1064. """
  1065. ),
  1066. },
  1067. ],
  1068. },
  1069. )
  1070. @self.base_endpoint
  1071. async def extract(
  1072. id: UUID = Path(
  1073. ...,
  1074. description="The ID of the document to extract entities and relationships from.",
  1075. ),
  1076. run_type: KGRunType = Query(
  1077. default=KGRunType.RUN,
  1078. description="Whether to return an estimate of the creation cost or to actually extract the document.",
  1079. ),
  1080. settings: Optional[KGCreationSettings] = Body(
  1081. default=None,
  1082. description="Settings for the entities and relationships extraction process.",
  1083. ),
  1084. run_with_orchestration: Optional[bool] = Query(
  1085. default=True,
  1086. description="Whether to run the entities and relationships extraction process with orchestration.",
  1087. ),
  1088. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1089. ):
  1090. """
  1091. Extracts entities and relationships from a document.
  1092. The entities and relationships extraction process involves:
  1093. 1. Parsing documents into semantic chunks
  1094. 2. Extracting entities and relationships using LLMs
  1095. """
  1096. await authorize_collection_action(
  1097. auth_user, id, CollectionAction.EDIT, self.services
  1098. )
  1099. settings = settings.dict() if settings else None # type: ignore
  1100. if not auth_user.is_superuser:
  1101. logger.warning("Implement permission checks here.")
  1102. # If no run type is provided, default to estimate
  1103. if not run_type:
  1104. run_type = KGRunType.ESTIMATE
  1105. # Apply runtime settings overrides
  1106. server_graph_creation_settings = (
  1107. self.providers.database.config.graph_creation_settings
  1108. )
  1109. if settings:
  1110. server_graph_creation_settings = update_settings_from_dict(
  1111. server_settings=server_graph_creation_settings,
  1112. settings_dict=settings, # type: ignore
  1113. )
  1114. if run_with_orchestration:
  1115. workflow_input = {
  1116. "collection_id": str(id),
  1117. "graph_creation_settings": server_graph_creation_settings.model_dump_json(),
  1118. "user": auth_user.json(),
  1119. }
  1120. return await self.providers.orchestration.run_workflow( # type: ignore
  1121. "extract-triples", {"request": workflow_input}, {}
  1122. )
  1123. else:
  1124. from core.main.orchestration import simple_kg_factory
  1125. logger.info("Running extract-triples without orchestration.")
  1126. simple_kg = simple_kg_factory(self.services.graph)
  1127. await simple_kg["extract-triples"](workflow_input) # type: ignore
  1128. return { # type: ignore
  1129. "message": "Graph created successfully.",
  1130. "task_id": None,
  1131. }