Skip to content

Commit

Permalink
Add geolocation
Browse files Browse the repository at this point in the history
  • Loading branch information
rudrankriyam committed Oct 30, 2024
1 parent 684843c commit cbf465e
Show file tree
Hide file tree
Showing 3 changed files with 315 additions and 8 deletions.
43 changes: 35 additions & 8 deletions Sources/AgniKit/AgniKit.swift
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ public struct AgniKit {
/// - timeout: An optional timeout in milliseconds for the request. Default is 30000.
/// - extract: An optional extraction configuration.
/// - actions: An optional array of actions to perform on the page before grabbing the content.
/// - location: An optional location configuration to apply to the request.
///
/// - Returns: A dictionary containing the scraped data.
///
Expand All @@ -63,7 +64,8 @@ public struct AgniKit {
waitFor: Int? = nil,
timeout: Int = 30000,
extract: [String: Any]? = nil,
actions: [[String: Any]]? = nil
actions: [[String: Any]]? = nil,
location: LocationConfig? = nil
) async throws -> ScrapeResponse {
var request = makeRequest(for: "v1/scrape")
request.httpMethod = "POST"
Expand All @@ -83,6 +85,13 @@ public struct AgniKit {
if let extract = extract { body["extract"] = extract }
if let actions = actions { body["actions"] = actions }

if let location = location {
body["location"] = [
"country": location.country.rawValue,
"languages": location.languages
].compactMapValues { $0 }
}

request.httpBody = try JSONSerialization.data(withJSONObject: body)

let (data, response) = try await URLSession.shared.data(for: request)
Expand Down Expand Up @@ -296,6 +305,7 @@ public struct AgniKit {
/// - formats: Array of desired output formats (e.g. ["markdown", "html"]). Default is ["markdown"]
/// - onlyMainContent: Whether to return only the main content. Default is true
/// - timeout: Timeout in milliseconds. Default is 30000
/// - location: An optional location configuration to apply to the request.
///
/// - Returns: A BatchScrapeResponse containing all scraped results
///
Expand All @@ -304,19 +314,27 @@ public struct AgniKit {
urls: [String],
formats: [String] = ["markdown"],
onlyMainContent: Bool = true,
timeout: Int = 30000
timeout: Int = 30000,
location: LocationConfig? = nil
) async throws -> BatchScrapeResponse {
var request = makeRequest(for: "v1/batch/scrape")
request.httpMethod = "POST"
request.setValue("application/json", forHTTPHeaderField: "Content-Type")

let body: [String: Any] = [
var body: [String: Any] = [
"urls": urls,
"formats": formats,
"onlyMainContent": onlyMainContent,
"timeout": timeout
"timeout": timeout,
]


if let location = location {
body["location"] = [
"country": location.country.rawValue,
"languages": location.languages
].compactMapValues { $0 }
}

request.httpBody = try JSONSerialization.data(withJSONObject: body)

let (data, response) = try await URLSession.shared.data(for: request)
Expand All @@ -343,26 +361,35 @@ public struct AgniKit {
/// - urls: Array of URLs to scrape
/// - formats: Array of desired output formats (e.g. ["markdown", "html"]). Default is ["markdown"]
/// - onlyMainContent: Whether to return only the main content. Default is true
/// - location: An optional location configuration to apply to the request.
///
/// - Returns: A BatchScrapeJobResponse containing the job ID and status URL
///
/// - Throws: An error if the request fails or if the response cannot be decoded
public func createBatchScrapeJob(
urls: [String],
formats: [String] = ["markdown"],
onlyMainContent: Bool = true
onlyMainContent: Bool = true,
location: LocationConfig? = nil
) async throws -> BatchScrapeJobResponse {
var request = makeRequest(for: "v1/batch/scrape")
request.httpMethod = "POST"
request.setValue("application/json", forHTTPHeaderField: "Content-Type")

let body: [String: Any] = [
var body: [String: Any] = [
"urls": urls,
"formats": formats,
"onlyMainContent": onlyMainContent,
"async": true
]

if let location = location {
body["location"] = [
"country": location.country.rawValue,
"languages": location.languages
].compactMapValues { $0 }
}

request.httpBody = try JSONSerialization.data(withJSONObject: body)

let (data, response) = try await URLSession.shared.data(for: request)
Expand Down Expand Up @@ -404,4 +431,4 @@ public struct AgniKit {
decoder.dateDecodingStrategy = .iso8601
return try decoder.decode(BatchScrapeResponse.self, from: data)
}
}
}
261 changes: 261 additions & 0 deletions Sources/AgniKit/CountryCode.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,261 @@
/// ISO 3166-1 alpha-2 country codes
///
/// Represents standardized two-letter country codes as defined by ISO 3166-1 alpha-2.
/// Used for specifying geographic locations in API requests.
///
/// This enum contains all official ISO 3166-1 alpha-2 country codes.
/// Each case represents a country or territory with its corresponding two-letter code.
///
/// - Note: The codes follow the official ISO 3166-1 alpha-2 standard maintained by the
/// International Organization for Standardization (ISO).
public enum CountryCode: String, Codable {
case afghanistan = "AF"
case alandIslands = "AX"
case albania = "AL"
case algeria = "DZ"
case americanSamoa = "AS"
case andorra = "AD"
case angola = "AO"
case anguilla = "AI"
case antarctica = "AQ"
case antiguaAndBarbuda = "AG"
case argentina = "AR"
case armenia = "AM"
case aruba = "AW"
case australia = "AU"
case austria = "AT"
case azerbaijan = "AZ"
case bahamas = "BS"
case bahrain = "BH"
case bangladesh = "BD"
case barbados = "BB"
case belarus = "BY"
case belgium = "BE"
case belize = "BZ"
case benin = "BJ"
case bermuda = "BM"
case bhutan = "BT"
case bolivia = "BO"
case bonaireStEustasiusSaba = "BQ"
case bosniaAndHerzegovina = "BA"
case botswana = "BW"
case bouvetIsland = "BV"
case brazil = "BR"
case britishIndianOceanTerritory = "IO"
case bruneiDarussalam = "BN"
case bulgaria = "BG"
case burkinaFaso = "BF"
case burundi = "BI"
case cambodia = "KH"
case cameroon = "CM"
case canada = "CA"
case capeVerde = "CV"
case caymanIslands = "KY"
case centralAfricanRepublic = "CF"
case chad = "TD"
case chile = "CL"
case china = "CN"
case christmasIsland = "CX"
case cocosIslands = "CC"
case colombia = "CO"
case comoros = "KM"
case congo = "CG"
case congoDemocraticRepublic = "CD"
case cookIslands = "CK"
case costaRica = "CR"
case coteDIvoire = "CI"
case croatia = "HR"
case cuba = "CU"
case curacao = "CW"
case cyprus = "CY"
case czechRepublic = "CZ"
case denmark = "DK"
case djibouti = "DJ"
case dominica = "DM"
case dominicanRepublic = "DO"
case ecuador = "EC"
case egypt = "EG"
case elSalvador = "SV"
case equatorialGuinea = "GQ"
case eritrea = "ER"
case estonia = "EE"
case ethiopia = "ET"
case falklandIslands = "FK"
case faroeIslands = "FO"
case fiji = "FJ"
case finland = "FI"
case france = "FR"
case frenchGuiana = "GF"
case frenchPolynesia = "PF"
case frenchSouthernTerritories = "TF"
case gabon = "GA"
case gambia = "GM"
case georgia = "GE"
case germany = "DE"
case ghana = "GH"
case gibraltar = "GI"
case greece = "GR"
case greenland = "GL"
case grenada = "GD"
case guadeloupe = "GP"
case guam = "GU"
case guatemala = "GT"
case guernsey = "GG"
case guinea = "GN"
case guineaBissau = "GW"
case guyana = "GY"
case haiti = "HT"
case heardIslandMcdonald = "HM"
case holySee = "VA"
case honduras = "HN"
case hongKong = "HK"
case hungary = "HU"
case iceland = "IS"
case india = "IN"
case indonesia = "ID"
case iran = "IR"
case iraq = "IQ"
case ireland = "IE"
case isleOfMan = "IM"
case israel = "IL"
case italy = "IT"
case jamaica = "JM"
case japan = "JP"
case jersey = "JE"
case jordan = "JO"
case kazakhstan = "KZ"
case kenya = "KE"
case kiribati = "KI"
case koreaNorth = "KP"
case southKorea = "KR"
case kuwait = "KW"
case kyrgyzstan = "KG"
case laos = "LA"
case latvia = "LV"
case lebanon = "LB"
case lesotho = "LS"
case liberia = "LR"
case libya = "LY"
case liechtenstein = "LI"
case lithuania = "LT"
case luxembourg = "LU"
case macao = "MO"
case macedonia = "MK"
case madagascar = "MG"
case malawi = "MW"
case malaysia = "MY"
case maldives = "MV"
case mali = "ML"
case malta = "MT"
case marshallIslands = "MH"
case martinique = "MQ"
case mauritania = "MR"
case mauritius = "MU"
case mayotte = "YT"
case mexico = "MX"
case micronesia = "FM"
case moldova = "MD"
case monaco = "MC"
case mongolia = "MN"
case montenegro = "ME"
case montserrat = "MS"
case morocco = "MA"
case mozambique = "MZ"
case myanmar = "MM"
case namibia = "NA"
case nauru = "NR"
case nepal = "NP"
case netherlands = "NL"
case newCaledonia = "NC"
case newZealand = "NZ"
case nicaragua = "NI"
case niger = "NE"
case nigeria = "NG"
case niue = "NU"
case norfolkIsland = "NF"
case northernMarianaIslands = "MP"
case norway = "NO"
case oman = "OM"
case pakistan = "PK"
case palau = "PW"
case palestine = "PS"
case panama = "PA"
case papuaNewGuinea = "PG"
case paraguay = "PY"
case peru = "PE"
case philippines = "PH"
case pitcairn = "PN"
case poland = "PL"
case portugal = "PT"
case puertoRico = "PR"
case qatar = "QA"
case reunion = "RE"
case romania = "RO"
case russianFederation = "RU"
case rwanda = "RW"
case saintBarthelemy = "BL"
case saintHelena = "SH"
case saintKittsAndNevis = "KN"
case saintLucia = "LC"
case saintMartin = "MF"
case saintPierreAndMiquelon = "PM"
case saintVincentGrenadines = "VC"
case samoa = "WS"
case sanMarino = "SM"
case saoTomeAndPrincipe = "ST"
case saudiArabia = "SA"
case senegal = "SN"
case serbia = "RS"
case seychelles = "SC"
case sierraLeone = "SL"
case singapore = "SG"
case sintMaarten = "SX"
case slovakia = "SK"
case slovenia = "SI"
case solomonIslands = "SB"
case somalia = "SO"
case southAfrica = "ZA"
case southGeorgia = "GS"
case southSudan = "SS"
case spain = "ES"
case sriLanka = "LK"
case sudan = "SD"
case suriname = "SR"
case svalbardJanMayen = "SJ"
case swaziland = "SZ"
case sweden = "SE"
case switzerland = "CH"
case syrianArabRepublic = "SY"
case taiwan = "TW"
case tajikistan = "TJ"
case tanzania = "TZ"
case thailand = "TH"
case timorLeste = "TL"
case togo = "TG"
case tokelau = "TK"
case tonga = "TO"
case trinidadAndTobago = "TT"
case tunisia = "TN"
case turkey = "TR"
case turkmenistan = "TM"
case turksAndCaicosIslands = "TC"
case tuvalu = "TV"
case uganda = "UG"
case ukraine = "UA"
case unitedArabEmirates = "AE"
case unitedKingdom = "GB"
case unitedStates = "US"
case unitedStatesOutlyingIslands = "UM"
case uruguay = "UY"
case uzbekistan = "UZ"
case vanuatu = "VU"
case venezuela = "VE"
case vietnam = "VN"
case virginIslandsBritish = "VG"
case virginIslandsUS = "VI"
case wallisAndFutuna = "WF"
case westernSahara = "EH"
case yemen = "YE"
case zambia = "ZM"
case zimbabwe = "ZW"
}
Loading

0 comments on commit cbf465e

Please sign in to comment.