Skip to content

allow passing context #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jackhopner wants to merge 1 commit into
base: main
Choose a base branch
from
Open

allow passing context #30

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 46 additions & 5 deletions firecrawl.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,9 @@ package firecrawl

import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"math"
Expand Down Expand Up @@ -271,6 +273,12 @@ func NewFirecrawlApp(apiKey, apiURL string, timeout ...time.Duration) (*Firecraw
// - *FirecrawlDocument or *FirecrawlDocumentV0: The scraped document data depending on the API version.
// - error: An error if the scrape request fails.
func (app *FirecrawlApp) ScrapeURL(url string, params *ScrapeParams) (*FirecrawlDocument, error) {
return app.ScrapeURLWithContext(context.Background(), url, params)
}

// ScrapeURLWithContext scrapes the content of the specified URL using the Firecrawl API.
// Accepts a context to control cancellation and deadlines.
func (app *FirecrawlApp) ScrapeURLWithContext(ctx context.Context, url string, params *ScrapeParams) (*FirecrawlDocument, error) {
headers := app.prepareHeaders(nil)
scrapeBody := map[string]any{"url": url}

Expand Down Expand Up @@ -326,6 +334,7 @@ func (app *FirecrawlApp) ScrapeURL(url string, params *ScrapeParams) (*Firecrawl
}

resp, err := app.makeRequest(
ctx,
http.MethodPost,
fmt.Sprintf("%s/v1/scrape", app.APIURL),
scrapeBody,
Expand Down Expand Up @@ -362,6 +371,11 @@ func (app *FirecrawlApp) ScrapeURL(url string, params *ScrapeParams) (*Firecrawl
// - CrawlStatusResponse: The crawl result if the job is completed.
// - error: An error if the crawl request fails.
func (app *FirecrawlApp) CrawlURL(url string, params *CrawlParams, idempotencyKey *string, pollInterval ...int) (*CrawlStatusResponse, error) {
return app.CrawlURLWithContext(context.Background(), url, params, idempotencyKey, pollInterval...)
}

// CrawlURLWithContext starts a crawl job for the specified URL using the Firecrawl API with context.
func (app *FirecrawlApp) CrawlURLWithContext(ctx context.Context, url string, params *CrawlParams, idempotencyKey *string, pollInterval ...int) (*CrawlStatusResponse, error) {
var key string
if idempotencyKey != nil {
key = *idempotencyKey
Expand Down Expand Up @@ -409,6 +423,7 @@ func (app *FirecrawlApp) CrawlURL(url string, params *CrawlParams, idempotencyKe
}

resp, err := app.makeRequest(
ctx,
http.MethodPost,
fmt.Sprintf("%s/v1/crawl", app.APIURL),
crawlBody,
Expand All @@ -427,7 +442,7 @@ func (app *FirecrawlApp) CrawlURL(url string, params *CrawlParams, idempotencyKe
return nil, err
}

return app.monitorJobStatus(crawlResponse.ID, headers, actualPollInterval)
return app.monitorJobStatus(ctx, crawlResponse.ID, headers, actualPollInterval)
}

// CrawlURL starts a crawl job for the specified URL using the Firecrawl API.
Expand All @@ -441,6 +456,11 @@ func (app *FirecrawlApp) CrawlURL(url string, params *CrawlParams, idempotencyKe
// - *CrawlResponse: The crawl response with id.
// - error: An error if the crawl request fails.
func (app *FirecrawlApp) AsyncCrawlURL(url string, params *CrawlParams, idempotencyKey *string) (*CrawlResponse, error) {
return app.AsyncCrawlURLWithContext(context.Background(), url, params, idempotencyKey)
}

// AsyncCrawlURLWithContext starts a crawl job asynchronously with context.
func (app *FirecrawlApp) AsyncCrawlURLWithContext(ctx context.Context, url string, params *CrawlParams, idempotencyKey *string) (*CrawlResponse, error) {
var key string
if idempotencyKey != nil {
key = *idempotencyKey
Expand Down Expand Up @@ -483,6 +503,7 @@ func (app *FirecrawlApp) AsyncCrawlURL(url string, params *CrawlParams, idempote
}

resp, err := app.makeRequest(
ctx,
http.MethodPost,
fmt.Sprintf("%s/v1/crawl", app.APIURL),
crawlBody,
Expand Down Expand Up @@ -518,10 +539,16 @@ func (app *FirecrawlApp) AsyncCrawlURL(url string, params *CrawlParams, idempote
// - *CrawlStatusResponse: The status of the crawl job.
// - error: An error if the crawl status check request fails.
func (app *FirecrawlApp) CheckCrawlStatus(ID string) (*CrawlStatusResponse, error) {
return app.CheckCrawlStatusWithContext(context.Background(), ID)
}

// CheckCrawlStatusWithContext checks the status of a crawl job using context.
func (app *FirecrawlApp) CheckCrawlStatusWithContext(ctx context.Context, ID string) (*CrawlStatusResponse, error) {
headers := app.prepareHeaders(nil)
apiURL := fmt.Sprintf("%s/v1/crawl/%s", app.APIURL, ID)

resp, err := app.makeRequest(
ctx,
http.MethodGet,
apiURL,
nil,
Expand Down Expand Up @@ -552,9 +579,15 @@ func (app *FirecrawlApp) CheckCrawlStatus(ID string) (*CrawlStatusResponse, erro
// - string: The status of the crawl job after cancellation.
// - error: An error if the crawl job cancellation request fails.
func (app *FirecrawlApp) CancelCrawlJob(ID string) (string, error) {
return app.CancelCrawlJobWithContext(context.Background(), ID)
}

// CancelCrawlJobWithContext cancels a crawl job using context.
func (app *FirecrawlApp) CancelCrawlJobWithContext(ctx context.Context, ID string) (string, error) {
headers := app.prepareHeaders(nil)
apiURL := fmt.Sprintf("%s/v1/crawl/%s", app.APIURL, ID)
resp, err := app.makeRequest(
ctx,
http.MethodDelete,
apiURL,
nil,
Expand Down Expand Up @@ -584,6 +617,11 @@ func (app *FirecrawlApp) CancelCrawlJob(ID string) (string, error) {
// - *MapResponse: The response from the mapping operation.
// - error: An error if the mapping request fails.
func (app *FirecrawlApp) MapURL(url string, params *MapParams) (*MapResponse, error) {
return app.MapURLWithContext(context.Background(), url, params)
}

// MapURLWithContext initiates a mapping operation for a URL using context.
func (app *FirecrawlApp) MapURLWithContext(ctx context.Context, url string, params *MapParams) (*MapResponse, error) {
headers := app.prepareHeaders(nil)
jsonData := map[string]any{"url": url}

Expand All @@ -603,6 +641,7 @@ func (app *FirecrawlApp) MapURL(url string, params *MapParams) (*MapResponse, er
}

resp, err := app.makeRequest(
ctx,
http.MethodPost,
fmt.Sprintf("%s/v1/map", app.APIURL),
jsonData,
Expand Down Expand Up @@ -670,7 +709,7 @@ func (app *FirecrawlApp) prepareHeaders(idempotencyKey *string) map[string]strin
// Returns:
// - []byte: The response body from the request.
// - error: An error if the request fails.
func (app *FirecrawlApp) makeRequest(method, url string, data map[string]any, headers map[string]string, action string, opts ...requestOption) ([]byte, error) {
func (app *FirecrawlApp) makeRequest(ctx context.Context, method, url string, data map[string]any, headers map[string]string, action string, opts ...requestOption) ([]byte, error) {
var body []byte
var err error
if data != nil {
Expand All @@ -680,7 +719,7 @@ func (app *FirecrawlApp) makeRequest(method, url string, data map[string]any, he
}
}

req, err := http.NewRequest(method, url, bytes.NewBuffer(body))
req, err := http.NewRequestWithContext(ctx, method, url, bytes.NewBuffer(body))
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -728,11 +767,12 @@ func (app *FirecrawlApp) makeRequest(method, url string, data map[string]any, he
// Returns:
// - *CrawlStatusResponse: The crawl result if the job is completed.
// - error: An error if the crawl status check request fails.
func (app *FirecrawlApp) monitorJobStatus(ID string, headers map[string]string, pollInterval int) (*CrawlStatusResponse, error) {
func (app *FirecrawlApp) monitorJobStatus(ctx context.Context, ID string, headers map[string]string, pollInterval int) (*CrawlStatusResponse, error) {
attempts := 3

for {
resp, err := app.makeRequest(
ctx,
http.MethodGet,
fmt.Sprintf("%s/v1/crawl/%s", app.APIURL, ID),
nil,
Expand Down Expand Up @@ -760,6 +800,7 @@ func (app *FirecrawlApp) monitorJobStatus(ID string, headers map[string]string,
allData := statusData.Data
for statusData.Next != nil {
resp, err := app.makeRequest(
ctx,
http.MethodGet,
*statusData.Next,
nil,
Expand Down Expand Up @@ -833,5 +874,5 @@ func (app *FirecrawlApp) handleError(statusCode int, body []byte, action string)
message = fmt.Sprintf("Unexpected error during %s: Status code %d. %s", action, statusCode, errorMessage)
}

return fmt.Errorf(message)
return errors.New(message)
}