writer 0.8.3rc4__py3-none-any.whl → 1.25.1rc1__py3-none-any.whl

writer/{ai.py → ai/__init__.py}

@@ -1,6 +1,8 @@
 import json
 import logging
+from contextvars import ContextVar
 from datetime import datetime
+from functools import wraps
 from typing import (
     Any,
     Callable,
@@ -12,40 +14,72 @@ from typing import (
     Optional,
     Set,
     TypedDict,
+    TypeVar,
     Union,
     cast,
 )
 from uuid import uuid4
 
 from httpx import Timeout
-from writerai import Writer
-from writerai._exceptions import WriterError
+from writerai import DefaultHttpxClient, Writer
+from writerai._exceptions import BadRequestError, WriterError
 from writerai._response import BinaryAPIResponse
 from writerai._streaming import Stream
 from writerai._types import Body, Headers, NotGiven, Query
 from writerai.resources import FilesResource, GraphsResource
 from writerai.types import (
-    Chat,
+    ApplicationListResponse,
+    ApplicationRetrieveResponse,
+    ChatCompletion,
     Completion,
+    CompletionChunk,
     FileDeleteResponse,
     GraphDeleteResponse,
     GraphRemoveFileFromGraphResponse,
     GraphUpdateResponse,
-    StreamingData,
 )
 from writerai.types import File as SDKFile
 from writerai.types import Graph as SDKGraph
 from writerai.types.application_generate_content_params import Input
-from writerai.types.chat import ChoiceMessage, ChoiceMessageGraphData, ChoiceMessageToolCall
+from writerai.types.applications import (
+    ApplicationGenerateAsyncResponse,
+    JobCreateResponse,
+    JobRetryResponse,
+)
+from writerai.types.applications.application_graphs_response import ApplicationGraphsResponse
+from writerai.types.chat_chat_params import (
+    GraphData,
+    MessageContentMixedContentImageFragment,
+    MessageContentMixedContentImageFragmentImageURL,
+    MessageContentMixedContentTextFragment,
+    ResponseFormat,
+    ToolChoice,
+)
 from writerai.types.chat_chat_params import Message as WriterAIMessage
-from writerai.types.chat_chat_params import MessageGraphData
-from writerai.types.chat_chat_params import ToolFunctionTool as SDKFunctionTool
-from writerai.types.chat_chat_params import ToolGraphTool as SDKGraphTool
+from writerai.types.chat_completion_message import ChatCompletionMessage
 from writerai.types.question import Question
 from writerai.types.question_response_chunk import QuestionResponseChunk
+from writerai.types.shared_params.tool_param import FunctionTool as SDKFunctionTool
+from writerai.types.shared_params.tool_param import GraphTool as SDKGraphTool
+from writerai.types.shared_params.tool_param import LlmTool as SDKLlmTool
+from writerai.types.shared_params.tool_param import WebSearchTool as SDKWebSearchTool
 
 from writer.core import get_app_process
 
+DEFAULT_CHAT_MODEL = "palmyra-x5"
+DEFAULT_COMPLETION_MODEL = "palmyra-x5"
+
+
+_ai_client: ContextVar[Optional[Writer]] = ContextVar(
+    "ai_client", default=None
+)
+
+
+class ExtendedWebSearchTool(TypedDict, total=False):
+    """Extended web search tool that includes all fields supported by the API"""
+    type: Literal["web_search"]
+    function: Dict[str, Any]  # Flexible to support include_raw_content and other fields
+
 
 class APIOptions(TypedDict, total=False):
     extra_headers: Optional[Headers]
@@ -56,7 +90,15 @@ class APIOptions(TypedDict, total=False):
 
 class ChatOptions(APIOptions, total=False):
     model: str
-    tools: Union[Iterable[Union[SDKGraphTool, SDKFunctionTool]], NotGiven]
+    tool_choice: ToolChoice
+    tools: Union[
+            Iterable[
+                Union[SDKGraphTool, SDKFunctionTool, SDKLlmTool, SDKWebSearchTool]
+                ],
+            NotGiven
+        ]
+    response_format: Union[ResponseFormat, NotGiven]
+    logprobs: Union[bool, NotGiven]
     max_tokens: Union[int, NotGiven]
     n: Union[int, NotGiven]
     stop: Union[List[str], str, NotGiven]
@@ -81,6 +123,15 @@ class APIListOptions(APIOptions, total=False):
     order: Union[Literal["asc", "desc"], NotGiven]
 
 
+class APIRetrieveJobsOptions(APIOptions, total=False):
+    limit: Union[int, NotGiven]
+    offset: Union[int, NotGiven]
+    status: Union[
+        Literal["completed", "failed", "in_progress"],
+        NotGiven
+        ]
+
+
 class Tool(TypedDict, total=False):
     type: str
 
@@ -88,6 +139,7 @@ class Tool(TypedDict, total=False):
 class GraphTool(Tool):
     graph_ids: List[str]
     subqueries: bool
+    description: Optional[str]
 
 
 class FunctionToolParameterMeta(TypedDict):
@@ -128,7 +180,18 @@ def create_function_tool(
     )
 
 
-def _process_completion_data_chunk(choice: StreamingData) -> str:
+class LLMTool(Tool):
+    model: str
+    description: str
+
+
+class WebSearchTool(Tool):
+    include_domains: Optional[List[str]]
+    exclude_domains: Optional[List[str]]
+    include_raw_content: Optional[bool]
+
+
+def _process_completion_data_chunk(choice: CompletionChunk) -> str:
     text = choice.value
     if not text:
         return ""
@@ -137,7 +200,7 @@ def _process_completion_data_chunk(choice: StreamingData) -> str:
     raise ValueError("Failed to retrieve text from completion stream")
 
 
-def _process_chat_data_chunk(chat_data: Chat) -> tuple[dict, dict]:
+def _process_chat_data_chunk(chat_data: ChatCompletion) -> tuple[dict, dict]:
     choices = chat_data.choices
     for entry in choices:
         dict_entry = entry.model_dump()
@@ -173,18 +236,8 @@ class WriterAIManager:
         :raises RuntimeError: If an API key was not provided to initialize
         SDK client properly.
         """
-        try:
-            self.client = Writer(
-                # This is the default and can be omitted
-                api_key=token,
-            )
-        except WriterError:
-            raise RuntimeError(
-                "Failed to acquire Writer API key. " +
-                "Provide it by either setting a WRITER_API_KEY" +
-                " environment variable, or by initializing the" +
-                " AI module explicitly: writer.ai.init(\"my-writer-api-key\")"
-                ) from None
+        self.token = token
+        from writer.core import get_app_process
         current_process = get_app_process()
         setattr(current_process, 'ai_manager', self)
 
@@ -218,7 +271,7 @@ class WriterAIManager:
         :param token: The new token to use for authentication.
         """
         instance = cls.acquire_instance()
-        instance.client = Writer(api_key=token)
+        instance.token = token
 
     @classmethod
     def use_chat_model(cls) -> str:
@@ -227,7 +280,7 @@ class WriterAIManager:
 
         :returns: Name for the chat model.
         """
-        return "palmyra-x-004"
+        return DEFAULT_CHAT_MODEL
 
     @classmethod
     def use_completion_model(cls) -> str:
@@ -236,12 +289,50 @@ class WriterAIManager:
 
         :returns: Name for the completion model.
         """
-        return "palmyra-x-004"
+        return DEFAULT_COMPLETION_MODEL
 
     @classmethod
-    def acquire_client(cls) -> Writer:
+    def acquire_client(
+        cls,
+        custom_httpx_client: Optional[DefaultHttpxClient] = None,
+        force_new_client: Optional[bool] = False
+    ) -> Writer:
+        from writer.core import get_session
         instance = cls.acquire_instance()
-        return instance.client
+
+        # Acquire header from session
+        # and set it to the client
+
+        current_session = get_session()
+        custom_headers = {}
+
+        if current_session:
+            headers = current_session.headers or {}
+            agent_token_header = headers.get("x-agent-token")
+            if agent_token_header:
+                custom_headers = {
+                        "X-Agent-Token": agent_token_header
+                    }
+
+        try:
+            context_client = _ai_client.get(None)
+            if force_new_client or not context_client:
+                client = Writer(
+                    api_key=instance.token,
+                    default_headers=custom_headers,
+                    http_client=custom_httpx_client
+                    )
+                _ai_client.set(client)
+                return client
+            else:
+                return context_client
+        except WriterError:
+            raise RuntimeError(
+                "Failed to acquire Writer API key. " +
+                "Provide it by either setting a WRITER_API_KEY" +
+                " environment variable, or by initializing the" +
+                " AI module explicitly: writer.ai.init(\"my-writer-api-key\")"
+                ) from None
 
 
 class SDKWrapper:
@@ -372,7 +463,7 @@ class Graph(SDKWrapper):
         - `extra_body` (Optional[Body]):
         Additional body parameters for the request.
         - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-        Timeout for the request.
+        Timeout for the request in seconds.
         """
         config = config or {}
 
@@ -385,7 +476,12 @@ class Graph(SDKWrapper):
         if description:
             payload["description"] = description
         graphs = self._retrieve_graphs_accessor()
-        response = graphs.update(self.id, **payload, **config)
+        response = graphs.update(
+            self.id,
+            name=payload.get("name", NotGiven()),
+            description=payload.get("description", NotGiven()),
+            **config
+        )
         Graph.stale_ids.add(self.id)
         return response
 
@@ -414,7 +510,7 @@ class Graph(SDKWrapper):
         - `extra_body` (Optional[Body]):
         Additional body parameters for the request.
         - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-        Timeout for the request.
+        Timeout for the request in seconds.
         """
         config = config or {}
         file_id = None
@@ -461,7 +557,7 @@ class Graph(SDKWrapper):
         - `extra_body` (Optional[Body]):
         Additional body parameters for the request.
         - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-        Timeout for the request.
+        Timeout for the request in seconds.
         """
         config = config or {}
         file_id = None
@@ -516,25 +612,30 @@ class Graph(SDKWrapper):
         """
         Streams response for a question posed to the graph.
 
-        This method returns incremental chunks of the response, ideal for long 
+        This method returns incremental chunks of the response, ideal for long
         responses or when reduced latency is needed.
 
         :param question: The query or question to be answered by the graph.
-        :param subqueries: Enables subquery generation if set to True, 
+        :param subqueries: Enables subquery generation if set to True,
         enhancing the result.
-        :param config: Optional dictionary for additional API 
-        configuration settings. 
-        The configuration can include:
-            - ``extra_headers`` (Optional[Headers]): Additional headers.
-            - ``extra_query`` (Optional[Query]): Extra query parameters.
-            - ``extra_body`` (Optional[Body]): Additional body data.
-            - ``timeout`` (Union[float, Timeout, None, NotGiven]): Request timeout.
+        :param config: Optional dictionary for additional API
+        configuration settings.
 
         :yields: Incremental chunks of the answer to the question.
 
         :raises ValueError: If an invalid graph or graph ID
         is provided in `graphs_or_graph_ids`.
 
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+
         **Example Usage**:
 
         >>> for chunk in graph.stream_ask(
@@ -543,7 +644,7 @@ class Graph(SDKWrapper):
         ...     print(chunk)
         ...
         """
-        
+
         response = cast(
                 Stream[QuestionResponseChunk],
                 self._question(
@@ -572,22 +673,27 @@ class Graph(SDKWrapper):
             config: Optional[APIOptions] = None
     ):
         """
-        Sends a question to the graph and retrieves 
+        Sends a question to the graph and retrieves
         a single response.
 
         :param question: The query or question to be answered by the graph.
-        :param subqueries: Enables subquery generation if set to True, 
+        :param subqueries: Enables subquery generation if set to True,
         enhancing the result.
-        :param config: Optional dictionary for additional API 
+        :param config: Optional dictionary for additional API
         configuration settings.
-        The configuration can include:
-            - ``extra_headers`` (Optional[Headers]): Additional headers.
-            - ``extra_query`` (Optional[Query]): Extra query parameters.
-            - ``extra_body`` (Optional[Body]): Additional body data.
-            - ``timeout`` (Union[float, Timeout, None, NotGiven]): Request timeout.
 
         :return: The answer to the question from the graph(s).
 
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+
         **Example Usage**:
 
         >>> response = graph.ask(
@@ -631,7 +737,7 @@ def create_graph(
     - `extra_body` (Optional[Body]):
     Additional body parameters for the request.
     - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-    Timeout for the request.
+    Timeout for the request in seconds.
     """
     config = config or {}
     graphs = Graph._retrieve_graphs_accessor()
@@ -667,7 +773,7 @@ def retrieve_graph(
     - `extra_body` (Optional[Body]):
     Additional body parameters for the request.
     - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-    Timeout for the request.
+    Timeout for the request in seconds.
     """
     config = config or {}
     graphs = Graph._retrieve_graphs_accessor()
@@ -693,7 +799,7 @@ def list_graphs(config: Optional[APIListOptions] = None) -> List[Graph]:
     - `extra_body` (Optional[Body]):
     Additional body parameters for the request.
     - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-    Timeout for the request.
+    Timeout for the request in seconds.
     - `after` (Union[str, NotGiven]):
     Filter to retrieve items created after a specific cursor.
     - `before` (Union[str, NotGiven]):
@@ -808,7 +914,7 @@ def retrieve_file(file_id: str, config: Optional[APIOptions] = None) -> File:
     - `extra_body` (Optional[Body]):
     Additional body parameters for the request.
     - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-    Timeout for the request.
+    Timeout for the request in seconds.
     """
     config = config or {}
     files = File._retrieve_files_accessor()
@@ -834,7 +940,7 @@ def list_files(config: Optional[APIListOptions] = None) -> List[File]:
     - `extra_body` (Optional[Body]):
     Additional body parameters for the request.
     - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-    Timeout for the request.
+    Timeout for the request in seconds.
     - `after` (Union[str, NotGiven]):
     Filter to retrieve items created after a specific cursor.
     - `before` (Union[str, NotGiven]):
@@ -878,7 +984,7 @@ def upload_file(
     - `extra_body` (Optional[Body]):
     Additional body parameters for the request.
     - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-    Timeout for the request.
+    Timeout for the request in seconds.
     """
     config = config or {}
     files = File._retrieve_files_accessor()
@@ -919,7 +1025,7 @@ def delete_file(
     - `extra_body` (Optional[Body]):
     Additional body parameters for the request.
     - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
-    Timeout for the request.
+    Timeout for the request in seconds.
     """
     config = config or {}
     file_id = None
@@ -937,6 +1043,39 @@ def delete_file(
     return files.delete(file_id, **config)
 
 
+class GuardrailError(Exception):
+    def __init__(self, name: str, message: str, *args):
+        super().__init__(f"{message}: {name}", *args)
+
+
+R = TypeVar("R")
+
+
+def catch_guardrail_error(func: Callable[..., R]) -> Callable[..., R]:
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> R:
+        try:
+            return func(*args, **kwargs)
+        except BadRequestError as e:
+            parsed = e.response.json()
+            errors = parsed.get("errors")
+            if not errors:
+                raise
+
+            extras = parsed.get("extras")
+            if extras is None:
+                raise
+            guardrail_info = extras.get("guardrail_info")
+            if guardrail_info is None:
+                raise
+
+            raise GuardrailError(
+                name=guardrail_info["guardrail_name"],
+                message=errors[0]["description"],
+            ) from None
+
+    return wrapper
+
 class Conversation:
     """
     Manages messages within a conversation flow with an AI system,
@@ -970,13 +1109,44 @@ class Conversation:
     :param prompt_or_history: Initial system prompt as a string, or
                               history of messages as a list, used to start
                               a new conversation or to load an existing one.
-    :param config: Configuration settings for the conversation. These settings
-                   can include parameters such as `max_tokens`, `temperature`,
-                   and `timeout`, which affect the behavior and performance
-                   of the conversation operations. This configuration provides
-                   a default context for all operations, but can be overridden
-                   or extended by additional configurations passed directly
-                   to specific methods.
+    :param config: Configuration settings for the conversation.
+
+    The `config` dictionary can include the following keys:
+    - `extra_headers` (Optional[Headers]):
+    Additional headers for the request.
+    - `extra_query` (Optional[Query]):
+    Additional query parameters for the request.
+    - `extra_body` (Optional[Body]):
+    Additional body parameters for the request.
+    - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+    Timeout for the request in seconds.
+    - `model` (str):
+    The model to use for completion.
+    - `tool_choice` (ToolChoice):
+    Configure how the model will call functions: `auto` will allow the model
+    to automatically choose the best tool, `none` disables tool calling.
+    You can also pass a specific previously defined function.
+    - `logprobs` (Union[bool, NotGiven]):
+    Specifies whether to return log probabilities of the output tokens.
+    - `tools` (Union[Iterable[Union[SDKGraphTool,
+    SDKFunctionTool, SDKLlmTool]], NotGiven]):
+    Tools available for the model to use.
+    - `max_tokens` (Union[int, NotGiven]):
+    Maximum number of tokens to generate.
+    - `n` (Union[int, NotGiven]):
+    Number of completions to generate.
+    - `stop` (Union[List[str], str, NotGiven]):
+    Sequences where the API will stop generating tokens.
+    - `temperature` (Union[float, NotGiven]):
+    Controls the randomness or creativity of the model's responses.
+    A higher temperature results in more varied and less predictable text,
+    while a lower temperature produces more deterministic
+    and conservative outputs.
+    - `top_p` (Union[float, NotGiven]):
+    Sets the threshold for "nucleus sampling," a technique to focus the model's
+    token generation on the most likely subset of tokens. Only tokens with
+    cumulative probability above this threshold are considered, controlling the
+    trade-off between creativity and coherence.
 
     **Configuration Example:**
 
@@ -990,7 +1160,7 @@ class Conversation:
     or extend the initial configuration:
 
     >>> response = conversation.complete(
-    ...             data={
+    ...             config={
     ...                 'max_tokens': 150,
     ...                 'temperature': 0.7
     ...                 }
@@ -1000,17 +1170,24 @@ class Conversation:
     the `temperature` to 0.7 for this specific call.
 
     """
+    class ContentFragment(TypedDict, total=False):
+        """Content fragment for messages that can contain text or images."""
+        type: Literal["text", "image_url"]
+        text: Optional[str]
+        image_url: Optional[Dict[str, str]]
+
     class Message(TypedDict, total=False):
         """
         Typed dictionary for conversation messages.
 
         :param role: Specifies the sender role.
-        :param content: Text content of the message.
+        :param content: Text content of the message or array of content fragments
+        for multimodal messages.
         :param actions: Optional dictionary containing actions
         related to the message.
         """
         role: Literal["system", "assistant", "user", "tool"]
-        content: str
+        content: Union[str, List['Conversation.ContentFragment']]
         actions: Optional[dict]
         name: Optional[str]
         tool_call_id: Optional[str]
@@ -1037,13 +1214,44 @@ class Conversation:
                 f"Improper message format to add to Conversation: {message}"
                 )
 
-        if not (
-            isinstance(message["content"], str)
-            or
-            message["content"] is None
+        content = message["content"]
+        if isinstance(content, list):
+            # Validate multimodal content structure
+            for fragment in content:
+                if not isinstance(fragment, dict):
+                    raise ValueError(
+                        f"Invalid content fragment in message: {message}. "
+                        f"Fragments must be dictionaries."
+                    )
+                fragment_type = fragment.get("type")
+                if fragment_type not in ["text", "image_url"]:
+                    raise ValueError(
+                        f"Invalid fragment type '{fragment_type}' in message: {message}. "
+                        f"Type must be 'text' or 'image_url'."
+                    )
+                if fragment_type == "text" and "text" not in fragment:
+                    raise ValueError(
+                        f"Text fragment missing 'text' field in message: {message}"
+                    )
+                if fragment_type == "image_url" and "image_url" not in fragment:
+                    raise ValueError(
+                        f"Image fragment missing 'image_url' field in message: {message}"
+                    )
+                if fragment_type == "image_url" and not isinstance(fragment.get("image_url"), dict):
+                    raise ValueError(
+                        f"Image fragment 'image_url' must be a dict in message: {message}"
+                    )
+                if fragment_type == "image_url" and "url" not in fragment.get("image_url", {}):
+                    raise ValueError(
+                        f"Image fragment missing 'url' in 'image_url' in message: {message}"
+                    )
+        elif not (
+            isinstance(content, str)
+            or content is None
         ):
             raise ValueError(
-                f"Non-string content in message cannot be added: {message}"
+                f"Invalid content format in message: {message}. "
+                f"Content must be a string, None, or array of content fragments."
                 )
 
         if message["role"] not in ["system", "assistant", "user", "tool"]:
@@ -1125,7 +1333,24 @@ class Conversation:
         clear_chunk = _clear_chunk_flag(raw_chunk)
         updated_last_message: 'Conversation.Message' = self.messages[-1]
         if "content" in clear_chunk:
-            updated_last_message["content"] += clear_chunk.pop("content") or ""
+            new_content = clear_chunk.pop("content") or ""
+            if isinstance(updated_last_message["content"], list):
+                # Handle list content (multimodal)
+                if isinstance(new_content, str) and new_content:
+                    # Find the last text fragment and append to it, or create new one
+                    last_text_fragment = None
+                    for i in range(len(updated_last_message["content"]) - 1, -1, -1):
+                        if updated_last_message["content"][i].get("type") == "text":
+                            last_text_fragment = updated_last_message["content"][i]
+                            break
+
+                    if last_text_fragment:
+                        last_text_fragment["text"] = (last_text_fragment.get("text") or "") + new_content
+                    else:
+                        updated_last_message["content"].append({"type": "text", "text": new_content})
+            else:
+                # Handle string content
+                updated_last_message["content"] = str(updated_last_message["content"]) + str(new_content)
 
         if "tool_calls" in clear_chunk:
             # Ensure 'tool_calls' exists in updated_last_message as list
@@ -1138,12 +1363,13 @@ class Conversation:
             new_tool_calls = clear_chunk.pop("tool_calls", []) or []
 
             for new_call in new_tool_calls:
-                index = new_call.get("index", 0)
                 last_message_tool_calls = \
                     cast(list, updated_last_message["tool_calls"])
-                if index < len(last_message_tool_calls):
-                    # If there's an existing call at this index, update it
-                    existing_call = last_message_tool_calls[index]
+                if not new_call.get("id"):
+                    # Received ID is an indicator
+                    # of a new call that needs to be added:
+                    # otherwise, modify the latest existing call
+                    existing_call = last_message_tool_calls[-1]
                     for key, value in new_call.items():
                         if key == "function":
                             if key not in existing_call:
@@ -1166,12 +1392,12 @@ class Conversation:
                                     # to existing arguments
                                     call_function[fkey] += \
                                         fvalue or ""
-                                elif fvalue is not None:
+                                elif fvalue is not None and fvalue != '':
                                     call_function[fkey] = fvalue
-                        elif value is not None:
+                        elif value is not None and value != '':
                             existing_call[key] = value
                 else:
-                    # If no existing call, append the new call
+                    # If ID was received, this is the new call and we add it
                     last_message_tool_calls.append(new_call)
 
         updated_last_message |= clear_chunk
@@ -1182,15 +1408,48 @@ class Conversation:
         Converts a message object stored in Conversation to a Writer AI SDK
         `Message` model, suitable for calls to API.
 
-        :param raw_chunk: The data to be merged into the last message.
+        :param message: The message to prepare.
         :raises ValueError: If there are no messages in the conversation
         to merge with.
         """
         if not ("role" in message and "content" in message):
             raise ValueError("Improper message format")
-        sdk_message = WriterAIMessage(
-            content=message.get("content", None) or "",
-            role=message["role"]
+
+        content = message.get("content")
+        if isinstance(content, list):
+            # Handle multimodal content (text + images)
+            # Convert our ContentFragment format to SDK format
+            sdk_content: List[Union[MessageContentMixedContentTextFragment, MessageContentMixedContentImageFragment]] = []
+            for fragment in content:
+                if fragment.get("type") == "text":
+                    sdk_content.append({
+                        "type": "text",
+                        "text": fragment.get("text") or ""
+                    })
+                elif fragment.get("type") == "image_url":
+                    image_url_data = fragment.get("image_url", {})
+                    if isinstance(image_url_data, dict):
+                        # Extract the URL from the dict structure
+                        url = image_url_data.get("url", "")
+                    else:
+                        # Assume it's already a URL string
+                        url = str(image_url_data) if image_url_data else ""
+
+                    image_url_obj: MessageContentMixedContentImageFragmentImageURL = {"url": url}
+                    sdk_content.append({
+                        "type": "image_url",
+                        "image_url": image_url_obj
+                    })
+
+            sdk_message = WriterAIMessage(
+                content=sdk_content,
+                role=message["role"]
+            )
+        else:
+            # Handle simple text content
+            sdk_message = WriterAIMessage(
+                content=content or "",
+                role=message["role"]
             )
         if msg_name := message.get("name"):
             sdk_message["name"] = cast(str, msg_name)
@@ -1200,7 +1459,7 @@ class Conversation:
             sdk_message["tool_calls"] = cast(list, msg_tool_calls)
         if msg_graph_data := message.get("graph_data"):
             sdk_message["graph_data"] = cast(
-                MessageGraphData,
+                GraphData,
                 msg_graph_data
                 )
         if msg_refusal := message.get("refusal"):
@@ -1275,8 +1534,8 @@ class Conversation:
 
     def _prepare_tool(
             self,
-            tool_instance: Union['Graph', GraphTool, FunctionTool]
-            ) -> Union[SDKGraphTool, SDKFunctionTool]:
+            tool_instance: Union['Graph', GraphTool, FunctionTool, LLMTool, WebSearchTool, dict]
+            ) -> Union[SDKGraphTool, SDKFunctionTool, SDKLlmTool, SDKWebSearchTool]:
         """
         Internal helper function to process a tool instance
         into the required format.
@@ -1310,15 +1569,18 @@ class Conversation:
                     raise ValueError(
                         f"'type' for parameter '{param_name}' must be a string"
                         )
-                
+
                 supported_types = {
-                    "string", "number", "integer", "float", 
+                    "string", "number", "integer", "float",
                     "boolean", "array", "object", "null"
                 }
                 if param_info["type"] not in supported_types:
-                    raise ValueError(
+                    logging.warning(
                         f"Unsupported type '{param_info['type']}' " +
-                        f"for parameter '{param_name}'"
+                        f"for parameter '{param_name}'. " +
+                        "Fallback to 'string' type will be used. " +
+                        "This may lead to unexpected results. " +
+                        f"Supported types are: {supported_types}"
                     )
 
                 # Optional 'description' validation (if provided)
@@ -1334,7 +1596,9 @@ class Conversation:
 
             return True
 
-        def prepare_parameters(parameters: Dict[str, FunctionToolParameterMeta]) -> Dict:
+        def prepare_parameters(
+                parameters: Dict[str, FunctionToolParameterMeta]
+        ) -> Dict:
             """
             Prepares the parameters dictionary for a function tool.
 
@@ -1388,13 +1652,14 @@ class Conversation:
                     "type": "graph",
                     "function": {
                         "graph_ids": [tool_instance.id],
-                        "subqueries": True
+                        "subqueries": True,
+                        "description": tool_instance.description
                     }
                 }
             )
 
         elif isinstance(tool_instance, dict):
-            # Handle a dictionary (either a graph or a function)
+            # Handle a dictionary (either a graph, a function or a LLM tool)
             if "type" not in tool_instance:
                 raise ValueError(
                     "Invalid tool definition: 'type' field is missing"
@@ -1406,6 +1671,13 @@ class Conversation:
                 tool_instance = cast(GraphTool, tool_instance)
                 if "graph_ids" not in tool_instance:
                     raise ValueError("Graph tool must include 'graph_ids'")
+                if "description" not in tool_instance:
+                    logging.warning(
+                        "No description provided for `graph` tool. " +
+                        "This may produce suboptimal results. " +
+                        "To increase output quality, provide a description " +
+                        "for the tool."
+                        )
                 # Return graph tool JSON
 
                 graph_ids_valid = \
@@ -1419,6 +1691,9 @@ class Conversation:
                                 "graph_ids": tool_instance["graph_ids"],
                                 "subqueries": tool_instance.get(
                                     "subqueries", False
+                                    ),
+                                "description": tool_instance.get(
+                                    "description", None
                                     )
                             }
                         }
@@ -1441,6 +1716,13 @@ class Conversation:
                     raise ValueError(
                         "Function tool must include 'name' and 'parameters'"
                         )
+                if "description" not in tool_instance:
+                    logging.warning(
+                        "No description provided for `function` tool. " +
+                        "This may produce suboptimal results. " +
+                        "To increase output quality, provide a description " +
+                        "for the tool."
+                        )
 
                 parameters_valid = \
                     validate_parameters(tool_instance["parameters"])
@@ -1460,7 +1742,8 @@ class Conversation:
                                 "parameters":
                                     prepare_parameters(
                                         tool_instance["parameters"]
-                                        )
+                                        ),
+                                "description": tool_instance.get("description")
                                 }
                         }
                     )
