diff --git a/server/api/views/assistant/agentic_loop.py b/server/api/views/assistant/agentic_loop.py new file mode 100644 index 00000000..53b4601d --- /dev/null +++ b/server/api/views/assistant/agentic_loop.py @@ -0,0 +1,112 @@ +import json +import logging + +logger = logging.getLogger(__name__) + + +def handle_tool_calls_with_reasoning( + response, client, model_defaults: dict, tools: list, user +) -> 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. + 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 + ------- + 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: + function_responses = invoke_functions_from_response(response, tools, user) + if len(function_responses) == 0: # We're done reasoning + logger.info("Reasoning completed") + final_response_output_text = response.output_text + final_response_id = response.id + logger.info(f"Final response: {final_response_output_text}") + return final_response_output_text, final_response_id + else: + logger.info("More reasoning required, continuing...") + response = client.responses.create( + input=function_responses, + previous_response_id=response.id, + **model_defaults, + ) + + +def invoke_functions_from_response( + response, tools: list, user +) -> list[dict]: + """Extract all function calls from the response, look up the corresponding tool(s) and execute them. + (This would be a good place to handle asynchroneous tool calls, or ones that take a while to execute.) + This returns a list of messages to be added to the conversation history. + + Parameters + ---------- + 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 + ------- + 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 + + # Index the tools by name so a model-supplied call name can be looked up. .get() + # returns None for an unknown name, handled explicitly below. + tools_by_name = {tool.name: tool for tool in tools} + intermediate_messages = [] + for response_item in response.output: + if response_item.type == "function_call": + target_tool = tools_by_name.get(response_item.name) + if target_tool is not None: + try: + arguments = json.loads(response_item.arguments) + logger.info( + f"Invoking tool: {response_item.name} with arguments: {arguments}" + ) + tool_output = target_tool.run(user=user, **arguments) + logger.info(f"Tool {response_item.name} completed successfully") + except Exception as e: + msg = f"Error executing function call: {response_item.name}: {e}" + tool_output = msg + logger.error(msg, exc_info=True) + else: + msg = f"ERROR - No tool registered for function call: {response_item.name}" + tool_output = msg + logger.error(msg) + intermediate_messages.append( + { + "type": "function_call_output", + "call_id": response_item.call_id, + "output": tool_output, + } + ) + elif response_item.type == "reasoning": + logger.info(f"Reasoning step: {response_item.summary}") + return intermediate_messages \ No newline at end of file diff --git a/server/api/views/assistant/assistant_services.py b/server/api/views/assistant/assistant_services.py index ac339b9f..e1cc8280 100644 --- a/server/api/views/assistant/assistant_services.py +++ b/server/api/views/assistant/assistant_services.py @@ -3,12 +3,9 @@ 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.agentic_loop import handle_tool_calls_with_reasoning logger = logging.getLogger(__name__) @@ -44,15 +41,12 @@ def run_assistant( "model": "gpt-5-nano", # 400,000 token context window # A summary of the reasoning performed by the model. This can be useful for debugging and understanding the model's reasoning process. "reasoning": {"effort": "low", "summary": None}, - "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 +63,6 @@ def run_assistant( **MODEL_DEFAULTS, ) - return handle_tool_calls_with_reasoning(response, client, MODEL_DEFAULTS, tool_mapping) + # Pass TOOLS and user through to the loop, which indexes tools by name and binds + # user into each tool call at dispatch time. + return handle_tool_calls_with_reasoning(response, client, MODEL_DEFAULTS, TOOLS, user) diff --git a/server/api/views/assistant/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..9c0a8dd1 100644 --- a/server/api/views/assistant/test_assistant_services.py +++ b/server/api/views/assistant/test_assistant_services.py @@ -1,10 +1,10 @@ # Tests for run_assistant (assistant_services.py): the orchestrator that wires the -# OpenAI client, the search tool mapping, and the agentic loop together. +# OpenAI client, the tool schemas, and the agentic loop together. # # The OpenAI client and handle_tool_calls_with_reasoning are mocked, so these # tests cover only logic run_assistant owns: how it builds the user input message, -# its decision to include vs. omit previous_response_id, and that it binds the -# request user into the search tool. No live OpenAI calls and no database. +# its decision to include vs. omit previous_response_id, and that it forwards the +# TOOLS and the request user to the loop. No live OpenAI calls and no database. from unittest.mock import MagicMock, patch @@ -68,24 +68,23 @@ def test_run_assistant_omits_previous_response_id_when_none(mock_openai_cls, moc assert "previous_response_id" not in call_kwargs -@patch("api.views.assistant.tool_services.search_documents") @patch("api.views.assistant.assistant_services.handle_tool_calls_with_reasoning") @patch("api.views.assistant.assistant_services.OpenAI") -def test_run_assistant_binds_user_to_search_documents(mock_openai_cls, mock_handle, mock_search): +def test_run_assistant_forwards_tools_and_user_to_loop(mock_openai_cls, mock_handle): mock_client = MagicMock() mock_openai_cls.return_value = mock_client mock_client.responses.create.return_value = _make_terminal_response() mock_handle.return_value = ("answer", "resp-1") from api.views.assistant.assistant_services import run_assistant + from api.views.assistant.tool_services import TOOLS user = MagicMock() run_assistant(message="query", user=user) - # Extract the tool_mapping passed to handle_tool_calls_with_reasoning - tool_mapping = mock_handle.call_args.kwargs.get("tool_mapping") or mock_handle.call_args.args[3] - bound_search = tool_mapping["search_documents"] - - # Calling the bound function should forward user to search_documents - bound_search(query="test query") - mock_search.assert_called_once_with("test query", user) + # run_assistant no longer binds user itself — it forwards TOOLS and the user to the + # loop, which binds user into each tool call at dispatch time. + # handle_tool_calls_with_reasoning(response, client, model_defaults, tools, user) + args = mock_handle.call_args.args + assert args[3] is TOOLS + assert args[4] is user diff --git a/server/api/views/assistant/test_tool_services.py b/server/api/views/assistant/test_tool_services.py index 86e57eed..036c04e0 100644 --- a/server/api/views/assistant/test_tool_services.py +++ b/server/api/views/assistant/test_tool_services.py @@ -1,13 +1,13 @@ -# Tests for tool_services.py: the retrieval tooling and the agentic reasoning loop. +# Tests for the assistant's tools and the agentic reasoning loop. # -# Covers the logic this module owns, with mocked tools (no DB, no OpenAI): -# - make_search_tool_mapping: the closure that binds the request user to -# search_documents, including per-call user independence. -# - invoke_functions_from_response: dispatching the model's function calls — -# the call/no-call branch, output shaping, and the unregistered-tool and -# tool-raises error paths. -# - handle_tool_calls_with_reasoning: the while-loop that keeps calling the -# model until it stops emitting tool calls, including loop continuity via +# Covers the logic these modules own, with mocked tools (no DB, no OpenAI): +# - Tool instances: SEARCH_TOOL.run forwards the request user; ASK_DATABASE_TOOL.run +# ignores it; schema() emits the flattened Responses-API shape. +# - invoke_functions_from_response: dispatching the model's function calls — the +# call/no-call branch, output shaping, and the unregistered-tool and tool-raises +# error paths. Tools are indexed by name and invoked as tool.run(user, **arguments). +# - handle_tool_calls_with_reasoning: the while-loop that keeps calling the model +# until it stops emitting tool calls, including loop continuity via # previous_response_id. import json @@ -19,45 +19,49 @@ # mocking those two (like the rest of the suite mocks collaborators) covers all # three paths as fast, DB-free unit tests. -from api.views.assistant.tool_services import ( +from api.views.assistant.agentic_loop import ( invoke_functions_from_response, handle_tool_calls_with_reasoning, - make_search_tool_mapping, ) +from api.views.assistant.tool_services import Tool, SEARCH_TOOL, ASK_DATABASE_TOOL, TOOLS # --------------------------------------------------------------------------- -# make_search_tool_mapping tests +# Tool instances # --------------------------------------------------------------------------- @patch("api.views.assistant.tool_services.search_documents") -def test_make_search_tool_mapping_bound_fn_forwards_user(mock_search): +def test_search_tool_run_forwards_query_and_user(mock_search): mock_search.return_value = "results" user = MagicMock() - mapping = make_search_tool_mapping(user) - mapping["search_documents"](query="lithium") + SEARCH_TOOL.run(user=user, query="lithium") mock_search.assert_called_once_with("lithium", user) -@patch("api.views.assistant.tool_services.search_documents") -def test_make_search_tool_mapping_different_users_are_independent(mock_search): - # Each call to make_search_tool_mapping should capture its own user, - # so two mappings created with different users do not share state. - user_a = MagicMock() - user_b = MagicMock() - mapping_a = make_search_tool_mapping(user_a) - mapping_b = make_search_tool_mapping(user_b) +@patch("api.views.assistant.tool_services.ask_database") +def test_ask_database_tool_run_ignores_user(mock_ask): + mock_ask.return_value = "rows" + + ASK_DATABASE_TOOL.run(user=MagicMock(), query="SELECT 1") + + # user is not forwarded — ask_database queries the shared medication table. + mock_ask.assert_called_once_with("SELECT 1") - mapping_a["search_documents"](query="q") - mapping_b["search_documents"](query="q") - # bound_search calls search_documents(query, user) positionally, so each - # recorded call is (args, kwargs) == (("q", user), {}). - calls = mock_search.call_args_list - assert calls[0] == (("q", user_a), {}) - assert calls[1] == (("q", user_b), {}) +def test_tool_schema_is_flattened_shape(): + schema = SEARCH_TOOL.schema() + assert schema["type"] == "function" + assert schema["name"] == "search_documents" + assert "parameters" in schema + # Flattened Responses-API shape — not nested under a "function" key. + assert "function" not in schema + + +def test_tools_registry_contains_both_tools(): + names = {tool.name for tool in TOOLS} + assert names == {"search_documents", "ask_database"} # --------------------------------------------------------------------------- @@ -86,22 +90,28 @@ def _make_response(output_items): return response +def _fake_tool(name, run): + """A Tool whose run is a mock; description/parameters are irrelevant to dispatch.""" + return Tool(name=name, description="", parameters={}, run=run) + + def test_invoke_returns_empty_list_when_no_function_calls(): response = _make_response([_make_reasoning_item()]) - result = invoke_functions_from_response(response, tool_mapping={}) + result = invoke_functions_from_response(response, tools=[], user=MagicMock()) assert result == [] def test_invoke_calls_tool_and_returns_output(): - mock_tool = MagicMock(return_value="search result") + mock_run = MagicMock(return_value="search result") + tool = _fake_tool("search_documents", mock_run) + user = MagicMock() item = _make_function_call_item("search_documents", {"query": "lithium"}, "call-1") response = _make_response([item]) - result = invoke_functions_from_response( - response, tool_mapping={"search_documents": mock_tool} - ) + result = invoke_functions_from_response(response, tools=[tool], user=user) - mock_tool.assert_called_once_with(query="lithium") + # The loop binds user at dispatch and forwards the model's arguments. + mock_run.assert_called_once_with(user=user, query="lithium") assert result == [ {"type": "function_call_output", "call_id": "call-1", "output": "search result"} ] @@ -111,38 +121,36 @@ def test_invoke_returns_error_message_when_tool_not_registered(): item = _make_function_call_item("unknown_tool", {"query": "x"}, "call-2") response = _make_response([item]) - result = invoke_functions_from_response(response, tool_mapping={}) + result = invoke_functions_from_response(response, tools=[], user=MagicMock()) assert result[0]["call_id"] == "call-2" assert "ERROR" in result[0]["output"] def test_invoke_returns_error_message_when_tool_raises(): - mock_tool = MagicMock(side_effect=Exception("tool exploded")) + mock_run = MagicMock(side_effect=Exception("tool exploded")) + tool = _fake_tool("search_documents", mock_run) item = _make_function_call_item("search_documents", {"query": "x"}, "call-3") response = _make_response([item]) - result = invoke_functions_from_response( - response, tool_mapping={"search_documents": mock_tool} - ) + result = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) assert "Error executing function call" in result[0]["output"] def test_invoke_handles_multiple_function_calls(): - mock_tool = MagicMock(return_value="result") + mock_run = MagicMock(return_value="result") + tool = _fake_tool("search_documents", mock_run) items = [ _make_function_call_item("search_documents", {"query": "q1"}, "call-4"), _make_function_call_item("search_documents", {"query": "q2"}, "call-5"), ] response = _make_response(items) - result = invoke_functions_from_response( - response, tool_mapping={"search_documents": mock_tool} - ) + result = invoke_functions_from_response(response, tools=[tool], user=MagicMock()) assert len(result) == 2 - assert mock_tool.call_count == 2 + assert mock_run.call_count == 2 # --------------------------------------------------------------------------- @@ -171,7 +179,7 @@ def test_handle_terminates_immediately_when_no_tool_calls(): client = MagicMock() text, resp_id = handle_tool_calls_with_reasoning( - response, client, model_defaults={}, tool_mapping={} + response, client, model_defaults={}, tools=[], user=MagicMock() ) assert text == "Final answer." @@ -180,27 +188,27 @@ def test_handle_terminates_immediately_when_no_tool_calls(): def test_handle_calls_tool_then_terminates(): - mock_search = MagicMock(return_value="doc content") + mock_run = MagicMock(return_value="doc content") + tool = _fake_tool("search_documents", mock_run) first_response = _make_tool_call_response("resp-1") second_response = _make_terminal_response("Final answer.", "resp-2") client = MagicMock() client.responses.create.return_value = second_response + user = MagicMock() text, resp_id = handle_tool_calls_with_reasoning( - first_response, - client, - model_defaults={}, - tool_mapping={"search_documents": mock_search}, + first_response, client, model_defaults={}, tools=[tool], user=user ) - mock_search.assert_called_once_with(query="lithium") + mock_run.assert_called_once_with(user=user, query="lithium") assert text == "Final answer." assert resp_id == "resp-2" def test_handle_passes_previous_response_id_on_followup(): - mock_search = MagicMock(return_value="doc content") + mock_run = MagicMock(return_value="doc content") + tool = _fake_tool("search_documents", mock_run) first_response = _make_tool_call_response("resp-1") second_response = _make_terminal_response("Done.", "resp-2") @@ -208,10 +216,7 @@ def test_handle_passes_previous_response_id_on_followup(): client.responses.create.return_value = second_response handle_tool_calls_with_reasoning( - first_response, - client, - model_defaults={}, - tool_mapping={"search_documents": mock_search}, + first_response, client, model_defaults={}, tools=[tool], user=MagicMock() ) call_kwargs = client.responses.create.call_args.kwargs diff --git a/server/api/views/assistant/tool_services.py b/server/api/views/assistant/tool_services.py index 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..d6bca95c 100644 --- a/server/api/views/assistant/views.py +++ b/server/api/views/assistant/views.py @@ -9,7 +9,7 @@ from drf_spectacular.utils import extend_schema, inline_serializer from rest_framework import serializers as drf_serializers -from .assistant_services import run_assistant +from api.views.assistant.assistant_services import run_assistant logger = logging.getLogger(__name__)