Merge pull request #2 from oxylabs/feature/ai-map

Add ai-map app

Author: andriuskuk

examples/map.py

@@ -0,0 +1,15 @@
+from oxylabs_ai_studio.apps.ai_map import AiMap
+
+
+map = AiMap(api_key="<API_KEY>")
+
+payload = {
+    "url": "https://career.oxylabs.io",
+    "user_prompt": "job ad pages",
+    "return_sources_limit": 10,
+    "max_depth": 1,
+    "geo_location": None,
+    "render_javascript": False,
+}
+result = map.map(**payload)
+print(result.data)

pyproject.toml

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

readme.md

@@ -130,5 +130,29 @@ print(result.data)
 - `return_content` (bool): Whether to return markdown contents in results (default: True)
 - `geo_location` (str): search proxy location in ISO2 format.
 
+### Map (`AiMap.map`)
+```python
+from oxylabs_ai_studio.apps.ai_map import AiMap
+
+
+map = AiMap(api_key="<API_KEY>")
+payload = {
+    "url": "https://career.oxylabs.io",
+    "user_prompt": "job ad pages",
+    "return_sources_limit": 10,
+    "max_depth": 1,
+    "geo_location": None,
+    "render_javascript": False,
+}
+result = map.map(**payload)
+print(result.data)
+```
+**Parameters:**
+- `url` (str): Starting URL to crawl (**required**)
+- `user_prompt` (str): Natural language prompt to guide extraction (**required**)
+- `render_javascript` (bool): Render JavaScript (default: False)
+- `return_sources_limit` (int): Max number of sources to return (default: 25)
+- `geo_location` (str): proxy location in ISO2 format.
+
 ---
 See the [examples](https://github.com/oxylabs/oxylabs-ai-studio-py/tree/main/examples) folder for usage examples of each method. Each method has corresponding async version.

src/oxylabs_ai_studio/apps/ai_map.py

@@ -0,0 +1,137 @@
+import asyncio
+import time
+from typing import Any
+
+import httpx
+from pydantic import BaseModel
+
+from oxylabs_ai_studio.client import OxyStudioAIClient
+from oxylabs_ai_studio.logger import get_logger
+
+POLL_MAX_ATTEMPTS = 10
+POLL_INTERVAL_SECONDS = 3
+
+logger = get_logger(__name__)
+
+
+class AiMapJob(BaseModel):
+    run_id: str
+    message: str | None = None
+    data: dict[str, Any] | str | None
+
+
+class AiMap(OxyStudioAIClient):
+    """AI Map app."""
+
+    def __init__(self, api_key: str):
+        super().__init__(api_key=api_key)
+
+    def map(
+        self,
+        url: str,
+        user_prompt: str,
+        return_sources_limit: int = 25,
+        max_depth: int = 3,
+        geo_location: str | None = None,
+        render_javascript: bool = False,
+    ) -> AiMapJob:
+        body = {
+            "url": url,
+            "user_prompt": user_prompt,
+            "return_sources_limit": return_sources_limit,
+            "max_depth": max_depth,
+            "geo_location": geo_location,
+            "render_html": render_javascript,
+        }
+        create_response = self.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}"
+            )
+        resp_body = create_response.json()
+        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})
+                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),
+                    )
+                if resp_body["status"] == "failed":
+                    raise Exception(f"Failed to map {url}.")
+                time.sleep(POLL_INTERVAL_SECONDS)
+        except KeyboardInterrupt:
+            logger.info("[Cancelled] Mapping was cancelled by user.")
+            raise KeyboardInterrupt from None
+        except Exception as e:
+            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})
+        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 {}
+
+    async def map_async(
+        self,
+        url: str,
+        user_prompt: str,
+        return_sources_limit: int = 25,
+        max_depth: int = 3,
+        geo_location: str | None = None,
+        render_javascript: bool = False,
+    ) -> AiMapJob:
+        body = {
+            "url": url,
+            "user_prompt": user_prompt,
+            "return_sources_limit": return_sources_limit,
+            "max_depth": max_depth,
+            "geo_location": geo_location,
+            "render_html": render_javascript,
+        }
+        async with self.async_client() as client:
+            create_response = await 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}"
+                )
+            resp_body = create_response.json()
+            run_id = resp_body["run_id"]
+            try:
+                for _ in range(POLL_MAX_ATTEMPTS):
+                    get_response = await 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":
+                        data = await self.get_data_async(client, run_id=run_id)
+                        return AiMapJob(
+                            run_id=run_id,
+                            message=resp_body.get("message", None),
+                            data=data,
+                        )
+                    if resp_body["status"] == "failed":
+                        raise Exception(f"Failed to map {url}.")
+                    await asyncio.sleep(POLL_INTERVAL_SECONDS)
+            except KeyboardInterrupt:
+                logger.info("[Cancelled] Mapping was cancelled by user.")
+                raise KeyboardInterrupt from None
+            except Exception as e:
+                raise e
+            raise TimeoutError(f"Failed to map {url}: timeout.")
+
+    async def get_data_async(
+        self, client: httpx.AsyncClient, run_id: str
+    ) -> dict[str, Any]:
+        get_response = await 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 {}