@@ -1470,6 +1753,54 @@ class Conversation:
                         f"`{tool_instance['name']}`"
                         )
 
+            elif tool_type == "llm":
+                tool_instance = cast(LLMTool, tool_instance)
+                if "model" not in tool_instance:
+                    raise ValueError("LLM tool must include 'model'")
+                if "description" not in tool_instance:
+                    logging.warning(
+                        "No description provided for `llm` tool. " +
+                        "This may produce suboptimal results. " +
+                        "To increase output quality, provide a description " +
+                        "for the tool."
+                        )
+                # Return LLM tool JSON
+                return cast(
+                    SDKLlmTool,
+                    {
+                        "type": "llm",
+                        "function": {
+                            "model": tool_instance["model"],
+                            "description": tool_instance.get("description")
+                        }
+                    }
+                )
+            elif tool_type == "web_search":
+                # Return web search tool JSON - SDK format
+                tool_instance = cast(WebSearchTool, tool_instance)
+                function_config: Dict[str, Any] = {}
+
+                # Add required parameters - SDK expects these as lists
+                if "include_domains" in tool_instance:
+                    include_domains = tool_instance["include_domains"]
+                    if include_domains:
+                        function_config["include_domains"] = include_domains
+                if "exclude_domains" in tool_instance:
+                    exclude_domains = tool_instance["exclude_domains"]
+                    if exclude_domains:
+                        function_config["exclude_domains"] = exclude_domains
+                if "include_raw_content" in tool_instance:
+                    include_raw_content = tool_instance["include_raw_content"]
+                    if include_raw_content is not None:
+                        function_config["include_raw_content"] = include_raw_content
+
+                # Return as ExtendedWebSearchTool but ensure SDK compatibility
+                result = ExtendedWebSearchTool({
+                    "type": "web_search",
+                    "function": function_config
+                })
+                # Cast to SDKWebSearchTool for SDK compatibility while preserving extra fields
+                return cast(SDKWebSearchTool, result)
             else:
                 raise ValueError(f"Unsupported tool type: {tool_type}")
 
