fix(braintrust): dont decorate generate and clean up unused code

[rguan72]

Description

Decorating the braintrust generator causes an unreasonable number of logs to be sent to braintrust.
To still decorate llm streams, I traced some higher level functions. I deleted unused code I found when looking through calls to stream llm.

How Has This Been Tested?

Tested locally

Additional Options

• [Optional] Override Linear Check

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Preview	Comments	Updated (UTC)
internal-search	Ready	Preview	Comment	Oct 7, 2025 9:31pm

[greptile-apps]

Greptile Overview

Summary

This PR addresses Braintrust tracing issues and removes unused orchestration code from the agent search system. The changes focus on two main areas: selective application of Braintrust `@traced` decorators and cleanup of deprecated agent orchestration nodes.

The Braintrust tracing changes are carefully targeted - the PR removes the @traced decorator from the stream method in the LLM interface (which was causing issues with iterator tracing) while adding it to specific streaming functions where it works properly. In shared_graph_utils/llm.py, the stream_llm_answer function gains tracing, and in dr_a0_clarification.py, a new wrapper function stream_and_process() is created with tracing around LLM streaming operations.

The cleanup portion removes several orchestration nodes that are no longer used in the current agent search architecture: choose_tool.py (354 lines of tool selection logic), call_tool.py (tool execution orchestration), and prepare_tool_input.py (tool input preparation). These files appear to be part of an older orchestration approach that has been superseded.

The PR also includes a minor dependency cleanup in process_message.py where ToolCallException is moved from an external import to a local class definition, reducing coupling between modules. Additionally, the entire query_validation.py file is removed, which was already marked as deprecated with a "No longer used" comment.

These changes align with the codebase's principles of avoiding dead code while improving observability through selective tracing of critical LLM operations.

Important Files Changed

Changed Files

Filename	Score	Overview
backend/onyx/agents/agent_search/shared_graph_utils/llm.py	5/5	Added Braintrust tracing to stream_llm_answer function for observability
backend/onyx/secondary_llm_flows/query_validation.py	4/5	Deleted deprecated query validation file that was marked as no longer used
backend/onyx/llm/interfaces.py	4/5	Removed Braintrust tracing decorator from stream method to fix iterator issues
backend/onyx/agents/agent_search/orchestration/nodes/choose_tool.py	2/5	Deleted entire tool selection orchestration node (354 lines of critical logic)
backend/onyx/agents/agent_search/dr/nodes/dr_a0_clarification.py	4/5	Added Braintrust tracing wrapper around LLM streaming operations
backend/onyx/chat/process_message.py	5/5	Replaced external ToolCallException import with local class definition
backend/onyx/agents/agent_search/orchestration/nodes/prepare_tool_input.py	2/5	Deleted tool input preparation function used in agent orchestration
backend/onyx/agents/agent_search/orchestration/nodes/call_tool.py	3/5	Deleted tool execution orchestration node and ToolCallException class

Confidence score: 3/5

• This PR has significant risks due to deletion of critical orchestration nodes that may still be referenced elsewhere in the codebase
• Score reflects concerns about removing substantial agent functionality without clear evidence it's truly unused
• Pay close attention to the deleted orchestration nodes (choose_tool.py, call_tool.py, prepare_tool_input.py) as they contain core agent functionality

