Spaces:
Sleeping
Sleeping
Update document_rag_router.py
Browse files- 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:
|