@@ -1522,12 +1853,43 @@ class Conversation:
         """
         self.__add__({"role": role, "content": message})
 
+    def add_with_images(
+            self,
+            role: Literal["system", "assistant", "user", "tool"],
+            text: Optional[str] = None,
+            image_urls: Optional[List[str]] = None
+            ):
+        """
+        Adds a new multimodal message with text and/or images.
+
+        :param role: The role of the message sender.
+        :param text: Optional text content.
+        :param image_urls: Optional list of image URLs (web URLs or base64 data URLs).
+        """
+        if not text and not image_urls:
+            raise ValueError("At least one of text or image_urls must be provided")
+
+        content: List['Conversation.ContentFragment'] = []
+
+        if text:
+            content.append({"type": "text", "text": text})
+
+        if image_urls:
+            for image_url in image_urls:
+                content.append({
+                    "type": "image_url",
+                    "image_url": {"url": image_url}
+                })
+
+        self.__add__({"role": role, "content": content})
+
+    @catch_guardrail_error
     def _send_chat_request(
             self,
             request_model: str,
             request_data: ChatOptions,
             stream: bool = False
-    ) -> Union[Stream, Chat]:
+    ) -> Union[Stream, ChatCompletion]:
         """
         Helper function to send a chat request to the LLM.
 
