Repository: jerryjliu/gpt_index
File: llama-index-integrations/llms/llama-index-llms-openai/llama_index/llms/openai/base.py (Python)

Code before:
                    ToolSelection(
                        tool_id=tool_call.id,
                        tool_name=tool_call.function.name,
                        tool_kwargs=argument_dict,
                    )
                )

            return tool_selections

    def _prepare_schema(
        self, llm_kwargs: Optional[Dict[str, Any]], output_cls: Type[Model]
    ) -> Dict[str, Any]:
        from openai.resources.chat.completions.completions import (
            _type_to_response_format,
        )

        llm_kwargs = llm_kwargs or {}
        llm_kwargs["response_format"] = _type_to_response_format(output_cls)
        if "tool_choice" in llm_kwargs:
            del llm_kwargs["tool_choice"]
        return llm_kwargs

    def _should_use_structure_outputs(self) -> bool:
        return (
            self.pydantic_program_mode == PydanticProgramMode.DEFAULT
            and is_json_schema_supported(self.model)
        )

    @dispatcher.span
    def structured_predict(
        self,
        output_cls: Type[Model],
        prompt: PromptTemplate,
        llm_kwargs: Optional[Dict[str, Any]] = None,
        **prompt_args: Any,
    ) -> Model:
        """Structured predict."""
        llm_kwargs = llm_kwargs or {}

        if self._should_use_structure_outputs():
            messages = self._extend_messages(prompt.format_messages(**prompt_args))
            llm_kwargs = self._prepare_schema(llm_kwargs, output_cls)
            response = self.chat(messages, **llm_kwargs)
            return output_cls.model_validate_json(str(response.message.content))

        # when uses function calling to extract structured outputs
        # here we force tool_choice to be required
        llm_kwargs["tool_choice"] = (
            "required" if "tool_choice" not in llm_kwargs else llm_kwargs["tool_choice"]
        )
        return super().structured_predict(

Review comment (security):
Not super sure, but probably should also sanitize `response_format["json_schema"]["schema"]["title"]`

Code after:
                tool_selections.append(
                    ToolSelection(
                        tool_id=tool_call.id,
                        tool_name=tool_call.function.name,
                        tool_kwargs=argument_dict,
                    )
                )

            return tool_selections

    def _prepare_schema(
        self, llm_kwargs: Optional[Dict[str, Any]], output_cls: Type[Model]
    ) -> Dict[str, Any]:
        from openai.resources.chat.completions.completions import (
            _type_to_response_format,
        )

        llm_kwargs = llm_kwargs or {}
        response_format = _type_to_response_format(output_cls)
        if isinstance(response_format, dict):
            json_schema = response_format.get("json_schema")
            if isinstance(json_schema, dict) and "name" in json_schema:
                json_schema["name"] = re.sub(
                    r"[^a-zA-Z0-9_-]", "_", str(json_schema["name"])
                )
        llm_kwargs["response_format"] = response_format
        if "tool_choice" in llm_kwargs:
            del llm_kwargs["tool_choice"]
        return llm_kwargs

    def _should_use_structure_outputs(self) -> bool:
        return (
            self.pydantic_program_mode == PydanticProgramMode.DEFAULT
            and is_json_schema_supported(self.model)
        )

    @dispatcher.span
    def structured_predict(
        self,
        output_cls: Type[Model],
        prompt: PromptTemplate,
        llm_kwargs: Optional[Dict[str, Any]] = None,
        **prompt_args: Any,
    ) -> Model:
        """Structured predict."""
        llm_kwargs = llm_kwargs or {}

        if self._should_use_structure_outputs():
            messages = self._extend_messages(prompt.format_messages(**prompt_args))
            llm_kwargs = self._prepare_schema(llm_kwargs, output_cls)
            response = self.chat(messages, **llm_kwargs)