diff --git a/server/api/views/assistant/agentic_loop.py b/server/api/views/assistant/agentic_loop.py new file mode 100644 index 00000000..218eec79 --- /dev/null +++ b/server/api/views/assistant/agentic_loop.py @@ -0,0 +1,238 @@ +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. + + TODO: capture token usage and turn count — the other axis, alongside tool + selection, for comparing strategies. The design is settled but unbuilt: + - Six defaulted int fields: input_tokens, cached_tokens, output_tokens, + reasoning_tokens, total_tokens, turn_count. All three existing construction + sites pass keywords, so adding them is inert — this is exactly the property + the dataclass was chosen for over a widened tuple. + - Accumulate at the top of the while body in handle_tool_calls_with_reasoning, + before invoke_functions_from_response. That counts the initial response + (created in run_assistant and passed in) and every continuation exactly once, + including the terminal turn before the return. turn_count is then simply the + number of responses.create calls the run made. + - Read through a helper that walks response.usage defensively rather than + type-checking only the leaf: reasoning_tokens lives at + usage.output_tokens_details.reasoning_tokens and cached_tokens at + usage.input_tokens_details.cached_tokens, so a None usage raises + AttributeError before any leaf check runs. The helper must also reject + non-ints — the loop tests build responses as bare MagicMocks, and MagicMock + implements __add__/__radd__, so mock values would accumulate silently into + the CSV rather than failing loudly. + - cached_tokens is not optional: every turn resends context via + previous_response_id, so a large share of input_tokens bills at the cached + rate. Without the split, a cost figure derived later from input_tokens + overstates spend and cannot be corrected from the CSV afterwards. + - Dollar cost stays out of this dataclass: it needs a price table keyed by + model *and* date, which goes stale and then lies. Derive it in pandas from + the token columns plus the CSV's model column. If a table is ever wanted, the + house pattern is PRICING_DOLLARS_PER_MILLION_TOKENS in + api/services/llm_services.py — which has no reasoning or cached tier yet. + + Known hole that work would widen: if client.responses.create raises mid-loop the + exception propagates out and every ToolCall collected so far is lost with it — + the eval row reads tool_call_count 0 despite real calls having run, and would + likewise read total_tokens 0 despite tokens having been billed. Closing it means + deciding what this dataclass describes: a successful run, or whatever actually + happened (a partial result returned with an error field, or carried on the + exception). + + Answer-quality scoring — ground truth, citation accuracy, LLM-as-judge — is a + separate layer above this one, not a field here; see the scoring TODO in + eval_assistant.py. + """ + + output_text: str + response_id: str + tool_calls: list[ToolCall] + + +def handle_tool_calls_with_reasoning( + response, client, model_defaults: dict, tools: list, user +) -> AssistantResult: + """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. + 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 + ------- + 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: + # 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 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=tool_output_messages, + previous_response_id=response.id, + **model_defaults, + ) + + +def invoke_functions_from_response( + response, tools: list, user +) -> 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.) + + Parameters + ---------- + response : OpenAI Response + The response object from OpenAI containing output items that may include function calls + 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 + ------- + 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 + # 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 = [] + 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) + logger.info( + f"Invoking tool: {response_item.name} with arguments: {arguments}" + ) + 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", + "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, tool_calls diff --git a/server/api/views/assistant/assistant_services.py b/server/api/views/assistant/assistant_services.py index ac339b9f..be70b414 100644 --- a/server/api/views/assistant/assistant_services.py +++ b/server/api/views/assistant/assistant_services.py @@ -3,21 +3,27 @@ from openai import OpenAI -from .assistant_prompts import INSTRUCTIONS -from .tool_services import ( - SEARCH_TOOLS_SCHEMA, - make_search_tool_mapping, +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, + AssistantResult, ) logger = logging.getLogger(__name__) +# The single source of truth for which model the assistant runs on. Module-level so +# eval_assistant.py can import it and label its CSV with the model that actually ran, +# rather than repeating the string and silently mislabelling results the first time +# this changes. +MODEL_NAME = "gpt-5-nano" # 400,000 token context window + 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 @@ -31,28 +37,23 @@ 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 = { "instructions": INSTRUCTIONS, - "model": "gpt-5-nano", # 400,000 token context window + "model": MODEL_NAME, # 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, + # 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], } - # 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) - if not previous_response_id: response = client.responses.create( input=[ @@ -69,4 +70,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/eval_assistant.py b/server/api/views/assistant/eval_assistant.py index b44a2174..fd6edd35 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 @@ -35,20 +38,45 @@ from django.contrib.auth import get_user_model -from api.views.assistant.assistant_services import run_assistant -# TODO: remove unused import or use INSTRUCTIONS to record an instructions_hash column +from api.views.assistant.assistant_services import run_assistant, MODEL_NAME +from api.views.assistant.agentic_loop import ToolCallStatus +# TODO: write INSTRUCTIONS to a sidecar file alongside the CSV in main(), named +# results/{branch}-{timestamp}.prompt.txt so the pairing cannot come apart: +# f.write(f"branch: {branch}\nmodel: {MODEL_NAME}\n\n{INSTRUCTIONS}") +# Two alternatives were considered and rejected: a full-text CSV column repeats +# ~2.1KB of multi-line prose in every row and buries a cross-branch CSV diff in +# prompt noise; logging it at run time leaves nothing behind in results/, which is +# exactly the failure this is meant to prevent (a CSV whose prompt is unrecoverable +# months later). Add an instructions_hash column *as well* only if runs are ever +# concatenated into one DataFrame, where a groupby-able key beats diffing sidecars. +# Until that lands, INSTRUCTIONS is imported but deliberately unused. from api.views.assistant.assistant_prompts import INSTRUCTIONS logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) -# Read model and INSTRUCTIONS from the source file or add a lightweight config endpoint to the backend - -# Read model and INSTRUCTIONS from the source file -# INSTRUCTIONS is imported from assistant_prompts.py -# MODEL is read from assistant_services.py MODEL_DEFAULTS -# TODO: import a shared MODEL_NAME constant from assistant_services instead of hardcoding -MODEL = "gpt-5-nano" +# Model and INSTRUCTIONS both come from their source of truth rather than being +# restated here: MODEL_NAME from assistant_services.py (imported above, and used for +# the CSV's model column), INSTRUCTIONS from assistant_prompts.py (see sidecar TODO). + +# TODO: add a scoring layer. This is the biggest remaining gap, and it needs a design +# pass rather than a patch. As it stands this file is a *generation* harness, not an +# eval: QUESTIONS below carries no ground truth, so the CSV records what the +# assistant said and — since the tool-call columns landed — which tools it chose, but +# nothing about whether the answer was right. Open questions for that pass: +# - Ground truth per question: expected medications/claims, expected source +# documents (which citations the answer should rest on), or both. +# - Grading method: deterministic assertions (does the answer cite doc X, name drug +# Y) vs LLM-as-judge for faithfulness. Likely both — assertions for retrieval +# correctness, judge for answer quality. +# - Citation accuracy is the cheapest real signal available: INSTRUCTIONS mandates +# the [Name {name}, Page {page_number}] format, so citations can be parsed out of +# response_output_text and checked against what search_documents actually +# returned — already captured in the tool_calls_json column. That catches +# fabricated citations, the failure mode that matters most clinically. +# - Where scoring runs: as a separate pass over an already-written CSV, not inside +# run_one, so scoring can be revised and re-run without paying for generation +# again. # Set of representative questions to evaluate the assistant QUESTIONS = [ @@ -92,22 +120,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, + "model": MODEL_NAME, "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, + "model": MODEL_NAME, "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), } @@ -120,11 +173,11 @@ def main(): if not user: raise RuntimeError("No superuser found. Create one with manage.py createsuperuser.") - logger.info(f"Starting evaluation: branch={branch}, model={MODEL}, questions={len(QUESTIONS)}") + logger.info(f"Starting evaluation: branch={branch}, model={MODEL_NAME}, questions={len(QUESTIONS)}") # ThreadPoolExecutor runs questions concurrently — see run_one docstring # for trade-off discussion vs asyncio.gather + await run_assistant. - # max_workers=5 stays safely under OpenAI rate limits for gpt-5-nano. + # max_workers=5 stays safely under OpenAI rate limits for MODEL_NAME. results = [] with ThreadPoolExecutor(max_workers=5) as pool: futures = { diff --git a/server/api/views/assistant/search_tool.py b/server/api/views/assistant/search_tool.py new file mode 100644 index 00000000..a0641e38 --- /dev/null +++ b/server/api/views/assistant/search_tool.py @@ -0,0 +1,50 @@ +from api.services.embedding_services import get_closest_embeddings +from api.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_assistant_services.py b/server/api/views/assistant/test_assistant_services.py index 9d911920..99ec3518 100644 --- a/server/api/views/assistant/test_assistant_services.py +++ b/server/api/views/assistant/test_assistant_services.py @@ -1,13 +1,15 @@ # 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 +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 @@ -68,24 +74,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") + mock_handle.return_value = _make_result() 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_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 86e57eed..86dc425f 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 @@ -19,45 +19,52 @@ # 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, + AssistantResult, + ToolCall, + ToolCallStatus, ) +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") + + +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 - 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_tools_registry_contains_both_tools(): + names = {tool.name for tool in TOOLS} + assert names == {"search_documents", "ask_database"} # --------------------------------------------------------------------------- @@ -86,63 +93,86 @@ def _make_response(output_items): return response -def test_invoke_returns_empty_list_when_no_function_calls(): +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_lists_when_no_function_calls(): response = _make_response([_make_reasoning_item()]) - result = invoke_functions_from_response(response, tool_mapping={}) - assert result == [] + messages, calls = invoke_functions_from_response(response, tools=[], user=MagicMock()) + assert messages == [] + assert calls == [] 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} - ) + messages, calls = invoke_functions_from_response(response, tools=[tool], user=user) - mock_tool.assert_called_once_with(query="lithium") - assert result == [ + # The loop binds user at dispatch and forwards the model's arguments. + mock_run.assert_called_once_with(user=user, query="lithium") + # 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, tool_mapping={}) + 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(): - mock_tool = MagicMock(side_effect=Exception("tool exploded")) +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, tool_mapping={"search_documents": mock_tool} - ) + 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(): - 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} - ) + messages, calls = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) - assert len(result) == 2 - assert mock_tool.call_count == 2 + assert len(messages) == 2 + assert len(calls) == 2 + assert mock_run.call_count == 2 # --------------------------------------------------------------------------- @@ -170,37 +200,63 @@ 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( - response, client, model_defaults={}, tool_mapping={} + 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() 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() + + 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 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] - text, resp_id = handle_tool_calls_with_reasoning( - first_response, - client, - model_defaults={}, - tool_mapping={"search_documents": mock_search}, + result = handle_tool_calls_with_reasoning( + first_response, client, model_defaults={}, tools=[tool], user=MagicMock() ) - mock_search.assert_called_once_with(query="lithium") - assert text == "Final answer." - assert resp_id == "resp-2" + # 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(): - 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 +264,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 0fb96cef..a7e46850 100644 --- a/server/api/views/assistant/tool_services.py +++ b/server/api/views/assistant/tool_services.py @@ -1,214 +1,141 @@ -import json -import logging +from dataclasses import dataclass from typing import Callable -from ...services.embedding_services import get_closest_embeddings -from ...services.conversions_services import convert_uuids - -logger = logging.getLogger(__name__) +# 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. 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 + + +@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. + + 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. + + 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. + """ -TOOL_DESCRIPTION = """ + 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, + } + + + +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. -""" - -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": TOOL_DESCRIPTION, - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": TOOL_PROPERTY_DESCRIPTION, - } - }, - "required": ["query"], +""", + } }, - } -] - - -# 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. - -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} - - -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." - - -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 - -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"], + }, + # search_documents needs the request user for document access control. + run=lambda user, query: search_documents(query, user), +) + + +# 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", + 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. +""", + 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), +) + + +# 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] 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..73082abc 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__) @@ -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, )