diff --git a/server/api/views/assistant/agentic_loop.py b/server/api/views/assistant/agentic_loop.py new file mode 100644 index 00000000..a1fbd645 --- /dev/null +++ b/server/api/views/assistant/agentic_loop.py @@ -0,0 +1,117 @@ +import json +import logging + +from api.views.assistant.assistant_types import ( + AgentResult, + ToolCallExecution, + ToolCallStatus, +) + +logger = logging.getLogger(__name__) + + +def run_agentic_loop( + response, client, model_defaults: dict, tools: list, user +) -> AgentResult: + + # Every tool call the agentic loop made before exiting + agentic_loop_tool_call_executions= [] + + while True: + # user is threaded through so tools that need it get it at dispatch time + + tool_output_schemas, tool_call_executions = handle_tool_calls(response, tools, user) + + # .extend splices every iteration's list of tools into one list + agentic_loop_tool_call_executions.extend(tool_call_executions) + + # Exit agentic loop when model response doesn't contain any tool calls + if not tool_output_schemas: + return AgentResult( + output_text=response.output_text, + response_id=response.id, + tool_calls=agentic_loop_tool_call_executions, + ) + + #TODO: Add error handling to collect partial AgentResult tool calls + response = client.responses.create( + input=tool_output_schemas, + previous_response_id=response.id, + **model_defaults, + ) + + +def handle_tool_calls( + response, tools: list, user +) -> tuple[list[dict], list[ToolCallExecution]]: + + # 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} + + tool_output_schemas = [] + tool_call_executions: list[ToolCallExecution] = [] + + for response_item in response.output: + if response_item.type == "reasoning": + #logger.info(f"Reasoning step: {response_item.summary}") + + elif response_item.type == "function_call": + + tool_output, tool_call_execution = _execute_function_call(response_item, tools_by_name, user) + + tool_output_schemas.append( + { + "type": "function_call_output", + "call_id": response_item.call_id, + "output": tool_output, + } + ) + + tool_call_executions.append(tool_call_execution) + + + return tool_output_schemas, tool_call_executions + + +def _execute_function_call( + response_item, tools_by_name: dict, user +) -> tuple[str, ToolCallExecution]: + + 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 None: + msg = f"ERROR - No tool registered for function call: {response_item.name}" + logger.error(msg) + return msg, ToolCallExecution( + name=response_item.name, + status=ToolCallStatus.UNREGISTERED, + error=msg, + ) + + 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") + return tool_output, ToolCallExecution( + 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}" + logger.error(msg, exc_info=True) + return msg, ToolCallExecution( + name=response_item.name, + status=ToolCallStatus.FAILED, + arguments=arguments, + error=str(e), + ) diff --git a/server/api/views/assistant/assistant_prompts.py b/server/api/views/assistant/assistant_prompts.py index 44bf9b9b..dde89eb8 100644 --- a/server/api/views/assistant/assistant_prompts.py +++ b/server/api/views/assistant/assistant_prompts.py @@ -1,3 +1,23 @@ +# TODO: rewrite the citation template below (RESPONSE FORMAT item 4) so the braces are not +# emitted literally. `[Name {name}, Page {page_number}]` is read by the model as required +# output *syntax* rather than as placeholders: the 20260807 eval returned +# [Pharmacological Treatment of Bipolar Depression: ... Options? {Pharmacological +# Treatment of Bipolar Depression: ... Options?}, Page 2] +# — the name filled in AND the braces kept, duplicating the title. Also observed: +# "Page: 3" (stray colon), "Page 4, Chunk 32" (extra field), "various pages", +# "multiple pages including 1-5". Show a filled-in example instead of a brace template, +# e.g. `[Name advancespharmaco.pdf, Page 9]`, and state that exactly one page number is +# cited per reference. +# +# This is one of two separable citation defects; the other is search_tool.py handing the +# model a UUID alongside the name (see the TODO there). Neither is cosmetic — citations +# are unparseable until both land, which blocks citation accuracy, the "cheapest real +# signal" the scoring TODO in eval_assistant.py is built on. +# +# Note both known importers pass this string through verbatim — assistant_services.py +# hands it to the API as `instructions`, eval_assistant.py imports it for a planned +# sidecar and does not use it — so no .format() reads the braces. They are inert to +# Python; the only thing interpreting them is the model. INSTRUCTIONS = """ You are an AI assistant that helps users find and understand information about bipolar disorder from your internal library of bipolar disorder research sources using semantic search. diff --git a/server/api/views/assistant/assistant_services.py b/server/api/views/assistant/assistant_services.py index ac339b9f..b7c4afad 100644 --- a/server/api/views/assistant/assistant_services.py +++ b/server/api/views/assistant/assistant_services.py @@ -3,65 +3,35 @@ 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 api.views.assistant.assistant_prompts import INSTRUCTIONS +from api.views.assistant.tool_services import TOOLS +from api.views.assistant.assistant_types import AgentResult +from api.views.assistant.agentic_loop import run_agentic_loop logger = logging.getLogger(__name__) +# Module-level so eval_assistant.py can import it and label its CSV with the model that actually ran +MODEL_NAME = "gpt-5-nano" + def run_assistant( - message: str, user, + message: str, previous_response_id: str | None = None, -) -> tuple[str, str]: - """Wire together the OpenAI client, retrieval, and the agentic reasoning loop. - - Parameters - ---------- - message : str - The user's input message. - user : User - The Django user object used for document access control in search_documents. - previous_response_id : str | None - ID of a prior response for multi-turn conversation continuity. - - Returns - ------- - tuple[str, str] - (final_response_output_text, final_response_id) - """ - # 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 +) -> AgentResult: client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY")) MODEL_DEFAULTS = { "instructions": INSTRUCTIONS, - "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. + "model": MODEL_NAME, + # TODO: Flip "summary" to "auto" once this org is confirmed verified with OpenAI "reasoning": {"effort": "low", "summary": None}, - "tools": SEARCH_TOOLS_SCHEMA, + "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=[ - {"type": "message", "role": "user", "content": str(message)} - ], - **MODEL_DEFAULTS, - ) - else: - response = client.responses.create( + if previous_response_id: + initial_response = client.responses.create( input=[ {"type": "message", "role": "user", "content": str(message)} ], @@ -69,4 +39,15 @@ def run_assistant( **MODEL_DEFAULTS, ) - return handle_tool_calls_with_reasoning(response, client, MODEL_DEFAULTS, tool_mapping) + # search_documents needs the request user for document access control + return run_agentic_loop(initial_response, client, MODEL_DEFAULTS, TOOLS, user) + + initial_response = client.responses.create( + input=[ + {"type": "message", "role": "user", "content": str(message)} + ], + **MODEL_DEFAULTS, + ) + + # search_documents needs the request user for document access control + return run_agentic_loop(initial_response, client, MODEL_DEFAULTS, TOOLS, user) diff --git a/server/api/views/assistant/assistant_types.py b/server/api/views/assistant/assistant_types.py new file mode 100644 index 00000000..6a04a3a4 --- /dev/null +++ b/server/api/views/assistant/assistant_types.py @@ -0,0 +1,71 @@ +from dataclasses import dataclass +from enum import Enum +from typing import Callable + +@dataclass(frozen=True) +class Tool: + """ + Instances are registered in tool_services.py's TOOLS list. + """ + name: str + description: str + parameters: dict + # Function we run: 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 + + # Schema that the model sees: Flattened Responses-API shape + def schema(self) -> dict: + return { + "type": "function", + "name": self.name, + "description": self.description, + "parameters": self.parameters, + } + + +class ToolCallStatus(str, Enum): + """ + Evaluate tool selection and distinguish between FAILED and UNREGISTERED + """ + + OK = "ok" + # Tool matched but raised an error + FAILED = "failed" + # No tool registered for the model's requested name: + # The model asked for a tool name we don't have + UNREGISTERED = "unregistered" + + +@dataclass(frozen=True) +class ToolCallExecution: + """ + A record of one tool call the model made + """ + + name: str + # `output` and `error` are disjoint by status + status: ToolCallStatus + # the query the model generated (the primary tool selection signal) + # None only when the model's argument JSON could not be parsed + arguments: dict | None = None + # the tool's result on success (retrieved content) + output: str | None = None + # the failure detail when status is not OK + error: str | None = None + + +@dataclass(frozen=True) +class AgentResult: + """ + # Built by the agentic loop as a run proceeds, + # and read by eval_assistant.py to fill the result CSV + """ + + # The model's final text + output_text: str + # The id of the final response (for multi-turn continuity) + response_id: str + # The ordered ToolCallExecutionrecords for every tool invocation across all loop iterations + tool_calls: list[ToolCallExecution] diff --git a/server/api/views/assistant/eval_assistant.py b/server/api/views/assistant/eval_assistant.py index 7584ae18..8e837f1d 100644 --- a/server/api/views/assistant/eval_assistant.py +++ b/server/api/views/assistant/eval_assistant.py @@ -1,33 +1,23 @@ -#!/usr/bin/env -S uv run --script -# /// script -# requires-python = "==3.11.11" -# dependencies = [ -# "pandas==2.2.3", -# "openai", -# "django", -# ] -# /// - -# uv script (or plain Python) to generate results to CSV, run from the terminal -# Run from inside the container (working dir is /usr/src/server): -# docker compose exec backend python api/views/assistant/eval_assistant.py -# - +# Generates eval results to CSV. Run from inside the container: +# docker compose exec backend python api/views/assistant/eval_assistant.py +# Writes to results/ next to this file, which the ./server bind mount surfaces on the host. import os import sys +import csv +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 -# NOTE: from api/views/assistant/, "../../../../" resolves four levels up to -# /usr/src (not /usr/src/server, where balancer_backend lives). So this insert -# alone does not put the settings package on sys.path — running the script -# relies on the container already having /usr/src/server on PYTHONPATH. Sanity- -# check this the first time the eval is run for real; the path depth may need -# adjusting (e.g. "../../../"). -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../"))) +# Django setup must come before any imports that touch the ORM. +# Three levels up from api/views/assistant/ is /usr/src/server, where the balancer_backend settings package lives. +# Running a script file puts the *script's* directory on sys.path[0], not the working +# directory, and the image sets no PYTHONPATH — so without it django.setup() below +# raises ModuleNotFoundError on balancer_backend.settings. +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))) os.environ.setdefault("DJANGO_SETTINGS_MODULE", "balancer_backend.settings") import django @@ -35,20 +25,32 @@ from django.contrib.auth import get_user_model # noqa: E402 -from api.views.assistant.assistant_services import run_assistant # noqa: E402 +from api.views.assistant.assistant_services import run_assistant, MODEL_NAME # noqa: E402 +from api.views.assistant.assistant_types import ToolCallStatus +# Imported to warm the embedding model in main() before the worker pool starts — +# see the call site for why this process needs it and the web path does not. +from api.services.sentencetTransformer_model import TransformerModel # noqa: E402 + 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" +FIELDNAMES = [ + "branch", + "model", + "question", + "response_output_text", + "response_id", + "tools_called", + "tool_call_count", + "tool_error_count", + "tool_calls_json", + "duration_s", + "error", +] # Set of representative questions to evaluate the assistant + QUESTIONS = [ "What medications are recommended for bipolar depression?", "What are the risks of lithium for patients with kidney disease?", @@ -61,51 +63,50 @@ def run_one(question: str, user, branch: str) -> dict: """Run the assistant for a single question and return a result row. - Uses ThreadPoolExecutor (not asyncio.gather + await run_assistant) for concurrency. - - Concurrency approach comparison: - - ThreadPoolExecutor (this implementation): - - run_assistant stays sync — views.py and the WSGI web app are unaffected - - Each question runs in a thread pool worker, blocking on OpenAI + DB I/O - - Django DB safe when run via `docker compose exec backend python eval_assistant.py`: - this is a synchronous Django process context. Each ThreadPoolExecutor worker - is a real OS thread with its own threading.local() storage, so each thread - gets its own DB connection created lazily on first use. There is no shared - event loop thread, so connections cannot clash or bleed between questions. - The connection isolation concern only arises in ASGI contexts where multiple - coroutines share one thread and therefore one threading.local() connection — - which is not the case here. - - Runtime: bottlenecked by OpenAI rate limits, not thread overhead - - asyncio.gather + await run_assistant (alternative): - - run_assistant becomes async — requires async def post in views.py, - AsyncOpenAI client, and async handle_tool_calls_with_reasoning - - Django DB unsafe if get_closest_embeddings is called directly in an async - context without wrapping: get_closest_embeddings is a sync function that - hits the ORM, so calling it on the event loop thread blocks all other - coroutines until the DB responds. The fix is sync_to_async(get_closest_embeddings), - which runs it in a dedicated worker thread with its own threading.local() - connection. Bare await does not work at all — Django ORM querysets are not - awaitables and raise TypeError immediately. - - Under WSGI (manage.py runserver), async views run in a new event loop - per request — adds overhead to every web request for no benefit - - Cleaner call site in eval_assistant.py but wrong trade-off given WSGI + Uses ThreadPoolExecutor for concurrency. + """ + # 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 AgentResult — 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), } @@ -118,11 +119,13 @@ 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)}") + + # Load the embedding model before starting any workers + TransformerModel.get_instance() - # 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. + # ThreadPoolExecutor runs questions concurrently + # max_workers=5 stays safely under OpenAI rate limits for MODEL_NAME. results = [] with ThreadPoolExecutor(max_workers=5) as pool: futures = { @@ -132,18 +135,17 @@ def main(): for future in as_completed(futures): results.append(future.result()) - # Import pandas here, not at module top, so that importing this module (e.g. - # run_one from test_eval_assistant.py) does not require pandas. It is only - # needed for the CSV output below, when this script is run directly. - import pandas as pd - - df = pd.DataFrame(results) results_dir = os.path.join(os.path.dirname(__file__), "results") os.makedirs(results_dir, exist_ok=True) timestamp = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%S") output_path = os.path.join(results_dir, f"{branch}-{timestamp}.csv") - df.to_csv(output_path, index=False) + + # pandas was never in the backend image's requirements.txt + with open(output_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=FIELDNAMES) + writer.writeheader() + writer.writerows(results) logger.info(f"Results saved to {output_path}") diff --git a/server/api/views/assistant/search_tool.py b/server/api/views/assistant/search_tool.py new file mode 100644 index 00000000..f187fdda --- /dev/null +++ b/server/api/views/assistant/search_tool.py @@ -0,0 +1,47 @@ +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, or a + message saying nothing matched. Matching nothing is a legitimate outcome, + not a failure, so it returns normally and the call is recorded as OK. + + Raises + ------ + Exception + If the embedding search fails. Deliberately not caught here. + + """ + + 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) diff --git a/server/api/views/assistant/test_assistant_services.py b/server/api/views/assistant/test_assistant_services.py index 9d911920..19c0d12a 100644 --- a/server/api/views/assistant/test_assistant_services.py +++ b/server/api/views/assistant/test_assistant_services.py @@ -1,13 +1,38 @@ # 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. +# The OpenAI client and run_agentic_loop are mocked, so what remains +# to test is the one decision run_assistant actually makes: whether to include +# previous_response_id in the call at all. Everything else it does is forwarding — a +# hardcoded message dict, TOOLS and user passed straight through to the loop — and +# the tests that asserted those forwards were removed as glue. The only bugs they +# could catch were renames and reorderings, and one of them (`args[3] is TOOLS`) was +# coupled to positional argument order, so it would have gone red on a harmless +# switch to keyword arguments. +# +# Coverage that leaves open, deliberately noted rather than silently dropped: +# - The user -> run_assistant -> loop leg is no longer asserted. It is a bare +# positional forward with no decision in it, and the legs on either side are +# still covered (test_invoke_calls_tool_and_returns_output asserts the loop +# dispatches run(user=user, ...); test_search_tool_run_forwards_query_and_user +# asserts the handoff into retrieval). +# - Nothing asserts that MODEL_DEFAULTS["tools"] == [tool.schema() for tool in +# TOOLS] reaches the model. That comprehension is a real transformation and is +# genuinely untested — but it is not what the deleted test checked either. from unittest.mock import MagicMock, patch +import pytest + +from api.views.assistant.assistant_types import AgentResult + +# Distinguishes "the kwarg was omitted" from "the kwarg was passed as None", which is +# the entire point of the test below. It cannot use dict.get()'s usual None default: +# a regression that sent previous_response_id=None explicitly would then be +# indistinguishable from correctly omitting the key, which is exactly the bug the +# omit-branch exists to prevent. +ABSENT = object() + def _make_terminal_response(output_text="Final answer.", response_id="resp-1"): response = MagicMock() @@ -16,76 +41,44 @@ def _make_terminal_response(output_text="Final answer.", response_id="resp-1"): response.id = response_id return response -@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") - - from api.views.assistant.assistant_services import run_assistant - - run_assistant(message="Tell me about valproate.", user=MagicMock()) - call_kwargs = mock_client.responses.create.call_args - input_messages = call_kwargs.kwargs.get("input") or call_kwargs.args[0] - assert any( - item.get("role") == "user" and "valproate" in item.get("content", "") - for item in input_messages - ) +def _make_result(output_text="answer", response_id="resp-1"): + return AgentResult(output_text=output_text, response_id=response_id, tool_calls=[]) -@patch("api.views.assistant.assistant_services.handle_tool_calls_with_reasoning") +@pytest.mark.parametrize( + "previous_response_id, expected", + [ + pytest.param("resp-1", "resp-1", id="forwarded-when-provided"), + pytest.param(None, ABSENT, id="omitted-entirely-when-none"), + ], +) +@patch("api.views.assistant.assistant_services.run_agentic_loop") @patch("api.views.assistant.assistant_services.OpenAI") -def test_run_assistant_passes_previous_response_id(mock_openai_cls, mock_handle): +def test_run_assistant_includes_previous_response_id_only_when_set( + mock_openai_cls, mock_loop, previous_response_id, expected +): + """run_assistant's `if not previous_response_id` branch, both ways. + + Parametrized rather than written twice: the two cases are the same call with one + input changed, and previously duplicated four lines of client/loop mock setup to + assert two halves of one decision. + + Asserting on call_args is the only way to see this decision — omitting a kwarg + has no return-value footprint, since both branches return the same loop result. + """ 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_loop.return_value = _make_result() from api.views.assistant.assistant_services import run_assistant - run_assistant(message="More info.", user=MagicMock(), previous_response_id="resp-1") - - call_kwargs = mock_client.responses.create.call_args.kwargs - assert call_kwargs.get("previous_response_id") == "resp-1" - - -@patch("api.views.assistant.assistant_services.handle_tool_calls_with_reasoning") -@patch("api.views.assistant.assistant_services.OpenAI") -def test_run_assistant_omits_previous_response_id_when_none(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 - - run_assistant(message="First message.", user=MagicMock(), previous_response_id=None) + run_assistant( + message="Tell me about valproate.", + user=MagicMock(), + previous_response_id=previous_response_id, + ) call_kwargs = mock_client.responses.create.call_args.kwargs - 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): - 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 - - 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) + assert call_kwargs.get("previous_response_id", ABSENT) == expected diff --git a/server/api/views/assistant/test_eval_assistant.py b/server/api/views/assistant/test_eval_assistant.py index 5853d340..860accc7 100644 --- a/server/api/views/assistant/test_eval_assistant.py +++ b/server/api/views/assistant/test_eval_assistant.py @@ -1,20 +1,97 @@ # Tests for run_one (eval_assistant.py): the helper that runs the assistant for a # single eval question and shapes the outcome into a result row. # -# run_assistant is mocked, so this covers the logic run_one owns — specifically -# that a raising question is captured as an error row (error text recorded, -# response left None) instead of aborting the whole eval batch. +# run_assistant is mocked, so this covers the logic run_one owns — the try/except +# that turns a raising question into an error row instead of aborting the batch, the +# tool columns derived from AgentResult.tool_calls, and the invariant that both +# paths emit every CSV column. from unittest.mock import MagicMock, patch -from api.views.assistant.eval_assistant import run_one +import pytest + +from api.views.assistant.assistant_types import ( + AgentResult, + ToolCallExecution, + ToolCallStatus, +) +from api.views.assistant.eval_assistant import FIELDNAMES, run_one # TODO: add coverage for main()'s CSV output. -@patch("api.views.assistant.eval_assistant.run_assistant", side_effect=Exception("boom")) +# The two run_assistant outcomes, as patch() kwargs so the same pair can drive both +# the per-path tests and the shared column invariant without restating either setup. +_SUCCEEDS = { + "return_value": AgentResult( + output_text="answer", + response_id="resp-1", + tool_calls=[ + ToolCallExecution( + name="search_documents", + status=ToolCallStatus.OK, + arguments={"query": "lithium"}, + output="docs", + ), + ToolCallExecution( + name="ask_database", + status=ToolCallStatus.FAILED, + arguments={"query": "SELECT"}, + error="bad sql", + ), + ], + ) +} +_RAISES = {"side_effect": Exception("boom")} + + +@pytest.mark.parametrize( + "run_assistant_behavior", + [pytest.param(_SUCCEEDS, id="success-row"), pytest.param(_RAISES, id="error-row")], +) +def test_run_one_row_carries_every_csv_column(run_assistant_behavior): + """One invariant over both code paths, so it is parametrized rather than restated. + + This is the only guard on it. csv.DictWriter raises on an *extra* key + (extrasaction="raise"), which is the direction the FIELDNAMES comment describes — + but a *missing* key is silently filled with restval (""). So a column added to + one row literal in run_one and forgotten in the other reaches the CSV as an empty + cell rather than an error, which is precisely the ragged-row failure FIELDNAMES + was introduced to prevent. + """ + with patch( + "api.views.assistant.eval_assistant.run_assistant", **run_assistant_behavior + ): + row = run_one("query", user=MagicMock(), branch="feature") + + assert set(row) == set(FIELDNAMES) + + +@patch("api.views.assistant.eval_assistant.run_assistant", **_RAISES) def test_run_one_captures_error(mock_run_assistant): row = run_one("query", user=MagicMock(), branch="feature") assert row["branch"] == "feature" assert row["response_output_text"] is None assert "boom" in row["error"] + # The error row *defaults* the tool columns rather than omitting them, and still + # records time-to-failure. That the columns are present at all is asserted above; + # these are their values. + assert row["tools_called"] == "" + assert row["tool_call_count"] == 0 + assert row["tool_error_count"] == 0 + assert row["tool_calls_json"] is None + assert row["duration_s"] > 0 + + +@patch("api.views.assistant.eval_assistant.run_assistant", **_SUCCEEDS) +def test_run_one_records_tool_calls(mock_run_assistant): + row = run_one("query", user=MagicMock(), branch="feature") + + assert row["tools_called"] == "search_documents|ask_database" + assert row["tool_call_count"] == 2 + # tool_error_count counts every non-OK status, so one FAILED call stays visible + # even though the run itself did not raise and `error` is None. That combination + # is the swallowed-failure hole this column exists to close — a run that reads + # clean at the row level while a retrieval underneath it broke. + 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..46197142 100644 --- a/server/api/views/assistant/test_tool_services.py +++ b/server/api/views/assistant/test_tool_services.py @@ -1,67 +1,65 @@ -# 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 -# previous_response_id. +# Covers the logic these modules own, with mocked collaborators (no DB, no OpenAI): +# - Tool instances: SEARCH_TOOL.run adapts the loop's uniform (user, **arguments) +# call into search_documents' own (query, user) signature; schema() emits the +# flattened Responses-API shape rather than the nested Chat-Completions one. +# - search_documents' error/empty contract: failing raises, matching nothing does not. +# - handle_tool_calls: dispatching the model's function calls — the call/no-call +# branch, output shaping, and both error outcomes. Tools are indexed by name and +# invoked as tool.run(user, **arguments). +# - run_agentic_loop: the while-loop that keeps calling the model until it stops +# emitting tool calls, including loop continuity via previous_response_id and the +# function_call_output payload fed back on each continuation. +# +# Two tests were removed as glue. test_ask_database_tool_run_ignores_user asserted a +# single-argument forward whose wrong version raises TypeError on first call, and +# test_tools_registry_contains_both_tools restated the TOOLS list literal — a +# change-detector that made "adding a tool is appending one Tool to TOOLS; nothing +# else changes" (tool_services.py) false, since the intended way to extend the code +# was also the way to break the test. +# +# Where two tests were the same test with one input changed, they are now one +# pytest.mark.parametrize case table. Tests whose assertions differ in kind are left +# separate on purpose: folding those together needs a column per optional assertion +# and a body full of conditionals, which costs more clarity than the duplication did. import json -from unittest.mock import MagicMock, patch - -# TODO: add coverage for search_documents itself (formatting of embeddings -# results, the empty-results message, and the exception path). No DB needed: -# search_documents only calls get_closest_embeddings and convert_uuids, so -# 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 ( - invoke_functions_from_response, - handle_tool_calls_with_reasoning, - make_search_tool_mapping, -) - +from unittest.mock import MagicMock, call, patch -# --------------------------------------------------------------------------- -# make_search_tool_mapping tests -# --------------------------------------------------------------------------- +import pytest -@patch("api.views.assistant.tool_services.search_documents") -def test_make_search_tool_mapping_bound_fn_forwards_user(mock_search): - mock_search.return_value = "results" - user = MagicMock() - mapping = make_search_tool_mapping(user) - - mapping["search_documents"](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) - - 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), {}) +# TODO: add coverage for search_documents' formatting of embeddings results — the +# [Document N - File: ..., Similarity: ...] shape and the multi-result join. No DB +# needed: search_documents only calls get_closest_embeddings and convert_uuids, so +# mocking those two (like the rest of the suite mocks collaborators) is enough. The +# empty-results and exception paths are covered below. +# +# Sequence this after the file_id removal queued in search_tool.py, not before: that +# format string is about to lose its `File: {file_id}` field, so a test written against +# today's shape would be red on arrival. Pinning the format is worth doing either way — +# the field is there because the model reads it, and a change-detector objection doesn't +# apply to output whose exact text is the contract with the model. + +from api.views.assistant.assistant_types import ( + AgentResult, + Tool, + ToolCallExecution, + ToolCallStatus, +) +from api.views.assistant.agentic_loop import ( + handle_tool_calls, + run_agentic_loop, +) +from api.views.assistant.search_tool import search_documents +from api.views.assistant.tool_services import SEARCH_TOOL # --------------------------------------------------------------------------- -# invoke_functions_from_response tests +# Response / tool builders +# +# Defined before the tests because pytest.mark.parametrize case tables are built at +# import time, so anything they construct must already exist. # --------------------------------------------------------------------------- def _make_function_call_item(name, arguments, call_id): @@ -86,133 +84,352 @@ def _make_response(output_items): return response -def test_invoke_returns_empty_list_when_no_function_calls(): - response = _make_response([_make_reasoning_item()]) - result = invoke_functions_from_response(response, tool_mapping={}) - assert result == [] +def _make_terminal_response(output_text, response_id): + """A response with no function calls — terminates the loop.""" + response = MagicMock() + response.output = [] + response.output_text = output_text + response.id = response_id + return response -def test_invoke_calls_tool_and_returns_output(): - mock_tool = MagicMock(return_value="search result") - item = _make_function_call_item("search_documents", {"query": "lithium"}, "call-1") - response = _make_response([item]) +def _make_tool_call_response(response_id, query="lithium"): + """A response with one function call — continues the loop.""" + response = MagicMock() + response.output = [_make_function_call_item("search_documents", {"query": query}, "call-loop")] + response.id = response_id + return response - result = invoke_functions_from_response( - response, tool_mapping={"search_documents": mock_tool} - ) - mock_tool.assert_called_once_with(query="lithium") - assert result == [ - {"type": "function_call_output", "call_id": "call-1", "output": "search result"} - ] +def _make_client(*responses): + """A client whose successive responses.create calls return `responses` in order. + side_effect rather than return_value on purpose: return_value would hand the same + terminal response back forever, so a loop that failed to terminate would hang or + silently pass. A list runs out, and the extra call raises StopIteration. + """ + client = MagicMock() + client.responses.create.side_effect = list(responses) + return client -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={}) +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) - assert result[0]["call_id"] == "call-2" - assert "ERROR" in result[0]["output"] +# --------------------------------------------------------------------------- +# Tool instances +# --------------------------------------------------------------------------- + +@patch("api.views.assistant.tool_services.search_documents") +def test_search_tool_run_forwards_query_and_user(mock_search): + """The adapter inverts the argument order, which is why this is worth asserting. + + The loop calls run(user=..., query=...); search_documents takes (query, user). + Getting the swap wrong searches with a User object as the query string and scopes + access control to a string — silent in both directions, and this is the leg where + document access control is actually enforced. + """ + mock_search.return_value = "results" + user = MagicMock() + + SEARCH_TOOL.run(user=user, query="lithium") + + mock_search.assert_called_once_with("lithium", user) + + +def test_tool_schema_is_flattened_shape(): + schema = SEARCH_TOOL.schema() + assert schema["type"] == "function" + assert schema["name"] == "search_documents" + assert "parameters" in schema + # The load-bearing assertion: this repo contains both tool-schema shapes, and + # services/tools/tools.py's create_tool_dict builds the nested Chat-Completions + # one. The Responses API needs the flattened form, so a copy-paste from there + # would be accepted by every other assertion here. + assert "function" not in schema + + +# --------------------------------------------------------------------------- +# search_documents error/empty contract +# +# These two lock in the distinction the tool's status reporting depends on: a +# retrieval that *fails* must raise (so the loop records FAILED), while a retrieval +# that legitimately *matches nothing* must return normally (so it stays OK). Both +# used to return a string, which made the two indistinguishable downstream. +# --------------------------------------------------------------------------- -def test_invoke_returns_error_message_when_tool_raises(): - mock_tool = MagicMock(side_effect=Exception("tool exploded")) - item = _make_function_call_item("search_documents", {"query": "x"}, "call-3") +@patch("api.views.assistant.search_tool.get_closest_embeddings") +def test_search_documents_raises_instead_of_returning_the_error(mock_get): + mock_get.side_effect = RuntimeError("embedding backend down") + + # Must propagate. Swallowing it here would report a failed retrieval as a + # successful tool call and leave ToolCallStatus.FAILED unreachable for this tool. + with pytest.raises(RuntimeError, match="embedding backend down"): + search_documents("lithium", user=MagicMock()) + + +@patch("api.views.assistant.search_tool.convert_uuids", return_value=[]) +@patch("api.views.assistant.search_tool.get_closest_embeddings", return_value=[]) +def test_search_documents_returns_message_when_nothing_matches(mock_get, mock_convert): + result = search_documents("lithium", user=MagicMock()) + + # No match is an outcome, not an error — returns normally so the call records OK. + assert "No relevant documents found" in result + + +@patch( + "api.views.assistant.search_tool.get_closest_embeddings", + side_effect=RuntimeError("embedding backend down"), +) +def test_failed_status_is_reachable_through_the_real_search_tool(mock_get): + """The two fixes composed: a real retrieval failure arrives at the eval as FAILED. + + Deliberately dispatches the *real* SEARCH_TOOL — only its embedding dependency is + mocked — rather than a fake tool that raises. A fake would exercise the identical + loop branch as test_handle_tool_calls_records_the_two_error_outcomes below and + prove nothing + extra; what is worth testing is that search_documents' decision not to swallow the + exception and the loop's decision to record FAILED actually meet, with the real + adapter between them. + """ + item = _make_function_call_item("search_documents", {"query": "lithium"}, "call-e2e") + + _, calls = handle_tool_calls( + _make_response([item]), tools=[SEARCH_TOOL], user=MagicMock() + ) + + assert calls[0].status is ToolCallStatus.FAILED + assert "embedding backend down" in calls[0].error + + +# --------------------------------------------------------------------------- +# handle_tool_calls tests +# --------------------------------------------------------------------------- + +def test_handle_tool_calls_returns_empty_lists_when_no_function_calls(): + response = _make_response([_make_reasoning_item()]) + messages, calls = handle_tool_calls(response, tools=[], user=MagicMock()) + assert messages == [] + assert calls == [] + + +def test_handle_tool_calls_dispatches_tool_and_returns_output(): + 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 = handle_tool_calls(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") + # The OpenAI payload (unchanged shape) is the first return value. + assert messages == [ + {"type": "function_call_output", "call_id": "call-1", "output": "search result"} + ] + # The ToolCallExecution record captures the outcome, the model's query, and output. + assert calls == [ + ToolCallExecution( + name="search_documents", + status=ToolCallStatus.OK, + arguments={"query": "lithium"}, + output="search result", + ) + ] + + +@pytest.mark.parametrize( + "tools, expected_output_fragment, expected_status, expected_error_fragment, expected_arguments", + [ + pytest.param( + [], + "ERROR - No tool registered", + ToolCallStatus.UNREGISTERED, + "No tool registered", + None, + id="model-named-a-tool-we-do-not-have", + ), + pytest.param( + [_fake_tool("search_documents", MagicMock(side_effect=Exception("tool exploded")))], + "Error executing function call", + ToolCallStatus.FAILED, + "tool exploded", + {"query": "x"}, + id="registered-tool-raised", + ), + ], +) +def test_handle_tool_calls_records_the_two_error_outcomes( + tools, + expected_output_fragment, + expected_status, + expected_error_fragment, + expected_arguments, +): + """FAILED vs UNREGISTERED, parametrized to keep the contrast readable. + + These are opposite diagnoses — a code or data fault on our side vs. the model + hallucinating a tool name — which is why ToolCallStatus is a three-state enum and + not a bool, and why a tool-selection eval has to tell them apart. + + Reading them as one table also surfaces a difference neither test stated when they + were separate: `arguments` is parsed inside the registered branch, so an + unregistered call records None while a raising tool still reports the query the + model generated. + """ + item = _make_function_call_item("search_documents", {"query": "x"}, "call-err") + + messages, calls = handle_tool_calls( + _make_response([item]), tools=tools, user=MagicMock() ) - assert "Error executing function call" in result[0]["output"] + # Either way the model still gets a message back, so it can retry or say it could + # not retrieve anything — the loop does not abandon the turn. + assert messages[0]["call_id"] == "call-err" + assert expected_output_fragment in messages[0]["output"] + assert calls[0].name == "search_documents" + assert calls[0].status is expected_status + assert expected_error_fragment in calls[0].error + assert calls[0].arguments == expected_arguments -def test_invoke_handles_multiple_function_calls(): - mock_tool = MagicMock(return_value="result") + +def test_handle_tool_calls_handles_multiple_calls_in_one_response(): + 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 = handle_tool_calls(response, tools=[tool], user=MagicMock()) - assert len(result) == 2 - assert mock_tool.call_count == 2 + # Two calls in one response accumulate rather than overwrite — distinct from the + # cross-iteration accumulation covered in the loop test below. + assert [m["call_id"] for m in messages] == ["call-4", "call-5"] + assert [c.arguments for c in calls] == [{"query": "q1"}, {"query": "q2"}] + assert mock_run.call_count == 2 # --------------------------------------------------------------------------- -# handle_tool_calls_with_reasoning tests +# run_agentic_loop tests # --------------------------------------------------------------------------- -def _make_terminal_response(output_text, response_id): - """A response with no function calls — terminates the loop.""" - response = MagicMock() - response.output = [] - response.output_text = output_text - response.id = response_id - return response - - -def _make_tool_call_response(response_id, query="lithium"): - """A response with one function call — continues the loop.""" - response = MagicMock() - response.output = [_make_function_call_item("search_documents", {"query": query}, "call-loop")] - response.id = response_id - return response - - -def test_handle_terminates_immediately_when_no_tool_calls(): +def test_run_agentic_loop_terminates_immediately_when_no_tool_calls(): response = _make_terminal_response("Final answer.", "resp-1") - client = MagicMock() + client = _make_client() - text, resp_id = handle_tool_calls_with_reasoning( - response, client, model_defaults={}, tool_mapping={} + result = run_agentic_loop( + response, client, model_defaults={}, tools=[], user=MagicMock() ) - assert text == "Final answer." - assert resp_id == "resp-1" + assert isinstance(result, AgentResult) + 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") - 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 +@pytest.mark.parametrize( + "queries", + [ + pytest.param(["lithium"], id="one-tool-turn"), + pytest.param(["q1", "q2"], id="two-tool-turns"), + ], +) +def test_run_agentic_loop_continues_until_the_model_stops_calling_tools(queries): + """The loop at one and two tool-calling turns. + + Three tests collapsed into this table — they were the same scenario at different + turn counts, asserting one facet each (that a tool runs then the loop terminates, + that ToolCallExecution records accumulate across iterations, that the follow-up chains + off previous_response_id). Asserting all three at every turn count is strictly + more coverage than the originals: continuity was previously only checked on the + first follow-up, so a loop that re-sent resp-1 forever would have passed. + """ + mock_run = MagicMock(return_value="doc content") + tool = _fake_tool("search_documents", mock_run) + user = MagicMock() - text, resp_id = handle_tool_calls_with_reasoning( - first_response, - client, - model_defaults={}, - tool_mapping={"search_documents": mock_search}, + # One tool-calling response per query, then a terminal one that ends the loop. + tool_turns = [ + _make_tool_call_response(f"resp-{i + 1}", query=q) for i, q in enumerate(queries) + ] + terminal_id = f"resp-{len(queries) + 1}" + # The first response is the one run_assistant creates and passes in; only the rest + # come back from the client. + client = _make_client( + *tool_turns[1:], _make_terminal_response("Final answer.", terminal_id) ) - mock_search.assert_called_once_with(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") - first_response = _make_tool_call_response("resp-1") - second_response = _make_terminal_response("Done.", "resp-2") + result = run_agentic_loop( + tool_turns[0], client, model_defaults={}, tools=[tool], user=user + ) - client = MagicMock() - client.responses.create.return_value = second_response + # The tool ran once per turn, with user bound at each dispatch. + assert mock_run.call_args_list == [call(user=user, query=q) for q in queries] + # ToolCallExecution records from every iteration accumulate into one flat list. + assert [c.arguments for c in result.tool_calls] == [{"query": q} for q in queries] + assert all(c.status is ToolCallStatus.OK for c in result.tool_calls) + # Loop continuity: each follow-up chains off the id of the response it answers, + # so the chain advances resp-1 -> resp-2 -> ... rather than repeating resp-1. + assert [ + c.kwargs["previous_response_id"] for c in client.responses.create.call_args_list + ] == [turn.id for turn in tool_turns] + # Terminating returns the *last* response's text and id, not the first. + assert result.output_text == "Final answer." + assert result.response_id == terminal_id + + +def test_run_agentic_loop_feeds_each_turns_tool_output_back_to_the_model(): + """The tool's result actually reaches the model on the following turn. + + Kept separate from the loop test above because that test never looks at `input`. + It asserts previous_response_id, the dispatch call args, the accumulated records + and the final text/id — every one of which still holds if the continuation sends + an empty or stale payload. So the single thing a tool-calling turn exists to do, + hand the tool's output back, was the one thing unasserted: the model would answer + from nothing while the whole suite stayed green. + + Asserted on every continuation rather than just the first, and with a different + output per turn, because the two failure modes are distinct. The loop rebuilds + tool_output_schemas per iteration, so a bug that re-sent turn 1's payload forever + is not the same as one that sent none at all, and identical outputs could not tell + them apart. + """ + # Distinct outputs per turn so a stale payload is visible, not just a missing one. + mock_run = MagicMock(side_effect=["first result", "second result"]) + tool = _fake_tool("search_documents", mock_run) + + tool_turns = [_make_tool_call_response("resp-1"), _make_tool_call_response("resp-2")] + client = _make_client( + tool_turns[1], _make_terminal_response("Final answer.", "resp-3") + ) - handle_tool_calls_with_reasoning( - first_response, - client, - model_defaults={}, - tool_mapping={"search_documents": mock_search}, + run_agentic_loop( + tool_turns[0], client, model_defaults={}, tools=[tool], user=MagicMock() ) - call_kwargs = client.responses.create.call_args.kwargs - assert call_kwargs["previous_response_id"] == "resp-1" + # The function_call_output payload the loop hands back on each continuation call. + # call_id must echo the model's own call_id or the API cannot pair the output with + # the call it answers. + assert [c.kwargs["input"] for c in client.responses.create.call_args_list] == [ + [ + { + "type": "function_call_output", + "call_id": "call-loop", + "output": "first result", + } + ], + [ + { + "type": "function_call_output", + "call_id": "call-loop", + "output": "second result", + } + ], + ] diff --git a/server/api/views/assistant/tool_services.py b/server/api/views/assistant/tool_services.py index 0fb96cef..02caa4fb 100644 --- a/server/api/views/assistant/tool_services.py +++ b/server/api/views/assistant/tool_services.py @@ -1,214 +1,81 @@ -import json -import logging -from typing import Callable +from api.views.assistant.assistant_types import Tool +from api.views.assistant.search_tool import search_documents +from api.services.tools.database import ask_database -from ...services.embedding_services import get_closest_embeddings -from ...services.conversions_services import convert_uuids -logger = logging.getLogger(__name__) - -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. -""" - -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." + "required": ["query"], + }, + # 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. + + # 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 rather than deriving it from +# Django's Model._meta becuase the table is 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"], + }, -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. + # 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. - 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 + # ask_database queries the shared medication table, so it ignores the request user. + run=lambda user, query: ask_database(query), +) -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. +# 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. - 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, - ) +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..5f988d86 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__) @@ -36,26 +36,21 @@ class Assistant(APIView): def post(self, request): try: user = request.user - - # TODO: validate message and return a 400 when it is omitted or blank. - # @extend_schema documents message as required, but that schema is not - # enforced at runtime, so a missing/empty message reaches run_assistant - # and becomes the literal string "None" (str(None)) in the model input — - # producing confusing model behavior. Add a 400 to the responses schema - # when implementing. + + # TODO: Missing/empty message reaches run_assistant and becomes the literal string "None" (str(None)) in the model input 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( - message=message, + result = run_assistant( user=user, + message=message, previous_response_id=previous_response_id, ) 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, )