⚡️ Speed up method MistralStreamedResponse._try_get_output_tool_from_text by 5%

[codeflash-ai]

📄 5% (0.05x) speedup for MistralStreamedResponse._try_get_output_tool_from_text in pydantic_ai_slim/pydantic_ai/models/mistral.py

⏱️ Runtime : 462 microseconds → 439 microseconds (best of 154 runs)

📝 Explanation and details

Here's an optimized rewrite of your code. The main improvements are.

• Avoid unnecessary local variable assignments (direct returns instead of storing before return).
• Short-circuit checks: Use early return False in loops.
• Avoid repeated dict access: Fetch and reuse properties in nested checks.
• Tighten loops using references: Reduce attribute accesses inside loops.
• Minor micro-optimizations for clarity and speed.
• No change to function signatures or core logic; output is unchanged.

Summary of changes:

• Remove unnecessary type annotations on local variables for speed.
• Combined guard/checks and used fast local variable lookups.
• Only access dicts when really needed, and cache lookups where possible.
• Added explicit return None at the end of _try_get_output_tool_from_text for safety.
• Everything else is strictly as in the original for correctness and API.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 43 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import json
# Patch pydantic_core.from_json for testing
import sys
import types
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Literal, Union

# imports
import pytest  # used for our unit tests
from pydantic_ai.models.mistral import MistralStreamedResponse

# --- Dummy/Minimal implementations for dependencies ---

class ToolCallPart:
    def __init__(self, tool_name: str, args: dict):
        self.tool_name = tool_name
        self.args = args

    def __eq__(self, other):
        if not isinstance(other, ToolCallPart):
            return False
        return self.tool_name == other.tool_name and self.args == other.args

class ToolDefinition:
    def __init__(self, name: str, parameters_json_schema: dict):
        self.name = name
        self.parameters_json_schema = parameters_json_schema

# --- Unit tests ---

# --- Basic Test Cases ---

def test_valid_single_toolcall_simple():
    """Test with a simple valid JSON and a matching tool schema."""
    schema = {
        "type": "object",
        "properties": {"foo": {"type": "string"}},
        "required": ["foo"]
    }
    tool = ToolDefinition(name="simple_tool", parameters_json_schema=schema)
    text = '{"foo": "bar"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"simple_tool": tool}); result = codeflash_output # 6.88μs -> 5.79μs (18.7% faster)

def test_valid_multiple_tools_only_one_matches():
    """Test with multiple tools, only one matches the JSON."""
    schema1 = {
        "type": "object",
        "properties": {"foo": {"type": "string"}},
        "required": ["foo"]
    }
    schema2 = {
        "type": "object",
        "properties": {"bar": {"type": "integer"}},
        "required": ["bar"]
    }
    tool1 = ToolDefinition(name="tool1", parameters_json_schema=schema1)
    tool2 = ToolDefinition(name="tool2", parameters_json_schema=schema2)
    text = '{"bar": 123}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"tool1": tool1, "tool2": tool2}); result = codeflash_output # 6.17μs -> 5.46μs (13.0% faster)

def test_valid_toolcall_with_extra_fields():
    """Test JSON with extra fields not in schema (should still match if required fields are present and correct)."""
    schema = {
        "type": "object",
        "properties": {"foo": {"type": "string"}},
        "required": ["foo"]
    }
    tool = ToolDefinition(name="tool", parameters_json_schema=schema)
    text = '{"foo": "bar", "extra": 123}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"tool": tool}); result = codeflash_output # 5.38μs -> 5.12μs (4.88% faster)

def test_invalid_json_returns_none():
    """Test with invalid JSON input (should return None)."""
    schema = {
        "type": "object",
        "properties": {"foo": {"type": "string"}},
        "required": ["foo"]
    }
    tool = ToolDefinition(name="tool", parameters_json_schema=schema)
    text = '{"foo": "bar"'  # Missing closing brace
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"tool": tool}); result = codeflash_output # 5.00μs -> 4.83μs (3.46% faster)

def test_json_missing_required_field_returns_none():
    """Test JSON missing required field (should return None)."""
    schema = {
        "type": "object",
        "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}},
        "required": ["foo", "bar"]
    }
    tool = ToolDefinition(name="tool", parameters_json_schema=schema)
    text = '{"foo": "baz"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"tool": tool}); result = codeflash_output # 1.25μs -> 1.25μs (0.000% faster)

# --- Edge Test Cases ---

