From b3669b082dd7dc3c682b392bedcb3fd6eb701930 Mon Sep 17 00:00:00 2001 From: Sahil D Shah Date: Wed, 15 Jul 2026 18:49:33 -0400 Subject: [PATCH 1/6] Add TODOs for adding tools --- server/api/views/assistant/tool_services.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/server/api/views/assistant/tool_services.py b/server/api/views/assistant/tool_services.py index 0fb96cef..d8c6647a 100644 --- a/server/api/views/assistant/tool_services.py +++ b/server/api/views/assistant/tool_services.py @@ -46,6 +46,11 @@ # that combine all tool schemas and mappings so assistant_services.py never needs # to change when a new tool is added — only tool_services.py does. +#TODO: Add existing tools from server/api/views/conversations/views.py + +#TODO: Add keyword tool if it doesn't overlap with semantic search because too +#many overlapping tools can return conflicting answers + def make_search_tool_mapping(user) -> dict[str, Callable]: # make_search_tool_mapping binds user to search_documents at call time. # user is a request-time value the model cannot generate, so it must be From b5710e966f7de2de48dbc8489fa7f1e4feed1fc5 Mon Sep 17 00:00:00 2001 From: Sahil D Shah Date: Thu, 16 Jul 2026 15:10:39 -0400 Subject: [PATCH 2/6] Add ask_database tool and split tool module 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. --- server/api/views/assistant/agentic_loop.py | 106 ++++++++ .../api/views/assistant/assistant_services.py | 22 +- server/api/views/assistant/search_tool.py | 50 ++++ .../api/views/assistant/test_tool_services.py | 4 +- server/api/views/assistant/tool_services.py | 254 +++++++----------- 5 files changed, 266 insertions(+), 170 deletions(-) create mode 100644 server/api/views/assistant/agentic_loop.py create mode 100644 server/api/views/assistant/search_tool.py diff --git a/server/api/views/assistant/agentic_loop.py b/server/api/views/assistant/agentic_loop.py new file mode 100644 index 00000000..d6af0d75 --- /dev/null +++ b/server/api/views/assistant/agentic_loop.py @@ -0,0 +1,106 @@ +import json +import logging +from typing import Callable + +logger = logging.getLogger(__name__) + + +def handle_tool_calls_with_reasoning( + response, client, model_defaults: dict, tool_mapping: dict[str, Callable] +) -> tuple[str, str]: + """Run the agentic loop until the model stops emitting function calls. + + Parameters + ---------- + response : OpenAI Response + The initial response from the model. + client : OpenAI + The OpenAI client instance. + model_defaults : dict + Keyword arguments forwarded to every client.responses.create call. + tool_mapping : dict[str, Callable] + Maps function names to their implementations. + + Returns + ------- + tuple[str, str] + (final_response_output_text, final_response_id) + """ + # Open AI Cookbook: Handling Function Calls with Reasoning Models + # https://cookbook.openai.com/examples/reasoning_function_calls + while True: + # Mapping of the tool names we tell the model about and the functions that implement them + function_responses = invoke_functions_from_response(response, tool_mapping) + if len(function_responses) == 0: # We're done reasoning + logger.info("Reasoning completed") + final_response_output_text = response.output_text + final_response_id = response.id + logger.info(f"Final response: {final_response_output_text}") + return final_response_output_text, final_response_id + else: + logger.info("More reasoning required, continuing...") + response = client.responses.create( + input=function_responses, + previous_response_id=response.id, + **model_defaults, + ) + + +def invoke_functions_from_response( + response, tool_mapping: dict[str, Callable] +) -> list[dict]: + """Extract all function calls from the response, look up the corresponding tool function(s) and execute them. + (This would be a good place to handle asynchroneous tool calls, or ones that take a while to execute.) + This returns a list of messages to be added to the conversation history. + + Parameters + ---------- + response : OpenAI Response + The response object from OpenAI containing output items that may include function calls + tool_mapping : dict[str, Callable] + A dictionary mapping function names (as strings) to their corresponding Python functions. + Keys should match the function names defined in the tools schema. + + Returns + ------- + list[dict] + List of function call output messages formatted for the OpenAI conversation. + Each message contains: + - type: "function_call_output" + - call_id: The unique identifier for the function call + - output: The result returned by the executed function (string or error message) + """ + + # Open AI Cookbook: Handling Function Calls with Reasoning Models + # https://cookbook.openai.com/examples/reasoning_function_calls + + intermediate_messages = [] + for response_item in response.output: + if response_item.type == "function_call": + target_tool = tool_mapping.get(response_item.name) + if target_tool: + try: + arguments = json.loads(response_item.arguments) + logger.info( + f"Invoking tool: {response_item.name} with arguments: {arguments}" + ) + tool_output = target_tool(**arguments) + logger.info(f"Tool {response_item.name} completed successfully") + except Exception as e: + msg = f"Error executing function call: {response_item.name}: {e}" + tool_output = msg + logger.error(msg, exc_info=True) + else: + msg = f"ERROR - No tool registered for function call: {response_item.name}" + tool_output = msg + logger.error(msg) + intermediate_messages.append( + { + "type": "function_call_output", + "call_id": response_item.call_id, + "output": tool_output, + } + ) + elif response_item.type == "reasoning": + logger.info(f"Reasoning step: {response_item.summary}") + return intermediate_messages \ No newline at end of file diff --git a/server/api/views/assistant/assistant_services.py b/server/api/views/assistant/assistant_services.py index ac339b9f..f730d901 100644 --- a/server/api/views/assistant/assistant_services.py +++ b/server/api/views/assistant/assistant_services.py @@ -4,11 +4,8 @@ from openai import OpenAI from .assistant_prompts import INSTRUCTIONS -from .tool_services import ( - SEARCH_TOOLS_SCHEMA, - make_search_tool_mapping, - handle_tool_calls_with_reasoning, -) +from .tool_services import get_tools_schema, make_tool_mapping +from .agentic_loop import handle_tool_calls_with_reasoning logger = logging.getLogger(__name__) @@ -44,14 +41,17 @@ def run_assistant( "model": "gpt-5-nano", # 400,000 token context window # A summary of the reasoning performed by the model. This can be useful for debugging and understanding the model's reasoning process. "reasoning": {"effort": "low", "summary": None}, - "tools": SEARCH_TOOLS_SCHEMA, + # get_tools_schema() / make_tool_mapping() are aggregators owned by + # tool_services.py — we deliberately call them here instead of naming individual + # tools, so adding a tool never requires editing this file. + "tools": get_tools_schema(), } - # TOOLS_SCHEMA tells the model what tools exist and what arguments to generate. - # tool_mapping wires those tool names to the Python functions that execute them. - # They are separate because the model generates arguments (schema concern) but - # cannot supply request-time values like user (mapping concern). - tool_mapping = make_search_tool_mapping(user) + # The schema (get_tools_schema) tells the model what tools exist and what arguments + # to generate; the tool_mapping wires those tool names to the Python functions that + # execute them. They are kept separate because the model generates arguments (schema + # concern) but cannot supply request-time values like `user` (mapping concern). + tool_mapping = make_tool_mapping(user) if not previous_response_id: response = client.responses.create( diff --git a/server/api/views/assistant/search_tool.py b/server/api/views/assistant/search_tool.py new file mode 100644 index 00000000..8724930e --- /dev/null +++ b/server/api/views/assistant/search_tool.py @@ -0,0 +1,50 @@ +from ...services.embedding_services import get_closest_embeddings +from ...services.conversions_services import convert_uuids + + +def search_documents(query: str, user) -> str: + """ + Search through user's uploaded documents using semantic similarity. + + This function performs vector similarity search against the user's document corpus + and returns formatted results with context information for the LLM to use. + + Parameters + ---------- + query : str + The search query string + user : User + The authenticated user whose documents to search + + Returns + ------- + str + Formatted search results containing document excerpts with metadata + + Raises + ------ + Exception + If embedding search fails + """ + + try: + embeddings_results = get_closest_embeddings( + user=user, message_data=query.strip() + ) + embeddings_results = convert_uuids(embeddings_results) + + if not embeddings_results: + return "No relevant documents found for your query. Please try different search terms or upload documents first." + + # Format results with clear structure and metadata + prompt_texts = [ + f"[Document {i + 1} - File: {obj['file_id']}, Name: {obj['name']}, Page: {obj['page_number']}, Chunk: {obj['chunk_number']}, Similarity: {1 - obj['distance']:.3f}]\n{obj['text']}\n[End Document {i + 1}]" + for i, obj in enumerate(embeddings_results) + ] + + return "\n\n".join(prompt_texts) + + except Exception as e: + return f"Error searching documents: {str(e)}. Please try again if the issue persists." + + \ No newline at end of file diff --git a/server/api/views/assistant/test_tool_services.py b/server/api/views/assistant/test_tool_services.py index 86e57eed..83524fa6 100644 --- a/server/api/views/assistant/test_tool_services.py +++ b/server/api/views/assistant/test_tool_services.py @@ -19,11 +19,11 @@ # mocking those two (like the rest of the suite mocks collaborators) covers all # three paths as fast, DB-free unit tests. -from api.views.assistant.tool_services import ( +from api.views.assistant.agentic_loop import ( invoke_functions_from_response, handle_tool_calls_with_reasoning, - make_search_tool_mapping, ) +from api.views.assistant.tool_services import make_search_tool_mapping # --------------------------------------------------------------------------- diff --git a/server/api/views/assistant/tool_services.py b/server/api/views/assistant/tool_services.py index d8c6647a..2159424c 100644 --- a/server/api/views/assistant/tool_services.py +++ b/server/api/views/assistant/tool_services.py @@ -1,20 +1,67 @@ -import json -import logging from typing import Callable -from ...services.embedding_services import get_closest_embeddings -from ...services.conversions_services import convert_uuids +# make_search_tool_mapping (below) calls search_documents, which now lives in its own +# module after the refactor — import it directly from there. +from .search_tool import search_documents +# Reuse the existing ask_database implementation from services/tools rather than +# reimplementing it here — it already enforces the SELECT-only and ALLOWED_TABLES +# guards. It takes no request-time secret (no `user`), so it is safe to import at +# module top level; database.py does no DB work at import time. +# +# IMPORT-TIME DB ACCESS: we import the ask_database *function* here (cheap, no DB), +# NOT database_schema_string from services/tools/tools.py. That module runs +# get_database_info(connection) at *its* import time to build the schema string, so +# importing it at module top level would fire a DB query just to import this file +# (breaking manage.py commands / any context where the DB isn't ready). +# _build_ask_database_schema() therefore defers that import to call time instead. +from ...services.tools.database import ask_database + +# TODO: Add keyword tool if it doesn't overlap with semantic search +# because too many overlapping tools can return conflicting answers + +# get_tools_schema / make_tool_mapping are the aggregation seam: assistant_services.py +# reads these two functions instead of naming individual tools. That way a new tool is +# added by editing ONLY this file (append its schema below, register its callable in the +# mapping) — assistant_services.py never has to change again. +def get_tools_schema() -> list[dict]: + """Return every tool schema the assistant exposes to the model. + + assistant_services.py reads this instead of naming individual schemas, so adding a + tool means appending here — not editing assistant_services.py. + """ + # OVERLAP RISK: this exposes a semantic document-search tool AND a SQL + # medication-lookup tool at the same time. For a question both could plausibly + # answer, the model chooses which to call, and they can return conflicting + # answers. If that becomes a problem, sharpen each tool's description to carve + # out when to prefer which (unstructured docs/citations vs. structured + # medication facts) rather than adding yet more overlapping tools — see the + # keyword-tool TODO at the top of this module. + return SEARCH_TOOLS_SCHEMA + [_build_ask_database_schema()] + + +def make_tool_mapping(user) -> dict[str, Callable]: + """Return the full name->callable mapping for every tool. + + ask_database needs no request-time binding (it queries the shared medication table, + not per-user data), so it is registered directly. search_documents is bound to the + user via make_search_tool_mapping. Reuses ask_database from services/tools. + """ + return { + **make_search_tool_mapping(user), + "ask_database": ask_database, + } -logger = logging.getLogger(__name__) -TOOL_DESCRIPTION = """ + + +SEARCH_TOOL_DESCRIPTION = """ Search the user's uploaded documents for information relevant to answering their question. Call this function when you need to find specific information from the user's documents to provide an accurate, citation-backed response. Always search before answering questions about document content. """ -TOOL_PROPERTY_DESCRIPTION = """ +SEARCH_TOOL_PROPERTY_DESCRIPTION = """ A specific search query to find relevant information in the user's documents. Use keywords, phrases, or questions related to what the user is asking about. Be specific rather than generic - use terms that would appear in the relevant documents. @@ -27,13 +74,13 @@ { "type": "function", "name": "search_documents", - "description": TOOL_DESCRIPTION, + "description": SEARCH_TOOL_DESCRIPTION, "parameters": { "type": "object", "properties": { "query": { "type": "string", - "description": TOOL_PROPERTY_DESCRIPTION, + "description": SEARCH_TOOL_PROPERTY_DESCRIPTION, } }, "required": ["query"], @@ -42,15 +89,6 @@ ] -# TODO: Add get_tools_schema() and make_tool_mapping(user) aggregation functions -# that combine all tool schemas and mappings so assistant_services.py never needs -# to change when a new tool is added — only tool_services.py does. - -#TODO: Add existing tools from server/api/views/conversations/views.py - -#TODO: Add keyword tool if it doesn't overlap with semantic search because too -#many overlapping tools can return conflicting answers - def make_search_tool_mapping(user) -> dict[str, Callable]: # make_search_tool_mapping binds user to search_documents at call time. # user is a request-time value the model cannot generate, so it must be @@ -73,147 +111,49 @@ def bound_search(query: str) -> str: return {"search_documents": bound_search} -def search_documents(query: str, user) -> str: - """ - Search through user's uploaded documents using semantic similarity. - - This function performs vector similarity search against the user's document corpus - and returns formatted results with context information for the LLM to use. - - Parameters - ---------- - query : str - The search query string - user : User - The authenticated user whose documents to search - - Returns - ------- - str - Formatted search results containing document excerpts with metadata - - Raises - ------ - Exception - If embedding search fails - """ - - try: - embeddings_results = get_closest_embeddings( - user=user, message_data=query.strip() - ) - embeddings_results = convert_uuids(embeddings_results) - - if not embeddings_results: - return "No relevant documents found for your query. Please try different search terms or upload documents first." - - # Format results with clear structure and metadata - prompt_texts = [ - f"[Document {i + 1} - File: {obj['file_id']}, Name: {obj['name']}, Page: {obj['page_number']}, Chunk: {obj['chunk_number']}, Similarity: {1 - obj['distance']:.3f}]\n{obj['text']}\n[End Document {i + 1}]" - for i, obj in enumerate(embeddings_results) - ] - - return "\n\n".join(prompt_texts) - - except Exception as e: - return f"Error searching documents: {str(e)}. Please try again if the issue persists." - +ASK_DATABASE_TOOL_DESCRIPTION = """ +Use this tool to answer questions about the medications in the Balancer database. +Medications are stored by their official generic names, not brand names, so convert +brand names to generic names first and match case-insensitively +(e.g. LOWER(name) = LOWER('lurasidone')). The input must be a single, fully-formed +SQL SELECT query. +""" -def invoke_functions_from_response( - response, tool_mapping: dict[str, Callable] -) -> list[dict]: - """Extract all function calls from the response, look up the corresponding tool function(s) and execute them. - (This would be a good place to handle asynchroneous tool calls, or ones that take a while to execute.) - This returns a list of messages to be added to the conversation history. - Parameters - ---------- - response : OpenAI Response - The response object from OpenAI containing output items that may include function calls - tool_mapping : dict[str, Callable] - A dictionary mapping function names (as strings) to their corresponding Python functions. - Keys should match the function names defined in the tools schema. +def _build_ask_database_schema() -> dict: + """Build the ask_database tool schema in the flattened Responses-API shape. - Returns - ------- - list[dict] - List of function call output messages formatted for the OpenAI conversation. - Each message contains: - - type: "function_call_output" - - call_id: The unique identifier for the function call - - output: The result returned by the executed function (string or error message) + The table/column names the model may query are injected from the live database + schema (reused from services/tools). The import is deferred to call time so that + merely importing this module never triggers a database query. """ - - # Open AI Cookbook: Handling Function Calls with Reasoning Models - # https://cookbook.openai.com/examples/reasoning_function_calls - - intermediate_messages = [] - for response_item in response.output: - if response_item.type == "function_call": - target_tool = tool_mapping.get(response_item.name) - if target_tool: - try: - arguments = json.loads(response_item.arguments) - logger.info( - f"Invoking tool: {response_item.name} with arguments: {arguments}" - ) - tool_output = target_tool(**arguments) - logger.info(f"Tool {response_item.name} completed successfully") - except Exception as e: - msg = f"Error executing function call: {response_item.name}: {e}" - tool_output = msg - logger.error(msg, exc_info=True) - else: - msg = f"ERROR - No tool registered for function call: {response_item.name}" - tool_output = msg - logger.error(msg) - intermediate_messages.append( - { - "type": "function_call_output", - "call_id": response_item.call_id, - "output": tool_output, + # Deferred (function-local) import on purpose: services/tools/tools.py builds + # database_schema_string by querying the DB at *its* import time. Importing it + # here at call time keeps that query out of this module's import, so loading the + # assistant (e.g. during manage.py commands or when the DB isn't ready) never + # triggers it. We reuse the prebuilt string instead of rebuilding the schema. + from ...services.tools.tools import database_schema_string + + # Flattened Responses-API shape: name/description/parameters live at the top level. + # This is intentionally NOT the nested {"function": {...}} shape that + # services/tools/tools.py's create_tool_dict produces for the Chat Completions API — + # which is also why create_tool_dict could not be reused here. + return { + "type": "function", + "name": "ask_database", + "description": ASK_DATABASE_TOOL_DESCRIPTION, + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": ( + "A plain-text SQL SELECT query answering the user's question, " + "written against this schema:\n" + f"{database_schema_string}" + ), } - ) - elif response_item.type == "reasoning": - logger.info(f"Reasoning step: {response_item.summary}") - return intermediate_messages - -def handle_tool_calls_with_reasoning( - response, client, model_defaults: dict, tool_mapping: dict[str, Callable] -) -> tuple[str, str]: - """Run the agentic loop until the model stops emitting function calls. - - Parameters - ---------- - response : OpenAI Response - The initial response from the model. - client : OpenAI - The OpenAI client instance. - model_defaults : dict - Keyword arguments forwarded to every client.responses.create call. - tool_mapping : dict[str, Callable] - Maps function names to their implementations. - - Returns - ------- - tuple[str, str] - (final_response_output_text, final_response_id) - """ - # Open AI Cookbook: Handling Function Calls with Reasoning Models - # https://cookbook.openai.com/examples/reasoning_function_calls - while True: - # Mapping of the tool names we tell the model about and the functions that implement them - function_responses = invoke_functions_from_response(response, tool_mapping) - if len(function_responses) == 0: # We're done reasoning - logger.info("Reasoning completed") - final_response_output_text = response.output_text - final_response_id = response.id - logger.info(f"Final response: {final_response_output_text}") - return final_response_output_text, final_response_id - else: - logger.info("More reasoning required, continuing...") - response = client.responses.create( - input=function_responses, - previous_response_id=response.id, - **model_defaults, - ) + }, + "required": ["query"], + }, + } From aba10f81c786a2b1021a1791802b408b6a9042ca Mon Sep 17 00:00:00 2001 From: Sahil D Shah Date: Fri, 17 Jul 2026 16:41:00 -0400 Subject: [PATCH 3/6] Refactor model tools as a Tool dataclass, drop the aggregators 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. --- server/api/views/assistant/agentic_loop.py | 38 +-- .../api/views/assistant/assistant_services.py | 20 +- .../assistant/test_assistant_services.py | 23 +- .../api/views/assistant/test_tool_services.py | 119 +++++----- server/api/views/assistant/tool_services.py | 219 ++++++++---------- 5 files changed, 194 insertions(+), 225 deletions(-) diff --git a/server/api/views/assistant/agentic_loop.py b/server/api/views/assistant/agentic_loop.py index d6af0d75..53b4601d 100644 --- a/server/api/views/assistant/agentic_loop.py +++ b/server/api/views/assistant/agentic_loop.py @@ -1,12 +1,11 @@ import json import logging -from typing import Callable logger = logging.getLogger(__name__) def handle_tool_calls_with_reasoning( - response, client, model_defaults: dict, tool_mapping: dict[str, Callable] + response, client, model_defaults: dict, tools: list, user ) -> tuple[str, str]: """Run the agentic loop until the model stops emitting function calls. @@ -18,8 +17,10 @@ def handle_tool_calls_with_reasoning( The OpenAI client instance. model_defaults : dict Keyword arguments forwarded to every client.responses.create call. - tool_mapping : dict[str, Callable] - Maps function names to their implementations. + tools : list[Tool] + The available tools; each has a `.name` and a `.run(user, **arguments)`. + user : User + The request user, bound into each tool call at dispatch time. Returns ------- @@ -29,8 +30,7 @@ def handle_tool_calls_with_reasoning( # Open AI Cookbook: Handling Function Calls with Reasoning Models # https://cookbook.openai.com/examples/reasoning_function_calls while True: - # Mapping of the tool names we tell the model about and the functions that implement them - function_responses = invoke_functions_from_response(response, tool_mapping) + function_responses = invoke_functions_from_response(response, tools, user) if len(function_responses) == 0: # We're done reasoning logger.info("Reasoning completed") final_response_output_text = response.output_text @@ -47,9 +47,9 @@ def handle_tool_calls_with_reasoning( def invoke_functions_from_response( - response, tool_mapping: dict[str, Callable] + response, tools: list, user ) -> list[dict]: - """Extract all function calls from the response, look up the corresponding tool function(s) and execute them. + """Extract all function calls from the response, look up the corresponding tool(s) and execute them. (This would be a good place to handle asynchroneous tool calls, or ones that take a while to execute.) This returns a list of messages to be added to the conversation history. @@ -57,9 +57,12 @@ def invoke_functions_from_response( ---------- response : OpenAI Response The response object from OpenAI containing output items that may include function calls - tool_mapping : dict[str, Callable] - A dictionary mapping function names (as strings) to their corresponding Python functions. - Keys should match the function names defined in the tools schema. + tools : list[Tool] + The available tools. Indexed by `.name` here to dispatch the model's calls; each + is invoked as `tool.run(user, **arguments)`. + user : User + The request user, forwarded to every tool call so tools that need it (e.g. + document access control) get it, and tools that don't simply ignore it. Returns ------- @@ -70,21 +73,24 @@ def invoke_functions_from_response( - call_id: The unique identifier for the function call - output: The result returned by the executed function (string or error message) """ - + # Open AI Cookbook: Handling Function Calls with Reasoning Models # https://cookbook.openai.com/examples/reasoning_function_calls - + + # Index the tools by name so a model-supplied call name can be looked up. .get() + # returns None for an unknown name, handled explicitly below. + tools_by_name = {tool.name: tool for tool in tools} intermediate_messages = [] for response_item in response.output: if response_item.type == "function_call": - target_tool = tool_mapping.get(response_item.name) - if target_tool: + target_tool = tools_by_name.get(response_item.name) + if target_tool is not None: try: arguments = json.loads(response_item.arguments) logger.info( f"Invoking tool: {response_item.name} with arguments: {arguments}" ) - tool_output = target_tool(**arguments) + tool_output = target_tool.run(user=user, **arguments) logger.info(f"Tool {response_item.name} completed successfully") except Exception as e: msg = f"Error executing function call: {response_item.name}: {e}" diff --git a/server/api/views/assistant/assistant_services.py b/server/api/views/assistant/assistant_services.py index f730d901..55a5b19a 100644 --- a/server/api/views/assistant/assistant_services.py +++ b/server/api/views/assistant/assistant_services.py @@ -4,7 +4,7 @@ from openai import OpenAI from .assistant_prompts import INSTRUCTIONS -from .tool_services import get_tools_schema, make_tool_mapping +from .tool_services import TOOLS from .agentic_loop import handle_tool_calls_with_reasoning logger = logging.getLogger(__name__) @@ -41,18 +41,12 @@ def run_assistant( "model": "gpt-5-nano", # 400,000 token context window # A summary of the reasoning performed by the model. This can be useful for debugging and understanding the model's reasoning process. "reasoning": {"effort": "low", "summary": None}, - # get_tools_schema() / make_tool_mapping() are aggregators owned by - # tool_services.py — we deliberately call them here instead of naming individual - # tools, so adding a tool never requires editing this file. - "tools": get_tools_schema(), + # The model only ever sees each tool's schema (name/description/parameters), + # derived from the single TOOLS list in tool_services.py. The request `user` is + # not part of the schema — it is bound into each call later, at dispatch time. + "tools": [tool.schema() for tool in TOOLS], } - # The schema (get_tools_schema) tells the model what tools exist and what arguments - # to generate; the tool_mapping wires those tool names to the Python functions that - # execute them. They are kept separate because the model generates arguments (schema - # concern) but cannot supply request-time values like `user` (mapping concern). - tool_mapping = make_tool_mapping(user) - if not previous_response_id: response = client.responses.create( input=[ @@ -69,4 +63,6 @@ def run_assistant( **MODEL_DEFAULTS, ) - return handle_tool_calls_with_reasoning(response, client, MODEL_DEFAULTS, tool_mapping) + # Pass TOOLS and user through to the loop, which indexes tools by name and binds + # user into each tool call at dispatch time. + return handle_tool_calls_with_reasoning(response, client, MODEL_DEFAULTS, TOOLS, user) diff --git a/server/api/views/assistant/test_assistant_services.py b/server/api/views/assistant/test_assistant_services.py index 9d911920..9c0a8dd1 100644 --- a/server/api/views/assistant/test_assistant_services.py +++ b/server/api/views/assistant/test_assistant_services.py @@ -1,10 +1,10 @@ # Tests for run_assistant (assistant_services.py): the orchestrator that wires the -# OpenAI client, the search tool mapping, and the agentic loop together. +# OpenAI client, the tool schemas, and the agentic loop together. # # The OpenAI client and handle_tool_calls_with_reasoning are mocked, so these # tests cover only logic run_assistant owns: how it builds the user input message, -# its decision to include vs. omit previous_response_id, and that it binds the -# request user into the search tool. No live OpenAI calls and no database. +# its decision to include vs. omit previous_response_id, and that it forwards the +# TOOLS and the request user to the loop. No live OpenAI calls and no database. from unittest.mock import MagicMock, patch @@ -68,24 +68,23 @@ def test_run_assistant_omits_previous_response_id_when_none(mock_openai_cls, moc assert "previous_response_id" not in call_kwargs -@patch("api.views.assistant.tool_services.search_documents") @patch("api.views.assistant.assistant_services.handle_tool_calls_with_reasoning") @patch("api.views.assistant.assistant_services.OpenAI") -def test_run_assistant_binds_user_to_search_documents(mock_openai_cls, mock_handle, mock_search): +def test_run_assistant_forwards_tools_and_user_to_loop(mock_openai_cls, mock_handle): mock_client = MagicMock() mock_openai_cls.return_value = mock_client mock_client.responses.create.return_value = _make_terminal_response() mock_handle.return_value = ("answer", "resp-1") from api.views.assistant.assistant_services import run_assistant + from api.views.assistant.tool_services import TOOLS user = MagicMock() run_assistant(message="query", user=user) - # Extract the tool_mapping passed to handle_tool_calls_with_reasoning - tool_mapping = mock_handle.call_args.kwargs.get("tool_mapping") or mock_handle.call_args.args[3] - bound_search = tool_mapping["search_documents"] - - # Calling the bound function should forward user to search_documents - bound_search(query="test query") - mock_search.assert_called_once_with("test query", user) + # run_assistant no longer binds user itself — it forwards TOOLS and the user to the + # loop, which binds user into each tool call at dispatch time. + # handle_tool_calls_with_reasoning(response, client, model_defaults, tools, user) + args = mock_handle.call_args.args + assert args[3] is TOOLS + assert args[4] is user diff --git a/server/api/views/assistant/test_tool_services.py b/server/api/views/assistant/test_tool_services.py index 83524fa6..036c04e0 100644 --- a/server/api/views/assistant/test_tool_services.py +++ b/server/api/views/assistant/test_tool_services.py @@ -1,13 +1,13 @@ -# Tests for tool_services.py: the retrieval tooling and the agentic reasoning loop. +# Tests for the assistant's tools and the agentic reasoning loop. # -# Covers the logic this module owns, with mocked tools (no DB, no OpenAI): -# - make_search_tool_mapping: the closure that binds the request user to -# search_documents, including per-call user independence. -# - invoke_functions_from_response: dispatching the model's function calls — -# the call/no-call branch, output shaping, and the unregistered-tool and -# tool-raises error paths. -# - handle_tool_calls_with_reasoning: the while-loop that keeps calling the -# model until it stops emitting tool calls, including loop continuity via +# Covers the logic these modules own, with mocked tools (no DB, no OpenAI): +# - Tool instances: SEARCH_TOOL.run forwards the request user; ASK_DATABASE_TOOL.run +# ignores it; schema() emits the flattened Responses-API shape. +# - invoke_functions_from_response: dispatching the model's function calls — the +# call/no-call branch, output shaping, and the unregistered-tool and tool-raises +# error paths. Tools are indexed by name and invoked as tool.run(user, **arguments). +# - handle_tool_calls_with_reasoning: the while-loop that keeps calling the model +# until it stops emitting tool calls, including loop continuity via # previous_response_id. import json @@ -23,41 +23,45 @@ invoke_functions_from_response, handle_tool_calls_with_reasoning, ) -from api.views.assistant.tool_services import make_search_tool_mapping +from api.views.assistant.tool_services import Tool, SEARCH_TOOL, ASK_DATABASE_TOOL, TOOLS # --------------------------------------------------------------------------- -# make_search_tool_mapping tests +# Tool instances # --------------------------------------------------------------------------- @patch("api.views.assistant.tool_services.search_documents") -def test_make_search_tool_mapping_bound_fn_forwards_user(mock_search): +def test_search_tool_run_forwards_query_and_user(mock_search): mock_search.return_value = "results" user = MagicMock() - mapping = make_search_tool_mapping(user) - mapping["search_documents"](query="lithium") + SEARCH_TOOL.run(user=user, query="lithium") mock_search.assert_called_once_with("lithium", user) -@patch("api.views.assistant.tool_services.search_documents") -def test_make_search_tool_mapping_different_users_are_independent(mock_search): - # Each call to make_search_tool_mapping should capture its own user, - # so two mappings created with different users do not share state. - user_a = MagicMock() - user_b = MagicMock() - mapping_a = make_search_tool_mapping(user_a) - mapping_b = make_search_tool_mapping(user_b) +@patch("api.views.assistant.tool_services.ask_database") +def test_ask_database_tool_run_ignores_user(mock_ask): + mock_ask.return_value = "rows" + + ASK_DATABASE_TOOL.run(user=MagicMock(), query="SELECT 1") + + # user is not forwarded — ask_database queries the shared medication table. + mock_ask.assert_called_once_with("SELECT 1") - mapping_a["search_documents"](query="q") - mapping_b["search_documents"](query="q") - # bound_search calls search_documents(query, user) positionally, so each - # recorded call is (args, kwargs) == (("q", user), {}). - calls = mock_search.call_args_list - assert calls[0] == (("q", user_a), {}) - assert calls[1] == (("q", user_b), {}) +def test_tool_schema_is_flattened_shape(): + schema = SEARCH_TOOL.schema() + assert schema["type"] == "function" + assert schema["name"] == "search_documents" + assert "parameters" in schema + # Flattened Responses-API shape — not nested under a "function" key. + assert "function" not in schema + + +def test_tools_registry_contains_both_tools(): + names = {tool.name for tool in TOOLS} + assert names == {"search_documents", "ask_database"} # --------------------------------------------------------------------------- @@ -86,22 +90,28 @@ def _make_response(output_items): return response +def _fake_tool(name, run): + """A Tool whose run is a mock; description/parameters are irrelevant to dispatch.""" + return Tool(name=name, description="", parameters={}, run=run) + + def test_invoke_returns_empty_list_when_no_function_calls(): response = _make_response([_make_reasoning_item()]) - result = invoke_functions_from_response(response, tool_mapping={}) + result = invoke_functions_from_response(response, tools=[], user=MagicMock()) assert result == [] def test_invoke_calls_tool_and_returns_output(): - mock_tool = MagicMock(return_value="search result") + mock_run = MagicMock(return_value="search result") + tool = _fake_tool("search_documents", mock_run) + user = MagicMock() item = _make_function_call_item("search_documents", {"query": "lithium"}, "call-1") response = _make_response([item]) - result = invoke_functions_from_response( - response, tool_mapping={"search_documents": mock_tool} - ) + result = invoke_functions_from_response(response, tools=[tool], user=user) - mock_tool.assert_called_once_with(query="lithium") + # The loop binds user at dispatch and forwards the model's arguments. + mock_run.assert_called_once_with(user=user, query="lithium") assert result == [ {"type": "function_call_output", "call_id": "call-1", "output": "search result"} ] @@ -111,38 +121,36 @@ def test_invoke_returns_error_message_when_tool_not_registered(): item = _make_function_call_item("unknown_tool", {"query": "x"}, "call-2") response = _make_response([item]) - result = invoke_functions_from_response(response, tool_mapping={}) + result = invoke_functions_from_response(response, tools=[], user=MagicMock()) assert result[0]["call_id"] == "call-2" assert "ERROR" in result[0]["output"] def test_invoke_returns_error_message_when_tool_raises(): - mock_tool = MagicMock(side_effect=Exception("tool exploded")) + mock_run = MagicMock(side_effect=Exception("tool exploded")) + tool = _fake_tool("search_documents", mock_run) item = _make_function_call_item("search_documents", {"query": "x"}, "call-3") response = _make_response([item]) - result = invoke_functions_from_response( - response, tool_mapping={"search_documents": mock_tool} - ) + result = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) assert "Error executing function call" in result[0]["output"] def test_invoke_handles_multiple_function_calls(): - mock_tool = MagicMock(return_value="result") + mock_run = MagicMock(return_value="result") + tool = _fake_tool("search_documents", mock_run) items = [ _make_function_call_item("search_documents", {"query": "q1"}, "call-4"), _make_function_call_item("search_documents", {"query": "q2"}, "call-5"), ] response = _make_response(items) - result = invoke_functions_from_response( - response, tool_mapping={"search_documents": mock_tool} - ) + result = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) assert len(result) == 2 - assert mock_tool.call_count == 2 + assert mock_run.call_count == 2 # --------------------------------------------------------------------------- @@ -171,7 +179,7 @@ def test_handle_terminates_immediately_when_no_tool_calls(): client = MagicMock() text, resp_id = handle_tool_calls_with_reasoning( - response, client, model_defaults={}, tool_mapping={} + response, client, model_defaults={}, tools=[], user=MagicMock() ) assert text == "Final answer." @@ -180,27 +188,27 @@ def test_handle_terminates_immediately_when_no_tool_calls(): def test_handle_calls_tool_then_terminates(): - mock_search = MagicMock(return_value="doc content") + mock_run = MagicMock(return_value="doc content") + tool = _fake_tool("search_documents", mock_run) first_response = _make_tool_call_response("resp-1") second_response = _make_terminal_response("Final answer.", "resp-2") client = MagicMock() client.responses.create.return_value = second_response + user = MagicMock() text, resp_id = handle_tool_calls_with_reasoning( - first_response, - client, - model_defaults={}, - tool_mapping={"search_documents": mock_search}, + first_response, client, model_defaults={}, tools=[tool], user=user ) - mock_search.assert_called_once_with(query="lithium") + mock_run.assert_called_once_with(user=user, query="lithium") assert text == "Final answer." assert resp_id == "resp-2" def test_handle_passes_previous_response_id_on_followup(): - mock_search = MagicMock(return_value="doc content") + mock_run = MagicMock(return_value="doc content") + tool = _fake_tool("search_documents", mock_run) first_response = _make_tool_call_response("resp-1") second_response = _make_terminal_response("Done.", "resp-2") @@ -208,10 +216,7 @@ def test_handle_passes_previous_response_id_on_followup(): client.responses.create.return_value = second_response handle_tool_calls_with_reasoning( - first_response, - client, - model_defaults={}, - tool_mapping={"search_documents": mock_search}, + first_response, client, model_defaults={}, tools=[tool], user=MagicMock() ) call_kwargs = client.responses.create.call_args.kwargs diff --git a/server/api/views/assistant/tool_services.py b/server/api/views/assistant/tool_services.py index 2159424c..dbd179ae 100644 --- a/server/api/views/assistant/tool_services.py +++ b/server/api/views/assistant/tool_services.py @@ -1,159 +1,122 @@ +from dataclasses import dataclass from typing import Callable -# make_search_tool_mapping (below) calls search_documents, which now lives in its own -# module after the refactor — import it directly from there. +# search_documents lives in its own module; imported here so SEARCH_TOOL.run can call +# it (and so tests can patch api.views.assistant.tool_services.search_documents). from .search_tool import search_documents # Reuse the existing ask_database implementation from services/tools rather than # reimplementing it here — it already enforces the SELECT-only and ALLOWED_TABLES -# guards. It takes no request-time secret (no `user`), so it is safe to import at -# module top level; database.py does no DB work at import time. -# -# IMPORT-TIME DB ACCESS: we import the ask_database *function* here (cheap, no DB), -# NOT database_schema_string from services/tools/tools.py. That module runs -# get_database_info(connection) at *its* import time to build the schema string, so -# importing it at module top level would fire a DB query just to import this file -# (breaking manage.py commands / any context where the DB isn't ready). -# _build_ask_database_schema() therefore defers that import to call time instead. +# guards, and does no DB work at import time. from ...services.tools.database import ask_database +# The Medication model is the source of truth for the queryable columns; we read them +# from its metadata (below) instead of introspecting the live database. +from ..listMeds.models import Medication -# TODO: Add keyword tool if it doesn't overlap with semantic search -# because too many overlapping tools can return conflicting answers -# get_tools_schema / make_tool_mapping are the aggregation seam: assistant_services.py -# reads these two functions instead of naming individual tools. That way a new tool is -# added by editing ONLY this file (append its schema below, register its callable in the -# mapping) — assistant_services.py never has to change again. -def get_tools_schema() -> list[dict]: - """Return every tool schema the assistant exposes to the model. +@dataclass(frozen=True) +class Tool: + """One assistant tool: the schema the model sees (data) and the function we run + (behavior), bundled together under a single name. - assistant_services.py reads this instead of naming individual schemas, so adding a - tool means appending here — not editing assistant_services.py. + Bundling name/description/parameters/run in one object means each tool is + registered in exactly one place — the TOOLS list at the bottom of this module — + so the schema sent to the model and the callable actually invoked can never drift + apart. Adding a tool is appending one Tool to TOOLS; nothing else changes. """ - # OVERLAP RISK: this exposes a semantic document-search tool AND a SQL - # medication-lookup tool at the same time. For a question both could plausibly - # answer, the model chooses which to call, and they can return conflicting - # answers. If that becomes a problem, sharpen each tool's description to carve - # out when to prefer which (unstructured docs/citations vs. structured - # medication facts) rather than adding yet more overlapping tools — see the - # keyword-tool TODO at the top of this module. - return SEARCH_TOOLS_SCHEMA + [_build_ask_database_schema()] - - -def make_tool_mapping(user) -> dict[str, Callable]: - """Return the full name->callable mapping for every tool. - ask_database needs no request-time binding (it queries the shared medication table, - not per-user data), so it is registered directly. search_documents is bound to the - user via make_search_tool_mapping. Reuses ask_database from services/tools. + name: str + description: str + parameters: dict + # run(user, **arguments) -> str. Every tool takes the request `user` so the dispatch + # loop can call them uniformly; a tool that doesn't need it simply ignores it. + run: Callable + + def schema(self) -> dict: + # Flattened Responses-API shape: name/description/parameters at the top level. + # This is intentionally NOT the nested {"function": {...}} shape that the Chat + # Completions API (and services/tools/tools.py's create_tool_dict) uses. + return { + "type": "function", + "name": self.name, + "description": self.description, + "parameters": self.parameters, + } + + +def _medication_schema_string() -> str: + """Describe the queryable medication table for the ask_database tool's prompt. + + The column list is read from the Medication model's metadata (``Model._meta``), + which Django populates from the class definition at import — so this needs no + database connection. That is why building the ask_database Tool below never + triggers a query (unlike introspecting information_schema over a live connection). """ - return { - **make_search_tool_mapping(user), - "ask_database": ask_database, - } - - + meta = Medication._meta + columns = ", ".join(field.column for field in meta.concrete_fields) + return f"Table: {meta.db_table}\nColumns: {columns}" -SEARCH_TOOL_DESCRIPTION = """ +SEARCH_TOOL = Tool( + name="search_documents", + description=""" Search the user's uploaded documents for information relevant to answering their question. Call this function when you need to find specific information from the user's documents to provide an accurate, citation-backed response. Always search before answering questions about document content. -""" - -SEARCH_TOOL_PROPERTY_DESCRIPTION = """ +""", + parameters={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": """ A specific search query to find relevant information in the user's documents. Use keywords, phrases, or questions related to what the user is asking about. Be specific rather than generic - use terms that would appear in the relevant documents. -""" - -# SEARCH_TOOLS_SCHEMA defines the search_documents tool for the OpenAI API. -# The model reads this schema to know what tools are available and what -# arguments to generate — it can only generate arguments declared here. -SEARCH_TOOLS_SCHEMA = [ - { - "type": "function", - "name": "search_documents", - "description": SEARCH_TOOL_DESCRIPTION, - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": SEARCH_TOOL_PROPERTY_DESCRIPTION, - } - }, - "required": ["query"], +""", + } }, - } -] - + "required": ["query"], + }, + # search_documents needs the request user for document access control. + run=lambda user, query: search_documents(query, user), +) -def make_search_tool_mapping(user) -> dict[str, Callable]: - # make_search_tool_mapping binds user to search_documents at call time. - # user is a request-time value the model cannot generate, so it must be - # captured here and kept out of the schema. - """Return a tool mapping with search_documents bound to the given user. - Parameters - ---------- - user : User - The Django user object used for document access control. - - Returns - ------- - dict[str, Callable] - Tool mapping ready to pass to invoke_functions_from_response. - """ - def bound_search(query: str) -> str: - return search_documents(query, user) - - return {"search_documents": bound_search} - - -ASK_DATABASE_TOOL_DESCRIPTION = """ +ASK_DATABASE_TOOL = Tool( + name="ask_database", + description=""" Use this tool to answer questions about the medications in the Balancer database. Medications are stored by their official generic names, not brand names, so convert brand names to generic names first and match case-insensitively (e.g. LOWER(name) = LOWER('lurasidone')). The input must be a single, fully-formed SQL SELECT query. -""" - - -def _build_ask_database_schema() -> dict: - """Build the ask_database tool schema in the flattened Responses-API shape. +""", + parameters={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": ( + "A plain-text SQL SELECT query answering the user's question, " + "written against this schema:\n" + f"{_medication_schema_string()}" + ), + } + }, + "required": ["query"], + }, + # ask_database queries the shared medication table, so it ignores the request user. + run=lambda user, query: ask_database(query), +) - The table/column names the model may query are injected from the live database - schema (reused from services/tools). The import is deferred to call time so that - merely importing this module never triggers a database query. - """ - # Deferred (function-local) import on purpose: services/tools/tools.py builds - # database_schema_string by querying the DB at *its* import time. Importing it - # here at call time keeps that query out of this module's import, so loading the - # assistant (e.g. during manage.py commands or when the DB isn't ready) never - # triggers it. We reuse the prebuilt string instead of rebuilding the schema. - from ...services.tools.tools import database_schema_string - # Flattened Responses-API shape: name/description/parameters live at the top level. - # This is intentionally NOT the nested {"function": {...}} shape that - # services/tools/tools.py's create_tool_dict produces for the Chat Completions API — - # which is also why create_tool_dict could not be reused here. - return { - "type": "function", - "name": "ask_database", - "description": ASK_DATABASE_TOOL_DESCRIPTION, - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": ( - "A plain-text SQL SELECT query answering the user's question, " - "written against this schema:\n" - f"{database_schema_string}" - ), - } - }, - "required": ["query"], - }, - } +# Single source of truth for the assistant's tools. assistant_services builds the +# schema list the model sees with [tool.schema() for tool in TOOLS]; the agentic loop +# indexes this by name to dispatch calls. Register a new tool by appending it here. +# +# OVERLAP RISK: this exposes a semantic document-search tool AND a SQL medication-lookup +# tool at once. For a question both could answer, the model chooses which to call and +# they can conflict. If that becomes a problem, sharpen each tool's description to carve +# out when to prefer which rather than adding more overlapping tools. +TOOLS = [SEARCH_TOOL, ASK_DATABASE_TOOL] From 76a78b0d3759fc38e55f1c315247a6275387cdd5 Mon Sep 17 00:00:00 2001 From: Sahil D Shah Date: Tue, 21 Jul 2026 12:20:31 -0400 Subject: [PATCH 4/6] Use absolute imports in assistant module; clarify test-patching contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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.), 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. --- .../api/views/assistant/assistant_services.py | 6 +-- server/api/views/assistant/search_tool.py | 4 +- server/api/views/assistant/tool_services.py | 42 +++++++++++-------- server/api/views/assistant/urls.py | 2 +- server/api/views/assistant/views.py | 2 +- 5 files changed, 31 insertions(+), 25 deletions(-) diff --git a/server/api/views/assistant/assistant_services.py b/server/api/views/assistant/assistant_services.py index 55a5b19a..e1cc8280 100644 --- a/server/api/views/assistant/assistant_services.py +++ b/server/api/views/assistant/assistant_services.py @@ -3,9 +3,9 @@ from openai import OpenAI -from .assistant_prompts import INSTRUCTIONS -from .tool_services import TOOLS -from .agentic_loop import handle_tool_calls_with_reasoning +from api.views.assistant.assistant_prompts import INSTRUCTIONS +from api.views.assistant.tool_services import TOOLS +from api.views.assistant.agentic_loop import handle_tool_calls_with_reasoning logger = logging.getLogger(__name__) diff --git a/server/api/views/assistant/search_tool.py b/server/api/views/assistant/search_tool.py index 8724930e..a0641e38 100644 --- a/server/api/views/assistant/search_tool.py +++ b/server/api/views/assistant/search_tool.py @@ -1,5 +1,5 @@ -from ...services.embedding_services import get_closest_embeddings -from ...services.conversions_services import convert_uuids +from api.services.embedding_services import get_closest_embeddings +from api.services.conversions_services import convert_uuids def search_documents(query: str, user) -> str: diff --git a/server/api/views/assistant/tool_services.py b/server/api/views/assistant/tool_services.py index dbd179ae..ef912ddb 100644 --- a/server/api/views/assistant/tool_services.py +++ b/server/api/views/assistant/tool_services.py @@ -1,16 +1,22 @@ from dataclasses import dataclass from typing import Callable -# search_documents lives in its own module; imported here so SEARCH_TOOL.run can call -# it (and so tests can patch api.views.assistant.tool_services.search_documents). -from .search_tool import search_documents +# search_documents is defined in search_tool.py; this `from ... import` binds a local +# name, api.views.assistant.tool_services.search_documents, that SEARCH_TOOL.run looks +# up at call time. Tests patch that local name — the use site here, NOT the definition +# site (search_tool.search_documents) — because mock.patch must rebind the reference the +# code actually resolves. Keep this as a bare-name import: rewriting SEARCH_TOOL.run to +# call search_tool.search_documents(...) would move the patch target and break the tests. +from api.views.assistant.search_tool import search_documents # Reuse the existing ask_database implementation from services/tools rather than # reimplementing it here — it already enforces the SELECT-only and ALLOWED_TABLES -# guards, and does no DB work at import time. -from ...services.tools.database import ask_database +# guards, and does no DB work at import time. Same use-site patching applies: tests patch +# api.views.assistant.tool_services.ask_database (the name bound here), not the definition +# in api.services.tools.database. +from api.services.tools.database import ask_database # The Medication model is the source of truth for the queryable columns; we read them # from its metadata (below) instead of introspecting the live database. -from ..listMeds.models import Medication +from api.views.listMeds.models import Medication @dataclass(frozen=True) @@ -43,18 +49,6 @@ def schema(self) -> dict: } -def _medication_schema_string() -> str: - """Describe the queryable medication table for the ask_database tool's prompt. - - The column list is read from the Medication model's metadata (``Model._meta``), - which Django populates from the class definition at import — so this needs no - database connection. That is why building the ask_database Tool below never - triggers a query (unlike introspecting information_schema over a live connection). - """ - meta = Medication._meta - columns = ", ".join(field.column for field in meta.concrete_fields) - return f"Table: {meta.db_table}\nColumns: {columns}" - SEARCH_TOOL = Tool( name="search_documents", @@ -83,6 +77,18 @@ def _medication_schema_string() -> str: ) +def _medication_schema_string() -> str: + """Describe the queryable medication table for the ask_database tool's prompt. + + The column list is read from the Medication model's metadata (``Model._meta``), + which Django populates from the class definition at import — so this needs no + database connection. That is why building the ask_database Tool below never + triggers a query (unlike introspecting information_schema over a live connection). + """ + meta = Medication._meta + columns = ", ".join(field.column for field in meta.concrete_fields) + return f"Table: {meta.db_table}\nColumns: {columns}" + ASK_DATABASE_TOOL = Tool( name="ask_database", description=""" diff --git a/server/api/views/assistant/urls.py b/server/api/views/assistant/urls.py index 4c68f952..53467803 100644 --- a/server/api/views/assistant/urls.py +++ b/server/api/views/assistant/urls.py @@ -1,5 +1,5 @@ from django.urls import path -from .views import Assistant +from api.views.assistant.views import Assistant urlpatterns = [path("v1/api/assistant", Assistant.as_view(), name="assistant")] diff --git a/server/api/views/assistant/views.py b/server/api/views/assistant/views.py index 74bee8f6..d6bca95c 100644 --- a/server/api/views/assistant/views.py +++ b/server/api/views/assistant/views.py @@ -9,7 +9,7 @@ from drf_spectacular.utils import extend_schema, inline_serializer from rest_framework import serializers as drf_serializers -from .assistant_services import run_assistant +from api.views.assistant.assistant_services import run_assistant logger = logging.getLogger(__name__) From 35a360dc505c34d705dbad14753ee8c41b0b32a6 Mon Sep 17 00:00:00 2001 From: Sahil D Shah Date: Tue, 21 Jul 2026 14:39:16 -0400 Subject: [PATCH 5/6] Hardcode ask_database schema string; document Tool composition choice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- server/api/views/assistant/tool_services.py | 43 ++++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/server/api/views/assistant/tool_services.py b/server/api/views/assistant/tool_services.py index ef912ddb..a7e46850 100644 --- a/server/api/views/assistant/tool_services.py +++ b/server/api/views/assistant/tool_services.py @@ -14,9 +14,6 @@ # api.views.assistant.tool_services.ask_database (the name bound here), not the definition # in api.services.tools.database. from api.services.tools.database import ask_database -# The Medication model is the source of truth for the queryable columns; we read them -# from its metadata (below) instead of introspecting the live database. -from api.views.listMeds.models import Medication @dataclass(frozen=True) @@ -28,6 +25,20 @@ class Tool: registered in exactly one place — the TOOLS list at the bottom of this module — so the schema sent to the model and the callable actually invoked can never drift apart. Adding a tool is appending one Tool to TOOLS; nothing else changes. + + Behavior is stored as the `run` field (composition) rather than a method on a + subclass because our tools differ only in *which* function runs — same schema() + machinery, same fields, just a different callable. They are instances of one + concept, not distinct kinds of thing. + + TODO: Flip to `Tool(ABC)` + one subclass per tool (with `run` as a method) if a + tool ever needs more than a swapped-in function — specifically when it: + - carries per-type state/setup (a client, connection, cache, validated config); + - overrides more than run (e.g. a custom schema() shape, or extra methods like + validate_arguments / cost_estimate); + - needs a per-type run signature or an @abstractmethod-enforced contract so a + tool with no behavior fails at class-definition time, not at call time. + Until then the callable field is lighter and keeps registration drift-proof. """ name: str @@ -77,17 +88,19 @@ def schema(self) -> dict: ) -def _medication_schema_string() -> str: - """Describe the queryable medication table for the ask_database tool's prompt. - - The column list is read from the Medication model's metadata (``Model._meta``), - which Django populates from the class definition at import — so this needs no - database connection. That is why building the ask_database Tool below never - triggers a query (unlike introspecting information_schema over a live connection). - """ - meta = Medication._meta - columns = ", ".join(field.column for field in meta.concrete_fields) - return f"Table: {meta.db_table}\nColumns: {columns}" +# The schema string describing the queryable medication table for ask_database's prompt. +# +# Kept in sync by hand with api.views.listMeds.models.Medication: if you add/rename a +# column there, update this string so ask_database's prompt matches the real table. +# +# Hand-writing the column list (rather than deriving it from Django's Model._meta) is a +# deliberate trade-off. _meta.concrete_fields would auto-sync with the model and needs no +# DB connection, but it dumps *every* column indiscriminately. A hand-written list lets us +# curate what the LLM sees — e.g. omit `id`, which the model never needs to filter on — and +# it drops the app-registry dependency (_meta requires the app registry loaded, so importing +# this module during app startup could raise AppRegistryNotReady). The cost is the manual +# update above, cheap for a table this small and stable. +_MEDICATION_SCHEMA_STRING = "Table: api_medication\nColumns: name, benefits, risks" ASK_DATABASE_TOOL = Tool( name="ask_database", @@ -106,7 +119,7 @@ def _medication_schema_string() -> str: "description": ( "A plain-text SQL SELECT query answering the user's question, " "written against this schema:\n" - f"{_medication_schema_string()}" + f"{_MEDICATION_SCHEMA_STRING}" ), } }, From 100ed122b1c1dcaff3463af5f351737a524e4a7a Mon Sep 17 00:00:00 2001 From: Sahil D Shah Date: Fri, 24 Jul 2026 16:48:17 -0400 Subject: [PATCH 6/6] Capture tool calls and duration in the eval CSV MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- server/api/views/assistant/agentic_loop.py | 122 +++++++++++++++--- .../api/views/assistant/assistant_services.py | 15 ++- server/api/views/assistant/eval_assistant.py | 33 ++++- .../assistant/test_assistant_services.py | 14 +- .../views/assistant/test_eval_assistant.py | 34 +++++ .../api/views/assistant/test_tool_services.py | 88 ++++++++++--- server/api/views/assistant/views.py | 8 +- 7 files changed, 262 insertions(+), 52 deletions(-) diff --git a/server/api/views/assistant/agentic_loop.py b/server/api/views/assistant/agentic_loop.py index 53b4601d..2d53f743 100644 --- a/server/api/views/assistant/agentic_loop.py +++ b/server/api/views/assistant/agentic_loop.py @@ -1,12 +1,63 @@ import json import logging +from dataclasses import dataclass +from enum import Enum logger = logging.getLogger(__name__) +class ToolCallStatus(str, Enum): + """The outcome of a single tool call. Three distinct states, deliberately not a + bool: FAILED (the tool matched but raised) and UNREGISTERED (the model asked for + a tool name we don't have) are opposite diagnoses — a code/data fault vs. the + model hallucinating a tool — and an eval on tool selection needs to tell them + apart. str-based so it serializes straight into CSV/JSON. + """ + + OK = "ok" + FAILED = "failed" # tool matched but raised + UNREGISTERED = "unregistered" # no tool registered for the model's requested name + + +@dataclass(frozen=True) +class ToolCall: + """A record of one tool call the model made — a complete unit for eval: + which tool, with what query (`arguments`), and what came back (`output`) or + broke (`error`). + + `output` and `error` are disjoint by status: `output` holds the retrieved + content on OK; `error` holds the detail when status is not OK. `arguments` + is the model-generated query (the primary tool-selection signal); it is None + only when the model's argument JSON could not be parsed. + """ + + name: str + status: ToolCallStatus + arguments: dict | None = None # the query the model generated (parsed) + output: str | None = None # the tool's result on success (retrieved content) + error: str | None = None # the failure detail when status is not OK + + +@dataclass(frozen=True) +class AssistantResult: + """What a full agentic run produced: the model's final text, the id of the final + response (for multi-turn continuity), and the ordered ToolCall records for every + tool invocation across all loop iterations. + + Deliberately deferred (add later as defaulted fields, no call-site churn): + token/cost usage (from response.usage) and turn count. Answer-quality scoring — + ground truth, citation accuracy, LLM-as-judge — is a separate layer above this + one, not a field here. + """ + + output_text: str + response_id: str + tool_calls: list[ToolCall] + + def handle_tool_calls_with_reasoning( response, client, model_defaults: dict, tools: list, user -) -> tuple[str, str]: +) -> AssistantResult: """Run the agentic loop until the model stops emitting function calls. Parameters @@ -24,23 +75,32 @@ def handle_tool_calls_with_reasoning( Returns ------- - tuple[str, str] - (final_response_output_text, final_response_id) + AssistantResult + The final response text and id, plus the ordered ToolCall records for every + tool invocation across all loop iterations (for eval / observability). """ # Open AI Cookbook: Handling Function Calls with Reasoning Models # https://cookbook.openai.com/examples/reasoning_function_calls + tool_calls: list[ToolCall] = [] while True: - function_responses = invoke_functions_from_response(response, tools, user) - if len(function_responses) == 0: # We're done reasoning + # user is threaded through so tools that need it (document access control) + # get it at dispatch time; tools that don't simply ignore it. + tool_output_messages, calls = invoke_functions_from_response(response, tools, user) + tool_calls.extend(calls) + if not tool_output_messages: # model emitted no tool calls this turn → done logger.info("Reasoning completed") final_response_output_text = response.output_text final_response_id = response.id logger.info(f"Final response: {final_response_output_text}") - return final_response_output_text, final_response_id + return AssistantResult( + output_text=final_response_output_text, + response_id=final_response_id, + tool_calls=tool_calls, + ) else: logger.info("More reasoning required, continuing...") response = client.responses.create( - input=function_responses, + input=tool_output_messages, previous_response_id=response.id, **model_defaults, ) @@ -48,10 +108,9 @@ def handle_tool_calls_with_reasoning( def invoke_functions_from_response( response, tools: list, user -) -> list[dict]: +) -> tuple[list[dict], list[ToolCall]]: """Extract all function calls from the response, look up the corresponding tool(s) and execute them. (This would be a good place to handle asynchroneous tool calls, or ones that take a while to execute.) - This returns a list of messages to be added to the conversation history. Parameters ---------- @@ -66,12 +125,16 @@ def invoke_functions_from_response( Returns ------- - list[dict] - List of function call output messages formatted for the OpenAI conversation. - Each message contains: - - type: "function_call_output" - - call_id: The unique identifier for the function call - - output: The result returned by the executed function (string or error message) + tuple[list[dict], list[ToolCall]] + - intermediate_messages: the function_call_output messages to append to the + conversation and feed back as the next turn's input. Each contains: + - type: "function_call_output" + - call_id: the unique identifier for the function call + - output: the result returned by the executed function (string or error message) + - tool_calls: one ToolCall record per function call — tool name, outcome status, + the model's arguments, and the output or error — kept distinct from the OpenAI + payload above so the run's tool usage is observable (see AssistantResult, + eval_assistant.py). """ # Open AI Cookbook: Handling Function Calls with Reasoning Models @@ -81,9 +144,13 @@ def invoke_functions_from_response( # returns None for an unknown name, handled explicitly below. tools_by_name = {tool.name: tool for tool in tools} intermediate_messages = [] + tool_calls: list[ToolCall] = [] for response_item in response.output: if response_item.type == "function_call": target_tool = tools_by_name.get(response_item.name) + # Parsed below; stays None if the model's argument JSON can't be parsed, + # so a FAILED record still reports whatever we managed to read. + arguments = None if target_tool is not None: try: arguments = json.loads(response_item.arguments) @@ -92,14 +159,37 @@ def invoke_functions_from_response( ) tool_output = target_tool.run(user=user, **arguments) logger.info(f"Tool {response_item.name} completed successfully") + tool_calls.append( + ToolCall( + name=response_item.name, + status=ToolCallStatus.OK, + arguments=arguments, + output=tool_output, + ) + ) except Exception as e: msg = f"Error executing function call: {response_item.name}: {e}" tool_output = msg logger.error(msg, exc_info=True) + tool_calls.append( + ToolCall( + name=response_item.name, + status=ToolCallStatus.FAILED, + arguments=arguments, + error=str(e), + ) + ) else: msg = f"ERROR - No tool registered for function call: {response_item.name}" tool_output = msg logger.error(msg) + tool_calls.append( + ToolCall( + name=response_item.name, + status=ToolCallStatus.UNREGISTERED, + error=msg, + ) + ) intermediate_messages.append( { "type": "function_call_output", @@ -109,4 +199,4 @@ def invoke_functions_from_response( ) elif response_item.type == "reasoning": logger.info(f"Reasoning step: {response_item.summary}") - return intermediate_messages \ No newline at end of file + return intermediate_messages, tool_calls diff --git a/server/api/views/assistant/assistant_services.py b/server/api/views/assistant/assistant_services.py index e1cc8280..4052cc7a 100644 --- a/server/api/views/assistant/assistant_services.py +++ b/server/api/views/assistant/assistant_services.py @@ -5,7 +5,10 @@ from api.views.assistant.assistant_prompts import INSTRUCTIONS from api.views.assistant.tool_services import TOOLS -from api.views.assistant.agentic_loop import handle_tool_calls_with_reasoning +from api.views.assistant.agentic_loop import ( + handle_tool_calls_with_reasoning, + AssistantResult, +) logger = logging.getLogger(__name__) @@ -14,7 +17,7 @@ def run_assistant( message: str, user, previous_response_id: str | None = None, -) -> tuple[str, str]: +) -> AssistantResult: """Wire together the OpenAI client, retrieval, and the agentic reasoning loop. Parameters @@ -28,12 +31,10 @@ def run_assistant( Returns ------- - tuple[str, str] - (final_response_output_text, final_response_id) + AssistantResult + The final response text and id, plus the ToolCall records made during the run. + Built by the loop and passed straight through — this function does not repack it. """ - # TODO: Track total duration, cost metrics, and tool_calls_made count - # and return them from run_assistant for use in eval_assistant.py CSV output - client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY")) MODEL_DEFAULTS = { diff --git a/server/api/views/assistant/eval_assistant.py b/server/api/views/assistant/eval_assistant.py index b44a2174..bed7a730 100644 --- a/server/api/views/assistant/eval_assistant.py +++ b/server/api/views/assistant/eval_assistant.py @@ -16,8 +16,11 @@ import os import sys +import json import logging import datetime +from dataclasses import asdict +from time import perf_counter from concurrent.futures import ThreadPoolExecutor, as_completed # Django setup must come before any imports that touch the ORM @@ -36,6 +39,7 @@ from django.contrib.auth import get_user_model from api.views.assistant.assistant_services import run_assistant +from api.views.assistant.agentic_loop import ToolCallStatus # TODO: remove unused import or use INSTRUCTIONS to record an instructions_hash column from api.views.assistant.assistant_prompts import INSTRUCTIONS @@ -92,22 +96,47 @@ def run_one(question: str, user, branch: str) -> dict: per request — adds overhead to every web request for no benefit - Cleaner call site in eval_assistant.py but wrong trade-off given WSGI """ + # Time the full run_assistant call here rather than inside it: run_one already + # owns the whole call, so wall-clock duration needs no plumbing through the + # production code path (see AssistantResult — duration is not carried). + start = perf_counter() try: - response_text, response_id = run_assistant(message=question, user=user) + result = run_assistant(message=question, user=user) + duration_s = perf_counter() - start + tool_error_count = sum( + 1 for c in result.tool_calls if c.status is not ToolCallStatus.OK + ) return { "branch": branch, "model": MODEL, "question": question, - "response_output_text": response_text, + "response_output_text": result.output_text, + "response_id": result.response_id, + # Flat summaries for at-a-glance scanning; the swallowed-failure hole this + # closes shows up as tool_error_count > 0 while error is None. + "tools_called": "|".join(c.name for c in result.tool_calls), + "tool_call_count": len(result.tool_calls), + "tool_error_count": tool_error_count, + # Full per-call detail — status, the model's arguments (query), output/error — + # for analysis that the flat columns can't hold. + "tool_calls_json": json.dumps([asdict(c) for c in result.tool_calls]), + "duration_s": duration_s, "error": None, } except Exception as e: + duration_s = perf_counter() - start logger.error(f"Error evaluating question '{question}': {e}") return { "branch": branch, "model": MODEL, "question": question, "response_output_text": None, + "response_id": None, + "tools_called": "", + "tool_call_count": 0, + "tool_error_count": 0, + "tool_calls_json": None, + "duration_s": duration_s, "error": str(e), } diff --git a/server/api/views/assistant/test_assistant_services.py b/server/api/views/assistant/test_assistant_services.py index 9c0a8dd1..99ec3518 100644 --- a/server/api/views/assistant/test_assistant_services.py +++ b/server/api/views/assistant/test_assistant_services.py @@ -8,6 +8,8 @@ from unittest.mock import MagicMock, patch +from api.views.assistant.agentic_loop import AssistantResult + def _make_terminal_response(output_text="Final answer.", response_id="resp-1"): response = MagicMock() @@ -16,13 +18,17 @@ def _make_terminal_response(output_text="Final answer.", response_id="resp-1"): response.id = response_id return response + +def _make_result(output_text="answer", response_id="resp-1"): + return AssistantResult(output_text=output_text, response_id=response_id, tool_calls=[]) + @patch("api.views.assistant.assistant_services.handle_tool_calls_with_reasoning") @patch("api.views.assistant.assistant_services.OpenAI") def test_run_assistant_sends_message_as_user_input(mock_openai_cls, mock_handle): mock_client = MagicMock() mock_openai_cls.return_value = mock_client mock_client.responses.create.return_value = _make_terminal_response() - mock_handle.return_value = ("answer", "resp-1") + mock_handle.return_value = _make_result() from api.views.assistant.assistant_services import run_assistant @@ -42,7 +48,7 @@ def test_run_assistant_passes_previous_response_id(mock_openai_cls, mock_handle) mock_client = MagicMock() mock_openai_cls.return_value = mock_client mock_client.responses.create.return_value = _make_terminal_response() - mock_handle.return_value = ("answer", "resp-2") + mock_handle.return_value = _make_result(response_id="resp-2") from api.views.assistant.assistant_services import run_assistant @@ -58,7 +64,7 @@ def test_run_assistant_omits_previous_response_id_when_none(mock_openai_cls, moc mock_client = MagicMock() mock_openai_cls.return_value = mock_client mock_client.responses.create.return_value = _make_terminal_response() - mock_handle.return_value = ("answer", "resp-1") + mock_handle.return_value = _make_result() from api.views.assistant.assistant_services import run_assistant @@ -74,7 +80,7 @@ def test_run_assistant_forwards_tools_and_user_to_loop(mock_openai_cls, mock_han mock_client = MagicMock() mock_openai_cls.return_value = mock_client mock_client.responses.create.return_value = _make_terminal_response() - mock_handle.return_value = ("answer", "resp-1") + mock_handle.return_value = _make_result() from api.views.assistant.assistant_services import run_assistant from api.views.assistant.tool_services import TOOLS diff --git a/server/api/views/assistant/test_eval_assistant.py b/server/api/views/assistant/test_eval_assistant.py index 5853d340..57d9e42e 100644 --- a/server/api/views/assistant/test_eval_assistant.py +++ b/server/api/views/assistant/test_eval_assistant.py @@ -18,3 +18,37 @@ def test_run_one_captures_error(mock_run_assistant): assert row["branch"] == "feature" assert row["response_output_text"] is None assert "boom" in row["error"] + # The error row carries the same tool/duration columns (defaulted) so the + # DataFrame is not ragged, and still records time-to-failure. + assert row["tools_called"] == "" + assert row["tool_call_count"] == 0 + assert row["tool_error_count"] == 0 + assert row["tool_calls_json"] is None + assert "duration_s" in row + + +@patch("api.views.assistant.eval_assistant.run_assistant") +def test_run_one_records_tool_calls(mock_run_assistant): + from api.views.assistant.agentic_loop import AssistantResult, ToolCall, ToolCallStatus + + mock_run_assistant.return_value = AssistantResult( + output_text="answer", + response_id="resp-1", + tool_calls=[ + ToolCall(name="search_documents", status=ToolCallStatus.OK, + arguments={"query": "lithium"}, output="docs"), + ToolCall(name="ask_database", status=ToolCallStatus.FAILED, + arguments={"query": "SELECT"}, error="bad sql"), + ], + ) + + row = run_one("query", user=MagicMock(), branch="feature") + + assert row["response_output_text"] == "answer" + assert row["response_id"] == "resp-1" + assert row["tools_called"] == "search_documents|ask_database" + assert row["tool_call_count"] == 2 + # One FAILED call is visible even though the run itself did not raise — + # the swallowed-failure hole this change closes. + assert row["tool_error_count"] == 1 + assert row["error"] is None diff --git a/server/api/views/assistant/test_tool_services.py b/server/api/views/assistant/test_tool_services.py index 036c04e0..86dc425f 100644 --- a/server/api/views/assistant/test_tool_services.py +++ b/server/api/views/assistant/test_tool_services.py @@ -22,6 +22,9 @@ from api.views.assistant.agentic_loop import ( invoke_functions_from_response, handle_tool_calls_with_reasoning, + AssistantResult, + ToolCall, + ToolCallStatus, ) from api.views.assistant.tool_services import Tool, SEARCH_TOOL, ASK_DATABASE_TOOL, TOOLS @@ -95,10 +98,11 @@ def _fake_tool(name, run): return Tool(name=name, description="", parameters={}, run=run) -def test_invoke_returns_empty_list_when_no_function_calls(): +def test_invoke_returns_empty_lists_when_no_function_calls(): response = _make_response([_make_reasoning_item()]) - result = invoke_functions_from_response(response, tools=[], user=MagicMock()) - assert result == [] + messages, calls = invoke_functions_from_response(response, tools=[], user=MagicMock()) + assert messages == [] + assert calls == [] def test_invoke_calls_tool_and_returns_output(): @@ -108,34 +112,51 @@ def test_invoke_calls_tool_and_returns_output(): item = _make_function_call_item("search_documents", {"query": "lithium"}, "call-1") response = _make_response([item]) - result = invoke_functions_from_response(response, tools=[tool], user=user) + messages, calls = invoke_functions_from_response(response, tools=[tool], user=user) # The loop binds user at dispatch and forwards the model's arguments. mock_run.assert_called_once_with(user=user, query="lithium") - assert result == [ + # The OpenAI payload (unchanged shape) is the first return value. + assert messages == [ {"type": "function_call_output", "call_id": "call-1", "output": "search result"} ] + # The ToolCall record captures the outcome, the model's query, and the output. + assert calls == [ + ToolCall( + name="search_documents", + status=ToolCallStatus.OK, + arguments={"query": "lithium"}, + output="search result", + ) + ] -def test_invoke_returns_error_message_when_tool_not_registered(): +def test_invoke_records_unregistered_when_tool_not_registered(): item = _make_function_call_item("unknown_tool", {"query": "x"}, "call-2") response = _make_response([item]) - result = invoke_functions_from_response(response, tools=[], user=MagicMock()) + messages, calls = invoke_functions_from_response(response, tools=[], user=MagicMock()) - assert result[0]["call_id"] == "call-2" - assert "ERROR" in result[0]["output"] + assert messages[0]["call_id"] == "call-2" + assert "ERROR" in messages[0]["output"] + assert calls[0].name == "unknown_tool" + assert calls[0].status is ToolCallStatus.UNREGISTERED + assert calls[0].error is not None -def test_invoke_returns_error_message_when_tool_raises(): +def test_invoke_records_failed_when_tool_raises(): mock_run = MagicMock(side_effect=Exception("tool exploded")) tool = _fake_tool("search_documents", mock_run) item = _make_function_call_item("search_documents", {"query": "x"}, "call-3") response = _make_response([item]) - result = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) + messages, calls = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) - assert "Error executing function call" in result[0]["output"] + assert "Error executing function call" in messages[0]["output"] + assert calls[0].status is ToolCallStatus.FAILED + assert "tool exploded" in calls[0].error + # arguments parsed before the tool raised, so they are still captured. + assert calls[0].arguments == {"query": "x"} def test_invoke_handles_multiple_function_calls(): @@ -147,9 +168,10 @@ def test_invoke_handles_multiple_function_calls(): ] response = _make_response(items) - result = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) + messages, calls = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) - assert len(result) == 2 + assert len(messages) == 2 + assert len(calls) == 2 assert mock_run.call_count == 2 @@ -178,12 +200,14 @@ def test_handle_terminates_immediately_when_no_tool_calls(): response = _make_terminal_response("Final answer.", "resp-1") client = MagicMock() - text, resp_id = handle_tool_calls_with_reasoning( + result = handle_tool_calls_with_reasoning( response, client, model_defaults={}, tools=[], user=MagicMock() ) - assert text == "Final answer." - assert resp_id == "resp-1" + assert isinstance(result, AssistantResult) + assert result.output_text == "Final answer." + assert result.response_id == "resp-1" + assert result.tool_calls == [] client.responses.create.assert_not_called() @@ -197,13 +221,37 @@ def test_handle_calls_tool_then_terminates(): client.responses.create.return_value = second_response user = MagicMock() - text, resp_id = handle_tool_calls_with_reasoning( + result = handle_tool_calls_with_reasoning( first_response, client, model_defaults={}, tools=[tool], user=user ) mock_run.assert_called_once_with(user=user, query="lithium") - assert text == "Final answer." - assert resp_id == "resp-2" + assert result.output_text == "Final answer." + assert result.response_id == "resp-2" + # The one tool call from the first turn is recorded on the result. + assert [c.name for c in result.tool_calls] == ["search_documents"] + assert result.tool_calls[0].status is ToolCallStatus.OK + + +def test_handle_accumulates_tool_calls_across_iterations(): + mock_run = MagicMock(return_value="doc content") + tool = _fake_tool("search_documents", mock_run) + # Two tool-calling turns, then a terminal one. + first_response = _make_tool_call_response("resp-1", query="q1") + second_response = _make_tool_call_response("resp-2", query="q2") + third_response = _make_terminal_response("Final answer.", "resp-3") + + client = MagicMock() + client.responses.create.side_effect = [second_response, third_response] + + result = handle_tool_calls_with_reasoning( + first_response, client, model_defaults={}, tools=[tool], user=MagicMock() + ) + + # Tool calls from every loop iteration are collected into one flat list. + assert len(result.tool_calls) == 2 + assert [c.arguments for c in result.tool_calls] == [{"query": "q1"}, {"query": "q2"}] + assert result.response_id == "resp-3" def test_handle_passes_previous_response_id_on_followup(): diff --git a/server/api/views/assistant/views.py b/server/api/views/assistant/views.py index d6bca95c..73082abc 100644 --- a/server/api/views/assistant/views.py +++ b/server/api/views/assistant/views.py @@ -46,16 +46,18 @@ def post(self, request): message = request.data.get("message", None) previous_response_id = request.data.get("previous_response_id", None) - final_response_output_text, final_response_id = run_assistant( + result = run_assistant( message=message, user=user, previous_response_id=previous_response_id, ) + # run_assistant now returns an AssistantResult; the JSON body is unchanged. + # tool_calls are captured for eval/observability and intentionally not exposed here. return Response( { - "response_output_text": final_response_output_text, - "final_response_id": final_response_id, + "response_output_text": result.output_text, + "final_response_id": result.response_id, }, status=status.HTTP_200_OK, )