@@ -1535,7 +1897,7 @@ class Conversation:
         :param request_data: Configuration settings for the chat request.
         :param stream: Whether to use streaming mode.
         :return: The response from the LLM, either as
-        a Stream or a Chat object.
+        a Stream or a ChatCompletion object.
         """
         client = WriterAIManager.acquire_client()
         prepared_messages = [
@@ -1547,11 +1909,20 @@ class Conversation:
             f"prepared messages – {prepared_messages}, " +
             f"request_data – {request_data}"
             )
+        tools = request_data.get('tools', NotGiven())
+        tool_choice: Union[ToolChoice, NotGiven]
+        if isinstance(tools, NotGiven):
+            tool_choice = NotGiven()
+        else:
+            tool_choice = request_data.get('tool_choice', cast(ToolChoice, 'auto'))
         return client.chat.chat(
             messages=prepared_messages,
             model=request_model,
             stream=stream,
-            tools=request_data.get('tools', NotGiven()),
+            logprobs=request_data.get('logprobs', NotGiven()),
+            tools=tools,
+            tool_choice=tool_choice,
+            response_format=request_data.get('response_format', NotGiven()),
             max_tokens=request_data.get('max_tokens', NotGiven()),
             n=request_data.get('n', NotGiven()),
             stop=request_data.get('stop', NotGiven()),
@@ -1602,7 +1973,11 @@ class Conversation:
         elif target_type == "null":
             return None
         else:
-            raise ValueError(f"Unsupported target type: {target_type}")
+            logging.warning(
+                f"Unsupported target type: {target_type}. " +
+                "Falling back to string type. " +
+                "This may lead to unexpected results.")
+            return str(value)
 
     def _check_if_arguments_are_required(self, function_name: str) -> bool:
         callable_entry = self._callable_registry.get(function_name)
@@ -1654,7 +2029,7 @@ class Conversation:
                                     value,
                                     target_type
                                     )
-                        else:
+                        elif param_info.get("required") is True:
                             raise ValueError(
                                 f"Missing required parameter: {param_name}"
                                 )
@@ -1723,28 +2098,17 @@ class Conversation:
             }
 
         # Capture `tool_call_id` from the message
-        if tool_call_id is not None:
+        if tool_call_id is not None and tool_call_id != '':
             self._ongoing_tool_calls[index]["tool_call_id"] = tool_call_id
 
         # Capture `name` for function call
-        if tool_call_name is not None:
+        if tool_call_name is not None and tool_call_name != '':
             self._ongoing_tool_calls[index]["name"] = tool_call_name
 
         # Accumulate arguments across chunks
-        if tool_call_arguments is not None:
-            if (
-                tool_call_arguments.startswith("{")
-                and tool_call_arguments.endswith("}")
-            ):
-                # For cases when LLM "bugs" and returns
-                # the whole arguments string as a last chunk
-                fixed_chunk = tool_call_arguments.rsplit("{")[-1]
-                self._ongoing_tool_calls[index]["arguments"] = \
-                    "{" + fixed_chunk
-            else:
-                # Process normally
-                self._ongoing_tool_calls[index]["arguments"] += \
-                    tool_call_arguments
+        if tool_call_arguments is not None and tool_call_arguments != '':
+            self._ongoing_tool_calls[index]["arguments"] += \
+                tool_call_arguments
 
         # Check if we have all necessary data to execute the function
         tool_call_id, tool_call_name, tool_call_arguments = \
@@ -1752,8 +2116,12 @@ class Conversation:
             self._ongoing_tool_calls[index]["name"], \
             self._ongoing_tool_calls[index]["arguments"]
 
-        tool_call_id_ready = tool_call_id is not None
-        tool_call_name_ready = tool_call_name is not None
+        tool_call_id_ready = \
+            tool_call_id is not None \
+            and tool_call_id != ''
+        tool_call_name_ready = \
+            tool_call_name is not None \
+            and tool_call_name != ''
 
         # Check whether the arguments are prepared properly -
         # either present in correct format
@@ -1794,7 +2162,7 @@ class Conversation:
             if follow_up_message:
                 self._ongoing_tool_calls[index]["res"] = follow_up_message
 
-    def _process_tool_calls(self, message: ChoiceMessage):
+    def _process_tool_calls(self, message: ChatCompletionMessage):
         if message.tool_calls:
             for helper_index, tool_call in enumerate(message.tool_calls):
                 index = tool_call.index or helper_index
@@ -1812,8 +2180,8 @@ class Conversation:
     def _process_streaming_tool_calls(self, chunk: Dict):
         tool_calls = chunk["tool_calls"]
         if isinstance(tool_calls, list):
-            for tool_call in tool_calls:
-                index = tool_call["index"]
+            for helper_index, tool_call in enumerate(tool_calls):
+                index = tool_call["index"] or helper_index
                 tool_call_id = tool_call["id"]
                 tool_call_name = tool_call["function"]["name"]
                 tool_call_arguments = tool_call["function"]["arguments"]
@@ -1825,7 +2193,10 @@ class Conversation:
                     tool_call_arguments
                 )
 
-    def _prepare_received_message_for_history(self, message: ChoiceMessage):
+    def _prepare_received_message_for_history(
+            self,
+            message: ChatCompletionMessage
+            ):
         """
         Prepares a received message for adding to the conversation history.
 
