| @@ -1,64 +0,0 @@ | |||
| # [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() | |||