Add new types

Author: andriuskuk

src/oxylabs_ai_studio/apps/ai_crawler.py

@@ -31,14 +31,16 @@ def crawl(
         self,
         url: str,
         user_prompt: str,
-        output_format: Literal["json", "markdown"] = "markdown",
+        output_format: Literal["json", "markdown", "csv"] = "markdown",
         schema: dict[str, Any] | None = None,
         render_javascript: bool = False,
         return_sources_limit: int = 25,
         geo_location: str | None = None,
     ) -> AiCrawlerJob:
-        if output_format == "json" and schema is None:
-            raise ValueError("openapi_schema is required when output_format is json")
+        if output_format in ["json", "csv"] and schema is None:
+            raise ValueError(
+                "openapi_schema is required when output_format is json or csv.",
+            )
 
         body = {
             "domain": url,
@@ -119,15 +121,17 @@ async def crawl_async(
         self,
         url: str,
         user_prompt: str = "",
-        output_format: Literal["json", "markdown"] = "markdown",
+        output_format: Literal["json", "markdown", "csv"] = "markdown",
         schema: dict[str, Any] | None = None,
         render_javascript: bool = False,
         return_sources_limit: int = 25,
         geo_location: str | None = None,
     ) -> AiCrawlerJob:
         """Async version of crawl."""
-        if output_format == "json" and schema is None:
-            raise ValueError("openapi_schema is required when output_format is json")
+        if output_format in ["json", "csv"] and schema is None:
+            raise ValueError(
+                "openapi_schema is required when output_format is json or csv.",
+            )
 
         body = {
             "domain": url,

src/oxylabs_ai_studio/apps/ai_scraper.py

@@ -31,13 +31,15 @@ def __init__(self, api_key: str | None = None):
     def scrape(
         self,
         url: str,
-        output_format: Literal["json", "markdown"] = "markdown",
+        output_format: Literal["json", "markdown", "csv", "screenshot"] = "markdown",
         schema: dict[str, Any] | None = None,
         render_javascript: bool = False,
         geo_location: str | None = None,
     ) -> AiScraperJob:
-        if output_format == "json" and schema is None:
-            raise ValueError("openapi_schema is required when output_format is json")
+        if output_format in ["json", "csv"] and schema is None:
+            raise ValueError(
+                "openapi_schema is required when output_format is json or csv.",
+            )
 
         body = {
             "url": url,
@@ -117,14 +119,16 @@ def generate_schema(self, prompt: str) -> dict[str, Any] | None:
     async def scrape_async(
         self,
         url: str,
-        output_format: Literal["json", "markdown"] = "markdown",
+        output_format: Literal["json", "markdown", "csv", "screenshot"] = "markdown",
         schema: dict[str, Any] | None = None,
         render_javascript: bool = False,
         geo_location: str | None = None,
     ) -> AiScraperJob:
         """Async version of scrape."""
-        if output_format == "json" and schema is None:
-            raise ValueError("openapi_schema is required when output_format is json")
+        if output_format in ["json", "csv"] and schema is None:
+            raise ValueError(
+                "openapi_schema is required when output_format is json or csv.",
+            )
 
         body = {
             "url": url,

src/oxylabs_ai_studio/apps/browser_agent.py

@@ -16,7 +16,7 @@
 
 
 class DataModel(BaseModel):
-    type: Literal["json", "markdown", "html", "screenshot"]
+    type: Literal["json", "markdown", "html", "screenshot", "csv"]
     content: dict[str, Any] | str | None
 
 
@@ -34,12 +34,16 @@ def run(
         self,
         url: str,
         user_prompt: str = "",
-        output_format: Literal["json", "markdown", "html", "screenshot"] = "markdown",
+        output_format: Literal[
+            "json", "markdown", "html", "screenshot", "csv"
+        ] = "markdown",
         schema: dict[str, Any] | None = None,
         geo_location: str | None = None,
     ) -> BrowserAgentJob:
-        if output_format == "json" and schema is None:
-            raise ValueError("openapi_schema is required when output_format is json")
+        if output_format in ["json", "csv"] and schema is None:
+            raise ValueError(
+                "openapi_schema is required when output_format is json or csv.",
+            )
 
         body = {
             "url": url,
@@ -115,13 +119,21 @@ async def run_async(
         self,
         url: str,
         user_prompt: str = "",
-        output_format: Literal["json", "markdown", "html", "screenshot"] = "markdown",
+        output_format: Literal[
+            "json",
+            "markdown",
+            "html",
+            "screenshot",
+            "csv",
+        ] = "markdown",
         schema: dict[str, Any] | None = None,
         geo_location: str | None = None,
     ) -> BrowserAgentJob:
         """Async version of run."""
-        if output_format == "json" and schema is None:
-            raise ValueError("openapi_schema is required when output_format is json")
+        if output_format in ["json", "csv"] and schema is None:
+            raise ValueError(
+                "openapi_schema is required when output_format is json or csv.",
+            )
 
         body = {
             "url": url,