@@ -1839,7 +2210,7 @@ class Conversation:
 
     def _process_response_data(
             self,
-            response_data: Chat,
+            response_data: ChatCompletion,
             request_model: str,
             request_data: ChatOptions,
             depth=1,
@@ -1865,7 +2236,7 @@ class Conversation:
                     # Send follow-up call to LLM
                     logging.debug("Sending a request to LLM")
                     follow_up_response = cast(
-                        Chat,
+                        ChatCompletion,
                         self._send_chat_request(
                             request_model=request_model,
                             request_data=request_data
@@ -1963,10 +2334,12 @@ class Conversation:
                     Graph,
                     GraphTool,
                     FunctionTool,
-                    List[Union[Graph, GraphTool, FunctionTool]]
+                    LLMTool,
+                    List[Union[Graph, GraphTool, FunctionTool, LLMTool]]
                     ]  # can be an instance of tool or a list of instances
                 ] = None,
             max_tool_depth: int = 5,
+            response_format: Optional[ResponseFormat] = None
             ) -> 'Conversation.Message':
         """
         Processes the conversation with the current messages and additional
@@ -1974,7 +2347,10 @@ class Conversation:
         Note: this method only produces AI model output and does not attach the
         result to the existing conversation history.
 
+        :param tools: Optional tools to use for processing.
         :param config: Optional parameters to pass for processing.
+        :param max_tool_depth: Maximum depth for tool calls processing.
+        :param response_format: Optional JSON schema used to format the model's output.
         :return: Generated message.
         :raises RuntimeError: If response data was not properly formatted
         to retrieve model text.
@@ -1991,11 +2367,18 @@ class Conversation:
         request_data: ChatOptions = {**config, **self.config}
         if prepared_tools:
             request_data |= {"tools": prepared_tools}
+        if response_format:
+            if not isinstance(response_format, dict):
+                raise ValueError(
+                    "Invalid schema for response_format: "
+                    f"dictionary required, got {type(response_format)}"
+                    )
+            request_data |= {"response_format": response_format}
         request_model = \
             request_data.get("model") or WriterAIManager.use_chat_model()
 
-        response_data: Chat = cast(
-                Chat,
+        response_data: ChatCompletion = cast(
+                ChatCompletion,
                 self._send_chat_request(
                     request_model=request_model,
                     request_data=request_data
@@ -2025,14 +2408,17 @@ class Conversation:
                     List[Union[Graph, GraphTool, FunctionTool]]
                     ]  # can be an instance of tool or a list of instances
                 ] = None,
-            max_tool_depth: int = 5
+            max_tool_depth: int = 5,
+            response_format: Optional[ResponseFormat] = None
             ) -> Generator[dict, None, None]:
         """
         Initiates a stream to receive chunks of the model's reply.
         Note: this method only produces AI model output and does not attach
         the result to the existing conversation history.
 
