dify/api/libs/bearer_data_source.py

# [REVIEW] Implement if Needed? Do we need a new type of data source
from abc import abstractmethod

import requests
from flask_login import current_user

from extensions.ext_database import db
from models.source import DataSourceBearerBinding


class BearerDataSource:
    def __init__(self, api_key: str, api_base_url: str):
        self.api_key = api_key
        self.api_base_url = api_base_url

    @abstractmethod
    def validate_bearer_data_source(self):
        """
        Validate the data source
        """


class FireCrawlDataSource(BearerDataSource):
    def validate_bearer_data_source(self):
        TEST_CRAWL_SITE_URL = "https://www.google.com"
        FIRECRAWL_API_VERSION = "v0"

        test_api_endpoint = self.api_base_url.rstrip("/") + f"/{FIRECRAWL_API_VERSION}/scrape"

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        data = {
            "url": TEST_CRAWL_SITE_URL,
        }

        response = requests.get(test_api_endpoint, headers=headers, json=data)

        return response.json().get("status") == "success"

    def save_credentials(self):
        # save data source binding
        data_source_binding = DataSourceBearerBinding.query.filter(
            db.and_(
                DataSourceBearerBinding.tenant_id == current_user.current_tenant_id,
                DataSourceBearerBinding.provider == "firecrawl",
                DataSourceBearerBinding.endpoint_url == self.api_base_url,
                DataSourceBearerBinding.bearer_key == self.api_key,
            )
        ).first()
        if data_source_binding:
            data_source_binding.disabled = False
            db.session.commit()
        else:
            new_data_source_binding = DataSourceBearerBinding(
                tenant_id=current_user.current_tenant_id,
                provider="firecrawl",
                endpoint_url=self.api_base_url,
                bearer_key=self.api_key,
            )
            db.session.add(new_data_source_binding)
            db.session.commit()