def test_empty_json_object():
    """Test with empty JSON object and schema with no required fields (should match)."""
    schema = {
        "type": "object",
        "properties": {},
        "required": []
    }
    tool = ToolDefinition(name="empty_tool", parameters_json_schema=schema)
    text = '{}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"empty_tool": tool}); result = codeflash_output # 375ns -> 333ns (12.6% faster)

def test_empty_json_object_with_required_fields():
    """Test with empty JSON object and schema with required fields (should not match)."""
    schema = {
        "type": "object",
        "properties": {"foo": {"type": "string"}},
        "required": ["foo"]
    }
    tool = ToolDefinition(name="tool", parameters_json_schema=schema)
    text = '{}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"tool": tool}); result = codeflash_output # 291ns -> 333ns (12.6% slower)

def test_json_with_array_type():
    """Test with JSON containing an array and schema expecting array of strings."""
    schema = {
        "type": "object",
        "properties": {"items": {"type": "array", "items": {"type": "string"}}},
        "required": ["items"]
    }
    tool = ToolDefinition(name="array_tool", parameters_json_schema=schema)
    text = '{"items": ["a", "b", "c"]}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"array_tool": tool}); result = codeflash_output # 5.88μs -> 5.75μs (2.17% faster)

def test_json_with_wrong_array_type():
    """Test with JSON array of wrong type (should not match schema)."""
    schema = {
        "type": "object",
        "properties": {"nums": {"type": "array", "items": {"type": "integer"}}},
        "required": ["nums"]
    }
    tool = ToolDefinition(name="num_tool", parameters_json_schema=schema)
    text = '{"nums": ["a", "b"]}'  # Should be integers
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"num_tool": tool}); result = codeflash_output # 1.54μs -> 1.54μs (0.000% faster)

def test_json_with_nested_object():
    """Test with JSON containing nested object and schema expecting nested object."""
    schema = {
        "type": "object",
        "properties": {
            "config": {
                "type": "object",
                "properties": {"enabled": {"type": "boolean"}},
                "required": ["enabled"]
            }
        },
        "required": ["config"]
    }
    tool = ToolDefinition(name="nested_tool", parameters_json_schema=schema)
    text = '{"config": {"enabled": true}}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"nested_tool": tool}); result = codeflash_output # 6.00μs -> 5.79μs (3.59% faster)

def test_json_with_missing_nested_required_field():
    """Test with JSON missing required field in nested object (should not match)."""
    schema = {
        "type": "object",
        "properties": {
            "config": {
                "type": "object",
                "properties": {"enabled": {"type": "boolean"}},
                "required": ["enabled"]
            }
        },
        "required": ["config"]
    }
    tool = ToolDefinition(name="nested_tool", parameters_json_schema=schema)
    text = '{"config": {}}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"nested_tool": tool}); result = codeflash_output # 1.54μs -> 1.46μs (5.69% faster)

def test_json_with_number_type_accepts_int_and_float():
    """Test JSON with number type accepts both int and float."""
    schema = {
        "type": "object",
        "properties": {"val": {"type": "number"}},
        "required": ["val"]
    }
    tool = ToolDefinition(name="num_tool", parameters_json_schema=schema)
    text1 = '{"val": 42}'
    text2 = '{"val": 3.14}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text1, {"num_tool": tool}); result1 = codeflash_output # 1.25μs -> 1.21μs (3.39% faster)
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text2, {"num_tool": tool}); result2 = codeflash_output # 4.92μs -> 4.71μs (4.42% faster)

def test_json_with_trailing_strings():
    """Test JSON with trailing non-JSON strings (should parse only the JSON part)."""
    schema = {
        "type": "object",
        "properties": {"foo": {"type": "string"}},
        "required": ["foo"]
    }
    tool = ToolDefinition(name="tool", parameters_json_schema=schema)
    text = '{"foo": "bar"} trailing text that is not JSON'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"tool": tool}); result = codeflash_output # 4.88μs -> 4.79μs (1.73% faster)

def test_json_with_null_type():
    """Test JSON with null type in schema."""
    schema = {
        "type": "object",
        "properties": {"foo": {"type": "null"}},
        "required": ["foo"]
    }
    tool = ToolDefinition(name="null_tool", parameters_json_schema=schema)
    text = '{"foo": null}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"null_tool": tool}); result = codeflash_output # 4.88μs -> 4.79μs (1.75% faster)

def test_json_with_boolean_type():
    """Test JSON with boolean type in schema."""
    schema = {
        "type": "object",
        "properties": {"flag": {"type": "boolean"}},
        "required": ["flag"]
    }
    tool = ToolDefinition(name="bool_tool", parameters_json_schema=schema)
    text = '{"flag": false}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"bool_tool": tool}); result = codeflash_output # 5.04μs -> 4.96μs (1.67% faster)

