참고소스 수정본

참고/instructor-main/instructor/providers/xai/utils.py

@@ -0,0 +1,185 @@
+"""xAI-specific utilities.
+
+This module contains utilities specific to the xAI provider,
+including reask functions, response handlers, and message formatting.
+"""
+
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from ...mode import Mode
+
+if TYPE_CHECKING:
+    from xai_sdk import chat as xchat
+else:
+    try:
+        from xai_sdk import chat as xchat
+    except ImportError:
+        xchat = None
+
+
+def _convert_messages(messages: list[dict[str, Any]]):
+    """Convert OpenAI-style messages to xAI format."""
+    if xchat is None:
+        from ...core.exceptions import ConfigurationError
+
+        raise ConfigurationError(
+            "The xAI provider needs the optional dependency `xai-sdk`. "
+            'Install it with `uv pip install "instructor[xai]"` (or `pip install "instructor[xai]"`). '
+            "Note: xai-sdk requires Python 3.10+."
+        ) from None
+
+    converted = []
+    for m in messages:
+        role = m["role"]
+        content = m.get("content", "")
+        if isinstance(content, str):
+            c = xchat.text(content)
+        else:
+            raise ValueError("Only string content supported for xAI provider")
+        if role == "user":
+            converted.append(xchat.user(c))
+        elif role == "assistant":
+            converted.append(xchat.assistant(c))
+        elif role == "system":
+            converted.append(xchat.system(c))
+        elif role == "tool":
+            converted.append(xchat.tool_result(content))
+        else:
+            raise ValueError(f"Unsupported role: {role}")
+    return converted
+
+
+def reask_xai_json(
+    kwargs: dict[str, Any],
+    response: Any,
+    exception: Exception,
+):
+    """
+    Handle reask for xAI JSON mode when validation fails.
+
+    Kwargs modifications:
+    - Modifies: "messages" (appends user message requesting correction)
+    """
+    kwargs = kwargs.copy()
+    reask_msg = {
+        "role": "user",
+        "content": f"Validation Errors found:\n{exception}\nRecall the function correctly, fix the errors found in the following attempt:\n{response}",
+    }
+    kwargs["messages"].append(reask_msg)
+    return kwargs
+
+
+def reask_xai_tools(
+    kwargs: dict[str, Any],
+    response: Any,
+    exception: Exception,
+):
+    """
+    Handle reask for xAI tools mode when validation fails.
+
+    Kwargs modifications:
+    - Modifies: "messages" (appends assistant and user messages for tool correction)
+    """
+    kwargs = kwargs.copy()
+
+    # Add assistant response to conversation history
+    assistant_msg = {
+        "role": "assistant",
+        "content": str(response),
+    }
+    kwargs["messages"].append(assistant_msg)
+
+    # Add user correction request
+    reask_msg = {
+        "role": "user",
+        "content": f"Validation Error found:\n{exception}\nRecall the function correctly, fix the errors",
+    }
+    kwargs["messages"].append(reask_msg)
+    return kwargs
+
+
+def handle_xai_json(
+    response_model: type[Any] | None, new_kwargs: dict[str, Any]
+) -> tuple[type[Any] | None, dict[str, Any]]:
+    """
+    Handle xAI JSON mode.
+
+    When response_model is None:
+        - Converts messages from OpenAI format to xAI format
+        - No schema is added to the request
+
+    When response_model is provided:
+        - Converts messages from OpenAI format to xAI format
+        - Sets up the model for JSON parsing mode
+
+    Kwargs modifications:
+    - Modifies: "messages" (converts from OpenAI to xAI format)
+    - Removes: instructor-specific kwargs (max_retries, validation_context, context, hooks)
+    """
+    # Convert messages to xAI format
+    messages = new_kwargs.get("messages", [])
+    new_kwargs["x_messages"] = _convert_messages(messages)
+
+    # Remove instructor-specific kwargs that xAI doesn't support
+    new_kwargs.pop("max_retries", None)
+    new_kwargs.pop("validation_context", None)
+    new_kwargs.pop("context", None)
+    new_kwargs.pop("hooks", None)
+
+    return response_model, new_kwargs
+
+
+def handle_xai_tools(
+    response_model: type[Any] | None, new_kwargs: dict[str, Any]
+) -> tuple[type[Any] | None, dict[str, Any]]:
+    """
+    Handle xAI tools mode.
+
+    When response_model is None:
+        - Converts messages from OpenAI format to xAI format
+        - No tools are configured
+
+    When response_model is provided:
+        - Converts messages from OpenAI format to xAI format
+        - Sets up tool schema from the response model
+        - Configures tool choice for automatic tool selection
+
+    Kwargs modifications:
+    - Modifies: "messages" (converts from OpenAI to xAI format)
+    - Adds: "tool" (xAI tool schema) - only when response_model provided
+    - Removes: instructor-specific kwargs (max_retries, validation_context, context, hooks)
+    """
+    # Convert messages to xAI format
+    messages = new_kwargs.get("messages", [])
+    new_kwargs["x_messages"] = _convert_messages(messages)
+
+    # Remove instructor-specific kwargs that xAI doesn't support
+    new_kwargs.pop("max_retries", None)
+    new_kwargs.pop("validation_context", None)
+    new_kwargs.pop("context", None)
+    new_kwargs.pop("hooks", None)
+
+    if response_model is not None and xchat is not None:
+        # Set up tool schema for structured output
+        new_kwargs["tool"] = xchat.tool(
+            name=response_model.__name__,
+            description=response_model.__doc__ or "",
+            parameters=response_model.model_json_schema(),
+        )
+
+    return response_model, new_kwargs
+
+
+# Handler registry for xAI
+XAI_HANDLERS = {
+    Mode.XAI_JSON: {
+        "reask": reask_xai_json,
+        "response": handle_xai_json,
+    },
+    Mode.XAI_TOOLS: {
+        "reask": reask_xai_tools,
+        "response": handle_xai_tools,
+    },
+}