update client settings

Author: KarolisKluonaitis

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "oxylabs-ai-studio"
-version = "0.2.12"
+version = "0.2.13"
 description = "Oxylabs studio python sdk"
 readme = "README.md"
 keywords = ["oxylabs", "ai", "studio"]

src/oxylabs_ai_studio/apps/ai_crawler.py

@@ -50,7 +50,8 @@ def crawl(
             "return_sources_limit": return_sources_limit,
             "geo_location": geo_location,
         }
-        create_response = self.client.post(url="/extract/run", json=body)
+        client = self.get_client()
+        create_response = client.post(url="/extract/run", json=body)
         if create_response.status_code != 200:
             raise Exception(
                 f"Failed to create crawl job for {url}: {create_response.text}"
@@ -60,7 +61,7 @@ def crawl(
         logger.info(f"Starting crawl for url: {url}. Job id: {run_id}.")
         try:
             for _ in range(POLL_MAX_ATTEMPTS):
-                get_response = self.client.get(
+                get_response = client.get(
                     "/extract/run/data", params={"run_id": run_id}
                 )
                 if get_response.status_code == 202:
@@ -89,7 +90,7 @@ def crawl(
     def generate_schema(self, prompt: str) -> dict[str, Any] | None:
         logger.info("Generating schema")
         body = {"user_prompt": prompt}
-        response = self.client.post(url="/extract/generate-params", json=body)
+        response = self.get_client().post(url="/extract/generate-params", json=body)
         if response.status_code != 200:
             raise Exception(f"Failed to generate schema: {response.text}")
         json_response: SchemaResponse = response.json()

src/oxylabs_ai_studio/apps/ai_map.py

@@ -24,7 +24,7 @@ class AiMapJob(BaseModel):
 class AiMap(OxyStudioAIClient):
     """AI Map app."""
 
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str | None = None):
         super().__init__(api_key=api_key)
 
     def map(
@@ -42,7 +42,8 @@ def map(
             "geo_location": geo_location,
             "render_html": render_javascript,
         }
-        create_response = self.client.post(url="/map", json=body)
+        client = self.get_client()
+        create_response = client.post(url="/map", json=body)
         if create_response.status_code != 200:
             raise Exception(
                 f"Failed to create map job for {url}: {create_response.text}"
@@ -51,15 +52,15 @@ def map(
         run_id = resp_body["run_id"]
         try:
             for _ in range(POLL_MAX_ATTEMPTS):
-                get_response = self.client.get("/map/run", params={"run_id": run_id})
+                get_response = client.get("/map/run", params={"run_id": run_id})
                 if get_response.status_code != 200:
                     raise Exception(f"Failed to map {url}: {get_response.text}")
                 resp_body = get_response.json()
                 if resp_body["status"] == "completed":
                     return AiMapJob(
                         run_id=run_id,
                         message=resp_body.get("message", None),
-                        data=self._get_data(run_id=run_id),
+                        data=self._get_data(client=client, run_id=run_id),
                     )
                 if resp_body["status"] == "failed":
                     raise Exception(f"Failed to map {url}.")
@@ -71,8 +72,8 @@ def map(
             raise e
         raise TimeoutError(f"Failed to map {url}: timeout.")
 
-    def _get_data(self, run_id: str) -> dict[str, Any]:
-        get_response = self.client.get("/map/run/data", params={"run_id": run_id})
+    def _get_data(self, client: httpx.Client, run_id: str) -> dict[str, Any]:
+        get_response = client.get("/map/run/data", params={"run_id": run_id})
         if get_response.status_code != 200:
             raise Exception(f"Failed to get data for run {run_id}: {get_response.text}")
         return get_response.json().get("data", {}) or {}

src/oxylabs_ai_studio/apps/ai_scraper.py

@@ -25,7 +25,7 @@ class AiScraperJob(BaseModel):
 class AiScraper(OxyStudioAIClient):
     """AI Scraper app."""
 
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str | None = None):
         super().__init__(api_key=api_key)
 
     def scrape(
@@ -46,7 +46,8 @@ def scrape(
             "render_html": render_javascript,
             "geo_location": geo_location,
         }
-        create_response = self.client.post(url="/scrape", json=body)
+        client = self.get_client()
+        create_response = client.post(url="/scrape", json=body)
         if create_response.status_code != 200:
             raise Exception(
                 f"Failed to create scrape job for {url}: {create_response.text}"
@@ -55,15 +56,15 @@ def scrape(
         run_id = resp_body["run_id"]
         try:
             for _ in range(POLL_MAX_ATTEMPTS):
-                get_response = self.client.get("/scrape/run", params={"run_id": run_id})
+                get_response = client.get("/scrape/run", params={"run_id": run_id})
                 if get_response.status_code != 200:
                     raise Exception(f"Failed to scrape {url}: {get_response.text}")
                 resp_body = get_response.json()
                 if resp_body["status"] == "completed":
                     return AiScraperJob(
                         run_id=run_id,
                         message=resp_body.get("message", None),
-                        data=self._get_data(run_id=run_id),
+                        data=self._get_data(client=client, run_id=run_id),
                     )
                 if resp_body["status"] == "failed":
                     raise Exception(f"Failed to scrape {url}.")
@@ -75,16 +76,16 @@ def scrape(
             raise e
         raise TimeoutError(f"Failed to scrape {url}: timeout.")
 
-    def _get_data(self, run_id: str) -> dict[str, Any]:
-        get_response = self.client.get("/scrape/run/data", params={"run_id": run_id})
+    def _get_data(self, client: httpx.Client, run_id: str) -> dict[str, Any]:
+        get_response = client.get("/scrape/run/data", params={"run_id": run_id})
         if get_response.status_code != 200:
             raise Exception(f"Failed to get data for run {run_id}: {get_response.text}")
         return get_response.json().get("data", {}) or {}
 
     def generate_schema(self, prompt: str) -> dict[str, Any] | None:
         logger.info("Generating schema")
         body = {"user_prompt": prompt}
-        response = self.client.post(url="/scrape/schema", json=body)
+        response = self.get_client().post(url="/scrape/schema", json=body)
         if response.status_code != 200:
             raise Exception(f"Failed to generate schema: {response.text}")
         json_response: SchemaResponse = response.json()

src/oxylabs_ai_studio/apps/ai_search.py

@@ -29,7 +29,7 @@ class AiSearchJob(BaseModel):
 class AiSearch(OxyStudioAIClient):
     """AI Search app."""
 
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str | None = None):
         super().__init__(api_key=api_key)
 
     def search(
@@ -50,16 +50,15 @@ def search(
             "return_content": return_content,
             "geo_location": geo_location,
         }
-        create_response = self.client.post(url="/search/run", json=body)
+        client = self.get_client()
+        create_response = client.post(url="/search/run", json=body)
         if create_response.status_code != 200:
             raise Exception(f"Failed to create search job: {create_response.text}")
         resp_body = create_response.json()
         run_id = resp_body["run_id"]
         try:
             for _ in range(POLL_MAX_ATTEMPTS):
-                get_response = self.client.get(
-                    "/search/run/data", params={"run_id": run_id}
-                )
+                get_response = client.get("/search/run/data", params={"run_id": run_id})
                 if get_response.status_code == 202:
                     time.sleep(POLL_INTERVAL_SECONDS)
                     continue

src/oxylabs_ai_studio/apps/browser_agent.py

@@ -27,7 +27,7 @@ class BrowserAgentJob(BaseModel):
 
 
 class BrowserAgent(OxyStudioAIClient):
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str | None = None):
         super().__init__(api_key=api_key)
 
     def run(
@@ -48,15 +48,16 @@ def run(
             "auxiliary_prompt": user_prompt,
             "geo_location": geo_location,
         }
-        create_response = self.client.post(url="/browser-agent/run", json=body)
+        client = self.get_client()
+        create_response = client.post(url="/browser-agent/run", json=body)
         if create_response.status_code != 200:
             raise Exception(f"Failed to launch browser agent: {create_response.text}")
         resp_body = create_response.json()
         run_id = resp_body["run_id"]
         logger.info(f"Starting browser agent run for url: {url}. Job id: {run_id}.")
         try:
             for _ in range(POLL_MAX_ATTEMPTS):
-                get_response = self.client.get(
+                get_response = client.get(
                     "/browser-agent/run/data", params={"run_id": run_id}
                 )
                 if get_response.status_code == 202:
@@ -85,7 +86,9 @@ def run(
     def generate_schema(self, prompt: str) -> dict[str, Any] | None:
         logger.info("Generating schema")
         body = {"user_prompt": prompt}
-        response = self.client.post(url="/browser-agent/generate-params", json=body)
+        response = self.get_client().post(
+            url="/browser-agent/generate-params", json=body
+        )
         if response.status_code != 200:
             raise Exception(f"Failed to generate schema: {response.text}")
         json_response: SchemaResponse = response.json()

src/oxylabs_ai_studio/client.py

@@ -8,6 +8,14 @@
 
 logger = get_logger(__name__)
 
+_UA_API: str | None = None
+
+
+def _resolve_ua() -> str:
+    return (
+        _UA_API.strip() if isinstance(_UA_API, str) and _UA_API.strip() else None
+    ) or "python-sdk"
+
 
 class OxyStudioAIClient:
     """Main client for interacting with the Oxy Studio AI API."""
@@ -25,15 +33,16 @@ def __init__(self, api_key: str | None = None, timeout: float = 30.0):
         self.api_key = resolved_key
         self.base_url = settings.OXYLABS_AI_STUDIO_API_URL
         self.timeout = timeout
-        # Initialize HTTP client with proper headers
-        self.client = httpx.Client(
+
+    def get_client(self) -> httpx.Client:
+        return httpx.Client(
             base_url=self.base_url,
             headers={
                 "x-api-key": self.api_key,
                 "Content-Type": "application/json",
-                "User-Agent": "python-sdk",
+                "User-Agent": _resolve_ua(),
             },
-            timeout=timeout,
+            timeout=self.timeout,
         )
 
     @asynccontextmanager
@@ -44,7 +53,7 @@ async def async_client(self) -> AsyncGenerator[httpx.AsyncClient, None]:
             headers={
                 "x-api-key": self.api_key,
                 "Content-Type": "application/json",
-                "User-Agent": "python-sdk",
+                "User-Agent": _resolve_ua(),
             },
             timeout=self.timeout,
         )