[WIP] [#521] Research Agent Tools - #524
Draft
sahilds1 wants to merge 11 commits into
Draft
Conversation
Add ask_database as a second tool for the assistant, alongside the existing semantic search_documents tool. - Reuse the SELECT-only, ALLOWED_TABLES-guarded ask_database implementation from services/tools/database.py rather than reimplementing the query guards in the assistant. - Add get_tools_schema() / make_tool_mapping() as an aggregation seam in tool_services.py so assistant_services.py no longer names individual tools; new tools are registered in one place. - Build the ask_database schema in the flattened Responses-API shape (not the nested Chat Completions shape from services/tools), and defer the database_schema_string import to call time so importing the module never triggers a DB query. - Split tool_services.py: move search_documents into search_tool.py and the agentic-loop helpers (handle_tool_calls_with_reasoning, invoke_functions_from_response) into agentic_loop.py, adding the imports each module needs. - Point importers (assistant_services.py, test_tool_services.py) at the defining module for each symbol instead of re-exporting through tool_services.py. - Document the search/SQL tool overlap risk and the import-time DB access caveat inline.
Replace the two parallel tool registries (get_tools_schema / make_tool_mapping) with a single Tool dataclass and a TOOLS list, so each tool's schema and callable live under one name and can't drift apart. - Add Tool(name, description, parameters, run) with a .schema() method; define SEARCH_TOOL and ASK_DATABASE_TOOL instances and a single TOOLS list as the source of truth. Adding a tool is appending one Tool. - Build the ask_database schema from Medication._meta (concrete_fields / db_table) instead of introspecting the live database, removing the import-time DB query and the deferred database_schema_string import. - Bind the request user at dispatch time: invoke_functions_from_response and handle_tool_calls_with_reasoning now take (tools, user), index tools by name, and call tool.run(user=user, **arguments). This drops make_tool_mapping / make_search_tool_mapping and their closures. - assistant_services builds the schema list with [tool.schema() for tool in TOOLS] and forwards TOOLS + user to the loop. - Update tests to cover the Tool instances and pass (tools, user) to the loop; the tool_services.search_documents / ask_database patch paths still resolve.
Convert every relative import in api/views/assistant to its absolute equivalent (assistant_services, search_tool, tool_services, urls, views). Multi-dot forms like `...services.tools.database` and `..listMeds.models` were error-prone to read and would break silently if a file's package depth changed; the absolute paths are move-safe and unambiguous. Expand the tool_services.py import comments to record why search_documents and ask_database are imported as bare names: the tests patch them at their use site (api.views.assistant.tool_services.<name>), not their definition site, so mock.patch rebinds the reference SEARCH_TOOL.run actually resolves at call time. Note the maintenance guard — qualifying those calls would move the patch target and break the tests. Move the _medication_schema_string helper to sit directly above its only caller, ASK_DATABASE_TOOL, instead of above SEARCH_TOOL. No behavior change: all bound names and patch targets are preserved.
Replace the _medication_schema_string() helper (which read columns from Medication._meta) with a hand-written _MEDICATION_SCHEMA_STRING constant, and drop the now-unused Medication import. The _meta approach auto-synced with the model but dumped every column and pulled in an app-registry dependency (AppRegistryNotReady if imported during app startup). For a 4-column, stable table, a curated constant is simpler, lets us hide columns from the LLM (omit `id`, which it never filters on), and removes the startup coupling — at the cost of a one-line manual update if the table's columns ever change, which the comment calls out. Note: this drops `id` from the schema the model sees (intentional curation, not just a port of the old behavior). Also document in the Tool docstring why behavior is a `run` field (composition) rather than a subclass method: the tools differ only in which function runs, so they are instances of one concept, not distinct types. Add a TODO listing the signals that would justify flipping to Tool(ABC) + per-tool subclasses (per-type state, overriding more than run, or a per-type/abstractmethod-enforced contract).
The eval showed what the assistant said but not how it chose tools. Tool-call info was produced in invoke_functions_from_response and dropped at every return boundary; caught tool exceptions were fed back to the model as strings, so a run where ask_database threw every question read as clean rows. Carry it up as return values (over a mutable out-param: honest domain data that grows via defaulted fields, no call-site churn): - agentic_loop: add ToolCallStatus (OK/FAILED/UNREGISTERED — a bool was both redundant with error and lossy), ToolCall(name, status, arguments, output, error), and AssistantResult(output_text, response_id, tool_calls). invoke_functions_from_response returns (messages, tool_calls); the loop accumulates across iterations and returns AssistantResult. - assistant_services: return AssistantResult (pass-through). Drops the TODO. - views: read result fields; JSON body unchanged. - eval_assistant: run_one times the call locally and adds tools_called, tool_call_count, tool_error_count, tool_calls_json, response_id, duration_s — making tool_error_count > 0 while error is None visible. Deferred: token-cost, turn count, correctness scoring. Tests updated.
Hoist MODEL_NAME into assistant_services, import it in eval_assistant, and drop the duplicated literal — the only functional change here. The rest is comments. Token usage and turn count, the scoring layer, and the INSTRUCTIONS sidecar were each designed in this pass and deliberately not built; the TODOs sit where the work will happen and carry the reasoning.
The eval had never been run. Every test mocks run_assistant, so a green suite proved run_one's row shaping but never that a CSV came out. Three defects stopped it running at all: the sys.path depth, pandas missing from the backend image, and the uv shebang and PEP 723 header, which never described a runnable configuration. A fourth is different in kind, and only the run could surface it. The cold-start race on the embedding model does not stop the eval — it corrupts it, completing normally and writing a CSV that looks clean. Warming the model before the pool addresses it here.
…run-unblocking fixes
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds a second tool to the assistant and makes the eval able to show which tool the
model chose — the thing this branch exists to change.
ask_databasetool — read-only SQL over the medication table, reusing theexisting function in
api/services/tools/database.py(SELECT-only +ALLOWED_TABLESguards not reimplemented). Schema built in the flattened Responses-API shape, which
differs from Chat Completions' nested one.
Tooldataclass + oneTOOLSlist — replaces two parallel registries(
get_tools_schema/make_tool_mapping) linked only by a duplicated, unenforcedtool-name string, where a half-registered tool failed at runtime.
useris now boundat dispatch time, which deletes the binder/closure machinery.
effectively frozen table). Also drops the
AppRegistryNotReadyrisk and lets us hideidfrom the model.run_assistantreturnsAssistantResult(output_text,response_id,tool_calls) instead of a tuple, carrying aToolCallper call.statusis a3-state enum, not a bool: a tool that raised and a tool name the model
hallucinated are opposite diagnoses. Dataclass over widened tuple so later metrics
are defaulted fields with zero call-site churn.
model as a string, so it writes fluent prose around a failed retrieval. Previously a
run where
ask_databasethrew on every question read as five clean rows; now it showsas
tool_error_count > 0whileerror is None.response_id,tools_called,tool_call_count,tool_error_count,tool_calls_json,duration_s.search_tool.pysplit out (withthe imports that split had left broken); one shared
MODEL_NAMEconstant so the CSVcan't mislabel the model.
API unchanged —
views.pyunpacks the dataclass; the JSON body is byte-identical.tool_callsare intentionally not exposed to the frontend.Related Issue
#521
Manual Tests
Automated Tests
19 unit tests, OpenAI + DB + tools all mocked — no live calls, no fixtures.
test_tool_services.pyTool.run/schema()/TOOLS; aToolCallrecorded on all three dispatch branches (OK / FAILED / UNREGISTERED); loop termination,previous_response_idcontinuity, tool calls accumulating across iterationstest_assistant_services.pyprevious_response_id, and thatTOOLS+userare forwarded to the looptest_eval_assistant.pytool_error_count == 1,error is None)Deliberately uncovered:
search_documents' own paths (TODO in-file — mockable, no DBneeded) and
main()'s CSV writing.Documentation
No user-facing or CLAUDE.md changes — internal to the assistant module, HTTP contract
unchanged. Design rationale is documented inline where the code lives (why behavior is a
runfield not a subclass method; why the schema is hardcoded; why tests patch usesites, with a maintenance guard).
Reviewers
Notes
Targets
develop.Before merge
green suite is not evidence a CSV comes out. First run also tests the self-flagged
sys.pathcaveat ineval_assistant.py("../../../../"→/usr/src, not/usr/src/server).idfrom the schema string is a behavior change, not a pure refactor —worth confirming no prompt or eval expected it.
Known, documented rather than fixed
client.responses.createraises mid-loop, everyToolCallso far is lost with theexception (
tool_call_count 0despite real calls). Fixing it means deciding whetherAssistantResultdescribes a successful run or whatever happened — a contractchange.
columns exist to make that observable; sharper descriptions may be needed.
run_onerow literals —as_completedisnondeterministic, so that's the only thing keeping CSV column order stable.
Deferred, design written down
cached_tokens— every turnresends context, so cost from
input_tokensalone overstates spend).INSTRUCTIONS→{branch}-{timestamp}.prompt.txtsidecar beside the CSV.eval_assistant.pyis still a generationharness: no ground truth, so the CSV shows what was said and which tools ran, but not
whether the answer was right. Citation accuracy is the cheapest real signal (the
mandated
[Name, Page]format can be checked against whatsearch_documentsreturned).