_{8 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[cubic-dev-ai]

No issues found across 8 files

[blacksmith-sh]

3 Jobs Failed:

Run Integration Tests v2 / integration-tests (tests/mcp, tests-mcp) failed on "Wait for service to be ready"

[...]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/requests/sessions.py", line 589, in request
    resp = self.send(prep, **send_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/requests/sessions.py", line 703, in send
    r = adapter.send(request, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/requests/adapters.py", line 677, in send
    raise ConnectionError(e, request=request)
requests.exceptions.ConnectionError: HTTPConnectionPool(host='inference_model_server', port=9000): Max retries exceeded with url: /api/gpu-status (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0xff973236a2d0>: Failed to establish a new connection: [Errno 111] Connection refused'))

ERROR:    Application startup failed. Exiting.
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
Timeout reached. Service did not become ready in 5 minutes.
Error: Process completed with exit code 1.

Run Integration Tests v2 / integration-tests (tests/projects, tests-projects) failed on "Wait for service to be ready"

[...]
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
INFO:     10/07/2025 09:46:29 PM                       utils.py  142: Initialized HuggingFaceTokenizer for: nomic-ai/nomic-embed-text-v1
INFO:     10/07/2025 09:46:30 PM                       setup.py  331: No existing docs or connectors found. Checking GPU availability for multipass indexing.
INFO:     10/07/2025 09:46:30 PM                       setup.py  335: GPU available: False
NOTICE:   10/07/2025 09:46:30 PM                       setup.py  339: Updating multipass indexing setting to: False
INFO:     10/07/2025 09:46:30 PM             search_settings.py  228: Current search settings updated successfully
NOTICE:   10/07/2025 09:46:30 PM                       setup.py  351: Updated settings with GPU availability: False
INFO:     10/07/2025 09:46:30 PM                   load_docs.py  150: Seeding initial documents
INFO:     10/07/2025 09:46:32 PM           indexing_pipeline.py  334: Upserted 18 changed docs out of 18 total docs into the DB
INFO:     10/07/2025 09:46:32 PM                   load_docs.py  230: Indexing seeding documents into Vespa (Vespa may take a few seconds to become ready after receiving the schema)
INFO:     10/07/2025 09:46:32 PM                       utils.py   88: Vespa: Readiness probe starting.
INFO:     10/07/2025 09:46:32 PM                       utils.py  110: Vespa: Readiness probe ongoing. elapsed=0.0 timeout=60.0
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
INFO:     10/07/2025 09:46:37 PM                       utils.py  110: Vespa: Readiness probe ongoing. elapsed=5.1 timeout=60.0
Service not ready yet (HTTP status 000curl_error). Retrying in 5 seconds...
INFO:     10/07/2025 09:46:42 PM                       utils.py  110: Vespa: Readiness probe ongoing. elapsed=10.1 timeout=60.0
Timeout reached. Service did not become ready in 5 minutes.
INFO:     10/07/2025 09:46:47 PM                       utils.py  110: Vespa: Readiness probe ongoing. elapsed=15.1 timeout=60.0
Error: Process completed with exit code 1.

Run Integration Tests v2 / required failed on "Run actions/github-script@v7"

[...]
  PRIVATE_REGISTRY: experimental-registry.blacksmith.sh:5000
  PRIVATE_REGISTRY_USERNAME: ***
  PRIVATE_REGISTRY_PASSWORD: ***
  OPENAI_API_KEY: ***
  SLACK_BOT_TOKEN: ***
  CONFLUENCE_TEST_SPACE_URL: ***
  CONFLUENCE_USER_NAME: ***
  CONFLUENCE_ACCESS_TOKEN: ***
  CONFLUENCE_ACCESS_TOKEN_SCOPED: ***
  JIRA_BASE_URL: ***
  JIRA_USER_EMAIL: ***
  JIRA_API_TOKEN: ***
  JIRA_API_TOKEN_SCOPED: ***
  PERM_SYNC_SHAREPOINT_CLIENT_ID: ***
  PERM_SYNC_SHAREPOINT_PRIVATE_KEY: ***
  PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD: ***
  PERM_SYNC_SHAREPOINT_DIRECTORY_ID: ***
  GITHUB_REPO_NAME: ***/onyx
Error: One or more upstream jobs failed or were cancelled.

---

Summary: 8 successful workflows, 1 failed workflow

• ✅ Python Checks (1 job succeeded)
• ✅ Helm - Lint and Test Charts (1 job succeeded)
• ✅ Connector Tests (1 job cancelled)
• ✅ Python Unit Tests (1 job succeeded)
• ❌ Run Integration Tests v2 (35 jobs succeeded, 3 jobs failed)
• ✅ Check Lazy Imports (1 job succeeded)
• ✅ External Dependency Unit Tests (6 jobs succeeded)
• ✅ Ensure PR references Linear (1 job succeeded)
• ✅ Quality Checks PR (1 job succeeded)

Last updated: 2025-10-08 03:15:05 UTC