+        :param tools: Optional tools to use for processing.
         :param config: Optional parameters to pass for processing.
+        :param max_tool_depth: Maximum depth for tool calls processing.
         :yields: Model response chunks as they arrive from the stream.
         """
         config = config or {}
@@ -2047,6 +2433,13 @@ class Conversation:
         request_data: ChatOptions = {**config, **self.config}
         if prepared_tools:
             request_data |= {"tools": prepared_tools}
+        if response_format:
+            if not isinstance(response_format, dict):
+                raise ValueError(
+                    "Invalid schema for response_format: "
+                    f"dictionary required, got {type(response_format)}"
+                    )
+            request_data |= {"response_format": response_format}
         request_model = \
             request_data.get("model") or WriterAIManager.use_chat_model()
 
@@ -2074,7 +2467,7 @@ class Conversation:
         """
         Function to verify whether the message should be serializable.
 
-        :return: Boolean indicating if the message meets 
+        :return: Boolean indicating if the message meets
         the criteria for serialization.
         """
         if message["role"] in ["system", "tool"]:
@@ -2114,21 +2507,135 @@ class Conversation:
 
 
 class Apps:
+    def list(
+            self,
+            config: Optional[APIOptions] = None
+            ) -> List[ApplicationListResponse]:
+        """
+        Lists all applications available to the user.
+
+        :param config: Optional dictionary containing parameters
+        for the list call.
+        :return: List of applications.
+
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+        """
+        client = WriterAIManager.acquire_client()
+        config = config or {}
+
+        response = client.applications.list(**config)
+        # Convert the response to list
+        # to collect all the apps available
+        result = list(response)
+
+        return result
+
+    def retrieve(
+            self,
+            application_id: str,
+            config: Optional[APIOptions] = None
+            ) -> ApplicationRetrieveResponse:
+        """
+        Retrieves all information about a specific application by its ID.
+
+        :param application_id: The ID of the application to retrieve data for.
+        :param config: Optional dictionary containing parameters
+        for the retrieve call.
+        :return: The application data.
+
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+        """
+        client = WriterAIManager.acquire_client()
+        config = config or {}
+
+        response_data = client.applications.retrieve(
+            application_id=application_id,
+            **config
+            )
+        return response_data
+
     def generate_content(
             self,
             application_id: str,
-            input_dict: Optional[Dict[str, str]] = None,
-            config: Optional[APIOptions] = None) -> str:
+            input_dict: Optional[Dict[str, Optional[Union[List[str], str]]]] = None,
+            async_job: Optional[bool] = False,
+            config: Optional[APIOptions] = None
+            ) -> Union[str, JobCreateResponse]:
         """
         Generates output based on an existing AI Studio no-code application.
 
         :param application_id: The id for the application, which can be
