Function upload_file

api/oss/src/routers/testset_router.py:80–157

Uploads a CSV or JSON file and saves its data to Postgres. Args: upload_type : Either a json or csv file. file (UploadFile): The CSV or JSON file to upload. testset_name (Optional): the name of the testset if provided. Returns: dict: The result of the uploa

(
    request: Request,
    upload_type: str = Form(None),
    file: UploadFile = File(...),
    testset_name: Optional[str] = File(None),
)

    "/upload/", response_model=TestSetSimpleResponse, operation_id="upload_file"
)
async def upload_file(
    request: Request,
    upload_type: str = Form(None),
    file: UploadFile = File(...),
    testset_name: Optional[str] = File(None),
):
    """
    Uploads a CSV or JSON file and saves its data to Postgres.

    Args:
    upload_type : Either a json or csv file.
        file (UploadFile): The CSV or JSON file to upload.
        testset_name (Optional): the name of the testset if provided.

    Returns:
        dict: The result of the upload process.
    """

    if is_ee():
        has_permission = await check_action_access(
            user_uid=request.state.user_id,
            project_id=request.state.project_id,
            permission=Permission.CREATE_TESTSET,
        )
        if not has_permission:
            error_msg = f"You do not have permission to perform this action. Please contact your organization admin."
            log.error(error_msg)
            return JSONResponse(
                {"detail": error_msg},
                status_code=403,
            )

    if file.size > TESTSETS_SIZE_LIMIT:  # Preemptively check file size
        raise TESTSETS_SIZE_EXCEPTION

    # Create a document
    document = {
        "name": testset_name if testset_name else file.filename,
        "csvdata": [],
    }

    if upload_type == "JSON":
        # Read and parse the JSON file
        json_data = await file.read()
        json_text = json_data.decode("utf-8")
        json_object = json.loads(json_text)

        # Populate the document with column names and values
        for i, row in enumerate(json_object):
            document["csvdata"].append(row)

    else:
        # Read and parse the CSV file
        csv_data = await file.read()
        csv_text = csv_data.decode("utf-8")

        # Use StringIO to create a file-like object from the string
        csv_file_like_object = io.StringIO(csv_text)