# --- Large Scale Test Cases ---

def test_large_number_of_tools():
    """Test performance and correctness with a large number of tools, only one matches."""
    # 999 tools with non-matching schema, 1 with matching schema
    tools = {}
    for i in range(999):
        schema = {
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
        tools[f"tool_{i}"] = ToolDefinition(name=f"tool_{i}", parameters_json_schema=schema)
    # Add the matching tool
    match_schema = {
        "type": "object",
        "properties": {"bar": {"type": "integer"}},
        "required": ["bar"]
    }
    tools["match_tool"] = ToolDefinition(name="match_tool", parameters_json_schema=match_schema)
    text = '{"bar": 123}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, tools); result = codeflash_output # 95.0μs -> 93.0μs (2.15% faster)

def test_large_json_object():
    """Test with a large JSON object and schema with many required fields."""
    N = 200
    properties = {f"field{i}": {"type": "integer"} for i in range(N)}
    required = [f"field{i}" for i in range(N)]
    schema = {
        "type": "object",
        "properties": properties,
        "required": required
    }
    tool = ToolDefinition(name="big_tool", parameters_json_schema=schema)
    json_obj = {f"field{i}": i for i in range(N)}
    text = json.dumps(json_obj)
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"big_tool": tool}); result = codeflash_output # 41.2μs -> 38.1μs (8.09% faster)

def test_large_array_in_json():
    """Test with a large array in JSON and schema expecting array of integers."""
    N = 500
    schema = {
        "type": "object",
        "properties": {"nums": {"type": "array", "items": {"type": "integer"}}},
        "required": ["nums"]
    }
    tool = ToolDefinition(name="array_tool", parameters_json_schema=schema)
    nums = list(range(N))
    text = json.dumps({"nums": nums})
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"array_tool": tool}); result = codeflash_output # 19.8μs -> 17.8μs (11.5% faster)

def test_large_nested_json():
    """Test with deeply nested JSON and schema."""
    schema = {
        "type": "object",
        "properties": {
            "outer": {
                "type": "object",
                "properties": {
                    "middle": {
                        "type": "object",
                        "properties": {
                            "inner": {"type": "string"}
                        },
                        "required": ["inner"]
                    }
                },
                "required": ["middle"]
            }
        },
        "required": ["outer"]
    }
    tool = ToolDefinition(name="nested_tool", parameters_json_schema=schema)
    text = '{"outer": {"middle": {"inner": "value"}}}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"nested_tool": tool}); result = codeflash_output # 6.04μs -> 5.92μs (2.11% faster)

def test_large_scale_no_match():
    """Test with many tools and a JSON that matches none (should return None)."""
    tools = {}
    for i in range(500):
        schema = {
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
        tools[f"tool_{i}"] = ToolDefinition(name=f"tool_{i}", parameters_json_schema=schema)
    text = '{"bar": 123}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, tools); result = codeflash_output # 46.7μs -> 45.6μs (2.38% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import json
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Literal, Union

# imports
import pytest  # used for our unit tests
from pydantic_ai.models.mistral import MistralStreamedResponse


# ToolDefinition stub
@dataclass
class ToolDefinition:
    name: str
    parameters_json_schema: dict[str, Any]

# --- Unit Tests ---

# --- Basic Test Cases ---

def test_simple_valid_toolcall():
    # Tool expects a string parameter called 'foo'
    tool = ToolDefinition(
        name='echo',
        parameters_json_schema={
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
    )
    text = '{"foo": "bar"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"echo": tool}); result = codeflash_output # 5.21μs -> 5.33μs (2.36% slower)

def test_multiple_tools_only_one_matches():
    # Two tools, only one matches the input
    tool1 = ToolDefinition(
        name='adder',
        parameters_json_schema={
            "type": "object",
            "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}},
            "required": ["x", "y"]
        }
    )
    tool2 = ToolDefinition(
        name='echo',
        parameters_json_schema={
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
    )
    text = '{"foo": "hello"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"adder": tool1, "echo": tool2}); result = codeflash_output # 5.25μs -> 5.25μs (0.000% faster)

def test_valid_toolcall_with_extra_fields():
    # Tool expects 'foo', but input also has 'bar'
    tool = ToolDefinition(
        name='echo',
        parameters_json_schema={
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
    )
    text = '{"foo": "bar", "bar": 123}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"echo": tool}); result = codeflash_output # 4.96μs -> 4.88μs (1.70% faster)

def test_no_matching_tool():
    # Tool expects 'foo', input has 'bar'
    tool = ToolDefinition(
        name='echo',
        parameters_json_schema={
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
    )
    text = '{"bar": "baz"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"echo": tool}); result = codeflash_output # 875ns -> 916ns (4.48% slower)

def test_empty_tools_dict():
    text = '{"foo": "bar"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {}); result = codeflash_output # 500ns -> 500ns (0.000% faster)

# --- Edge Test Cases ---

def test_invalid_json_input():
    tool = ToolDefinition(
        name='echo',
        parameters_json_schema={
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
    )
    text = '{"foo": "bar"'  # missing closing }
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"echo": tool}); result = codeflash_output # 4.88μs -> 4.88μs (0.000% faster)


def test_null_json_input():
    tool = ToolDefinition(
        name='echo',
        parameters_json_schema={
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
    )
    text = 'null'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"echo": tool}); result = codeflash_output # 500ns -> 500ns (0.000% faster)

def test_required_field_wrong_type():
    tool = ToolDefinition(
        name='echo',
        parameters_json_schema={
            "type": "object",
            "properties": {"foo": {"type": "integer"}},
            "required": ["foo"]
        }
    )
    text = '{"foo": "not_an_int"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"echo": tool}); result = codeflash_output # 1.58μs -> 1.38μs (15.1% faster)

def test_nested_required_object():
    tool = ToolDefinition(
        name='nest',
        parameters_json_schema={
            "type": "object",
            "properties": {
                "outer": {
                    "type": "object",
                    "properties": {
                        "inner": {"type": "string"}
                    },
                    "required": ["inner"]
                }
            },
            "required": ["outer"]
        }
    )
    text = '{"outer": {"inner": "value"}}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"nest": tool}); result = codeflash_output # 7.38μs -> 6.38μs (15.7% faster)

def test_nested_required_object_missing_inner():
    tool = ToolDefinition(
        name='nest',
        parameters_json_schema={
            "type": "object",
            "properties": {
                "outer": {
                    "type": "object",
                    "properties": {
                        "inner": {"type": "string"}
                    },
                    "required": ["inner"]
                }
            },
            "required": ["outer"]
        }
    )
    text = '{"outer": {}}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"nest": tool}); result = codeflash_output # 1.58μs -> 1.46μs (8.64% faster)

def test_array_type_enforcement():
    tool = ToolDefinition(
        name='arraytool',
        parameters_json_schema={
            "type": "object",
            "properties": {
                "nums": {
                    "type": "array",
                    "items": {"type": "integer"}
                }
            },
            "required": ["nums"]
        }
    )
    text = '{"nums": [1, 2, 3]}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"arraytool": tool}); result = codeflash_output # 6.33μs -> 6.00μs (5.55% faster)

def test_array_type_wrong_inner_type():
    tool = ToolDefinition(
        name='arraytool',
        parameters_json_schema={
            "type": "object",
            "properties": {
                "nums": {
                    "type": "array",
                    "items": {"type": "integer"}
                }
            },
            "required": ["nums"]
        }
    )
    text = '{"nums": [1, "2", 3]}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"arraytool": tool}); result = codeflash_output # 1.71μs -> 1.67μs (2.58% faster)

def test_object_type_field_is_not_object():
    tool = ToolDefinition(
        name='objtool',
        parameters_json_schema={
            "type": "object",
            "properties": {
                "opt": {
                    "type": "object",
                    "properties": {"a": {"type": "string"}},
                    "required": ["a"]
                }
            },
            "required": ["opt"]
        }
    )
    text = '{"opt": "not_an_object"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"objtool": tool}); result = codeflash_output # 1.25μs -> 1.25μs (0.000% faster)

def test_required_field_null():
    tool = ToolDefinition(
        name='nulltool',
        parameters_json_schema={
            "type": "object",
            "properties": {"foo": {"type": "string"}},
            "required": ["foo"]
        }
    )
    text = '{"foo": null}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"nulltool": tool}); result = codeflash_output # 1.17μs -> 1.21μs (3.39% slower)

def test_required_field_boolean_type():
    tool = ToolDefinition(
        name='booltool',
        parameters_json_schema={
            "type": "object",
            "properties": {"flag": {"type": "boolean"}},
            "required": ["flag"]
        }
    )
    text = '{"flag": true}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"booltool": tool}); result = codeflash_output # 5.62μs -> 5.38μs (4.65% faster)

def test_required_field_boolean_type_wrong():
    tool = ToolDefinition(
        name='booltool',
        parameters_json_schema={
            "type": "object",
            "properties": {"flag": {"type": "boolean"}},
            "required": ["flag"]
        }
    )
    text = '{"flag": "true"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"booltool": tool}); result = codeflash_output # 1.12μs -> 1.12μs (0.000% faster)

# --- Large Scale Test Cases ---

def test_large_number_of_tools_one_matches():
    # 500 tools, only the last one matches
    tools = {}
    for i in range(499):
        tools[f"tool_{i}"] = ToolDefinition(
            name=f"tool_{i}",
            parameters_json_schema={
                "type": "object",
                "properties": {"foo": {"type": "integer"}},
                "required": ["foo"]
            }
        )
    tools["tool_match"] = ToolDefinition(
        name="tool_match",
        parameters_json_schema={
            "type": "object",
            "properties": {"bar": {"type": "string"}},
            "required": ["bar"]
        }
    )
    text = '{"bar": "found"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, tools); result = codeflash_output # 51.3μs -> 48.3μs (6.30% faster)

def test_large_json_payload():
    # Tool expects a list of 500 integers
    tool = ToolDefinition(
        name='bigarray',
        parameters_json_schema={
            "type": "object",
            "properties": {
                "nums": {
                    "type": "array",
                    "items": {"type": "integer"}
                }
            },
            "required": ["nums"]
        }
    )
    nums = list(range(500))
    text = json.dumps({"nums": nums})
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"bigarray": tool}); result = codeflash_output # 20.2μs -> 17.7μs (13.9% faster)

def test_large_json_payload_wrong_type():
    # Tool expects a list of 500 integers, but one is a string
    tool = ToolDefinition(
        name='bigarray',
        parameters_json_schema={
            "type": "object",
            "properties": {
                "nums": {
                    "type": "array",
                    "items": {"type": "integer"}
                }
            },
            "required": ["nums"]
        }
    )
    nums = list(range(499)) + ["oops"]
    text = json.dumps({"nums": nums})
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"bigarray": tool}); result = codeflash_output # 16.4μs -> 14.3μs (14.5% faster)

