pvanand commited on
Commit
c182518
·
verified ·
1 Parent(s): 8f209c2

Update document_rag_router.py

Browse files
Files changed (1) hide show
  1. document_rag_router.py +22 -0
document_rag_router.py CHANGED
@@ -372,6 +372,28 @@ async def delete_collection(collection_id: str, user_id: str):
372
  detail=f"Unexpected error: {str(e)}"
373
  )
374
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
375
  @router.post("/query_collection_tool")
376
  async def query_collection_tool(input_data: QueryInput):
377
  try:
 
372
  detail=f"Unexpected error: {str(e)}"
373
  )
374
 
375
+ @router.post("/get_collection_files")
376
+ def get_collection_files(collection_id: str, user_id: str) -> str:
377
+ """Get list of files in the specified collection"""
378
+ try:
379
+ # Get the full collection name
380
+ collection_name = f"{user_id}_{collection_id}"
381
+
382
+ # Open the table and convert to pandas
383
+ table = db.open_table(collection_name)
384
+ df = table.to_pandas()
385
+ logging.info(f"fetched chunks {str(df.head())}")
386
+
387
+ # Get unique file names
388
+ unique_files = df['file_name'].unique()
389
+
390
+ # Join the file names into a string
391
+ return ", ".join(unique_files)
392
+ except Exception as e:
393
+ logging.error(f"Error getting collection files: {str(e)}")
394
+ return f"Error getting files: {str(e)}"
395
+
396
+
397
  @router.post("/query_collection_tool")
398
  async def query_collection_tool(input_data: QueryInput):
399
  try: