graph_router.py 72 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886
  1. import logging
  2. import textwrap
  3. from typing import Optional
  4. from uuid import UUID
  5. from fastapi import Body, Depends, Path, Query
  6. from core.base import KGEnrichmentStatus, R2RException, Workflow
  7. from core.base.abstractions import KGRunType
  8. from core.base.api.models import (
  9. GenericBooleanResponse,
  10. WrappedBooleanResponse,
  11. WrappedCommunitiesResponse,
  12. WrappedCommunityResponse,
  13. WrappedEntitiesResponse,
  14. WrappedEntityResponse,
  15. WrappedGraphResponse,
  16. WrappedGraphsResponse,
  17. WrappedRelationshipResponse,
  18. WrappedRelationshipsResponse,
  19. )
  20. from core.utils import (
  21. generate_default_user_collection_id,
  22. update_settings_from_dict,
  23. )
  24. from ...abstractions import R2RProviders, R2RServices
  25. from .base_router import BaseRouterV3
  26. logger = logging.getLogger()
  27. class GraphRouter(BaseRouterV3):
  28. def __init__(
  29. self,
  30. providers: R2RProviders,
  31. services: R2RServices,
  32. ):
  33. super().__init__(providers, services)
  34. self._register_workflows()
  35. def _register_workflows(self):
  36. workflow_messages = {}
  37. if self.providers.orchestration.config.provider == "hatchet":
  38. workflow_messages["extract-triples"] = (
  39. "Document extraction task queued successfully."
  40. )
  41. workflow_messages["build-communities"] = (
  42. "Graph enrichment task queued successfully."
  43. )
  44. workflow_messages["entity-deduplication"] = (
  45. "KG Entity Deduplication task queued successfully."
  46. )
  47. else:
  48. workflow_messages["extract-triples"] = (
  49. "Document entities and relationships extracted successfully."
  50. )
  51. workflow_messages["build-communities"] = (
  52. "Graph communities created successfully."
  53. )
  54. workflow_messages["entity-deduplication"] = (
  55. "KG Entity Deduplication completed successfully."
  56. )
  57. self.providers.orchestration.register_workflows(
  58. Workflow.KG,
  59. self.services.graph,
  60. workflow_messages,
  61. )
  62. async def _deduplicate_entities(
  63. self,
  64. collection_id: UUID,
  65. settings,
  66. run_type: Optional[KGRunType] = KGRunType.ESTIMATE,
  67. run_with_orchestration: bool = True,
  68. auth_user=None,
  69. ):
  70. """Deduplicates entities in the knowledge graph using LLM-based analysis.
  71. The deduplication process:
  72. 1. Groups potentially duplicate entities by name/type
  73. 2. Uses LLM analysis to determine if entities refer to same thing
  74. 3. Merges duplicate entities while preserving relationships
  75. 4. Updates all references to use canonical entity IDs
  76. Args:
  77. id (UUID): Graph containing the entities
  78. settings (dict, optional): Deduplication settings including:
  79. - graph_entity_deduplication_type (str): Deduplication method (e.g. "by_name")
  80. - graph_entity_deduplication_prompt (str): Custom prompt for analysis
  81. - max_description_input_length (int): Max chars for entity descriptions
  82. - generation_config (dict): LLM generation parameters
  83. run_type (KGRunType): Whether to estimate cost or run deduplication
  84. run_with_orchestration (bool): Whether to run async with task queue
  85. auth_user: Authenticated user making request
  86. Returns:
  87. Result containing:
  88. message (str): Status message
  89. task_id (UUID): Async task ID if run with orchestration
  90. Raises:
  91. R2RException: If user unauthorized or deduplication fails
  92. """
  93. if not auth_user.is_superuser:
  94. raise R2RException(
  95. "Only superusers can deduplicate a graphs entities", 403
  96. )
  97. server_settings = (
  98. self.providers.database.config.graph_entity_deduplication_settings
  99. )
  100. if settings:
  101. server_settings = update_settings_from_dict(
  102. server_settings, settings
  103. )
  104. # Return cost estimate if requested
  105. if run_type == KGRunType.ESTIMATE:
  106. return await self.services.graph.get_deduplication_estimate(
  107. collection_id, server_settings
  108. )
  109. workflow_input = {
  110. "graph_id": str(collection_id),
  111. "graph_entity_deduplication_settings": server_settings.model_dump_json(),
  112. "user": auth_user.model_dump_json(),
  113. }
  114. if run_with_orchestration:
  115. return await self.providers.orchestration.run_workflow( # type: ignore
  116. "entity-deduplication", {"request": workflow_input}, {}
  117. )
  118. else:
  119. from core.main.orchestration import simple_kg_factory
  120. simple_kg = simple_kg_factory(self.services.graph)
  121. await simple_kg["entity-deduplication"](workflow_input)
  122. return { # type: ignore
  123. "message": "Entity deduplication completed successfully.",
  124. "task_id": None,
  125. }
  126. async def _get_collection_id(
  127. self, collection_id: Optional[UUID], auth_user
  128. ) -> UUID:
  129. """Helper method to get collection ID, using default if none provided"""
  130. if collection_id is None:
  131. return generate_default_user_collection_id(auth_user.id)
  132. return collection_id
  133. def _setup_routes(self):
  134. @self.router.get(
  135. "/graphs",
  136. dependencies=[Depends(self.rate_limit_dependency)],
  137. summary="List graphs",
  138. openapi_extra={
  139. "x-codeSamples": [
  140. { # TODO: Verify
  141. "lang": "Python",
  142. "source": textwrap.dedent(
  143. """
  144. from r2r import R2RClient
  145. client = R2RClient()
  146. # when using auth, do client.login(...)
  147. response = client.graphs.list()
  148. """
  149. ),
  150. },
  151. {
  152. "lang": "JavaScript",
  153. "source": textwrap.dedent(
  154. """
  155. const { r2rClient } = require("r2r-js");
  156. const client = new r2rClient();
  157. function main() {
  158. const response = await client.graphs.list({});
  159. }
  160. main();
  161. """
  162. ),
  163. },
  164. ]
  165. },
  166. )
  167. @self.base_endpoint
  168. async def list_graphs(
  169. collection_ids: list[str] = Query(
  170. [],
  171. description="A list of graph IDs to retrieve. If not provided, all graphs will be returned.",
  172. ),
  173. offset: int = Query(
  174. 0,
  175. ge=0,
  176. description="Specifies the number of objects to skip. Defaults to 0.",
  177. ),
  178. limit: int = Query(
  179. 100,
  180. ge=1,
  181. le=1000,
  182. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  183. ),
  184. auth_user=Depends(self.providers.auth.auth_wrapper()),
  185. ) -> WrappedGraphsResponse:
  186. """
  187. Returns a paginated list of graphs the authenticated user has access to.
  188. Results can be filtered by providing specific graph IDs. Regular users will only see
  189. graphs they own or have access to. Superusers can see all graphs.
  190. The graphs are returned in order of last modification, with most recent first.
  191. """
  192. requesting_user_id = (
  193. None if auth_user.is_superuser else [auth_user.id]
  194. )
  195. graph_uuids = [UUID(graph_id) for graph_id in collection_ids]
  196. list_graphs_response = await self.services.graph.list_graphs(
  197. # user_ids=requesting_user_id,
  198. graph_ids=graph_uuids,
  199. offset=offset,
  200. limit=limit,
  201. )
  202. return ( # type: ignore
  203. list_graphs_response["results"],
  204. {"total_entries": list_graphs_response["total_entries"]},
  205. )
  206. @self.router.get(
  207. "/graphs/{collection_id}",
  208. dependencies=[Depends(self.rate_limit_dependency)],
  209. summary="Retrieve graph details",
  210. openapi_extra={
  211. "x-codeSamples": [
  212. {
  213. "lang": "Python",
  214. "source": textwrap.dedent(
  215. """
  216. from r2r import R2RClient
  217. client = R2RClient()
  218. # when using auth, do client.login(...)
  219. response = client.graphs.get(
  220. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  221. )"""
  222. ),
  223. },
  224. {
  225. "lang": "JavaScript",
  226. "source": textwrap.dedent(
  227. """
  228. const { r2rClient } = require("r2r-js");
  229. const client = new r2rClient();
  230. function main() {
  231. const response = await client.graphs.retrieve({
  232. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  233. });
  234. }
  235. main();
  236. """
  237. ),
  238. },
  239. {
  240. "lang": "cURL",
  241. "source": textwrap.dedent(
  242. """
  243. curl -X GET "https://api.example.com/v3/graphs/d09dedb1-b2ab-48a5-b950-6e1f464d83e7" \\
  244. -H "Authorization: Bearer YOUR_API_KEY" """
  245. ),
  246. },
  247. ]
  248. },
  249. )
  250. @self.base_endpoint
  251. async def get_graph(
  252. collection_id: UUID = Path(...),
  253. auth_user=Depends(self.providers.auth.auth_wrapper()),
  254. ) -> WrappedGraphResponse:
  255. """
  256. Retrieves detailed information about a specific graph by ID.
  257. """
  258. if (
  259. # not auth_user.is_superuser
  260. collection_id
  261. not in auth_user.collection_ids
  262. ):
  263. raise R2RException(
  264. "The currently authenticated user does not have access to the specified collection associated with the given graph.",
  265. 403,
  266. )
  267. list_graphs_response = await self.services.graph.list_graphs(
  268. # user_ids=None,
  269. graph_ids=[collection_id],
  270. offset=0,
  271. limit=1,
  272. )
  273. return list_graphs_response["results"][0]
  274. @self.router.post(
  275. "/graphs/{collection_id}/communities/build",
  276. dependencies=[Depends(self.rate_limit_dependency)],
  277. )
  278. @self.base_endpoint
  279. async def build_communities(
  280. collection_id: UUID = Path(
  281. ..., description="The unique identifier of the collection"
  282. ),
  283. run_type: Optional[KGRunType] = Body(
  284. default=KGRunType.ESTIMATE,
  285. description="Run type for the graph enrichment process.",
  286. ),
  287. graph_enrichment_settings: Optional[dict] = Body(
  288. default=None,
  289. description="Settings for the graph enrichment process.",
  290. ),
  291. run_with_orchestration: Optional[bool] = Body(True),
  292. auth_user=Depends(self.providers.auth.auth_wrapper()),
  293. ):
  294. """
  295. Creates communities in the graph by analyzing entity relationships and similarities.
  296. Communities are created through the following process:
  297. 1. Analyzes entity relationships and metadata to build a similarity graph
  298. 2. Applies advanced community detection algorithms (e.g. Leiden) to identify densely connected groups
  299. 3. Creates hierarchical community structure with multiple granularity levels
  300. 4. Generates natural language summaries and statistical insights for each community
  301. The resulting communities can be used to:
  302. - Understand high-level graph structure and organization
  303. - Identify key entity groupings and their relationships
  304. - Navigate and explore the graph at different levels of detail
  305. - Generate insights about entity clusters and their characteristics
  306. The community detection process is configurable through settings like:
  307. - Community detection algorithm parameters
  308. - Summary generation prompt
  309. """
  310. if not auth_user.is_superuser:
  311. raise R2RException(
  312. "Only superusers can build communities", 403
  313. )
  314. if (
  315. # not auth_user.is_superuser
  316. collection_id
  317. not in auth_user.collection_ids
  318. ):
  319. raise R2RException(
  320. "The currently authenticated user does not have access to the collection associated with the given graph.",
  321. 403,
  322. )
  323. # If no collection ID is provided, use the default user collection
  324. # id = generate_default_user_collection_id(auth_user.id)
  325. # If no run type is provided, default to estimate
  326. if not run_type:
  327. run_type = KGRunType.ESTIMATE
  328. # Apply runtime settings overrides
  329. server_graph_enrichment_settings = (
  330. self.providers.database.config.graph_enrichment_settings
  331. )
  332. if graph_enrichment_settings:
  333. server_graph_enrichment_settings = update_settings_from_dict(
  334. server_graph_enrichment_settings, graph_enrichment_settings
  335. )
  336. workflow_input = {
  337. "collection_id": str(collection_id),
  338. "graph_enrichment_settings": server_graph_enrichment_settings.model_dump_json(),
  339. "user": auth_user.json(),
  340. }
  341. if run_with_orchestration:
  342. return await self.providers.orchestration.run_workflow( # type: ignore
  343. "build-communities", {"request": workflow_input}, {}
  344. )
  345. else:
  346. from core.main.orchestration import simple_kg_factory
  347. logger.info("Running build-communities without orchestration.")
  348. simple_kg = simple_kg_factory(self.services.graph)
  349. await simple_kg["build-communities"](workflow_input)
  350. return {
  351. "message": "Graph communities created successfully.",
  352. "task_id": None,
  353. }
  354. @self.router.post(
  355. "/graphs/{collection_id}/reset",
  356. dependencies=[Depends(self.rate_limit_dependency)],
  357. summary="Reset a graph back to the initial state.",
  358. openapi_extra={
  359. "x-codeSamples": [
  360. {
  361. "lang": "Python",
  362. "source": textwrap.dedent(
  363. """
  364. from r2r import R2RClient
  365. client = R2RClient()
  366. # when using auth, do client.login(...)
  367. response = client.graphs.reset(
  368. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  369. )"""
  370. ),
  371. },
  372. {
  373. "lang": "JavaScript",
  374. "source": textwrap.dedent(
  375. """
  376. const { r2rClient } = require("r2r-js");
  377. const client = new r2rClient();
  378. function main() {
  379. const response = await client.graphs.reset({
  380. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  381. });
  382. }
  383. main();
  384. """
  385. ),
  386. },
  387. {
  388. "lang": "cURL",
  389. "source": textwrap.dedent(
  390. """
  391. curl -X POST "https://api.example.com/v3/graphs/d09dedb1-b2ab-48a5-b950-6e1f464d83e7/reset" \\
  392. -H "Authorization: Bearer YOUR_API_KEY" """
  393. ),
  394. },
  395. ]
  396. },
  397. )
  398. @self.base_endpoint
  399. async def reset(
  400. collection_id: UUID = Path(...),
  401. auth_user=Depends(self.providers.auth.auth_wrapper()),
  402. ) -> WrappedBooleanResponse:
  403. """
  404. Deletes a graph and all its associated data.
  405. This endpoint permanently removes the specified graph along with all
  406. entities and relationships that belong to only this graph.
  407. The original source entities and relationships extracted from underlying documents are not deleted
  408. and are managed through the document lifecycle.
  409. """
  410. if not auth_user.is_superuser:
  411. raise R2RException("Only superusers can reset a graph", 403)
  412. if (
  413. # not auth_user.is_superuser
  414. collection_id
  415. not in auth_user.collection_ids
  416. ):
  417. raise R2RException(
  418. "The currently authenticated user does not have access to the collection associated with the given graph.",
  419. 403,
  420. )
  421. await self.services.graph.reset_graph_v3(id=collection_id)
  422. # await _pull(collection_id, auth_user)
  423. return GenericBooleanResponse(success=True) # type: ignore
  424. # update graph
  425. @self.router.post(
  426. "/graphs/{collection_id}",
  427. dependencies=[Depends(self.rate_limit_dependency)],
  428. summary="Update graph",
  429. openapi_extra={
  430. "x-codeSamples": [
  431. {
  432. "lang": "Python",
  433. "source": textwrap.dedent(
  434. """
  435. from r2r import R2RClient
  436. client = R2RClient()
  437. # when using auth, do client.login(...)
  438. response = client.graphs.update(
  439. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  440. graph={
  441. "name": "New Name",
  442. "description": "New Description"
  443. }
  444. )"""
  445. ),
  446. },
  447. {
  448. "lang": "JavaScript",
  449. "source": textwrap.dedent(
  450. """
  451. const { r2rClient } = require("r2r-js");
  452. const client = new r2rClient();
  453. function main() {
  454. const response = await client.graphs.update({
  455. collection_id: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  456. name: "New Name",
  457. description: "New Description",
  458. });
  459. }
  460. main();
  461. """
  462. ),
  463. },
  464. ]
  465. },
  466. )
  467. @self.base_endpoint
  468. async def update_graph(
  469. collection_id: UUID = Path(
  470. ...,
  471. description="The collection ID corresponding to the graph to update",
  472. ),
  473. name: Optional[str] = Body(
  474. None, description="The name of the graph"
  475. ),
  476. description: Optional[str] = Body(
  477. None, description="An optional description of the graph"
  478. ),
  479. auth_user=Depends(self.providers.auth.auth_wrapper()),
  480. ):
  481. """
  482. Update an existing graphs's configuration.
  483. This endpoint allows updating the name and description of an existing collection.
  484. The user must have appropriate permissions to modify the collection.
  485. """
  486. if not auth_user.is_superuser:
  487. raise R2RException(
  488. "Only superusers can update graph details", 403
  489. )
  490. if (
  491. not auth_user.is_superuser
  492. and id not in auth_user.collection_ids
  493. ):
  494. raise R2RException(
  495. "The currently authenticated user does not have access to the collection associated with the given graph.",
  496. 403,
  497. )
  498. return await self.services.graph.update_graph( # type: ignore
  499. collection_id,
  500. name=name,
  501. description=description,
  502. )
  503. @self.router.get(
  504. "/graphs/{collection_id}/entities",
  505. dependencies=[Depends(self.rate_limit_dependency)],
  506. openapi_extra={
  507. "x-codeSamples": [
  508. {
  509. "lang": "Python",
  510. "source": textwrap.dedent(
  511. """
  512. from r2r import R2RClient
  513. client = R2RClient()
  514. # when using auth, do client.login(...)
  515. response = client.graphs.list_entities(collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7")
  516. """
  517. ),
  518. },
  519. {
  520. "lang": "JavaScript",
  521. "source": textwrap.dedent(
  522. """
  523. const { r2rClient } = require("r2r-js");
  524. const client = new r2rClient();
  525. function main() {
  526. const response = await client.graphs.listEntities({
  527. collection_id: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  528. });
  529. }
  530. main();
  531. """
  532. ),
  533. },
  534. ],
  535. },
  536. )
  537. @self.base_endpoint
  538. async def get_entities(
  539. collection_id: UUID = Path(
  540. ...,
  541. description="The collection ID corresponding to the graph to list entities from.",
  542. ),
  543. offset: int = Query(
  544. 0,
  545. ge=0,
  546. description="Specifies the number of objects to skip. Defaults to 0.",
  547. ),
  548. limit: int = Query(
  549. 100,
  550. ge=1,
  551. le=1000,
  552. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  553. ),
  554. auth_user=Depends(self.providers.auth.auth_wrapper()),
  555. ) -> WrappedEntitiesResponse:
  556. """Lists all entities in the graph with pagination support."""
  557. if (
  558. # not auth_user.is_superuser
  559. collection_id
  560. not in auth_user.collection_ids
  561. ):
  562. raise R2RException(
  563. "The currently authenticated user does not have access to the collection associated with the given graph.",
  564. 403,
  565. )
  566. entities, count = await self.services.graph.get_entities(
  567. parent_id=collection_id,
  568. offset=offset,
  569. limit=limit,
  570. )
  571. return entities, { # type: ignore
  572. "total_entries": count,
  573. }
  574. @self.router.post(
  575. "/graphs/{collection_id}/entities",
  576. dependencies=[Depends(self.rate_limit_dependency)],
  577. )
  578. @self.base_endpoint
  579. async def create_entity(
  580. collection_id: UUID = Path(
  581. ...,
  582. description="The collection ID corresponding to the graph to add the entity to.",
  583. ),
  584. name: str = Body(
  585. ..., description="The name of the entity to create."
  586. ),
  587. description: str = Body(
  588. ..., description="The description of the entity to create."
  589. ),
  590. category: Optional[str] = Body(
  591. None, description="The category of the entity to create."
  592. ),
  593. metadata: Optional[dict] = Body(
  594. None, description="The metadata of the entity to create."
  595. ),
  596. auth_user=Depends(self.providers.auth.auth_wrapper()),
  597. ) -> WrappedEntityResponse:
  598. """Creates a new entity in the graph."""
  599. if (
  600. # not auth_user.is_superuser
  601. collection_id
  602. not in auth_user.collection_ids
  603. ):
  604. raise R2RException(
  605. "The currently authenticated user does not have access to the collection associated with the given graph.",
  606. 403,
  607. )
  608. return await self.services.graph.create_entity(
  609. name=name,
  610. description=description,
  611. parent_id=collection_id,
  612. category=category,
  613. metadata=metadata,
  614. )
  615. @self.router.post(
  616. "/graphs/{collection_id}/relationships",
  617. dependencies=[Depends(self.rate_limit_dependency)],
  618. )
  619. @self.base_endpoint
  620. async def create_relationship(
  621. collection_id: UUID = Path(
  622. ...,
  623. description="The collection ID corresponding to the graph to add the relationship to.",
  624. ),
  625. subject: str = Body(
  626. ..., description="The subject of the relationship to create."
  627. ),
  628. subject_id: UUID = Body(
  629. ...,
  630. description="The ID of the subject of the relationship to create.",
  631. ),
  632. predicate: str = Body(
  633. ..., description="The predicate of the relationship to create."
  634. ),
  635. object: str = Body(
  636. ..., description="The object of the relationship to create."
  637. ),
  638. object_id: UUID = Body(
  639. ...,
  640. description="The ID of the object of the relationship to create.",
  641. ),
  642. description: str = Body(
  643. ...,
  644. description="The description of the relationship to create.",
  645. ),
  646. weight: float = Body(
  647. 1.0, description="The weight of the relationship to create."
  648. ),
  649. metadata: Optional[dict] = Body(
  650. None, description="The metadata of the relationship to create."
  651. ),
  652. auth_user=Depends(self.providers.auth.auth_wrapper()),
  653. ) -> WrappedRelationshipResponse:
  654. """Creates a new relationship in the graph."""
  655. if not auth_user.is_superuser:
  656. raise R2RException(
  657. "Only superusers can create relationships.", 403
  658. )
  659. if (
  660. # not auth_user.is_superuser
  661. collection_id
  662. not in auth_user.collection_ids
  663. ):
  664. raise R2RException(
  665. "The currently authenticated user does not have access to the collection associated with the given graph.",
  666. 403,
  667. )
  668. return await self.services.graph.create_relationship(
  669. subject=subject,
  670. subject_id=subject_id,
  671. predicate=predicate,
  672. object=object,
  673. object_id=object_id,
  674. description=description,
  675. weight=weight,
  676. metadata=metadata,
  677. parent_id=collection_id,
  678. )
  679. @self.router.get(
  680. "/graphs/{collection_id}/entities/{entity_id}",
  681. dependencies=[Depends(self.rate_limit_dependency)],
  682. openapi_extra={
  683. "x-codeSamples": [
  684. {
  685. "lang": "Python",
  686. "source": textwrap.dedent(
  687. """
  688. from r2r import R2RClient
  689. client = R2RClient()
  690. # when using auth, do client.login(...)
  691. response = client.graphs.get_entity(
  692. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  693. entity_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  694. )
  695. """
  696. ),
  697. },
  698. {
  699. "lang": "JavaScript",
  700. "source": textwrap.dedent(
  701. """
  702. const { r2rClient } = require("r2r-js");
  703. const client = new r2rClient();
  704. function main() {
  705. const response = await client.graphs.get_entity({
  706. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  707. entityId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  708. });
  709. }
  710. main();
  711. """
  712. ),
  713. },
  714. ]
  715. },
  716. )
  717. @self.base_endpoint
  718. async def get_entity(
  719. collection_id: UUID = Path(
  720. ...,
  721. description="The collection ID corresponding to the graph containing the entity.",
  722. ),
  723. entity_id: UUID = Path(
  724. ..., description="The ID of the entity to retrieve."
  725. ),
  726. auth_user=Depends(self.providers.auth.auth_wrapper()),
  727. ) -> WrappedEntityResponse:
  728. """Retrieves a specific entity by its ID."""
  729. if (
  730. # not auth_user.is_superuser
  731. collection_id
  732. not in auth_user.collection_ids
  733. ):
  734. raise R2RException(
  735. "The currently authenticated user does not have access to the collection associated with the given graph.",
  736. 403,
  737. )
  738. result = await self.providers.database.graphs_handler.entities.get(
  739. parent_id=collection_id,
  740. store_type="graphs",
  741. offset=0,
  742. limit=1,
  743. entity_ids=[entity_id],
  744. )
  745. if len(result) == 0 or len(result[0]) == 0:
  746. raise R2RException("Entity not found", 404)
  747. return result[0][0]
  748. @self.router.post(
  749. "/graphs/{collection_id}/entities/{entity_id}",
  750. dependencies=[Depends(self.rate_limit_dependency)],
  751. )
  752. @self.base_endpoint
  753. async def update_entity(
  754. collection_id: UUID = Path(
  755. ...,
  756. description="The collection ID corresponding to the graph containing the entity.",
  757. ),
  758. entity_id: UUID = Path(
  759. ..., description="The ID of the entity to update."
  760. ),
  761. name: Optional[str] = Body(
  762. ..., description="The updated name of the entity."
  763. ),
  764. description: Optional[str] = Body(
  765. None, description="The updated description of the entity."
  766. ),
  767. category: Optional[str] = Body(
  768. None, description="The updated category of the entity."
  769. ),
  770. metadata: Optional[dict] = Body(
  771. None, description="The updated metadata of the entity."
  772. ),
  773. auth_user=Depends(self.providers.auth.auth_wrapper()),
  774. ) -> WrappedEntityResponse:
  775. """Updates an existing entity in the graph."""
  776. if not auth_user.is_superuser:
  777. raise R2RException(
  778. "Only superusers can update graph entities.", 403
  779. )
  780. if (
  781. # not auth_user.is_superuser
  782. collection_id
  783. not in auth_user.collection_ids
  784. ):
  785. raise R2RException(
  786. "The currently authenticated user does not have access to the collection associated with the given graph.",
  787. 403,
  788. )
  789. return await self.services.graph.update_entity(
  790. entity_id=entity_id,
  791. name=name,
  792. category=category,
  793. description=description,
  794. metadata=metadata,
  795. )
  796. @self.router.delete(
  797. "/graphs/{collection_id}/entities/{entity_id}",
  798. dependencies=[Depends(self.rate_limit_dependency)],
  799. summary="Remove an entity",
  800. openapi_extra={
  801. "x-codeSamples": [
  802. {
  803. "lang": "Python",
  804. "source": textwrap.dedent(
  805. """
  806. from r2r import R2RClient
  807. client = R2RClient()
  808. # when using auth, do client.login(...)
  809. response = client.graphs.remove_entity(
  810. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  811. entity_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  812. )
  813. """
  814. ),
  815. },
  816. {
  817. "lang": "JavaScript",
  818. "source": textwrap.dedent(
  819. """
  820. const { r2rClient } = require("r2r-js");
  821. const client = new r2rClient();
  822. function main() {
  823. const response = await client.graphs.removeEntity({
  824. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  825. entityId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  826. });
  827. }
  828. main();
  829. """
  830. ),
  831. },
  832. ]
  833. },
  834. )
  835. @self.base_endpoint
  836. async def delete_entity(
  837. collection_id: UUID = Path(
  838. ...,
  839. description="The collection ID corresponding to the graph to remove the entity from.",
  840. ),
  841. entity_id: UUID = Path(
  842. ...,
  843. description="The ID of the entity to remove from the graph.",
  844. ),
  845. auth_user=Depends(self.providers.auth.auth_wrapper()),
  846. ) -> WrappedBooleanResponse:
  847. """Removes an entity from the graph."""
  848. if not auth_user.is_superuser:
  849. raise R2RException(
  850. "Only superusers can delete graph details.", 403
  851. )
  852. if (
  853. # not auth_user.is_superuser
  854. collection_id
  855. not in auth_user.collection_ids
  856. ):
  857. raise R2RException(
  858. "The currently authenticated user does not have access to the collection associated with the given graph.",
  859. 403,
  860. )
  861. await self.services.graph.delete_entity(
  862. parent_id=collection_id,
  863. entity_id=entity_id,
  864. )
  865. return GenericBooleanResponse(success=True) # type: ignore
  866. @self.router.get(
  867. "/graphs/{collection_id}/relationships",
  868. dependencies=[Depends(self.rate_limit_dependency)],
  869. description="Lists all relationships in the graph with pagination support.",
  870. openapi_extra={
  871. "x-codeSamples": [
  872. {
  873. "lang": "Python",
  874. "source": textwrap.dedent(
  875. """
  876. from r2r import R2RClient
  877. client = R2RClient()
  878. # when using auth, do client.login(...)
  879. response = client.graphs.list_relationships(collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7")
  880. """
  881. ),
  882. },
  883. {
  884. "lang": "JavaScript",
  885. "source": textwrap.dedent(
  886. """
  887. const { r2rClient } = require("r2r-js");
  888. const client = new r2rClient();
  889. function main() {
  890. const response = await client.graphs.listRelationships({
  891. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  892. });
  893. }
  894. main();
  895. """
  896. ),
  897. },
  898. ],
  899. },
  900. )
  901. @self.base_endpoint
  902. async def get_relationships(
  903. collection_id: UUID = Path(
  904. ...,
  905. description="The collection ID corresponding to the graph to list relationships from.",
  906. ),
  907. offset: int = Query(
  908. 0,
  909. ge=0,
  910. description="Specifies the number of objects to skip. Defaults to 0.",
  911. ),
  912. limit: int = Query(
  913. 100,
  914. ge=1,
  915. le=1000,
  916. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  917. ),
  918. auth_user=Depends(self.providers.auth.auth_wrapper()),
  919. ) -> WrappedRelationshipsResponse:
  920. """
  921. Lists all relationships in the graph with pagination support.
  922. """
  923. if (
  924. # not auth_user.is_superuser
  925. collection_id
  926. not in auth_user.collection_ids
  927. ):
  928. raise R2RException(
  929. "The currently authenticated user does not have access to the collection associated with the given graph.",
  930. 403,
  931. )
  932. relationships, count = await self.services.graph.get_relationships(
  933. parent_id=collection_id,
  934. offset=offset,
  935. limit=limit,
  936. )
  937. return relationships, { # type: ignore
  938. "total_entries": count,
  939. }
  940. @self.router.get(
  941. "/graphs/{collection_id}/relationships/{relationship_id}",
  942. dependencies=[Depends(self.rate_limit_dependency)],
  943. description="Retrieves a specific relationship by its ID.",
  944. openapi_extra={
  945. "x-codeSamples": [
  946. {
  947. "lang": "Python",
  948. "source": textwrap.dedent(
  949. """
  950. from r2r import R2RClient
  951. client = R2RClient()
  952. # when using auth, do client.login(...)
  953. response = client.graphs.get_relationship(
  954. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  955. relationship_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  956. )
  957. """
  958. ),
  959. },
  960. {
  961. "lang": "JavaScript",
  962. "source": textwrap.dedent(
  963. """
  964. const { r2rClient } = require("r2r-js");
  965. const client = new r2rClient();
  966. function main() {
  967. const response = await client.graphs.getRelationship({
  968. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  969. relationshipId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  970. });
  971. }
  972. main();
  973. """
  974. ),
  975. },
  976. ],
  977. },
  978. )
  979. @self.base_endpoint
  980. async def get_relationship(
  981. collection_id: UUID = Path(
  982. ...,
  983. description="The collection ID corresponding to the graph containing the relationship.",
  984. ),
  985. relationship_id: UUID = Path(
  986. ..., description="The ID of the relationship to retrieve."
  987. ),
  988. auth_user=Depends(self.providers.auth.auth_wrapper()),
  989. ) -> WrappedRelationshipResponse:
  990. """Retrieves a specific relationship by its ID."""
  991. if (
  992. # not auth_user.is_superuser
  993. collection_id
  994. not in auth_user.collection_ids
  995. ):
  996. raise R2RException(
  997. "The currently authenticated user does not have access to the collection associated with the given graph.",
  998. 403,
  999. )
  1000. results = (
  1001. await self.providers.database.graphs_handler.relationships.get(
  1002. parent_id=collection_id,
  1003. store_type="graphs",
  1004. offset=0,
  1005. limit=1,
  1006. relationship_ids=[relationship_id],
  1007. )
  1008. )
  1009. if len(results) == 0 or len(results[0]) == 0:
  1010. raise R2RException("Relationship not found", 404)
  1011. return results[0][0]
  1012. @self.router.post(
  1013. "/graphs/{collection_id}/relationships/{relationship_id}",
  1014. dependencies=[Depends(self.rate_limit_dependency)],
  1015. )
  1016. @self.base_endpoint
  1017. async def update_relationship(
  1018. collection_id: UUID = Path(
  1019. ...,
  1020. description="The collection ID corresponding to the graph containing the relationship.",
  1021. ),
  1022. relationship_id: UUID = Path(
  1023. ..., description="The ID of the relationship to update."
  1024. ),
  1025. subject: Optional[str] = Body(
  1026. ..., description="The updated subject of the relationship."
  1027. ),
  1028. subject_id: Optional[UUID] = Body(
  1029. ..., description="The updated subject ID of the relationship."
  1030. ),
  1031. predicate: Optional[str] = Body(
  1032. ..., description="The updated predicate of the relationship."
  1033. ),
  1034. object: Optional[str] = Body(
  1035. ..., description="The updated object of the relationship."
  1036. ),
  1037. object_id: Optional[UUID] = Body(
  1038. ..., description="The updated object ID of the relationship."
  1039. ),
  1040. description: Optional[str] = Body(
  1041. None,
  1042. description="The updated description of the relationship.",
  1043. ),
  1044. weight: Optional[float] = Body(
  1045. None, description="The updated weight of the relationship."
  1046. ),
  1047. metadata: Optional[dict] = Body(
  1048. None, description="The updated metadata of the relationship."
  1049. ),
  1050. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1051. ) -> WrappedRelationshipResponse:
  1052. """Updates an existing relationship in the graph."""
  1053. if not auth_user.is_superuser:
  1054. raise R2RException(
  1055. "Only superusers can update graph details", 403
  1056. )
  1057. if (
  1058. # not auth_user.is_superuser
  1059. collection_id
  1060. not in auth_user.collection_ids
  1061. ):
  1062. raise R2RException(
  1063. "The currently authenticated user does not have access to the collection associated with the given graph.",
  1064. 403,
  1065. )
  1066. return await self.services.graph.update_relationship(
  1067. relationship_id=relationship_id,
  1068. subject=subject,
  1069. subject_id=subject_id,
  1070. predicate=predicate,
  1071. object=object,
  1072. object_id=object_id,
  1073. description=description,
  1074. weight=weight,
  1075. metadata=metadata,
  1076. )
  1077. @self.router.delete(
  1078. "/graphs/{collection_id}/relationships/{relationship_id}",
  1079. dependencies=[Depends(self.rate_limit_dependency)],
  1080. description="Removes a relationship from the graph.",
  1081. openapi_extra={
  1082. "x-codeSamples": [
  1083. {
  1084. "lang": "Python",
  1085. "source": textwrap.dedent(
  1086. """
  1087. from r2r import R2RClient
  1088. client = R2RClient()
  1089. # when using auth, do client.login(...)
  1090. response = client.graphs.delete_relationship(
  1091. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  1092. relationship_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  1093. )
  1094. """
  1095. ),
  1096. },
  1097. {
  1098. "lang": "JavaScript",
  1099. "source": textwrap.dedent(
  1100. """
  1101. const { r2rClient } = require("r2r-js");
  1102. const client = new r2rClient();
  1103. function main() {
  1104. const response = await client.graphs.deleteRelationship({
  1105. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  1106. relationshipId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  1107. });
  1108. }
  1109. main();
  1110. """
  1111. ),
  1112. },
  1113. ],
  1114. },
  1115. )
  1116. @self.base_endpoint
  1117. async def delete_relationship(
  1118. collection_id: UUID = Path(
  1119. ...,
  1120. description="The collection ID corresponding to the graph to remove the relationship from.",
  1121. ),
  1122. relationship_id: UUID = Path(
  1123. ...,
  1124. description="The ID of the relationship to remove from the graph.",
  1125. ),
  1126. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1127. ) -> WrappedBooleanResponse:
  1128. """Removes a relationship from the graph."""
  1129. if not auth_user.is_superuser:
  1130. raise R2RException(
  1131. "Only superusers can delete a relationship.", 403
  1132. )
  1133. if (
  1134. not auth_user.is_superuser
  1135. and collection_id not in auth_user.collection_ids
  1136. ):
  1137. raise R2RException(
  1138. "The currently authenticated user does not have access to the collection associated with the given graph.",
  1139. 403,
  1140. )
  1141. await self.services.graph.delete_relationship(
  1142. parent_id=collection_id,
  1143. relationship_id=relationship_id,
  1144. )
  1145. return GenericBooleanResponse(success=True) # type: ignore
  1146. @self.router.post(
  1147. "/graphs/{collection_id}/communities",
  1148. dependencies=[Depends(self.rate_limit_dependency)],
  1149. summary="Create a new community",
  1150. openapi_extra={
  1151. "x-codeSamples": [
  1152. {
  1153. "lang": "Python",
  1154. "source": textwrap.dedent(
  1155. """
  1156. from r2r import R2RClient
  1157. client = R2RClient()
  1158. # when using auth, do client.login(...)
  1159. response = client.graphs.create_community(
  1160. collection_id="9fbe403b-c11c-5aae-8ade-ef22980c3ad1",
  1161. name="My Community",
  1162. summary="A summary of the community",
  1163. findings=["Finding 1", "Finding 2"],
  1164. rating=5,
  1165. rating_explanation="This is a rating explanation",
  1166. )
  1167. """
  1168. ),
  1169. },
  1170. {
  1171. "lang": "JavaScript",
  1172. "source": textwrap.dedent(
  1173. """
  1174. const { r2rClient } = require("r2r-js");
  1175. const client = new r2rClient();
  1176. function main() {
  1177. const response = await client.graphs.createCommunity({
  1178. collectionId: "9fbe403b-c11c-5aae-8ade-ef22980c3ad1",
  1179. name: "My Community",
  1180. summary: "A summary of the community",
  1181. findings: ["Finding 1", "Finding 2"],
  1182. rating: 5,
  1183. ratingExplanation: "This is a rating explanation",
  1184. });
  1185. }
  1186. main();
  1187. """
  1188. ),
  1189. },
  1190. ]
  1191. },
  1192. )
  1193. @self.base_endpoint
  1194. async def create_community(
  1195. collection_id: UUID = Path(
  1196. ...,
  1197. description="The collection ID corresponding to the graph to create the community in.",
  1198. ),
  1199. name: str = Body(..., description="The name of the community"),
  1200. summary: str = Body(..., description="A summary of the community"),
  1201. findings: Optional[list[str]] = Body(
  1202. default=[], description="Findings about the community"
  1203. ),
  1204. rating: Optional[float] = Body(
  1205. default=5, ge=1, le=10, description="Rating between 1 and 10"
  1206. ),
  1207. rating_explanation: Optional[str] = Body(
  1208. default="", description="Explanation for the rating"
  1209. ),
  1210. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1211. ) -> WrappedCommunityResponse:
  1212. """
  1213. Creates a new community in the graph.
  1214. While communities are typically built automatically via the /graphs/{id}/communities/build endpoint,
  1215. this endpoint allows you to manually create your own communities.
  1216. This can be useful when you want to:
  1217. - Define custom groupings of entities based on domain knowledge
  1218. - Add communities that weren't detected by the automatic process
  1219. - Create hierarchical organization structures
  1220. - Tag groups of entities with specific metadata
  1221. The created communities will be integrated with any existing automatically detected communities
  1222. in the graph's community structure.
  1223. """
  1224. if not auth_user.is_superuser:
  1225. raise R2RException(
  1226. "Only superusers can create a community.", 403
  1227. )
  1228. if (
  1229. not auth_user.is_superuser
  1230. and collection_id not in auth_user.collection_ids
  1231. ):
  1232. raise R2RException(
  1233. "The currently authenticated user does not have access to the collection associated with the given graph.",
  1234. 403,
  1235. )
  1236. return await self.services.graph.create_community(
  1237. parent_id=collection_id,
  1238. name=name,
  1239. summary=summary,
  1240. findings=findings,
  1241. rating=rating,
  1242. rating_explanation=rating_explanation,
  1243. )
  1244. @self.router.get(
  1245. "/graphs/{collection_id}/communities",
  1246. dependencies=[Depends(self.rate_limit_dependency)],
  1247. summary="List communities",
  1248. openapi_extra={
  1249. "x-codeSamples": [
  1250. {
  1251. "lang": "Python",
  1252. "source": textwrap.dedent(
  1253. """
  1254. from r2r import R2RClient
  1255. client = R2RClient()
  1256. # when using auth, do client.login(...)
  1257. response = client.graphs.list_communities(collection_id="9fbe403b-c11c-5aae-8ade-ef22980c3ad1")
  1258. """
  1259. ),
  1260. },
  1261. {
  1262. "lang": "JavaScript",
  1263. "source": textwrap.dedent(
  1264. """
  1265. const { r2rClient } = require("r2r-js");
  1266. const client = new r2rClient();
  1267. function main() {
  1268. const response = await client.graphs.listCommunities({
  1269. collectionId: "9fbe403b-c11c-5aae-8ade-ef22980c3ad1",
  1270. });
  1271. }
  1272. main();
  1273. """
  1274. ),
  1275. },
  1276. ]
  1277. },
  1278. )
  1279. @self.base_endpoint
  1280. async def get_communities(
  1281. collection_id: UUID = Path(
  1282. ...,
  1283. description="The collection ID corresponding to the graph to get communities for.",
  1284. ),
  1285. offset: int = Query(
  1286. 0,
  1287. ge=0,
  1288. description="Specifies the number of objects to skip. Defaults to 0.",
  1289. ),
  1290. limit: int = Query(
  1291. 100,
  1292. ge=1,
  1293. le=1000,
  1294. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  1295. ),
  1296. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1297. ) -> WrappedCommunitiesResponse:
  1298. """
  1299. Lists all communities in the graph with pagination support.
  1300. """
  1301. if (
  1302. # not auth_user.is_superuser
  1303. collection_id
  1304. not in auth_user.collection_ids
  1305. ):
  1306. raise R2RException(
  1307. "The currently authenticated user does not have access to the collection associated with the given graph.",
  1308. 403,
  1309. )
  1310. communities, count = await self.services.graph.get_communities(
  1311. parent_id=collection_id,
  1312. offset=offset,
  1313. limit=limit,
  1314. )
  1315. return communities, { # type: ignore
  1316. "total_entries": count,
  1317. }
  1318. @self.router.get(
  1319. "/graphs/{collection_id}/communities/{community_id}",
  1320. dependencies=[Depends(self.rate_limit_dependency)],
  1321. summary="Retrieve a community",
  1322. openapi_extra={
  1323. "x-codeSamples": [
  1324. {
  1325. "lang": "Python",
  1326. "source": textwrap.dedent(
  1327. """
  1328. from r2r import R2RClient
  1329. client = R2RClient()
  1330. # when using auth, do client.login(...)
  1331. response = client.graphs.get_community(collection_id="9fbe403b-c11c-5aae-8ade-ef22980c3ad1")
  1332. """
  1333. ),
  1334. },
  1335. {
  1336. "lang": "JavaScript",
  1337. "source": textwrap.dedent(
  1338. """
  1339. const { r2rClient } = require("r2r-js");
  1340. const client = new r2rClient();
  1341. function main() {
  1342. const response = await client.graphs.getCommunity({
  1343. collectionId: "9fbe403b-c11c-5aae-8ade-ef22980c3ad1",
  1344. });
  1345. }
  1346. main();
  1347. """
  1348. ),
  1349. },
  1350. ]
  1351. },
  1352. )
  1353. @self.base_endpoint
  1354. async def get_community(
  1355. collection_id: UUID = Path(
  1356. ...,
  1357. description="The ID of the collection to get communities for.",
  1358. ),
  1359. community_id: UUID = Path(
  1360. ...,
  1361. description="The ID of the community to get.",
  1362. ),
  1363. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1364. ) -> WrappedCommunityResponse:
  1365. """
  1366. Retrieves a specific community by its ID.
  1367. """
  1368. if (
  1369. # not auth_user.is_superuser
  1370. collection_id
  1371. not in auth_user.collection_ids
  1372. ):
  1373. raise R2RException(
  1374. "The currently authenticated user does not have access to the collection associated with the given graph.",
  1375. 403,
  1376. )
  1377. results = (
  1378. await self.providers.database.graphs_handler.communities.get(
  1379. parent_id=collection_id,
  1380. community_ids=[community_id],
  1381. store_type="graphs",
  1382. offset=0,
  1383. limit=1,
  1384. )
  1385. )
  1386. if len(results) == 0 or len(results[0]) == 0:
  1387. raise R2RException("Community not found", 404)
  1388. return results[0][0]
  1389. @self.router.delete(
  1390. "/graphs/{collection_id}/communities/{community_id}",
  1391. dependencies=[Depends(self.rate_limit_dependency)],
  1392. summary="Delete a community",
  1393. openapi_extra={
  1394. "x-codeSamples": [
  1395. {
  1396. "lang": "Python",
  1397. "source": textwrap.dedent(
  1398. """
  1399. from r2r import R2RClient
  1400. client = R2RClient()
  1401. # when using auth, do client.login(...)
  1402. response = client.graphs.delete_community(
  1403. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  1404. community_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  1405. )
  1406. """
  1407. ),
  1408. },
  1409. {
  1410. "lang": "JavaScript",
  1411. "source": textwrap.dedent(
  1412. """
  1413. const { r2rClient } = require("r2r-js");
  1414. const client = new r2rClient();
  1415. function main() {
  1416. const response = await client.graphs.deleteCommunity({
  1417. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  1418. communityId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  1419. });
  1420. }
  1421. main();
  1422. """
  1423. ),
  1424. },
  1425. ]
  1426. },
  1427. )
  1428. @self.base_endpoint
  1429. async def delete_community(
  1430. collection_id: UUID = Path(
  1431. ...,
  1432. description="The collection ID corresponding to the graph to delete the community from.",
  1433. ),
  1434. community_id: UUID = Path(
  1435. ...,
  1436. description="The ID of the community to delete.",
  1437. ),
  1438. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1439. ):
  1440. if (
  1441. not auth_user.is_superuser
  1442. and collection_id not in auth_user.graph_ids
  1443. ):
  1444. raise R2RException(
  1445. "Only superusers can delete communities", 403
  1446. )
  1447. if (
  1448. # not auth_user.is_superuser
  1449. collection_id
  1450. not in auth_user.collection_ids
  1451. ):
  1452. raise R2RException(
  1453. "The currently authenticated user does not have access to the collection associated with the given graph.",
  1454. 403,
  1455. )
  1456. await self.services.graph.delete_community(
  1457. parent_id=collection_id,
  1458. community_id=community_id,
  1459. )
  1460. return GenericBooleanResponse(success=True) # type: ignore
  1461. @self.router.post(
  1462. "/graphs/{collection_id}/communities/{community_id}",
  1463. dependencies=[Depends(self.rate_limit_dependency)],
  1464. summary="Update community",
  1465. openapi_extra={
  1466. "x-codeSamples": [
  1467. {
  1468. "lang": "Python",
  1469. "source": textwrap.dedent(
  1470. """
  1471. from r2r import R2RClient
  1472. client = R2RClient()
  1473. # when using auth, do client.login(...)
  1474. response = client.graphs.update_community(
  1475. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  1476. community_update={
  1477. "metadata": {
  1478. "topic": "Technology",
  1479. "description": "Tech companies and products"
  1480. }
  1481. }
  1482. )"""
  1483. ),
  1484. },
  1485. {
  1486. "lang": "JavaScript",
  1487. "source": textwrap.dedent(
  1488. """
  1489. const { r2rClient } = require("r2r-js");
  1490. const client = new r2rClient();
  1491. async function main() {
  1492. const response = await client.graphs.updateCommunity({
  1493. collectionId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  1494. communityId: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7",
  1495. communityUpdate: {
  1496. metadata: {
  1497. topic: "Technology",
  1498. description: "Tech companies and products"
  1499. }
  1500. }
  1501. });
  1502. }
  1503. main();
  1504. """
  1505. ),
  1506. },
  1507. ]
  1508. },
  1509. )
  1510. @self.base_endpoint
  1511. async def update_community(
  1512. collection_id: UUID = Path(...),
  1513. community_id: UUID = Path(...),
  1514. name: Optional[str] = Body(None),
  1515. summary: Optional[str] = Body(None),
  1516. findings: Optional[list[str]] = Body(None),
  1517. rating: Optional[float] = Body(default=None, ge=1, le=10),
  1518. rating_explanation: Optional[str] = Body(None),
  1519. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1520. ) -> WrappedCommunityResponse:
  1521. """
  1522. Updates an existing community in the graph.
  1523. """
  1524. if (
  1525. not auth_user.is_superuser
  1526. and collection_id not in auth_user.graph_ids
  1527. ):
  1528. raise R2RException(
  1529. "Only superusers can update communities.", 403
  1530. )
  1531. if (
  1532. # not auth_user.is_superuser
  1533. collection_id
  1534. not in auth_user.collection_ids
  1535. ):
  1536. raise R2RException(
  1537. "The currently authenticated user does not have access to the collection associated with the given graph.",
  1538. 403,
  1539. )
  1540. return await self.services.graph.update_community(
  1541. community_id=community_id,
  1542. name=name,
  1543. summary=summary,
  1544. findings=findings,
  1545. rating=rating,
  1546. rating_explanation=rating_explanation,
  1547. )
  1548. @self.router.post(
  1549. "/graphs/{collection_id}/pull",
  1550. dependencies=[Depends(self.rate_limit_dependency)],
  1551. summary="Pull latest entities to the graph",
  1552. openapi_extra={
  1553. "x-codeSamples": [
  1554. {
  1555. "lang": "Python",
  1556. "source": textwrap.dedent(
  1557. """
  1558. from r2r import R2RClient
  1559. client = R2RClient()
  1560. # when using auth, do client.login(...)
  1561. response = client.graphs.pull(
  1562. collection_id="d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  1563. )"""
  1564. ),
  1565. },
  1566. {
  1567. "lang": "JavaScript",
  1568. "source": textwrap.dedent(
  1569. """
  1570. const { r2rClient } = require("r2r-js");
  1571. const client = new r2rClient();
  1572. async function main() {
  1573. const response = await client.graphs.pull({
  1574. collection_id: "d09dedb1-b2ab-48a5-b950-6e1f464d83e7"
  1575. });
  1576. }
  1577. main();
  1578. """
  1579. ),
  1580. },
  1581. ]
  1582. },
  1583. )
  1584. @self.base_endpoint
  1585. async def pull(
  1586. collection_id: UUID = Path(
  1587. ..., description="The ID of the graph to initialize."
  1588. ),
  1589. force: Optional[bool] = Body(
  1590. False,
  1591. description="If true, forces a re-pull of all entities and relationships.",
  1592. ),
  1593. # document_ids: list[UUID] = Body(
  1594. # ..., description="List of document IDs to add to the graph."
  1595. # ),
  1596. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1597. ) -> WrappedBooleanResponse:
  1598. """
  1599. Adds documents to a graph by copying their entities and relationships.
  1600. This endpoint:
  1601. 1. Copies document entities to the graphs_entities table
  1602. 2. Copies document relationships to the graphs_relationships table
  1603. 3. Associates the documents with the graph
  1604. When a document is added:
  1605. - Its entities and relationships are copied to graph-specific tables
  1606. - Existing entities/relationships are updated by merging their properties
  1607. - The document ID is recorded in the graph's document_ids array
  1608. Documents added to a graph will contribute their knowledge to:
  1609. - Graph analysis and querying
  1610. - Community detection
  1611. - Knowledge graph enrichment
  1612. The user must have access to both the graph and the documents being added.
  1613. """
  1614. # Check user permissions for graph
  1615. if not auth_user.is_superuser:
  1616. raise R2RException("Only superusers can `pull` a graph.", 403)
  1617. if (
  1618. # not auth_user.is_superuser
  1619. collection_id
  1620. not in auth_user.collection_ids
  1621. ):
  1622. raise R2RException(
  1623. "The currently authenticated user does not have access to the collection associated with the given graph.",
  1624. 403,
  1625. )
  1626. list_graphs_response = await self.services.graph.list_graphs(
  1627. # user_ids=None,
  1628. graph_ids=[collection_id],
  1629. offset=0,
  1630. limit=1,
  1631. )
  1632. if len(list_graphs_response["results"]) == 0:
  1633. raise R2RException("Graph not found", 404)
  1634. collection_id = list_graphs_response["results"][0].collection_id
  1635. documents = []
  1636. document_req = (
  1637. await self.providers.database.collections_handler.documents_in_collection(
  1638. collection_id, offset=0, limit=100
  1639. )
  1640. )["results"]
  1641. documents.extend(document_req)
  1642. while len(document_req) == 100:
  1643. document_req = (
  1644. await self.providers.database.collections_handler.documents_in_collection(
  1645. collection_id, offset=len(documents), limit=100
  1646. )
  1647. )["results"]
  1648. documents.extend(document_req)
  1649. success = False
  1650. for document in documents:
  1651. # TODO - Add better checks for user permissions
  1652. if (
  1653. not auth_user.is_superuser
  1654. and document.id
  1655. not in auth_user.document_ids # TODO - extend to include checks on collections
  1656. ):
  1657. raise R2RException(
  1658. f"The currently authenticated user does not have access to document {document.id}",
  1659. 403,
  1660. )
  1661. entities = (
  1662. await self.providers.database.graphs_handler.entities.get(
  1663. parent_id=document.id,
  1664. store_type="documents",
  1665. offset=0,
  1666. limit=100,
  1667. )
  1668. )
  1669. has_document = (
  1670. await self.providers.database.graphs_handler.has_document(
  1671. collection_id, document.id
  1672. )
  1673. )
  1674. if has_document:
  1675. logger.info(
  1676. f"Document {document.id} is already in graph {collection_id}, skipping."
  1677. )
  1678. continue
  1679. if len(entities[0]) == 0:
  1680. if not force:
  1681. logger.warning(
  1682. f"Document {document.id} has no entities, extraction may not have been called, skipping."
  1683. )
  1684. continue
  1685. else:
  1686. logger.warning(
  1687. f"Document {document.id} has no entities, but force=True, continuing."
  1688. )
  1689. success = (
  1690. await self.providers.database.graphs_handler.add_documents(
  1691. id=collection_id,
  1692. document_ids=[document.id],
  1693. )
  1694. )
  1695. if not success:
  1696. logger.warning(
  1697. f"No documents were added to graph {collection_id}, marking as failed."
  1698. )
  1699. if success:
  1700. await self.providers.database.documents_handler.set_workflow_status(
  1701. id=collection_id,
  1702. status_type="graph_sync_status",
  1703. status=KGEnrichmentStatus.SUCCESS,
  1704. )
  1705. return GenericBooleanResponse(success=success) # type: ignore