-        obtained on AI Studio.
+            obtained on AI Studio.
         :param input_dict: Optional dictionary containing parameters for
-        the generation call.
-        :return: The generated text.
+            the generation call.
+        :param async_job: Optional. If True, the function initiates
+            an asynchronous job and returns job details instead of
+            waiting for the immediate output.
+        :return: The generated text
+            or the information about new asynchronous job.
         :raises RuntimeError: If response data was not properly formatted
         to retrieve model text.
+
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+
+        Examples:
+
+        **Synchronous Call (Immediate Output)**
+
+        >>> response = writer.ai.apps.generate_content(
+        ...     application_id="app_123",
+        ...     input_dict={"topic": "Climate Change"},
+        ...     async_job=False
+        ... )
+        >>> print(response)
+        "Climate change refers to long-term shifts in temperatures and weather patterns..."
+
+        **Asynchronous Call (Job Creation)**
+
+        >>> response = writer.ai.apps.generate_content(
+        ...     application_id="app_123",
+        ...     input_dict={"topic": "Climate Change"},
+        ...     async_job=True
+        ... )
+        >>> print(response)
+        JobCreateResponse(
+            id="job_456",
+            created_at=datetime(2025, 2, 24, 12, 30, 45),
+            status="in_progress"
+        )
+        >>> result = writer.ai.apps.retrieve_job(job_id=response.id)
+        >>> if result.status == "completed":
+        ...     print(result.data)
+        {
+            "title": "output",
+            "suggestion": "Climate change refers to long-term shifts in "
+                         "temperatures and weather patterns..."
+        }
+
         """
 
         client = WriterAIManager.acquire_client()
@@ -2137,26 +2644,211 @@ class Apps:
         inputs = []
 
         for k, v in input_dict.items():
+            # Convert None/empty to []
+            # to avoid API 400 errors on optional inputs
+            if v is None or v == "":
+                value = []
+            else:
+                value = v if isinstance(v, list) else [v]
+
             inputs.append(Input({
                 "id": k,
-                "value": v if isinstance(v, list) else [v]
+                "value": value
             }))
 
-        response_data = client.applications.generate_content(
-            application_id=application_id,
-            inputs=inputs,
+        if not async_job:
+            response_data = client.applications.generate_content(
+                application_id=application_id,
+                inputs=inputs,
+                **config
+                )
+
+            text = response_data.suggestion
+            if text:
+                return text
+
+            raise RuntimeError(
+                "Failed to acquire proper response " +
+                "for completion from data: " +
+                f"{response_data}"
+            )
+
+        else:
+            async_response_data = client.applications.jobs.create(
+                application_id=application_id,
+                inputs=inputs,
+                **config
+            )
+
+            return async_response_data
+
+    def retry_job(
+            self,
+            job_id: str,
+            config: Optional[APIOptions] = None
+            ) -> JobRetryResponse:
+        """
+        Retries a specific asynchronous job execution.
+
+        :param job_id: The unique identifier of the job to retry.
+        :param config: Optional API configuration options for
+        the retry request.
+        :return: The response data from retrying the job.
+
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+        """
+        client = WriterAIManager.acquire_client()
+        config = config or {}
+
+        response_data = client.applications.jobs.retry(
+            job_id=job_id,
             **config
             )
 
-        text = response_data.suggestion
-        if text:
-            return text
+        return response_data
 
-        raise RuntimeError(
-            "Failed to acquire proper response for completion from data: " +
-            f"{response_data}"
+    def retrieve_jobs(
+            self,
+            application_id: str,
+            config: Optional[APIRetrieveJobsOptions] = None
+            ) -> List[ApplicationGenerateAsyncResponse]:
+        """
+        Retrieves a list of jobs for a specific application.
+
+        :param application_id: The unique identifier of the application.
+        :param config: Optional configuration parameters for the API request.
+        :return: A list of job responses associated with
+        the specified application.
+
+        The `config` dictionary can include:
+        - `limit` (int): The pagination limit for retrieving the jobs.
+        - `offset` (int): The pagination offset for retrieving the jobs.
+        - `status` (Literal['in_progress', 'failed', 'completed']):
+        Filter jobs by the provided status.
+        - `extra_headers` (Optional[Headers]): Additional headers.
+        - `extra_query` (Optional[Query]): Extra query parameters.
+        - `extra_body` (Optional[Body]): Additional body data.
+        - `timeout` (Union[float, Timeout, None, NotGiven]):
+        Request timeout in seconds.
+        """
+        client = WriterAIManager.acquire_client()
+        config = config or {}
+
+        jobs = client.applications.jobs.list(
+            application_id=application_id,
+            **config
+        )
+
+        return list(jobs)
+
+    def retrieve_job(
+            self,
+            job_id: str,
+            config: Optional[APIOptions] = None
+            ) -> ApplicationGenerateAsyncResponse:
+        """
+        Retrieves an asynchronous job from the Writer AI API based on its ID.
+
+        :param job_id: The unique identifier of the job to retrieve.
+        :param config: Additional API configuration options.
+        :returns: The retrieved job object from the API.
+
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+        """
+        client = WriterAIManager.acquire_client()
+        config = config or {}
+
+        job = client.applications.jobs.retrieve(
+            job_id=job_id,
+            **config
+        )
+
+        return job
+
+    def retrieve_graphs(
+            self,
+            application_id: str,
+            config: Optional[APIOptions] = None
+            ) -> List[str]:
+        """
+        Retrieves a list of graph IDs for a specific application.
+
+        :param application_id: The unique identifier of the application.
+        :param config: Optional configuration parameters for the API request.
+        :return: A list of graph IDs associated with the specified application.
+
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+        """
+        client = WriterAIManager.acquire_client()
+        config = config or {}
+
+        graphs = client.applications.graphs.list(
+            application_id=application_id,
+            **config
+        )
+
+        return graphs.graph_ids
+
+    def associate_graphs(
+            self,
+            application_id: str,
+            graph_ids: List[str],
+            config: Optional[APIOptions] = None
+            ) -> ApplicationGraphsResponse:
+        """
+        Associates a list of graph IDs with a specific application.
+
+        :param application_id: The unique identifier of the application.
+        :param graph_ids: A list of graph IDs to associate with
+        the application.
+        :param config: Optional configuration parameters for the API request.
+        :return: The response data from associating the graphs.
+
+        The `config` dictionary can include the following keys:
+        - `extra_headers` (Optional[Headers]):
+        Additional headers for the request.
+        - `extra_query` (Optional[Query]):
+        Additional query parameters for the request.
+        - `extra_body` (Optional[Body]):
+        Additional body parameters for the request.
+        - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+        Timeout for the request in seconds.
+        """
+        client = WriterAIManager.acquire_client()
+        config = config or {}
+
+        response_data = client.applications.graphs.update(
+            application_id=application_id,
+            graph_ids=graph_ids,
+            **config
             )
 
+        return response_data
+
 
 class Tools:
     SplittingStrategy = Union[
@@ -2262,6 +2954,7 @@ class Tools:
         return result.entities
 
 
+@catch_guardrail_error
 def complete(
         initial_text: str,
         config: Optional['CreateOptions'] = None
@@ -2308,6 +3001,7 @@ def complete(
         f"{response_data}")
 
 
+@catch_guardrail_error
 def stream_complete(
         initial_text: str,
         config: Optional['CreateOptions'] = None
@@ -2373,21 +3067,16 @@ def ask(
     config: Optional[APIOptions] = None
 ):
     """
