documents_router.py 93 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363
  1. import base64
  2. import logging
  3. import mimetypes
  4. import textwrap
  5. from datetime import datetime
  6. from io import BytesIO
  7. from typing import Any, Optional
  8. from urllib.parse import quote
  9. from uuid import UUID
  10. from fastapi import Body, Depends, File, Form, Path, Query, UploadFile
  11. from fastapi.background import BackgroundTasks
  12. from fastapi.responses import FileResponse, StreamingResponse
  13. from pydantic import Json
  14. from core.base import (
  15. IngestionConfig,
  16. R2RException,
  17. SearchMode,
  18. SearchSettings,
  19. UnprocessedChunk,
  20. Workflow,
  21. generate_document_id,
  22. generate_id,
  23. select_search_filters,
  24. )
  25. from core.base.abstractions import GraphCreationSettings, StoreType
  26. from core.base.api.models import (
  27. GenericBooleanResponse,
  28. WrappedBooleanResponse,
  29. WrappedChunksResponse,
  30. WrappedCollectionsResponse,
  31. WrappedDocumentResponse,
  32. WrappedDocumentSearchResponse,
  33. WrappedDocumentsResponse,
  34. WrappedEntitiesResponse,
  35. WrappedGenericMessageResponse,
  36. WrappedIngestionResponse,
  37. WrappedRelationshipsResponse,
  38. )
  39. from core.utils import update_settings_from_dict
  40. from shared.abstractions import IngestionMode
  41. from ...abstractions import R2RProviders, R2RServices
  42. from ...config import R2RConfig
  43. from .base_router import BaseRouterV3
  44. logger = logging.getLogger()
  45. MAX_CHUNKS_PER_REQUEST = 1024 * 100
  46. def merge_search_settings(
  47. base: SearchSettings, overrides: SearchSettings
  48. ) -> SearchSettings:
  49. # Convert both to dict
  50. base_dict = base.model_dump()
  51. overrides_dict = overrides.model_dump(exclude_unset=True)
  52. # Update base_dict with values from overrides_dict
  53. # This ensures that any field set in overrides takes precedence
  54. for k, v in overrides_dict.items():
  55. base_dict[k] = v
  56. # Construct a new SearchSettings from the merged dict
  57. return SearchSettings(**base_dict)
  58. def merge_ingestion_config(
  59. base: IngestionConfig, overrides: IngestionConfig
  60. ) -> IngestionConfig:
  61. base_dict = base.model_dump()
  62. overrides_dict = overrides.model_dump(exclude_unset=True)
  63. for k, v in overrides_dict.items():
  64. base_dict[k] = v
  65. return IngestionConfig(**base_dict)
  66. class DocumentsRouter(BaseRouterV3):
  67. def __init__(
  68. self,
  69. providers: R2RProviders,
  70. services: R2RServices,
  71. config: R2RConfig,
  72. ):
  73. logging.info("Initializing DocumentsRouter")
  74. super().__init__(providers, services, config)
  75. self._register_workflows()
  76. def _prepare_search_settings(
  77. self,
  78. auth_user: Any,
  79. search_mode: SearchMode,
  80. search_settings: Optional[SearchSettings],
  81. ) -> SearchSettings:
  82. """Prepare the effective search settings based on the provided
  83. search_mode, optional user-overrides in search_settings, and applied
  84. filters."""
  85. if search_mode != SearchMode.custom:
  86. # Start from mode defaults
  87. effective_settings = SearchSettings.get_default(search_mode.value)
  88. if search_settings:
  89. # Merge user-provided overrides
  90. effective_settings = merge_search_settings(
  91. effective_settings, search_settings
  92. )
  93. else:
  94. # Custom mode: use provided settings or defaults
  95. effective_settings = search_settings or SearchSettings()
  96. # Apply user-specific filters
  97. effective_settings.filters = select_search_filters(
  98. auth_user, effective_settings
  99. )
  100. return effective_settings
  101. # TODO - Remove this legacy method
  102. def _register_workflows(self):
  103. self.providers.orchestration.register_workflows(
  104. Workflow.INGESTION,
  105. self.services.ingestion,
  106. {
  107. "ingest-files": (
  108. "Ingest files task queued successfully."
  109. if self.providers.orchestration.config.provider != "simple"
  110. else "Document created and ingested successfully."
  111. ),
  112. "ingest-chunks": (
  113. "Ingest chunks task queued successfully."
  114. if self.providers.orchestration.config.provider != "simple"
  115. else "Document created and ingested successfully."
  116. ),
  117. "update-chunk": (
  118. "Update chunk task queued successfully."
  119. if self.providers.orchestration.config.provider != "simple"
  120. else "Chunk update completed successfully."
  121. ),
  122. "create-vector-index": (
  123. "Vector index creation task queued successfully."
  124. if self.providers.orchestration.config.provider != "simple"
  125. else "Vector index creation task completed successfully."
  126. ),
  127. "delete-vector-index": (
  128. "Vector index deletion task queued successfully."
  129. if self.providers.orchestration.config.provider != "simple"
  130. else "Vector index deletion task completed successfully."
  131. ),
  132. "select-vector-index": (
  133. "Vector index selection task queued successfully."
  134. if self.providers.orchestration.config.provider != "simple"
  135. else "Vector index selection task completed successfully."
  136. ),
  137. },
  138. )
  139. def _prepare_ingestion_config(
  140. self,
  141. ingestion_mode: IngestionMode,
  142. ingestion_config: Optional[IngestionConfig],
  143. ) -> IngestionConfig:
  144. # If not custom, start from defaults
  145. if ingestion_mode != IngestionMode.custom:
  146. effective_config = IngestionConfig.get_default(
  147. ingestion_mode.value, app=self.providers.auth.config.app
  148. )
  149. if ingestion_config:
  150. effective_config = merge_ingestion_config(
  151. effective_config, ingestion_config
  152. )
  153. else:
  154. effective_config = ingestion_config or IngestionConfig(
  155. app=self.providers.auth.config.app
  156. )
  157. effective_config.validate_config()
  158. return effective_config
  159. def _setup_routes(self):
  160. @self.router.post(
  161. "/documents",
  162. dependencies=[Depends(self.rate_limit_dependency)],
  163. status_code=202,
  164. summary="Create a new document",
  165. openapi_extra={
  166. "x-codeSamples": [
  167. {
  168. "lang": "Python",
  169. "source": textwrap.dedent("""
  170. from r2r import R2RClient
  171. client = R2RClient()
  172. # when using auth, do client.login(...)
  173. response = client.documents.create(
  174. file_path="pg_essay_1.html",
  175. metadata={"metadata_1":"some random metadata"},
  176. id=None
  177. )
  178. """),
  179. },
  180. {
  181. "lang": "JavaScript",
  182. "source": textwrap.dedent("""
  183. const { r2rClient } = require("r2r-js");
  184. const client = new r2rClient();
  185. function main() {
  186. const response = await client.documents.create({
  187. file: { path: "examples/data/marmeladov.txt", name: "marmeladov.txt" },
  188. metadata: { title: "marmeladov.txt" },
  189. });
  190. }
  191. main();
  192. """),
  193. },
  194. {
  195. "lang": "cURL",
  196. "source": textwrap.dedent("""
  197. curl -X POST "https://api.example.com/v3/documents" \\
  198. -H "Content-Type: multipart/form-data" \\
  199. -H "Authorization: Bearer YOUR_API_KEY" \\
  200. -F "file=@pg_essay_1.html;type=text/html" \\
  201. -F 'metadata={}' \\
  202. -F 'id=null'
  203. """),
  204. },
  205. ]
  206. },
  207. )
  208. @self.base_endpoint
  209. async def create_document(
  210. file: Optional[UploadFile] = File(
  211. None,
  212. description="The file to ingest. Exactly one of file, raw_text, or chunks must be provided.",
  213. ),
  214. raw_text: Optional[str] = Form(
  215. None,
  216. description="Raw text content to ingest. Exactly one of file, raw_text, or chunks must be provided.",
  217. ),
  218. chunks: Optional[Json[list[str]]] = Form(
  219. None,
  220. description="Pre-processed text chunks to ingest. Exactly one of file, raw_text, or chunks must be provided.",
  221. ),
  222. id: Optional[UUID] = Form(
  223. None,
  224. description="The ID of the document. If not provided, a new ID will be generated.",
  225. ),
  226. collection_ids: Optional[Json[list[UUID]]] = Form(
  227. None,
  228. description="Collection IDs to associate with the document. If none are provided, the document will be assigned to the user's default collection.",
  229. ),
  230. metadata: Optional[Json[dict]] = Form(
  231. None,
  232. description="Metadata to associate with the document, such as title, description, or custom fields.",
  233. ),
  234. ingestion_mode: IngestionMode = Form(
  235. default=IngestionMode.custom,
  236. description=(
  237. "Ingestion modes:\n"
  238. "- `hi-res`: Thorough ingestion with full summaries and enrichment.\n"
  239. "- `ocr`: OCR via Mistral and full summaries.\n"
  240. "- `fast`: Quick ingestion with minimal enrichment and no summaries.\n"
  241. "- `custom`: Full control via `ingestion_config`.\n\n"
  242. "If `filters` or `limit` (in `ingestion_config`) are provided alongside `hi-res` or `fast`, "
  243. "they will override the default settings for that mode."
  244. ),
  245. ),
  246. ingestion_config: Optional[Json[IngestionConfig]] = Form(
  247. None,
  248. description="An optional dictionary to override the default chunking configuration for the ingestion process. If not provided, the system will use the default server-side chunking configuration.",
  249. ),
  250. run_with_orchestration: Optional[bool] = Form(
  251. True,
  252. description="Whether or not ingestion runs with orchestration, default is `True`. When set to `False`, the ingestion process will run synchronous and directly return the result.",
  253. ),
  254. auth_user=Depends(self.providers.auth.auth_wrapper()),
  255. ) -> WrappedIngestionResponse:
  256. """
  257. Creates a new Document object from an input file, text content, or chunks. The chosen `ingestion_mode` determines
  258. how the ingestion process is configured:
  259. **Ingestion Modes:**
  260. - `hi-res`: Comprehensive parsing and enrichment, including summaries and possibly more thorough parsing.
  261. - `fast`: Speed-focused ingestion that skips certain enrichment steps like summaries.
  262. - `custom`: Provide a full `ingestion_config` to customize the entire ingestion process.
  263. Either a file or text content must be provided, but not both. Documents are shared through `Collections` which allow for tightly specified cross-user interactions.
  264. The ingestion process runs asynchronously and its progress can be tracked using the returned
  265. task_id.
  266. """
  267. if not auth_user.is_superuser:
  268. user_document_count = 0
  269. '''
  270. (
  271. await self.services.management.documents_overview(
  272. user_ids=[auth_user.id],
  273. offset=0,
  274. limit=1,
  275. )
  276. )["total_entries"]
  277. '''
  278. user_max_documents = (
  279. await self.services.management.get_user_max_documents(
  280. auth_user.id
  281. )
  282. )
  283. if user_document_count >= user_max_documents:
  284. raise R2RException(
  285. status_code=403,
  286. message=f"User has reached the maximum number of documents allowed ({user_max_documents}).",
  287. )
  288. # Get chunks using the vector handler's list_chunks method
  289. user_chunk_count = 0
  290. '''
  291. (
  292. await self.services.ingestion.list_chunks(
  293. filters={"owner_id": {"$eq": str(auth_user.id)}},
  294. offset=0,
  295. limit=1,
  296. )
  297. )["total_entries"]
  298. '''
  299. user_max_chunks = (
  300. await self.services.management.get_user_max_chunks(
  301. auth_user.id
  302. )
  303. )
  304. if user_chunk_count >= user_max_chunks:
  305. raise R2RException(
  306. status_code=403,
  307. message=f"User has reached the maximum number of chunks allowed ({user_max_chunks}).",
  308. )
  309. user_collections_count = 0
  310. '''
  311. (
  312. await self.services.management.collections_overview(
  313. user_ids=[auth_user.id],
  314. offset=0,
  315. limit=1,
  316. )
  317. )["total_entries"]
  318. '''
  319. user_max_collections = (
  320. await self.services.management.get_user_max_collections(
  321. auth_user.id
  322. )
  323. )
  324. if user_collections_count >= user_max_collections: # type: ignore
  325. raise R2RException(
  326. status_code=403,
  327. message=f"User has reached the maximum number of collections allowed ({user_max_collections}).",
  328. )
  329. effective_ingestion_config = self._prepare_ingestion_config(
  330. ingestion_mode=ingestion_mode,
  331. ingestion_config=ingestion_config,
  332. )
  333. if not file and not raw_text and not chunks:
  334. raise R2RException(
  335. status_code=422,
  336. message="Either a `file`, `raw_text`, or `chunks` must be provided.",
  337. )
  338. if (
  339. (file and raw_text)
  340. or (file and chunks)
  341. or (raw_text and chunks)
  342. ):
  343. raise R2RException(
  344. status_code=422,
  345. message="Only one of `file`, `raw_text`, or `chunks` may be provided.",
  346. )
  347. # Check if the user is a superuser
  348. metadata = metadata or {}
  349. if chunks:
  350. if len(chunks) == 0:
  351. raise R2RException("Empty list of chunks provided", 400)
  352. if len(chunks) > MAX_CHUNKS_PER_REQUEST:
  353. raise R2RException(
  354. f"Maximum of {MAX_CHUNKS_PER_REQUEST} chunks per request",
  355. 400,
  356. )
  357. document_id = id or generate_document_id(
  358. "".join(chunks), auth_user.id
  359. )
  360. # FIXME: Metadata doesn't seem to be getting passed through
  361. raw_chunks_for_doc = [
  362. UnprocessedChunk(
  363. text=chunk,
  364. metadata=metadata,
  365. id=generate_id(),
  366. )
  367. for chunk in chunks
  368. ]
  369. # Prepare workflow input
  370. workflow_input = {
  371. "document_id": str(document_id),
  372. "chunks": [
  373. chunk.model_dump(mode="json")
  374. for chunk in raw_chunks_for_doc
  375. ],
  376. "collection_ids": (
  377. [str(cid) for cid in collection_ids]
  378. if collection_ids
  379. else None
  380. ),
  381. "metadata": metadata, # Base metadata for the document
  382. "user": auth_user.model_dump_json(),
  383. "ingestion_config": effective_ingestion_config.model_dump(
  384. mode="json"
  385. ),
  386. }
  387. if run_with_orchestration:
  388. try:
  389. # Run ingestion with orchestration
  390. raw_message = (
  391. await self.providers.orchestration.run_workflow(
  392. "ingest-chunks",
  393. {"request": workflow_input},
  394. options={
  395. "additional_metadata": {
  396. "document_id": str(document_id),
  397. }
  398. },
  399. )
  400. )
  401. raw_message["document_id"] = str(document_id)
  402. return raw_message # type: ignore
  403. except Exception as e: # TODO: Need to find specific errors that we should be excepting (gRPC most likely?)
  404. logger.error(
  405. f"Error running orchestrated ingestion: {e} \n\nAttempting to run without orchestration."
  406. )
  407. logger.info("Running chunk ingestion without orchestration.")
  408. from core.main.orchestration import simple_ingestion_factory
  409. simple_ingestor = simple_ingestion_factory(
  410. self.services.ingestion
  411. )
  412. await simple_ingestor["ingest-chunks"](workflow_input)
  413. return { # type: ignore
  414. "message": "Document created and ingested successfully.",
  415. "document_id": str(document_id),
  416. "task_id": None,
  417. }
  418. else:
  419. if file:
  420. file_data = await self._process_file(file)
  421. if metadata.get("title"):
  422. file_data["filename"] = metadata["title"]
  423. if not file_data["filename"]:
  424. raise R2RException(
  425. status_code=422,
  426. message="Uploaded file must have a filename.",
  427. )
  428. file_ext = file_data["filename"].split(".")[
  429. -1
  430. ] # e.g. "pdf", "txt"
  431. max_allowed_size = await self.services.management.get_max_upload_size_by_type(
  432. user_id=auth_user.id, file_type_or_ext=file_ext
  433. )
  434. content_length = file_data["content_length"]
  435. if content_length > max_allowed_size:
  436. raise R2RException(
  437. status_code=413, # HTTP 413: Payload Too Large
  438. message=(
  439. f"File size exceeds maximum of {max_allowed_size} bytes "
  440. f"for extension '{file_ext}'."
  441. ),
  442. )
  443. file_content = BytesIO(
  444. base64.b64decode(file_data["content"])
  445. )
  446. file_data.pop("content", None)
  447. document_id = id or generate_document_id(
  448. file_data["filename"], auth_user.id
  449. )
  450. elif raw_text:
  451. content_length = len(raw_text)
  452. file_content = BytesIO(raw_text.encode("utf-8"))
  453. document_id = id or generate_document_id(
  454. raw_text, auth_user.id
  455. )
  456. title = metadata.get("title", None)
  457. title = title + ".txt" if title else None
  458. file_data = {
  459. "filename": title or "N/A",
  460. "content_type": "text/plain",
  461. }
  462. else:
  463. raise R2RException(
  464. status_code=422,
  465. message="Either a file or content must be provided.",
  466. )
  467. workflow_input = {
  468. "file_data": file_data,
  469. "document_id": str(document_id),
  470. "collection_ids": (
  471. [str(cid) for cid in collection_ids]
  472. if collection_ids
  473. else None
  474. ),
  475. "metadata": metadata,
  476. "ingestion_config": effective_ingestion_config.model_dump(
  477. mode="json"
  478. ),
  479. "user": auth_user.model_dump_json(),
  480. "size_in_bytes": content_length,
  481. "version": "v0",
  482. }
  483. file_name = file_data["filename"]
  484. await self.providers.file.store_file(
  485. document_id,
  486. file_name,
  487. file_content,
  488. file_data["content_type"],
  489. )
  490. await self.services.ingestion.ingest_file_ingress(
  491. file_data=workflow_input["file_data"],
  492. user=auth_user,
  493. document_id=workflow_input["document_id"],
  494. size_in_bytes=workflow_input["size_in_bytes"],
  495. metadata=workflow_input["metadata"],
  496. version=workflow_input["version"],
  497. )
  498. if run_with_orchestration:
  499. try:
  500. # TODO - Modify create_chunks so that we can add chunks to existing document
  501. workflow_result: dict[
  502. str, str | None
  503. ] = await self.providers.orchestration.run_workflow( # type: ignore
  504. "ingest-files",
  505. {"request": workflow_input},
  506. options={
  507. "additional_metadata": {
  508. "document_id": str(document_id),
  509. }
  510. },
  511. )
  512. workflow_result["document_id"] = str(document_id)
  513. return workflow_result # type: ignore
  514. except Exception as e: # TODO: Need to find specific error (gRPC most likely?)
  515. logger.error(
  516. f"Error running orchestrated ingestion: {e} \n\nAttempting to run without orchestration."
  517. )
  518. logger.info(
  519. f"Running ingestion without orchestration for file {file_name} and document_id {document_id}."
  520. )
  521. # TODO - Clean up implementation logic here to be more explicitly `synchronous`
  522. from core.main.orchestration import simple_ingestion_factory
  523. simple_ingestor = simple_ingestion_factory(self.services.ingestion)
  524. await simple_ingestor["ingest-files"](workflow_input)
  525. return { # type: ignore
  526. "message": "Document created and ingested successfully.",
  527. "document_id": str(document_id),
  528. "task_id": None,
  529. }
  530. @self.router.patch(
  531. "/documents/{id}/metadata",
  532. dependencies=[Depends(self.rate_limit_dependency)],
  533. summary="Append metadata to a document",
  534. openapi_extra={
  535. "x-codeSamples": [
  536. {
  537. "lang": "Python",
  538. "source": textwrap.dedent("""
  539. from r2r import R2RClient
  540. client = R2RClient()
  541. # when using auth, do client.login(...)
  542. response = client.documents.append_metadata(
  543. id="9fbe403b-c11c-5aae-8ade-ef22980c3ad1",
  544. metadata=[{"key": "new_key", "value": "new_value"}]
  545. )
  546. """),
  547. },
  548. {
  549. "lang": "JavaScript",
  550. "source": textwrap.dedent("""
  551. const { r2rClient } = require("r2r-js");
  552. const client = new r2rClient();
  553. function main() {
  554. const response = await client.documents.appendMetadata({
  555. id: "9fbe403b-c11c-5aae-8ade-ef22980c3ad1",
  556. metadata: [{ key: "new_key", value: "new_value" }],
  557. });
  558. }
  559. main();
  560. """),
  561. },
  562. ]
  563. },
  564. )
  565. @self.base_endpoint
  566. async def patch_metadata(
  567. id: UUID = Path(
  568. ...,
  569. description="The ID of the document to append metadata to.",
  570. ),
  571. metadata: list[dict] = Body(
  572. ...,
  573. description="Metadata to append to the document.",
  574. ),
  575. auth_user=Depends(self.providers.auth.auth_wrapper()),
  576. ) -> WrappedDocumentResponse:
  577. """Appends metadata to a document. This endpoint allows adding new metadata fields or updating existing ones."""
  578. request_user_ids = (
  579. None if auth_user.is_superuser else [auth_user.id]
  580. )
  581. documents_overview_response = (
  582. await self.services.management.documents_overview(
  583. user_ids=request_user_ids,
  584. document_ids=[id],
  585. offset=0,
  586. limit=1,
  587. )
  588. )
  589. results = documents_overview_response["results"]
  590. if len(results) == 0:
  591. raise R2RException("Document not found.", 404)
  592. return await self.services.management.update_document_metadata(
  593. document_id=id,
  594. metadata=metadata,
  595. overwrite=False,
  596. )
  597. @self.router.put(
  598. "/documents/{id}/metadata",
  599. dependencies=[Depends(self.rate_limit_dependency)],
  600. summary="Replace metadata of a document",
  601. openapi_extra={
  602. "x-codeSamples": [
  603. {
  604. "lang": "Python",
  605. "source": textwrap.dedent("""
  606. from r2r import R2RClient
  607. client = R2RClient()
  608. # when using auth, do client.login(...)
  609. response = client.documents.replace_metadata(
  610. id="9fbe403b-c11c-5aae-8ade-ef22980c3ad1",
  611. metadata=[{"key": "new_key", "value": "new_value"}]
  612. )
  613. """),
  614. },
  615. {
  616. "lang": "JavaScript",
  617. "source": textwrap.dedent("""
  618. const { r2rClient } = require("r2r-js");
  619. const client = new r2rClient();
  620. function main() {
  621. const response = await client.documents.replaceMetadata({
  622. id: "9fbe403b-c11c-5aae-8ade-ef22980c3ad1",
  623. metadata: [{ key: "new_key", value: "new_value" }],
  624. });
  625. }
  626. main();
  627. """),
  628. },
  629. ]
  630. },
  631. )
  632. @self.base_endpoint
  633. async def put_metadata(
  634. id: UUID = Path(
  635. ...,
  636. description="The ID of the document to append metadata to.",
  637. ),
  638. metadata: list[dict] = Body(
  639. ...,
  640. description="Metadata to append to the document.",
  641. ),
  642. auth_user=Depends(self.providers.auth.auth_wrapper()),
  643. ) -> WrappedDocumentResponse:
  644. """Replaces metadata in a document. This endpoint allows overwriting existing metadata fields."""
  645. request_user_ids = (
  646. None if auth_user.is_superuser else [auth_user.id]
  647. )
  648. documents_overview_response = (
  649. await self.services.management.documents_overview(
  650. user_ids=request_user_ids,
  651. document_ids=[id],
  652. offset=0,
  653. limit=1,
  654. )
  655. )
  656. results = documents_overview_response["results"]
  657. if len(results) == 0:
  658. raise R2RException("Document not found.", 404)
  659. return await self.services.management.update_document_metadata(
  660. document_id=id,
  661. metadata=metadata,
  662. overwrite=True,
  663. )
  664. @self.router.post(
  665. "/documents/export",
  666. summary="Export documents to CSV",
  667. dependencies=[Depends(self.rate_limit_dependency)],
  668. openapi_extra={
  669. "x-codeSamples": [
  670. {
  671. "lang": "Python",
  672. "source": textwrap.dedent("""
  673. from r2r import R2RClient
  674. client = R2RClient("http://localhost:7272")
  675. # when using auth, do client.login(...)
  676. response = client.documents.export(
  677. output_path="export.csv",
  678. columns=["id", "title", "created_at"],
  679. include_header=True,
  680. )
  681. """),
  682. },
  683. {
  684. "lang": "JavaScript",
  685. "source": textwrap.dedent("""
  686. const { r2rClient } = require("r2r-js");
  687. const client = new r2rClient("http://localhost:7272");
  688. function main() {
  689. await client.documents.export({
  690. outputPath: "export.csv",
  691. columns: ["id", "title", "created_at"],
  692. includeHeader: true,
  693. });
  694. }
  695. main();
  696. """),
  697. },
  698. {
  699. "lang": "cURL",
  700. "source": textwrap.dedent("""
  701. curl -X POST "http://127.0.0.1:7272/v3/documents/export" \
  702. -H "Authorization: Bearer YOUR_API_KEY" \
  703. -H "Content-Type: application/json" \
  704. -H "Accept: text/csv" \
  705. -d '{ "columns": ["id", "title", "created_at"], "include_header": true }' \
  706. --output export.csv
  707. """),
  708. },
  709. ]
  710. },
  711. )
  712. @self.base_endpoint
  713. async def export_documents(
  714. background_tasks: BackgroundTasks,
  715. columns: Optional[list[str]] = Body(
  716. None, description="Specific columns to export"
  717. ),
  718. filters: Optional[dict] = Body(
  719. None, description="Filters to apply to the export"
  720. ),
  721. include_header: Optional[bool] = Body(
  722. True, description="Whether to include column headers"
  723. ),
  724. auth_user=Depends(self.providers.auth.auth_wrapper()),
  725. ) -> FileResponse:
  726. """Export documents as a downloadable CSV file."""
  727. if not auth_user.is_superuser:
  728. raise R2RException(
  729. "Only a superuser can export data.",
  730. 403,
  731. )
  732. (
  733. csv_file_path,
  734. temp_file,
  735. ) = await self.services.management.export_documents(
  736. columns=columns,
  737. filters=filters,
  738. include_header=include_header
  739. if include_header is not None
  740. else True,
  741. )
  742. background_tasks.add_task(temp_file.close)
  743. return FileResponse(
  744. path=csv_file_path,
  745. media_type="text/csv",
  746. filename="documents_export.csv",
  747. )
  748. @self.router.get(
  749. "/documents/download_zip",
  750. dependencies=[Depends(self.rate_limit_dependency)],
  751. response_class=StreamingResponse,
  752. summary="Export multiple documents as zip",
  753. openapi_extra={
  754. "x-codeSamples": [
  755. {
  756. "lang": "Python",
  757. "source": textwrap.dedent("""
  758. client.documents.download_zip(
  759. document_ids=["uuid1", "uuid2"],
  760. start_date="2024-01-01",
  761. end_date="2024-12-31"
  762. )
  763. """),
  764. },
  765. {
  766. "lang": "cURL",
  767. "source": textwrap.dedent("""
  768. curl -X GET "https://api.example.com/v3/documents/download_zip?document_ids=uuid1,uuid2&start_date=2024-01-01&end_date=2024-12-31" \\
  769. -H "Authorization: Bearer YOUR_API_KEY"
  770. """),
  771. },
  772. ]
  773. },
  774. )
  775. @self.base_endpoint
  776. async def export_files(
  777. document_ids: Optional[list[UUID]] = Query(
  778. None,
  779. description="List of document IDs to include in the export. If not provided, all accessible documents will be included.",
  780. ),
  781. start_date: Optional[datetime] = Query(
  782. None,
  783. description="Filter documents created on or after this date.",
  784. ),
  785. end_date: Optional[datetime] = Query(
  786. None,
  787. description="Filter documents created before this date.",
  788. ),
  789. auth_user=Depends(self.providers.auth.auth_wrapper()),
  790. ) -> StreamingResponse:
  791. """Export multiple documents as a zip file. Documents can be
  792. filtered by IDs and/or date range.
  793. The endpoint allows downloading:
  794. - Specific documents by providing their IDs
  795. - Documents within a date range
  796. - All accessible documents if no filters are provided
  797. Files are streamed as a zip archive to handle potentially large downloads efficiently.
  798. """
  799. if not auth_user.is_superuser:
  800. # For non-superusers, verify access to requested documents
  801. if document_ids:
  802. documents_overview = (
  803. await self.services.management.documents_overview(
  804. user_ids=[auth_user.id],
  805. document_ids=document_ids,
  806. offset=0,
  807. limit=len(document_ids),
  808. )
  809. )
  810. if len(documents_overview["results"]) != len(document_ids):
  811. raise R2RException(
  812. status_code=403,
  813. message="You don't have access to one or more requested documents.",
  814. )
  815. if not document_ids:
  816. raise R2RException(
  817. status_code=403,
  818. message="Non-superusers must provide document IDs to export.",
  819. )
  820. (
  821. zip_name,
  822. zip_content,
  823. zip_size,
  824. ) = await self.services.management.export_files(
  825. document_ids=document_ids,
  826. start_date=start_date,
  827. end_date=end_date,
  828. )
  829. encoded_filename = quote(zip_name)
  830. async def stream_file():
  831. yield zip_content.getvalue()
  832. return StreamingResponse(
  833. stream_file(),
  834. media_type="application/zip",
  835. headers={
  836. "Content-Disposition": f"attachment; filename*=UTF-8''{encoded_filename}",
  837. "Content-Length": str(zip_size),
  838. },
  839. )
  840. @self.router.get(
  841. "/documents",
  842. dependencies=[Depends(self.rate_limit_dependency)],
  843. summary="List documents",
  844. openapi_extra={
  845. "x-codeSamples": [
  846. {
  847. "lang": "Python",
  848. "source": textwrap.dedent("""
  849. from r2r import R2RClient
  850. client = R2RClient()
  851. # when using auth, do client.login(...)
  852. response = client.documents.list(
  853. limit=10,
  854. offset=0
  855. )
  856. """),
  857. },
  858. {
  859. "lang": "JavaScript",
  860. "source": textwrap.dedent("""
  861. const { r2rClient } = require("r2r-js");
  862. const client = new r2rClient();
  863. function main() {
  864. const response = await client.documents.list({
  865. limit: 10,
  866. offset: 0,
  867. });
  868. }
  869. main();
  870. """),
  871. },
  872. {
  873. "lang": "cURL",
  874. "source": textwrap.dedent("""
  875. curl -X GET "https://api.example.com/v3/documents" \\
  876. -H "Authorization: Bearer YOUR_API_KEY"
  877. """),
  878. },
  879. ]
  880. },
  881. )
  882. @self.base_endpoint
  883. async def get_documents(
  884. ids: list[str] = Query(
  885. [],
  886. description="A list of document IDs to retrieve. If not provided, all documents will be returned.",
  887. ),
  888. offset: int = Query(
  889. 0,
  890. ge=0,
  891. description="Specifies the number of objects to skip. Defaults to 0.",
  892. ),
  893. limit: int = Query(
  894. 100,
  895. ge=1,
  896. le=1000,
  897. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  898. ),
  899. include_summary_embeddings: bool = Query(
  900. False,
  901. description="Specifies whether or not to include embeddings of each document summary.",
  902. ),
  903. owner_only: bool = Query(
  904. False,
  905. description="If true, only returns documents owned by the user, not all accessible documents.",
  906. ),
  907. auth_user=Depends(self.providers.auth.auth_wrapper()),
  908. ) -> WrappedDocumentsResponse:
  909. """Returns a paginated list of documents the authenticated user has
  910. access to.
  911. Results can be filtered by providing specific document IDs. Regular
  912. users will only see documents they own or have access to through
  913. collections. Superusers can see all documents.
  914. The documents are returned in order of last modification, with most
  915. recent first.
  916. """
  917. if auth_user.is_superuser:
  918. requesting_user_id = [auth_user.id] if owner_only else None
  919. filter_collection_ids = None
  920. else:
  921. requesting_user_id = [auth_user.id]
  922. filter_collection_ids = auth_user.collection_ids
  923. document_uuids = [
  924. UUID(document_id)
  925. for document_id in ids
  926. if document_id and document_id.strip()
  927. ]
  928. documents_overview_response = (
  929. await self.services.management.documents_overview(
  930. user_ids=requesting_user_id,
  931. collection_ids=filter_collection_ids,
  932. document_ids=document_uuids,
  933. offset=offset,
  934. limit=limit,
  935. owner_only=owner_only,
  936. )
  937. )
  938. if not include_summary_embeddings:
  939. for document in documents_overview_response["results"]:
  940. document.summary_embedding = None
  941. return ( # type: ignore
  942. documents_overview_response["results"],
  943. {
  944. "total_entries": documents_overview_response[
  945. "total_entries"
  946. ]
  947. },
  948. )
  949. @self.router.get(
  950. "/documents/{id}",
  951. dependencies=[Depends(self.rate_limit_dependency)],
  952. summary="Retrieve a document",
  953. openapi_extra={
  954. "x-codeSamples": [
  955. {
  956. "lang": "Python",
  957. "source": textwrap.dedent("""
  958. from r2r import R2RClient
  959. client = R2RClient()
  960. # when using auth, do client.login(...)
  961. response = client.documents.retrieve(
  962. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa"
  963. )
  964. """),
  965. },
  966. {
  967. "lang": "JavaScript",
  968. "source": textwrap.dedent("""
  969. const { r2rClient } = require("r2r-js");
  970. const client = new r2rClient();
  971. function main() {
  972. const response = await client.documents.retrieve({
  973. id: "b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  974. });
  975. }
  976. main();
  977. """),
  978. },
  979. {
  980. "lang": "cURL",
  981. "source": textwrap.dedent("""
  982. curl -X GET "https://api.example.com/v3/documents/b4ac4dd6-5f27-596e-a55b-7cf242ca30aa" \\
  983. -H "Authorization: Bearer YOUR_API_KEY"
  984. """),
  985. },
  986. ]
  987. },
  988. )
  989. @self.base_endpoint
  990. async def get_document(
  991. id: UUID = Path(
  992. ...,
  993. description="The ID of the document to retrieve.",
  994. ),
  995. auth_user=Depends(self.providers.auth.auth_wrapper()),
  996. ) -> WrappedDocumentResponse:
  997. """Retrieves detailed information about a specific document by its
  998. ID.
  999. This endpoint returns the document's metadata, status, and system information. It does not
  1000. return the document's content - use the `/documents/{id}/download` endpoint for that.
  1001. Users can only retrieve documents they own or have access to through collections.
  1002. Superusers can retrieve any document.
  1003. """
  1004. request_user_ids = (
  1005. None if auth_user.is_superuser else [auth_user.id]
  1006. )
  1007. filter_collection_ids = (
  1008. None if auth_user.is_superuser else auth_user.collection_ids
  1009. )
  1010. documents_overview_response = await self.services.management.documents_overview( # FIXME: This was using the pagination defaults from before... We need to review if this is as intended.
  1011. user_ids=request_user_ids,
  1012. collection_ids=filter_collection_ids,
  1013. document_ids=[id],
  1014. offset=0,
  1015. limit=100,
  1016. )
  1017. results = documents_overview_response["results"]
  1018. if len(results) == 0:
  1019. raise R2RException("Document not found.", 404)
  1020. return results[0]
  1021. @self.router.get(
  1022. "/documents/{id}/chunks",
  1023. dependencies=[Depends(self.rate_limit_dependency)],
  1024. summary="List document chunks",
  1025. openapi_extra={
  1026. "x-codeSamples": [
  1027. {
  1028. "lang": "Python",
  1029. "source": textwrap.dedent("""
  1030. from r2r import R2RClient
  1031. client = R2RClient()
  1032. # when using auth, do client.login(...)
  1033. response = client.documents.list_chunks(
  1034. id="32b6a70f-a995-5c51-85d2-834f06283a1e"
  1035. )
  1036. """),
  1037. },
  1038. {
  1039. "lang": "JavaScript",
  1040. "source": textwrap.dedent("""
  1041. const { r2rClient } = require("r2r-js");
  1042. const client = new r2rClient();
  1043. function main() {
  1044. const response = await client.documents.listChunks({
  1045. id: "32b6a70f-a995-5c51-85d2-834f06283a1e",
  1046. });
  1047. }
  1048. main();
  1049. """),
  1050. },
  1051. {
  1052. "lang": "cURL",
  1053. "source": textwrap.dedent("""
  1054. curl -X GET "https://api.example.com/v3/documents/b4ac4dd6-5f27-596e-a55b-7cf242ca30aa/chunks" \\
  1055. -H "Authorization: Bearer YOUR_API_KEY"\
  1056. """),
  1057. },
  1058. ]
  1059. },
  1060. )
  1061. @self.base_endpoint
  1062. async def list_chunks(
  1063. id: UUID = Path(
  1064. ...,
  1065. description="The ID of the document to retrieve chunks for.",
  1066. ),
  1067. offset: int = Query(
  1068. 0,
  1069. ge=0,
  1070. description="Specifies the number of objects to skip. Defaults to 0.",
  1071. ),
  1072. limit: int = Query(
  1073. 100,
  1074. ge=1,
  1075. le=1000,
  1076. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  1077. ),
  1078. include_vectors: Optional[bool] = Query(
  1079. False,
  1080. description="Whether to include vector embeddings in the response.",
  1081. ),
  1082. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1083. ) -> WrappedChunksResponse:
  1084. """Retrieves the text chunks that were generated from a document
  1085. during ingestion. Chunks represent semantic sections of the
  1086. document and are used for retrieval and analysis.
  1087. Users can only access chunks from documents they own or have access
  1088. to through collections. Vector embeddings are only included if
  1089. specifically requested.
  1090. Results are returned in chunk sequence order, representing their
  1091. position in the original document.
  1092. """
  1093. list_document_chunks = (
  1094. await self.services.management.list_document_chunks(
  1095. document_id=id,
  1096. offset=offset,
  1097. limit=limit,
  1098. include_vectors=include_vectors or False,
  1099. )
  1100. )
  1101. if not list_document_chunks["results"]:
  1102. raise R2RException(
  1103. "No chunks found for the given document ID.", 404
  1104. )
  1105. is_owner = str(
  1106. list_document_chunks["results"][0].get("owner_id")
  1107. ) == str(auth_user.id)
  1108. document_collections = (
  1109. await self.services.management.collections_overview(
  1110. offset=0,
  1111. limit=-1,
  1112. document_ids=[id],
  1113. )
  1114. )
  1115. user_has_access = (
  1116. is_owner
  1117. or set(auth_user.collection_ids).intersection(
  1118. {ele.id for ele in document_collections["results"]} # type: ignore
  1119. )
  1120. != set()
  1121. )
  1122. if not user_has_access and not auth_user.is_superuser:
  1123. raise R2RException(
  1124. "Not authorized to access this document's chunks.", 403
  1125. )
  1126. return ( # type: ignore
  1127. list_document_chunks["results"],
  1128. {"total_entries": list_document_chunks["total_entries"]},
  1129. )
  1130. @self.router.get(
  1131. "/documents/{id}/download",
  1132. dependencies=[Depends(self.rate_limit_dependency)],
  1133. response_class=StreamingResponse,
  1134. summary="Download document content",
  1135. openapi_extra={
  1136. "x-codeSamples": [
  1137. {
  1138. "lang": "Python",
  1139. "source": textwrap.dedent("""
  1140. from r2r import R2RClient
  1141. client = R2RClient()
  1142. # when using auth, do client.login(...)
  1143. response = client.documents.download(
  1144. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa"
  1145. )
  1146. """),
  1147. },
  1148. {
  1149. "lang": "JavaScript",
  1150. "source": textwrap.dedent("""
  1151. const { r2rClient } = require("r2r-js");
  1152. const client = new r2rClient();
  1153. function main() {
  1154. const response = await client.documents.download({
  1155. id: "b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1156. });
  1157. }
  1158. main();
  1159. """),
  1160. },
  1161. {
  1162. "lang": "cURL",
  1163. "source": textwrap.dedent("""
  1164. curl -X GET "https://api.example.com/v3/documents/b4ac4dd6-5f27-596e-a55b-7cf242ca30aa/download" \\
  1165. -H "Authorization: Bearer YOUR_API_KEY"
  1166. """),
  1167. },
  1168. ]
  1169. },
  1170. )
  1171. @self.base_endpoint
  1172. async def get_document_file(
  1173. id: str = Path(..., description="Document ID"),
  1174. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1175. ) -> StreamingResponse:
  1176. """Downloads the original file content of a document.
  1177. For uploaded files, returns the original file with its proper MIME
  1178. type. For text-only documents, returns the content as plain text.
  1179. Users can only download documents they own or have access to
  1180. through collections.
  1181. """
  1182. try:
  1183. document_uuid = UUID(id)
  1184. except ValueError:
  1185. raise R2RException(
  1186. status_code=422, message="Invalid document ID format."
  1187. ) from None
  1188. # Retrieve the document's information
  1189. documents_overview_response = (
  1190. await self.services.management.documents_overview(
  1191. user_ids=None,
  1192. collection_ids=None,
  1193. document_ids=[document_uuid],
  1194. offset=0,
  1195. limit=1,
  1196. )
  1197. )
  1198. if not documents_overview_response["results"]:
  1199. raise R2RException("Document not found.", 404)
  1200. document = documents_overview_response["results"][0]
  1201. is_owner = str(document.owner_id) == str(auth_user.id)
  1202. if not auth_user.is_superuser and not is_owner:
  1203. document_collections = (
  1204. await self.services.management.collections_overview(
  1205. offset=0,
  1206. limit=-1,
  1207. document_ids=[document_uuid],
  1208. )
  1209. )
  1210. document_collection_ids = {
  1211. str(ele.id)
  1212. for ele in document_collections["results"] # type: ignore
  1213. }
  1214. user_collection_ids = {
  1215. str(cid) for cid in auth_user.collection_ids
  1216. }
  1217. has_collection_access = user_collection_ids.intersection(
  1218. document_collection_ids
  1219. )
  1220. if not has_collection_access:
  1221. raise R2RException(
  1222. "Not authorized to access this document.", 403
  1223. )
  1224. file_tuple = await self.services.management.download_file(
  1225. document_uuid
  1226. )
  1227. if not file_tuple:
  1228. raise R2RException(status_code=404, message="File not found.")
  1229. file_name, file_content, file_size = file_tuple
  1230. encoded_filename = quote(file_name)
  1231. mime_type, _ = mimetypes.guess_type(file_name)
  1232. if not mime_type:
  1233. mime_type = "application/octet-stream"
  1234. async def file_stream():
  1235. chunk_size = 1024 * 1024 # 1MB
  1236. while True:
  1237. data = file_content.read(chunk_size)
  1238. if not data:
  1239. break
  1240. yield data
  1241. return StreamingResponse(
  1242. file_stream(),
  1243. media_type=mime_type,
  1244. headers={
  1245. "Content-Disposition": f"inline; filename*=UTF-8''{encoded_filename}",
  1246. "Content-Length": str(file_size),
  1247. },
  1248. )
  1249. @self.router.delete(
  1250. "/documents/by-filter",
  1251. dependencies=[Depends(self.rate_limit_dependency)],
  1252. summary="Delete documents by filter",
  1253. openapi_extra={
  1254. "x-codeSamples": [
  1255. {
  1256. "lang": "Python",
  1257. "source": textwrap.dedent("""
  1258. from r2r import R2RClient
  1259. client = R2RClient()
  1260. # when using auth, do client.login(...)
  1261. response = client.documents.delete_by_filter(
  1262. filters={"document_type": {"$eq": "txt"}}
  1263. )
  1264. """),
  1265. },
  1266. {
  1267. "lang": "cURL",
  1268. "source": textwrap.dedent("""
  1269. curl -X DELETE "https://api.example.com/v3/documents/by-filter?filters=%7B%22document_type%22%3A%7B%22%24eq%22%3A%22text%22%7D%2C%22created_at%22%3A%7B%22%24lt%22%3A%222023-01-01T00%3A00%3A00Z%22%7D%7D" \\
  1270. -H "Authorization: Bearer YOUR_API_KEY"
  1271. """),
  1272. },
  1273. ]
  1274. },
  1275. )
  1276. @self.base_endpoint
  1277. async def delete_document_by_filter(
  1278. filters: Json[dict] = Body(
  1279. ..., description="JSON-encoded filters"
  1280. ),
  1281. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1282. ) -> WrappedBooleanResponse:
  1283. """Delete documents based on provided filters.
  1284. Allowed operators
  1285. include: `eq`, `neq`, `gt`, `gte`, `lt`, `lte`, `like`,
  1286. `ilike`, `in`, and `nin`. Deletion requests are limited to a
  1287. user's own documents.
  1288. """
  1289. filters_dict = {
  1290. "$and": [{"owner_id": {"$eq": str(auth_user.id)}}, filters]
  1291. }
  1292. await (
  1293. self.services.management.delete_documents_and_chunks_by_filter(
  1294. filters=filters_dict
  1295. )
  1296. )
  1297. return GenericBooleanResponse(success=True) # type: ignore
  1298. @self.router.delete(
  1299. "/documents/{id}",
  1300. dependencies=[Depends(self.rate_limit_dependency)],
  1301. summary="Delete a document",
  1302. openapi_extra={
  1303. "x-codeSamples": [
  1304. {
  1305. "lang": "Python",
  1306. "source": textwrap.dedent("""
  1307. from r2r import R2RClient
  1308. client = R2RClient()
  1309. # when using auth, do client.login(...)
  1310. response = client.documents.delete(
  1311. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa"
  1312. )
  1313. """),
  1314. },
  1315. {
  1316. "lang": "JavaScript",
  1317. "source": textwrap.dedent("""
  1318. const { r2rClient } = require("r2r-js");
  1319. const client = new r2rClient();
  1320. function main() {
  1321. const response = await client.documents.delete({
  1322. id: "b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1323. });
  1324. }
  1325. main();
  1326. """),
  1327. },
  1328. {
  1329. "lang": "cURL",
  1330. "source": textwrap.dedent("""
  1331. curl -X DELETE "https://api.example.com/v3/documents/b4ac4dd6-5f27-596e-a55b-7cf242ca30aa" \\
  1332. -H "Authorization: Bearer YOUR_API_KEY"
  1333. """),
  1334. },
  1335. ]
  1336. },
  1337. )
  1338. @self.base_endpoint
  1339. async def delete_document_by_id(
  1340. id: UUID = Path(..., description="Document ID"),
  1341. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1342. ) -> WrappedBooleanResponse:
  1343. """Delete a specific document. All chunks corresponding to the
  1344. document are deleted, and all other references to the document are
  1345. removed.
  1346. NOTE - Deletions do not yet impact the knowledge graph or other derived data. This feature is planned for a future release.
  1347. """
  1348. filters: dict[str, Any] = {"document_id": {"$eq": str(id)}}
  1349. if not auth_user.is_superuser:
  1350. filters = {
  1351. "$and": [
  1352. {"owner_id": {"$eq": str(auth_user.id)}},
  1353. {"document_id": {"$eq": str(id)}},
  1354. ]
  1355. }
  1356. await (
  1357. self.services.management.delete_documents_and_chunks_by_filter(
  1358. filters=filters
  1359. )
  1360. )
  1361. return GenericBooleanResponse(success=True) # type: ignore
  1362. @self.router.get(
  1363. "/documents/{id}/collections",
  1364. dependencies=[Depends(self.rate_limit_dependency)],
  1365. summary="List document collections",
  1366. openapi_extra={
  1367. "x-codeSamples": [
  1368. {
  1369. "lang": "Python",
  1370. "source": textwrap.dedent("""
  1371. from r2r import R2RClient
  1372. client = R2RClient()
  1373. # when using auth, do client.login(...)
  1374. response = client.documents.list_collections(
  1375. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa", offset=0, limit=10
  1376. )
  1377. """),
  1378. },
  1379. {
  1380. "lang": "JavaScript",
  1381. "source": textwrap.dedent("""
  1382. const { r2rClient } = require("r2r-js");
  1383. const client = new r2rClient();
  1384. function main() {
  1385. const response = await client.documents.listCollections({
  1386. id: "b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1387. });
  1388. }
  1389. main();
  1390. """),
  1391. },
  1392. {
  1393. "lang": "cURL",
  1394. "source": textwrap.dedent("""
  1395. curl -X GET "https://api.example.com/v3/documents/b4ac4dd6-5f27-596e-a55b-7cf242ca30aa/collections" \\
  1396. -H "Authorization: Bearer YOUR_API_KEY"
  1397. """),
  1398. },
  1399. ]
  1400. },
  1401. )
  1402. @self.base_endpoint
  1403. async def get_document_collections(
  1404. id: str = Path(..., description="Document ID"),
  1405. offset: int = Query(
  1406. 0,
  1407. ge=0,
  1408. description="Specifies the number of objects to skip. Defaults to 0.",
  1409. ),
  1410. limit: int = Query(
  1411. 100,
  1412. ge=1,
  1413. le=1000,
  1414. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  1415. ),
  1416. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1417. ) -> WrappedCollectionsResponse:
  1418. """Retrieves all collections that contain the specified document.
  1419. This endpoint is restricted to superusers only and provides a
  1420. system-wide view of document organization.
  1421. Collections are used to organize documents and manage access control. A document can belong
  1422. to multiple collections, and users can access documents through collection membership.
  1423. The results are paginated and ordered by collection creation date, with the most recently
  1424. created collections appearing first.
  1425. NOTE - This endpoint is only available to superusers, it will be extended to regular users in a future release.
  1426. """
  1427. if not auth_user.is_superuser:
  1428. raise R2RException(
  1429. "Only a superuser can get the collections belonging to a document.",
  1430. 403,
  1431. )
  1432. collections_response = (
  1433. await self.services.management.collections_overview(
  1434. offset=offset,
  1435. limit=limit,
  1436. document_ids=[UUID(id)], # Convert string ID to UUID
  1437. )
  1438. )
  1439. return collections_response["results"], { # type: ignore
  1440. "total_entries": collections_response["total_entries"]
  1441. }
  1442. @self.router.post(
  1443. "/documents/{id}/extract",
  1444. dependencies=[Depends(self.rate_limit_dependency)],
  1445. summary="Extract entities and relationships",
  1446. openapi_extra={
  1447. "x-codeSamples": [
  1448. {
  1449. "lang": "Python",
  1450. "source": textwrap.dedent("""
  1451. from r2r import R2RClient
  1452. client = R2RClient()
  1453. # when using auth, do client.login(...)
  1454. response = client.documents.extract(
  1455. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa"
  1456. )
  1457. """),
  1458. },
  1459. ],
  1460. },
  1461. )
  1462. @self.base_endpoint
  1463. async def extract(
  1464. id: UUID = Path(
  1465. ...,
  1466. description="The ID of the document to extract entities and relationships from.",
  1467. ),
  1468. settings: Optional[GraphCreationSettings] = Body(
  1469. default=None,
  1470. description="Settings for the entities and relationships extraction process.",
  1471. ),
  1472. run_with_orchestration: Optional[bool] = Body(
  1473. default=True,
  1474. description="Whether to run the entities and relationships extraction process with orchestration.",
  1475. ),
  1476. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1477. ) -> WrappedGenericMessageResponse:
  1478. """Extracts entities and relationships from a document.
  1479. The entities and relationships extraction process involves:
  1480. 1. Parsing documents into semantic chunks
  1481. 2. Extracting entities and relationships using LLMs
  1482. 3. Storing the created entities and relationships in the knowledge graph
  1483. 4. Preserving the document's metadata and content, and associating the elements with collections the document belongs to
  1484. """
  1485. settings = settings.dict() if settings else None # type: ignore
  1486. documents_overview_response = (
  1487. await self.services.management.documents_overview(
  1488. user_ids=(
  1489. None if auth_user.is_superuser else [auth_user.id]
  1490. ),
  1491. collection_ids=(
  1492. None
  1493. if auth_user.is_superuser
  1494. else auth_user.collection_ids
  1495. ),
  1496. document_ids=[id],
  1497. offset=0,
  1498. limit=1,
  1499. )
  1500. )["results"]
  1501. if len(documents_overview_response) == 0:
  1502. raise R2RException("Document not found.", 404)
  1503. if (
  1504. not auth_user.is_superuser
  1505. and auth_user.id != documents_overview_response[0].owner_id
  1506. ):
  1507. raise R2RException(
  1508. "Only a superuser can extract entities and relationships from a document they do not own.",
  1509. 403,
  1510. )
  1511. # Apply runtime settings overrides
  1512. server_graph_creation_settings = (
  1513. self.providers.database.config.graph_creation_settings
  1514. )
  1515. if settings:
  1516. server_graph_creation_settings = update_settings_from_dict(
  1517. server_settings=server_graph_creation_settings,
  1518. settings_dict=settings, # type: ignore
  1519. )
  1520. workflow_input = {
  1521. "document_id": str(id),
  1522. "graph_creation_settings": server_graph_creation_settings.model_dump_json(),
  1523. "user": auth_user.json(),
  1524. }
  1525. if run_with_orchestration:
  1526. try:
  1527. return await self.providers.orchestration.run_workflow( # type: ignore
  1528. "graph-extraction", {"request": workflow_input}, {}
  1529. )
  1530. except Exception as e: # TODO: Need to find specific errors that we should be excepting (gRPC most likely?)
  1531. logger.error(
  1532. f"Error running orchestrated extraction: {e} \n\nAttempting to run without orchestration."
  1533. )
  1534. from core.main.orchestration import (
  1535. simple_graph_search_results_factory,
  1536. )
  1537. logger.info("Running extract-triples without orchestration.")
  1538. simple_graph_search_results = simple_graph_search_results_factory(
  1539. self.services.graph
  1540. )
  1541. await simple_graph_search_results["graph-extraction"](
  1542. workflow_input
  1543. )
  1544. return { # type: ignore
  1545. "message": "Graph created successfully.",
  1546. "task_id": None,
  1547. }
  1548. @self.router.post(
  1549. "/documents/{id}/deduplicate",
  1550. dependencies=[Depends(self.rate_limit_dependency)],
  1551. summary="Deduplicate entities",
  1552. openapi_extra={
  1553. "x-codeSamples": [
  1554. {
  1555. "lang": "Python",
  1556. "source": textwrap.dedent("""
  1557. from r2r import R2RClient
  1558. client = R2RClient()
  1559. response = client.documents.deduplicate(
  1560. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa"
  1561. )
  1562. """),
  1563. },
  1564. {
  1565. "lang": "JavaScript",
  1566. "source": textwrap.dedent("""
  1567. const { r2rClient } = require("r2r-js");
  1568. const client = new r2rClient();
  1569. function main() {
  1570. const response = await client.documents.deduplicate({
  1571. id: "b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1572. });
  1573. }
  1574. main();
  1575. """),
  1576. },
  1577. {
  1578. "lang": "cURL",
  1579. "source": textwrap.dedent("""
  1580. curl -X POST "https://api.example.com/v3/documents/b4ac4dd6-5f27-596e-a55b-7cf242ca30aa/deduplicate" \\
  1581. -H "Authorization: Bearer YOUR_API_KEY"
  1582. """),
  1583. },
  1584. ],
  1585. },
  1586. )
  1587. @self.base_endpoint
  1588. async def deduplicate(
  1589. id: UUID = Path(
  1590. ...,
  1591. description="The ID of the document to extract entities and relationships from.",
  1592. ),
  1593. settings: Optional[GraphCreationSettings] = Body(
  1594. default=None,
  1595. description="Settings for the entities and relationships extraction process.",
  1596. ),
  1597. run_with_orchestration: Optional[bool] = Body(
  1598. default=True,
  1599. description="Whether to run the entities and relationships extraction process with orchestration.",
  1600. ),
  1601. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1602. ) -> WrappedGenericMessageResponse:
  1603. """Deduplicates entities from a document."""
  1604. settings = settings.model_dump() if settings else None # type: ignore
  1605. documents_overview_response = (
  1606. await self.services.management.documents_overview(
  1607. user_ids=(
  1608. None if auth_user.is_superuser else [auth_user.id]
  1609. ),
  1610. collection_ids=(
  1611. None
  1612. if auth_user.is_superuser
  1613. else auth_user.collection_ids
  1614. ),
  1615. document_ids=[id],
  1616. offset=0,
  1617. limit=1,
  1618. )
  1619. )["results"]
  1620. if len(documents_overview_response) == 0:
  1621. raise R2RException("Document not found.", 404)
  1622. if (
  1623. not auth_user.is_superuser
  1624. and auth_user.id != documents_overview_response[0].owner_id
  1625. ):
  1626. raise R2RException(
  1627. "Only a superuser can run deduplication on a document they do not own.",
  1628. 403,
  1629. )
  1630. # Apply runtime settings overrides
  1631. server_graph_creation_settings = (
  1632. self.providers.database.config.graph_creation_settings
  1633. )
  1634. if settings:
  1635. server_graph_creation_settings = update_settings_from_dict(
  1636. server_settings=server_graph_creation_settings,
  1637. settings_dict=settings, # type: ignore
  1638. )
  1639. if run_with_orchestration:
  1640. try:
  1641. workflow_input = {
  1642. "document_id": str(id),
  1643. }
  1644. return await self.providers.orchestration.run_workflow( # type: ignore
  1645. "graph-deduplication",
  1646. {"request": workflow_input},
  1647. {},
  1648. )
  1649. except Exception as e: # TODO: Need to find specific errors that we should be excepting (gRPC most likely?)
  1650. logger.error(
  1651. f"Error running orchestrated deduplication: {e} \n\nAttempting to run without orchestration."
  1652. )
  1653. from core.main.orchestration import (
  1654. simple_graph_search_results_factory,
  1655. )
  1656. logger.info(
  1657. "Running deduplicate-document-entities without orchestration."
  1658. )
  1659. simple_graph_search_results = simple_graph_search_results_factory(
  1660. self.services.graph
  1661. )
  1662. await simple_graph_search_results["graph-deduplication"](
  1663. workflow_input
  1664. )
  1665. return { # type: ignore
  1666. "message": "Graph created successfully.",
  1667. "task_id": None,
  1668. }
  1669. @self.router.get(
  1670. "/documents/{id}/entities",
  1671. dependencies=[Depends(self.rate_limit_dependency)],
  1672. summary="Lists the entities from the document",
  1673. openapi_extra={
  1674. "x-codeSamples": [
  1675. {
  1676. "lang": "Python",
  1677. "source": textwrap.dedent("""
  1678. from r2r import R2RClient
  1679. client = R2RClient()
  1680. # when using auth, do client.login(...)
  1681. response = client.documents.extract(
  1682. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa"
  1683. )
  1684. """),
  1685. },
  1686. ],
  1687. },
  1688. )
  1689. @self.base_endpoint
  1690. async def get_entities(
  1691. id: UUID = Path(
  1692. ...,
  1693. description="The ID of the document to retrieve entities from.",
  1694. ),
  1695. offset: int = Query(
  1696. 0,
  1697. ge=0,
  1698. description="Specifies the number of objects to skip. Defaults to 0.",
  1699. ),
  1700. limit: int = Query(
  1701. 100,
  1702. ge=1,
  1703. le=1000,
  1704. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  1705. ),
  1706. include_embeddings: Optional[bool] = Query(
  1707. False,
  1708. description="Whether to include vector embeddings in the response.",
  1709. ),
  1710. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1711. ) -> WrappedEntitiesResponse:
  1712. """Retrieves the entities that were extracted from a document.
  1713. These represent important semantic elements like people, places,
  1714. organizations, concepts, etc.
  1715. Users can only access entities from documents they own or have
  1716. access to through collections. Entity embeddings are only included
  1717. if specifically requested.
  1718. Results are returned in the order they were extracted from the
  1719. document.
  1720. """
  1721. # if (
  1722. # not auth_user.is_superuser
  1723. # and id not in auth_user.collection_ids
  1724. # ):
  1725. # raise R2RException(
  1726. # "The currently authenticated user does not have access to the specified collection.",
  1727. # 403,
  1728. # )
  1729. # First check if the document exists and user has access
  1730. documents_overview_response = (
  1731. await self.services.management.documents_overview(
  1732. user_ids=(
  1733. None if auth_user.is_superuser else [auth_user.id]
  1734. ),
  1735. collection_ids=(
  1736. None
  1737. if auth_user.is_superuser
  1738. else auth_user.collection_ids
  1739. ),
  1740. document_ids=[id],
  1741. offset=0,
  1742. limit=1,
  1743. )
  1744. )
  1745. if not documents_overview_response["results"]:
  1746. raise R2RException("Document not found.", 404)
  1747. # Get all entities for this document from the document_entity table
  1748. (
  1749. entities,
  1750. count,
  1751. ) = await self.providers.database.graphs_handler.entities.get(
  1752. parent_id=id,
  1753. store_type=StoreType.DOCUMENTS,
  1754. offset=offset,
  1755. limit=limit,
  1756. include_embeddings=include_embeddings or False,
  1757. )
  1758. return entities, {"total_entries": count} # type: ignore
  1759. @self.router.post(
  1760. "/documents/{id}/entities/export",
  1761. summary="Export document entities to CSV",
  1762. dependencies=[Depends(self.rate_limit_dependency)],
  1763. openapi_extra={
  1764. "x-codeSamples": [
  1765. {
  1766. "lang": "Python",
  1767. "source": textwrap.dedent("""
  1768. from r2r import R2RClient
  1769. client = R2RClient("http://localhost:7272")
  1770. # when using auth, do client.login(...)
  1771. response = client.documents.export_entities(
  1772. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1773. output_path="export.csv",
  1774. columns=["id", "title", "created_at"],
  1775. include_header=True,
  1776. )
  1777. """),
  1778. },
  1779. {
  1780. "lang": "JavaScript",
  1781. "source": textwrap.dedent("""
  1782. const { r2rClient } = require("r2r-js");
  1783. const client = new r2rClient("http://localhost:7272");
  1784. function main() {
  1785. await client.documents.exportEntities({
  1786. id: "b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1787. outputPath: "export.csv",
  1788. columns: ["id", "title", "created_at"],
  1789. includeHeader: true,
  1790. });
  1791. }
  1792. main();
  1793. """),
  1794. },
  1795. {
  1796. "lang": "cURL",
  1797. "source": textwrap.dedent("""
  1798. curl -X POST "http://127.0.0.1:7272/v3/documents/export_entities" \
  1799. -H "Authorization: Bearer YOUR_API_KEY" \
  1800. -H "Content-Type: application/json" \
  1801. -H "Accept: text/csv" \
  1802. -d '{ "columns": ["id", "title", "created_at"], "include_header": true }' \
  1803. --output export.csv
  1804. """),
  1805. },
  1806. ]
  1807. },
  1808. )
  1809. @self.base_endpoint
  1810. async def export_entities(
  1811. background_tasks: BackgroundTasks,
  1812. id: UUID = Path(
  1813. ...,
  1814. description="The ID of the document to export entities from.",
  1815. ),
  1816. columns: Optional[list[str]] = Body(
  1817. None, description="Specific columns to export"
  1818. ),
  1819. filters: Optional[dict] = Body(
  1820. None, description="Filters to apply to the export"
  1821. ),
  1822. include_header: Optional[bool] = Body(
  1823. True, description="Whether to include column headers"
  1824. ),
  1825. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1826. ) -> FileResponse:
  1827. """Export documents as a downloadable CSV file."""
  1828. if not auth_user.is_superuser:
  1829. raise R2RException(
  1830. "Only a superuser can export data.",
  1831. 403,
  1832. )
  1833. (
  1834. csv_file_path,
  1835. temp_file,
  1836. ) = await self.services.management.export_document_entities(
  1837. id=id,
  1838. columns=columns,
  1839. filters=filters,
  1840. include_header=include_header
  1841. if include_header is not None
  1842. else True,
  1843. )
  1844. background_tasks.add_task(temp_file.close)
  1845. return FileResponse(
  1846. path=csv_file_path,
  1847. media_type="text/csv",
  1848. filename="documents_export.csv",
  1849. )
  1850. @self.router.get(
  1851. "/documents/{id}/relationships",
  1852. dependencies=[Depends(self.rate_limit_dependency)],
  1853. summary="List document relationships",
  1854. openapi_extra={
  1855. "x-codeSamples": [
  1856. {
  1857. "lang": "Python",
  1858. "source": textwrap.dedent("""
  1859. from r2r import R2RClient
  1860. client = R2RClient()
  1861. # when using auth, do client.login(...)
  1862. response = client.documents.list_relationships(
  1863. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1864. offset=0,
  1865. limit=100
  1866. )
  1867. """),
  1868. },
  1869. {
  1870. "lang": "JavaScript",
  1871. "source": textwrap.dedent("""
  1872. const { r2rClient } = require("r2r-js");
  1873. const client = new r2rClient();
  1874. function main() {
  1875. const response = await client.documents.listRelationships({
  1876. id: "b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1877. offset: 0,
  1878. limit: 100,
  1879. });
  1880. }
  1881. main();
  1882. """),
  1883. },
  1884. {
  1885. "lang": "cURL",
  1886. "source": textwrap.dedent("""
  1887. curl -X GET "https://api.example.com/v3/documents/b4ac4dd6-5f27-596e-a55b-7cf242ca30aa/relationships" \\
  1888. -H "Authorization: Bearer YOUR_API_KEY"
  1889. """),
  1890. },
  1891. ]
  1892. },
  1893. )
  1894. @self.base_endpoint
  1895. async def get_relationships(
  1896. id: UUID = Path(
  1897. ...,
  1898. description="The ID of the document to retrieve relationships for.",
  1899. ),
  1900. offset: int = Query(
  1901. 0,
  1902. ge=0,
  1903. description="Specifies the number of objects to skip. Defaults to 0.",
  1904. ),
  1905. limit: int = Query(
  1906. 100,
  1907. ge=1,
  1908. le=1000,
  1909. description="Specifies a limit on the number of objects to return, ranging between 1 and 100. Defaults to 100.",
  1910. ),
  1911. entity_names: Optional[list[str]] = Query(
  1912. None,
  1913. description="Filter relationships by specific entity names.",
  1914. ),
  1915. relationship_types: Optional[list[str]] = Query(
  1916. None,
  1917. description="Filter relationships by specific relationship types.",
  1918. ),
  1919. auth_user=Depends(self.providers.auth.auth_wrapper()),
  1920. ) -> WrappedRelationshipsResponse:
  1921. """Retrieves the relationships between entities that were extracted
  1922. from a document. These represent connections and interactions
  1923. between entities found in the text.
  1924. Users can only access relationships from documents they own or have
  1925. access to through collections. Results can be filtered by entity
  1926. names and relationship types.
  1927. Results are returned in the order they were extracted from the
  1928. document.
  1929. """
  1930. # if (
  1931. # not auth_user.is_superuser
  1932. # and id not in auth_user.collection_ids
  1933. # ):
  1934. # raise R2RException(
  1935. # "The currently authenticated user does not have access to the specified collection.",
  1936. # 403,
  1937. # )
  1938. # First check if the document exists and user has access
  1939. documents_overview_response = (
  1940. await self.services.management.documents_overview(
  1941. user_ids=(
  1942. None if auth_user.is_superuser else [auth_user.id]
  1943. ),
  1944. collection_ids=(
  1945. None
  1946. if auth_user.is_superuser
  1947. else auth_user.collection_ids
  1948. ),
  1949. document_ids=[id],
  1950. offset=0,
  1951. limit=1,
  1952. )
  1953. )
  1954. if not documents_overview_response["results"]:
  1955. raise R2RException("Document not found.", 404)
  1956. # Get relationships for this document
  1957. (
  1958. relationships,
  1959. count,
  1960. ) = await self.providers.database.graphs_handler.relationships.get(
  1961. parent_id=id,
  1962. store_type=StoreType.DOCUMENTS,
  1963. entity_names=entity_names,
  1964. relationship_types=relationship_types,
  1965. offset=offset,
  1966. limit=limit,
  1967. )
  1968. return relationships, {"total_entries": count} # type: ignore
  1969. @self.router.post(
  1970. "/documents/{id}/relationships/export",
  1971. summary="Export document relationships to CSV",
  1972. dependencies=[Depends(self.rate_limit_dependency)],
  1973. openapi_extra={
  1974. "x-codeSamples": [
  1975. {
  1976. "lang": "Python",
  1977. "source": textwrap.dedent("""
  1978. from r2r import R2RClient
  1979. client = R2RClient("http://localhost:7272")
  1980. # when using auth, do client.login(...)
  1981. response = client.documents.export_entities(
  1982. id="b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1983. output_path="export.csv",
  1984. columns=["id", "title", "created_at"],
  1985. include_header=True,
  1986. )
  1987. """),
  1988. },
  1989. {
  1990. "lang": "JavaScript",
  1991. "source": textwrap.dedent("""
  1992. const { r2rClient } = require("r2r-js");
  1993. const client = new r2rClient("http://localhost:7272");
  1994. function main() {
  1995. await client.documents.exportEntities({
  1996. id: "b4ac4dd6-5f27-596e-a55b-7cf242ca30aa",
  1997. outputPath: "export.csv",
  1998. columns: ["id", "title", "created_at"],
  1999. includeHeader: true,
  2000. });
  2001. }
  2002. main();
  2003. """),
  2004. },
  2005. {
  2006. "lang": "cURL",
  2007. "source": textwrap.dedent("""
  2008. curl -X POST "http://127.0.0.1:7272/v3/documents/export_entities" \
  2009. -H "Authorization: Bearer YOUR_API_KEY" \
  2010. -H "Content-Type: application/json" \
  2011. -H "Accept: text/csv" \
  2012. -d '{ "columns": ["id", "title", "created_at"], "include_header": true }' \
  2013. --output export.csv
  2014. """),
  2015. },
  2016. ]
  2017. },
  2018. )
  2019. @self.base_endpoint
  2020. async def export_relationships(
  2021. background_tasks: BackgroundTasks,
  2022. id: UUID = Path(
  2023. ...,
  2024. description="The ID of the document to export entities from.",
  2025. ),
  2026. columns: Optional[list[str]] = Body(
  2027. None, description="Specific columns to export"
  2028. ),
  2029. filters: Optional[dict] = Body(
  2030. None, description="Filters to apply to the export"
  2031. ),
  2032. include_header: Optional[bool] = Body(
  2033. True, description="Whether to include column headers"
  2034. ),
  2035. auth_user=Depends(self.providers.auth.auth_wrapper()),
  2036. ) -> FileResponse:
  2037. """Export documents as a downloadable CSV file."""
  2038. if not auth_user.is_superuser:
  2039. raise R2RException(
  2040. "Only a superuser can export data.",
  2041. 403,
  2042. )
  2043. (
  2044. csv_file_path,
  2045. temp_file,
  2046. ) = await self.services.management.export_document_relationships(
  2047. id=id,
  2048. columns=columns,
  2049. filters=filters,
  2050. include_header=include_header
  2051. if include_header is not None
  2052. else True,
  2053. )
  2054. background_tasks.add_task(temp_file.close)
  2055. return FileResponse(
  2056. path=csv_file_path,
  2057. media_type="text/csv",
  2058. filename="documents_export.csv",
  2059. )
  2060. @self.router.post(
  2061. "/documents/search",
  2062. dependencies=[Depends(self.rate_limit_dependency)],
  2063. summary="Search document summaries",
  2064. )
  2065. @self.base_endpoint
  2066. async def search_documents(
  2067. query: str = Body(
  2068. ...,
  2069. description="The search query to perform.",
  2070. ),
  2071. search_mode: SearchMode = Body(
  2072. default=SearchMode.custom,
  2073. description=(
  2074. "Default value of `custom` allows full control over search settings.\n\n"
  2075. "Pre-configured search modes:\n"
  2076. "`basic`: A simple semantic-based search.\n"
  2077. "`advanced`: A more powerful hybrid search combining semantic and full-text.\n"
  2078. "`custom`: Full control via `search_settings`.\n\n"
  2079. "If `filters` or `limit` are provided alongside `basic` or `advanced`, "
  2080. "they will override the default settings for that mode."
  2081. ),
  2082. ),
  2083. search_settings: SearchSettings = Body(
  2084. default_factory=SearchSettings,
  2085. description="Settings for document search",
  2086. ),
  2087. auth_user=Depends(self.providers.auth.auth_wrapper()),
  2088. ) -> WrappedDocumentSearchResponse:
  2089. """Perform a search query on the automatically generated document
  2090. summaries in the system.
  2091. This endpoint allows for complex filtering of search results using PostgreSQL-based queries.
  2092. Filters can be applied to various fields such as document_id, and internal metadata values.
  2093. Allowed operators include `eq`, `neq`, `gt`, `gte`, `lt`, `lte`, `like`, `ilike`, `in`, and `nin`.
  2094. """
  2095. effective_settings = self._prepare_search_settings(
  2096. auth_user, search_mode, search_settings
  2097. )
  2098. query_embedding = (
  2099. await self.providers.embedding.async_get_embedding(query)
  2100. )
  2101. results = await self.services.retrieval.search_documents(
  2102. query=query,
  2103. query_embedding=query_embedding,
  2104. settings=effective_settings,
  2105. )
  2106. return results # type: ignore
  2107. @staticmethod
  2108. async def _process_file(file):
  2109. import base64
  2110. content = await file.read()
  2111. return {
  2112. "filename": file.filename,
  2113. "content": base64.b64encode(content).decode("utf-8"),
  2114. "content_type": file.content_type,
  2115. "content_length": len(content),
  2116. }