Skip to content

feat: implements base optimization logic in the optimization sdk#116

Open
andrewklatzke wants to merge 11 commits intomainfrom
aklatzke/AIC-1990/optimize-method
Open

feat: implements base optimization logic in the optimization sdk#116
andrewklatzke wants to merge 11 commits intomainfrom
aklatzke/AIC-1990/optimize-method

Conversation

@andrewklatzke
Copy link
Copy Markdown
Contributor

@andrewklatzke andrewklatzke commented Mar 31, 2026

Requirements

  • I have added test coverage for new or changed functionality
  • I have followed the repository's pull request submission guidelines
  • I have validated my changes against all supported platform versions

Related issues

This specifically only affects the currently empty and 0.0.0 version of the optimization package.

Pulls in the optimize method from the moonshot branch, updates it to be more production-ready, adds tests, logically splits up code into more manageable chunks.

Describe the solution you've provided

This is the initial implementation of the optimization method that we're pulling into this SDK. Right now this services the same surface area as the moonshot - namely, it implements the optimize_from_options() method while leaving the optimize_from_config() method unimplemented. This also does not handle additional features we'll be adding, such as the ability to compare to ground-truth responses or the ability to post back to LaunchDarkly.

The logs are set to debug level, so enabling them will allow you to trace along with the progress.

Using the manual/passing options implementation of this looks like:

    tool_handlers = {
        "user_preferences_lookup": get_preferences,
    }

    def resolve_tools(
        config: AIAgentConfig,
        provided_tool_handlers: dict[str, Callable[[Any], Any]], 
    ) -> Sequence[Tool]:
        tools: list[Tool] = []
        # ... user implementation detail ...
        return tools

    async def handle_agent_call(
        key: str,
        config: AIAgentConfig,
        context: OptimizationContext,
        provided_tool_handlers: dict[str, Callable[[Any], Any]], # tools we provide, such as the required structured output tool, or the tool to validate a newly created variation generation
    ) -> str:
        model = config.model.get_parameter("name") if config.model else None
        root = Agent(
            name=key,
            instructions=config.instructions,
            handoffs=[],
            tools=resolve_tools(config, provided_tool_handlers),
            model=model,
        )
        response = await Runner.run(root, context.user_input or "")
        return response.final_output

    async def handle_judge_call( # separate method, as this can also be run as a completion by accessing `config.messages` and doesn't need to be distinctly agent. A user may also want to capture or log intermediary data here.
        judge_key: str,
        config: AIAgentConfig,
        context: OptimizationJudgeContext,
        provided_tool_handlers: dict[str, Callable[[Any], Any]],
    ) -> str:
        model = config.model.get_parameter("name") if config.model else None
        root = Agent(
            name=judge_key,
            instructions=config.instructions,
            handoffs=[],
            tools=resolve_tools(config, provided_tool_handlers),
            model=model,
        )
        response = await Runner.run(root, context.user_input or "")
        return response.final_output
        
     # Everything below this is the actual

    options = OptimizationOptions(
        judges={
            "acceptance": OptimizationJudge(
                acceptance_statement="The orchestrator should appropriately fetch the user preferences and route to the correct sub-agent, carrying through any relevant information from the users' query. The orchestrator should not provide any answers itself, just pass to the correct sub-agent. Inability to fetch user preferences or mentions of missing data should be automatic failures. If preferences are not included, that should be an automatic failure.",
                threshold=0.95,
            ),
        },
        context_choices=[
            context_builder("user-123"),
        ],
        max_attempts=5,
        model_choices=["gpt-5", "gpt-5.1", "gpt-5.4", "gpt-5.4-mini"],
        judge_model="gpt-5.4-mini",
        variable_choices=[
            {
                "user_id": "user-123",
                "trip_purpose": "business",
            },
            {
                "user_id": "user-125",
                "trip_purpose": "personal",
            },
        ],
        user_input_options=[
            "I'm going to austin next week, where should I stay?"
        ],
        handle_agent_call=handle_agent_call,
        handle_judge_call=handle_judge_call,
    )

    client = OptimizationClient(ld_ai_client) 
    result = await client.optimize_from_options("travel-agent-orchestrator", options) # distinct step so that optimization options can be re-used