-    Sends a question to the specified graph(s) and retrieves 
+    Sends a question to the specified graph(s) and retrieves
     a single response.
 
     :param question: The query or question to be answered by the graph(s).
-    :param graphs_or_graph_ids: A list of `Graph` objects or graph IDs that 
+    :param graphs_or_graph_ids: A list of `Graph` objects or graph IDs that
     should be queried.
-    :param subqueries: Enables subquery generation if set to True, 
+    :param subqueries: Enables subquery generation if set to True,
     enhancing the result.
-    :param config: Optional dictionary for additional API 
+    :param config: Optional dictionary for additional API
     configuration settings.
-    The configuration can include:
-        - ``extra_headers`` (Optional[Headers]): Additional headers.
-        - ``extra_query`` (Optional[Query]): Extra query parameters.
-        - ``extra_body`` (Optional[Body]): Additional body data.
-        - ``timeout`` (Union[float, Timeout, None, NotGiven]): Request timeout.
 
     :return: The answer to the question from the graph(s).
 
@@ -2396,6 +3085,16 @@ def ask(
     :raises RuntimeError: If the API response is improperly formatted
     or the answer cannot be retrieved.
 
+    The `config` dictionary can include the following keys:
+    - `extra_headers` (Optional[Headers]):
+    Additional headers for the request.
+    - `extra_query` (Optional[Query]):
+    Additional query parameters for the request.
+    - `extra_body` (Optional[Body]):
+    Additional body parameters for the request.
+    - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+    Timeout for the request in seconds.
+
     **Example Usage**:
 
     >>> response = ask(
@@ -2430,27 +3129,32 @@ def stream_ask(
     """
     Streams response for a question posed to the specified graph(s).
 
-    This method returns incremental chunks of the response, ideal for long 
+    This method returns incremental chunks of the response, ideal for long
     responses or when reduced latency is needed.
 
     :param question: The query or question to be answered by the graph(s).
-    :param graphs_or_graph_ids: A list of Graph objects or graph IDs that 
+    :param graphs_or_graph_ids: A list of Graph objects or graph IDs that
     should be queried.
-    :param subqueries: Enables subquery generation if set to True, 
+    :param subqueries: Enables subquery generation if set to True,
     enhancing the result.
-    :param config: Optional dictionary for additional API 
-    configuration settings. 
-    The configuration can include:
-        - ``extra_headers`` (Optional[Headers]): Additional headers.
-        - ``extra_query`` (Optional[Query]): Extra query parameters.
-        - ``extra_body`` (Optional[Body]): Additional body data.
-        - ``timeout`` (Union[float, Timeout, None, NotGiven]): Request timeout.
+    :param config: Optional dictionary for additional API
+    configuration settings.
 
     :yields: Incremental chunks of the answer to the question.
 
     :raises ValueError: If an invalid graph or graph ID
     is provided in `graphs_or_graph_ids`.
 
+    The `config` dictionary can include the following keys:
+    - `extra_headers` (Optional[Headers]):
+    Additional headers for the request.
+    - `extra_query` (Optional[Query]):
+    Additional query parameters for the request.
+    - `extra_body` (Optional[Body]):
+    Additional body parameters for the request.
+    - `timeout` (Union[float, httpx.Timeout, None, NotGiven]):
+    Timeout for the request in seconds.
+
     **Example Usage**:
 
     >>> for chunk in stream_ask(