def test_large_number_of_tools_none_match():
    # 500 tools, none match the input
    tools = {}
    for i in range(500):
        tools[f"tool_{i}"] = ToolDefinition(
            name=f"tool_{i}",
            parameters_json_schema={
                "type": "object",
                "properties": {"foo": {"type": "integer"}},
                "required": ["foo"]
            }
        )
    text = '{"bar": "notfound"}'
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, tools); result = codeflash_output # 45.7μs -> 44.5μs (2.62% faster)

def test_large_nested_object():
    # Tool expects a deeply nested object
    schema = {
        "type": "object",
        "properties": {
            "a": {
                "type": "object",
                "properties": {
                    "b": {
                        "type": "object",
                        "properties": {
                            "c": {
                                "type": "object",
                                "properties": {
                                    "d": {"type": "string"}
                                },
                                "required": ["d"]
                            }
                        },
                        "required": ["c"]
                    }
                },
                "required": ["b"]
            }
        },
        "required": ["a"]
    }
    tool = ToolDefinition(name="deep", parameters_json_schema=schema)
    text = json.dumps({"a": {"b": {"c": {"d": "deep_value"}}}})
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"deep": tool}); result = codeflash_output # 6.62μs -> 6.17μs (7.43% faster)

def test_large_nested_object_missing_deep_field():
    # Tool expects a deeply nested object, but missing the deepest required field
    schema = {
        "type": "object",
        "properties": {
            "a": {
                "type": "object",
                "properties": {
                    "b": {
                        "type": "object",
                        "properties": {
                            "c": {
                                "type": "object",
                                "properties": {
                                    "d": {"type": "string"}
                                },
                                "required": ["d"]
                            }
                        },
                        "required": ["c"]
                    }
                },
                "required": ["b"]
            }
        },
        "required": ["a"]
    }
    tool = ToolDefinition(name="deep", parameters_json_schema=schema)
    text = json.dumps({"a": {"b": {"c": {}}}})
    codeflash_output = MistralStreamedResponse._try_get_output_tool_from_text(text, {"deep": tool}); result = codeflash_output # 2.08μs -> 1.88μs (11.1% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from pydantic_ai.models.mistral import MistralStreamedResponse
from pydantic_ai.tools import ToolDefinition
import pytest

def test_MistralStreamedResponse__try_get_output_tool_from_text():
    with pytest.raises(TypeError, match='Expected\\ bytes,\\ bytearray\\ or\\ str'):
        MistralStreamedResponse._try_get_output_tool_from_text('', {'': ToolDefinition('', {}, description=None, outer_typed_dict_key=None, strict=None)})

To edit these changes git checkout codeflash/optimize-MistralStreamedResponse._try_get_output_tool_from_text-mddt3yvi and push.