Note

Medium Risk
Medium risk due to a large new async control-flow surface (agent calls, judge calls, variation generation, tool handling, JSON extraction) that can affect correctness and error handling, though it doesn’t change auth/data persistence paths.

Overview
Adds a real OptimizationClient (replacing the placeholder ApiAgentOptimizationClient) that runs an optimization loop: call the agent with per-iteration variable interpolation, evaluate the response via one or more judges (LD judge_config-backed or inline acceptance statements), and, on failure, ask an LLM to return a structured improved configuration for the next attempt.

Introduces new public dataclasses (OptimizationOptions, OptimizationContext, OptimizationJudge, AIJudgeCallConfig, ToolDefinition, etc.) plus utility helpers for structured-output tools and robust JSON extraction, and wires status/result callbacks throughout the loop.

Adds extensive test coverage for tool handling, judge evaluation paths, variation prompt generation, and end-to-end success/failure behavior; updates package exports and smoke tests to reflect the new API.

Written by Cursor Bugbot for commit e2ff561. This will update automatically on new commits. Configure here.

@andrewklatzke andrewklatzke requested a review from a team as a code owner March 31, 2026 01:22
@andrewklatzke andrewklatzke changed the title Aklatzke/aic 1990/optimize method feat: implements base optimization logic in the optimization sdk Mar 31, 2026
start_idx = response_str.find('{', start_idx + 1)
if start_idx == -1:
break
brace_count = 0
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Balanced-brace scanner retry uses stale start index

Medium Severity

The balanced-brace scanning fallback in extract_json_from_response has broken retry logic. When a balanced {…} block fails json.loads, start_idx is updated to the next { after the original start (which is before i), and brace_count is reset to 0, but the for loop continues from i + 1 — already past the new start_idx. This means the scanner never re-processes characters from the updated start position. On the next time brace_count hits 0, response_str[start_idx:i + 1] spans from the stale inner start_idx to the new closing brace, producing a garbage substring that won't parse. The retry path effectively never works.

Fix in Cursor Fix in Web

},
"required": ["passed", "rationale"],
},
)
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unused create_boolean_tool function is dead code

Low Severity

create_boolean_tool is defined but never called, imported, or exported anywhere in the codebase. It appears to be leftover scaffolding that was superseded by create_evaluation_tool, which is the tool actually used for judge evaluations.

Fix in Cursor Fix in Web

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Placeholder for now; we want to allow both boolean/range evals

if self.judges is None and self.on_turn is None:
raise ValueError("Either judges or on_turn must be provided")
if self.judge_model is None:
raise ValueError("judge_model must be provided")
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing validation for empty variable_choices list

Medium Severity

__post_init__ validates that context_choices and model_choices each have at least one element, but no equivalent check exists for variable_choices. An empty list passes validation but causes random.choice() to raise an IndexError at runtime in _run_optimization and _create_optimization_context.

Fix in Cursor Fix in Web

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should not be validated; it's valid for there to be no variables

)
return judge_results

async def _evaluate_config_judge(
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason we are not using the built in create_judge method from the ai sdk?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes an underlying call to:

    def _judge_config(
        self,
        judge_key: str,
        context: Context,
        default: AIJudgeConfigDefault,
        variables: Dict[str, Any],
    ) -> AIJudgeConfig:
        """
        Fetch a judge configuration from the LaunchDarkly client.

        Thin wrapper around LDAIClient.judge_config so callers do not need a
        direct reference to the client.

        :param judge_key: The key for the judge configuration in LaunchDarkly
        :param context: The evaluation context
        :param default: Fallback config when the flag is disabled or unreachable
        :param variables: Template variables for instruction interpolation
        :return: The resolved AIJudgeConfig
        """
        return self._ldClient.judge_config(judge_key, context, default, variables)

Where we get the judge config directly. This method is doing a lot of setup/manipulation of the messages. We don't want the auto-evaluating judges as they defer down to the user to execute. Figured using just the config directly here made more sense as we don't want any of the "auto" functionality

Copy link
Copy Markdown

@cursor cursor bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 4 total unresolved issues (including 3 from previous reviews).

Fix All in Cursor

Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants