diff --git a/jigsawstack/embedding_v2.py b/jigsawstack/embedding_v2.py
index 685cd52..cefcbfe 100644
--- a/jigsawstack/embedding_v2.py
+++ b/jigsawstack/embedding_v2.py
@@ -15,6 +15,9 @@ class EmbeddingV2Params(TypedDict):
     url: NotRequired[str]
     file_store_key: NotRequired[str]
     token_overflow_mode: NotRequired[Literal["truncate", "error"]]
+    dimensions: NotRequired[int]
+    instruction: NotRequired[str]
+    query: NotRequired[bool]
     speaker_fingerprint: NotRequired[bool]
 
 
diff --git a/jigsawstack/prediction.py b/jigsawstack/prediction.py
index 00bd3cf..65ea514 100644
--- a/jigsawstack/prediction.py
+++ b/jigsawstack/prediction.py
@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Union, cast
 
-from typing_extensions import TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 from ._config import ClientConfig
 from ._types import BaseResponse
@@ -25,9 +25,9 @@ class PredictionParams(TypedDict):
     """
     The dataset to make predictions on. This is an array of object with keys date and value. See example below for more information.
     """
-    steps: int
+    steps: NotRequired[int]
     """
-    The number of predictions to make. The default is 5.
+    The number of predictions to make. Min: 1, Max: 500. Default: 5.
     """
 
 
diff --git a/jigsawstack/prompt_engine.py b/jigsawstack/prompt_engine.py
index c264db9..9963343 100644
--- a/jigsawstack/prompt_engine.py
+++ b/jigsawstack/prompt_engine.py
@@ -48,6 +48,7 @@ class PromptEngineRunResponse(TypedDict):
 
 
 class PromptEngineCreateParams(TypedDict):
+    name: NotRequired[str]
     prompt: str
     inputs: NotRequired[List[object]]
     return_prompt: Union[str, List[object], Dict[str, str]]
diff --git a/jigsawstack/search.py b/jigsawstack/search.py
index 7898f8b..0b2af13 100644
--- a/jigsawstack/search.py
+++ b/jigsawstack/search.py
@@ -193,6 +193,11 @@ class SearchParams(TypedDict):
     Whether to perform spell checking on the query. Defaults to True.
     """
 
+    max_results: NotRequired[int]
+    """
+    Maximum number of search results to return.
+    """
+
     safe_search: NotRequired[Literal["strict", "moderate", "off"]]
     """
     Safe search filtering level. Can be 'strict', 'moderate', or 'off'
@@ -241,7 +246,7 @@ def search(self, params: SearchParams) -> SearchResponse:
         safe_search = params.get("safe_search", "moderate")
         spell_check = params.get("spell_check", "True")
 
-        body = {
+        body: Dict[str, Any] = {
             "byo_urls": params.get("byo_urls", []),
             "query": query,
             "ai_overview": ai_overview,
@@ -249,6 +254,15 @@ def search(self, params: SearchParams) -> SearchResponse:
             "spell_check": spell_check,
         }
 
+        if "max_results" in params:
+            body["max_results"] = params["max_results"]
+
+        if "country_code" in params:
+            body["country_code"] = params["country_code"]
+
+        if "auto_scrape" in params:
+            body["auto_scrape"] = params["auto_scrape"]
+
         path = "/web/search"
         resp = Request(
             config=self.config,
@@ -304,13 +318,23 @@ async def search(self, params: SearchParams) -> SearchResponse:
         safe_search = params.get("safe_search", "moderate")
         spell_check = params.get("spell_check", "True")
 
-        body = {
+        body: Dict[str, Any] = {
             "byo_urls": params.get("byo_urls", []),
             "query": query,
             "ai_overview": ai_overview,
             "safe_search": safe_search,
             "spell_check": spell_check,
         }
+
+        if "max_results" in params:
+            body["max_results"] = params["max_results"]
+
+        if "country_code" in params:
+            body["country_code"] = params["country_code"]
+
+        if "auto_scrape" in params:
+            body["auto_scrape"] = params["auto_scrape"]
+
         resp = await AsyncRequest(
             config=self.config,
             path=path,
diff --git a/jigsawstack/version.py b/jigsawstack/version.py
index b8d2d3c..c9225b9 100644
--- a/jigsawstack/version.py
+++ b/jigsawstack/version.py
@@ -1,4 +1,4 @@
-__version__ = "0.3.9"
+__version__ = "0.4.0"
 
 
 def get_version() -> str:
diff --git a/jigsawstack/vision.py b/jigsawstack/vision.py
index 0b5a483..2778cf3 100644
--- a/jigsawstack/vision.py
+++ b/jigsawstack/vision.py
@@ -150,6 +150,10 @@ class VOCRParams(TypedDict):
     url: NotRequired[str]
     file_store_key: NotRequired[str]
     page_range: NotRequired[List[int]]
+    fine_grained: NotRequired[bool]
+    """
+    High fidelity word-level bounding boxes within complex documents. Default: false.
+    """
 
 
 class Word(TypedDict):
diff --git a/jigsawstack/web.py b/jigsawstack/web.py
index d432c25..c6ffe4d 100644
--- a/jigsawstack/web.py
+++ b/jigsawstack/web.py
@@ -128,7 +128,11 @@ class BaseAIScrapeParams(TypedDict):
 
 
 class AIScrapeParams(BaseAIScrapeParams):
-    element_prompts: NotRequired[List[str]]
+    element_prompts: NotRequired[Union[List[str], Dict[str, str]]]
+    """
+    List of prompts or a dictionary of key-value prompts for element extraction.
+    Max 5 items. If dict, max 50 chars per key and max 500 chars per prompt value.
+    """
     root_element_selector: NotRequired[str]
     page_position: NotRequired[int]
 
diff --git a/setup.py b/setup.py
index 8a9a443..aaa15df 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,7 @@
 
 setup(
     name="jigsawstack",
-    version="0.3.9",
+    version="0.4.0",
     description="JigsawStack - The AI SDK for Python",
     long_description=open("README.md", encoding="utf8").read(),
     long_description_content_type="text/markdown",
diff --git a/tests/test_ai_scrape.py b/tests/test_ai_scrape.py
index 1f53453..e7e8e30 100644
--- a/tests/test_ai_scrape.py
+++ b/tests/test_ai_scrape.py
@@ -77,16 +77,16 @@
             "is_mobile": True,
         },
     },
-    {
-        "name": "scrape_with_cookies",
-        "params": {
-            "url": URL,
-            "element_prompts": ["user data"],
-            "cookies": [
-                {"name": "session", "value": "test123", "domain": "example.com"}
-            ],
-        },
-    },
+    # {
+    #     "name": "scrape_with_cookies",
+    #     "params": {
+    #         "url": URL,
+    #         "element_prompts": ["user data"],
+    #         "cookies": [
+    #             {"name": "session", "value": "test123", "domain": "example.com"}
+    #         ],
+    #     },
+    # },
     {
         "name": "scrape_with_advance_config",
         "params": {