chore(dr): finer grained tracing for clarification step, research plan step, and orchestration step

[rohoswagger]

Description

add finer grainer tracing for 3 diff planning steps for deep research

How Has This Been Tested?

locally

Additional Options

• [Optional] Override Linear Check

[cubic-dev-ai]

No issues found across 1 file

[greptile-apps]

Greptile Overview

Greptile Summary

This PR adds finer-grained tracing instrumentation to the deep research loop by wrapping three major steps in function_span context managers:

1. Clarification Step (lines 223-272): Wraps the optional clarification question logic that determines if the LLM needs more information from the user before proceeding.

2. Research Plan Step (lines 277-352): Wraps the research plan generation logic where the LLM creates a plan for the research tasks. This span properly sets the span_data.output to capture the generated research plan.

3. Research Execution Step (lines 357-670): Wraps the main orchestration loop that executes the research plan through multiple cycles, handling tool calls for thinking, research agents, and final report generation.

Changes Overview

The changes are purely additive - wrapping existing code blocks with tracing spans. The core logic remains unchanged. All control flow statements (return, break, continue) work correctly within the span contexts due to Python's context manager protocol properly handling cleanup.

Improvements for Observability

While the code is functionally correct, the tracing could be more informative by:

• Setting span.span_data.input for the clarification and execution steps to provide context (similar to existing patterns in generate_report at line 109)
• Setting span.span_data.output in the clarification step when it returns early

These additions would align with the tracing patterns already established in the codebase and make debugging/monitoring easier.

Confidence Score: 4/5

• This PR is safe to merge with minimal risk - it adds tracing instrumentation without changing core logic.
• The changes are purely additive tracing instrumentation that wraps existing code blocks in function_span context managers. All control flow (return, break, continue statements) works correctly with Python's context manager protocol. The code is functionally identical to before, just with better observability. The only improvements suggested are style enhancements to make trace data more informative by setting span input/output fields consistently with existing patterns in the codebase.
• No files require special attention - the single changed file has been thoroughly reviewed and contains only non-breaking tracing additions.

Important Files Changed

File Analysis

Filename	Score	Overview
backend/onyx/deep_research/dr_loop.py	4/5	Adds function span tracing around clarification, research plan, and research execution steps. Code is functionally correct but could benefit from setting span input/output data for better observability.

Sequence Diagram

sequenceDiagram
    participant Client
    participant DrLoop as run_deep_research_llm_loop
    participant ClarificationSpan as clarification_step
    participant PlanSpan as research_plan_step
    participant ExecutionSpan as research_execution_step
    participant LLM
    
    Client->>DrLoop: Start deep research
    
    alt skip_clarification = false
        DrLoop->>ClarificationSpan: Enter span
        ClarificationSpan->>LLM: run_llm_step with clarification tools
        alt No tool calls (clarification needed)
            LLM-->>ClarificationSpan: No tool calls
            ClarificationSpan->>Client: Return (wait for user input)
        else Tool calls present
            LLM-->>ClarificationSpan: Tool calls
            ClarificationSpan->>DrLoop: Continue
        end
    end
    
    DrLoop->>PlanSpan: Enter span
    PlanSpan->>LLM: run_llm_step_pkt_generator for plan
    loop Stream packets
        LLM-->>PlanSpan: AgentResponseStart/Delta
        PlanSpan->>Client: DeepResearchPlanStart/Delta
    end
    LLM-->>PlanSpan: research_plan
    PlanSpan->>DrLoop: Set span output
    
    DrLoop->>ExecutionSpan: Enter span
    loop orchestrator cycles (up to 8 or 4 for reasoning models)
        ExecutionSpan->>LLM: run_llm_step with orchestrator tools
        alt generate_report tool called
            LLM-->>ExecutionSpan: generate_report
            ExecutionSpan->>LLM: generate_final_report
            LLM-->>ExecutionSpan: final report
            ExecutionSpan->>DrLoop: Exit loop
        else think_tool called
            LLM-->>ExecutionSpan: think_tool
            ExecutionSpan->>ExecutionSpan: Save reasoning, continue
        else research_agent tools called
            LLM-->>ExecutionSpan: research_agent calls
            ExecutionSpan->>LLM: run_research_agent_calls
            LLM-->>ExecutionSpan: intermediate reports
            ExecutionSpan->>ExecutionSpan: Update history, continue
        end
    end
    
    DrLoop->>Client: Emit OverallStop

Loading

[greptile-apps]

_{3 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Consider adding input data to the clarification_step span for better traceability. Following the pattern used in other spans in the codebase (e.g., generate_report at line 109), you could add:

Suggested change

	with function_span("clarification_step") as span:
	clarification_prompt = CLARIFICATION_PROMPT.format(
	with function_span("clarification_step") as span:
	span.span_data.input = f"skip_clarification={skip_clarification}, history_length={len(simple_chat_history)}"
	clarification_prompt = CLARIFICATION_PROMPT.format(

This would make the trace data more informative when debugging or monitoring deep research operations.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

Prompt To Fix With AI

This is a comment left during a code review.
Path: backend/onyx/deep_research/dr_loop.py
Line: 223:224

Comment:
Consider adding input data to the clarification_step span for better traceability. Following the pattern used in other spans in the codebase (e.g., `generate_report` at line 109), you could add:

```suggestion
            with function_span("clarification_step") as span:
                span.span_data.input = f"skip_clarification={skip_clarification}, history_length={len(simple_chat_history)}"
                clarification_prompt = CLARIFICATION_PROMPT.format(
```

This would make the trace data more informative when debugging or monitoring deep research operations.

<sub>Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!</sub>

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

Consider adding input data to the research_execution_step span for better traceability. Following the pattern used in other spans:

Suggested change

	with function_span("research_execution_step") as span:
	is_reasoning_model = model_is_reasoning_model(
	with function_span("research_execution_step") as span:
	span.span_data.input = f"research_plan_length={len(research_plan) if research_plan else 0}, orchestrator_start_turn_index={orchestrator_start_turn_index}"
	is_reasoning_model = model_is_reasoning_model(

This would make it easier to understand the context when analyzing traces, especially when debugging issues in the research execution phase.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

Prompt To Fix With AI

This is a comment left during a code review.
Path: backend/onyx/deep_research/dr_loop.py
Line: 357:358

Comment:
Consider adding input data to the research_execution_step span for better traceability. Following the pattern used in other spans:

```suggestion
        with function_span("research_execution_step") as span:
            span.span_data.input = f"research_plan_length={len(research_plan) if research_plan else 0}, orchestrator_start_turn_index={orchestrator_start_turn_index}"
            is_reasoning_model = model_is_reasoning_model(
```

This would make it easier to understand the context when analyzing traces, especially when debugging issues in the research execution phase.

<sub>Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!</sub>

How can I resolve this? If you propose a fix, please make it concise.