Skip to content

Commit 3ad5a1e

Browse files
authored
Merge pull request #499 from sahilds1/490-improve-ai-evals
This PR refactors the assistant endpoint into smaller service modules (assistant orchestration + tool loop + prompts) and adds an evaluation script plus unit tests Changes: - Extracted assistant orchestration into assistant_services.py and tool-loop/retrieval logic into tool_services.py, updating the DRF view to call run_assistant. - Added an eval_assistant.py script (CSV output) and a small notebook for side-by-side response review. - Added focused unit tests for the new service modules and adjusted an existing uploadFile test import.
2 parents c798f73 + ecd9a72 commit 3ad5a1e

10 files changed

Lines changed: 879 additions & 299 deletions
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
INSTRUCTIONS = """
2+
You are an AI assistant that helps users find and understand information about bipolar disorder
3+
from your internal library of bipolar disorder research sources using semantic search.
4+
5+
IMPORTANT CONTEXT:
6+
- You have access to a library of sources that the user CANNOT see
7+
- The user did not upload these sources and doesn't know about them
8+
- You must explain what information exists in your sources and provide clear references
9+
10+
TOPIC RESTRICTIONS:
11+
When a prompt is received that is unrelated to bipolar disorder, mental health treatment,
12+
or psychiatric medications, respond by saying you are limited to bipolar-specific conversations.
13+
14+
SEMANTIC SEARCH STRATEGY:
15+
- Always perform semantic search using the search_documents function when users ask questions
16+
- Use conceptually related terms and synonyms, not just exact keyword matches
17+
- Search for the meaning and context of the user's question, not just literal words
18+
- Consider medical terminology, lay terms, and related conditions when searching
19+
20+
FUNCTION USAGE:
21+
- When a user asks about information that might be in your source library, ALWAYS use the search_documents function first
22+
- Perform semantic searches using concepts, symptoms, treatments, and related terms from the user's question
23+
- Only provide answers based on information found through your source searches
24+
25+
RESPONSE FORMAT:
26+
After gathering information through semantic searches, provide responses that:
27+
1. Answer the user's question directly using only the found information
28+
2. Structure responses with clear sections and paragraphs
29+
3. Explain what information you found in your sources and provide context
30+
4. Include citations using this exact format: [Name {name}, Page {page_number}]
31+
5. Only cite information that directly supports your statements
32+
33+
If no relevant information is found in your source library, clearly state that the information
34+
is not available in your current sources.
35+
36+
REMEMBER: You are working with an internal library of bipolar disorder sources that the user
37+
cannot see. Always search these sources first, explain what you found, and provide proper citations.
38+
"""
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import os
2+
import logging
3+
4+
from openai import OpenAI
5+
6+
from .assistant_prompts import INSTRUCTIONS
7+
from .tool_services import (
8+
SEARCH_TOOLS_SCHEMA,
9+
make_search_tool_mapping,
10+
handle_tool_calls_with_reasoning,
11+
)
12+
13+
logger = logging.getLogger(__name__)
14+
15+
16+
def run_assistant(
17+
message: str,
18+
user,
19+
previous_response_id: str | None = None,
20+
) -> tuple[str, str]:
21+
"""Wire together the OpenAI client, retrieval, and the agentic reasoning loop.
22+
23+
Parameters
24+
----------
25+
message : str
26+
The user's input message.
27+
user : User
28+
The Django user object used for document access control in search_documents.
29+
previous_response_id : str | None
30+
ID of a prior response for multi-turn conversation continuity.
31+
32+
Returns
33+
-------
34+
tuple[str, str]
35+
(final_response_output_text, final_response_id)
36+
"""
37+
# TODO: Track total duration, cost metrics, and tool_calls_made count
38+
# and return them from run_assistant for use in eval_assistant.py CSV output
39+
40+
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
41+
42+
MODEL_DEFAULTS = {
43+
"instructions": INSTRUCTIONS,
44+
"model": "gpt-5-nano", # 400,000 token context window
45+
# A summary of the reasoning performed by the model. This can be useful for debugging and understanding the model's reasoning process.
46+
"reasoning": {"effort": "low", "summary": None},
47+
"tools": SEARCH_TOOLS_SCHEMA,
48+
}
49+
50+
# TOOLS_SCHEMA tells the model what tools exist and what arguments to generate.
51+
# tool_mapping wires those tool names to the Python functions that execute them.
52+
# They are separate because the model generates arguments (schema concern) but
53+
# cannot supply request-time values like user (mapping concern).
54+
tool_mapping = make_search_tool_mapping(user)
55+
56+
if not previous_response_id:
57+
response = client.responses.create(
58+
input=[
59+
{"type": "message", "role": "user", "content": str(message)}
60+
],
61+
**MODEL_DEFAULTS,
62+
)
63+
else:
64+
response = client.responses.create(
65+
input=[
66+
{"type": "message", "role": "user", "content": str(message)}
67+
],
68+
previous_response_id=str(previous_response_id),
69+
**MODEL_DEFAULTS,
70+
)
71+
72+
return handle_tool_calls_with_reasoning(response, client, MODEL_DEFAULTS, tool_mapping)
Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
#!/usr/bin/env -S uv run --script
2+
# /// script
3+
# requires-python = "==3.11.11"
4+
# dependencies = [
5+
# "pandas==2.2.3",
6+
# "openai",
7+
# "django",
8+
# ]
9+
# ///
10+
11+
# uv script (or plain Python) to generate results to CSV, run from the terminal
12+
# Run from inside the container (working dir is /usr/src/server):
13+
# docker compose exec backend python api/views/assistant/eval_assistant.py
14+
#
15+
16+
17+
import os
18+
import sys
19+
import logging
20+
import datetime
21+
from concurrent.futures import ThreadPoolExecutor, as_completed
22+
23+
# Django setup must come before any imports that touch the ORM
24+
# NOTE: from api/views/assistant/, "../../../../" resolves four levels up to
25+
# /usr/src (not /usr/src/server, where balancer_backend lives). So this insert
26+
# alone does not put the settings package on sys.path — running the script
27+
# relies on the container already having /usr/src/server on PYTHONPATH. Sanity-
28+
# check this the first time the eval is run for real; the path depth may need
29+
# adjusting (e.g. "../../../").
30+
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../")))
31+
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "balancer_backend.settings")
32+
33+
import django
34+
django.setup()
35+
36+
from django.contrib.auth import get_user_model
37+
38+
from api.views.assistant.assistant_services import run_assistant
39+
# TODO: remove unused import or use INSTRUCTIONS to record an instructions_hash column
40+
from api.views.assistant.assistant_prompts import INSTRUCTIONS
41+
42+
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
43+
logger = logging.getLogger(__name__)
44+
45+
# Read model and INSTRUCTIONS from the source file or add a lightweight config endpoint to the backend
46+
47+
# Read model and INSTRUCTIONS from the source file
48+
# INSTRUCTIONS is imported from assistant_prompts.py
49+
# MODEL is read from assistant_services.py MODEL_DEFAULTS
50+
# TODO: import a shared MODEL_NAME constant from assistant_services instead of hardcoding
51+
MODEL = "gpt-5-nano"
52+
53+
# Set of representative questions to evaluate the assistant
54+
QUESTIONS = [
55+
"What medications are recommended for bipolar depression?",
56+
"What are the risks of lithium for patients with kidney disease?",
57+
"Which mood stabilizers are safe during pregnancy?",
58+
"What is the evidence for quetiapine in bipolar disorder?",
59+
"How does valproate compare to lithium for mania?",
60+
]
61+
62+
63+
def run_one(question: str, user, branch: str) -> dict:
64+
"""Run the assistant for a single question and return a result row.
65+
66+
Uses ThreadPoolExecutor (not asyncio.gather + await run_assistant) for concurrency.
67+
68+
Concurrency approach comparison:
69+
- ThreadPoolExecutor (this implementation):
70+
- run_assistant stays sync — views.py and the WSGI web app are unaffected
71+
- Each question runs in a thread pool worker, blocking on OpenAI + DB I/O
72+
- Django DB safe when run via `docker compose exec backend python eval_assistant.py`:
73+
this is a synchronous Django process context. Each ThreadPoolExecutor worker
74+
is a real OS thread with its own threading.local() storage, so each thread
75+
gets its own DB connection created lazily on first use. There is no shared
76+
event loop thread, so connections cannot clash or bleed between questions.
77+
The connection isolation concern only arises in ASGI contexts where multiple
78+
coroutines share one thread and therefore one threading.local() connection —
79+
which is not the case here.
80+
- Runtime: bottlenecked by OpenAI rate limits, not thread overhead
81+
- asyncio.gather + await run_assistant (alternative):
82+
- run_assistant becomes async — requires async def post in views.py,
83+
AsyncOpenAI client, and async handle_tool_calls_with_reasoning
84+
- Django DB unsafe if get_closest_embeddings is called directly in an async
85+
context without wrapping: get_closest_embeddings is a sync function that
86+
hits the ORM, so calling it on the event loop thread blocks all other
87+
coroutines until the DB responds. The fix is sync_to_async(get_closest_embeddings),
88+
which runs it in a dedicated worker thread with its own threading.local()
89+
connection. Bare await does not work at all — Django ORM querysets are not
90+
awaitables and raise TypeError immediately.
91+
- Under WSGI (manage.py runserver), async views run in a new event loop
92+
per request — adds overhead to every web request for no benefit
93+
- Cleaner call site in eval_assistant.py but wrong trade-off given WSGI
94+
"""
95+
try:
96+
response_text, response_id = run_assistant(message=question, user=user)
97+
return {
98+
"branch": branch,
99+
"model": MODEL,
100+
"question": question,
101+
"response_output_text": response_text,
102+
"error": None,
103+
}
104+
except Exception as e:
105+
logger.error(f"Error evaluating question '{question}': {e}")
106+
return {
107+
"branch": branch,
108+
"model": MODEL,
109+
"question": question,
110+
"response_output_text": None,
111+
"error": str(e),
112+
}
113+
114+
115+
def main():
116+
branch = os.environ.get("EVAL_BRANCH", "develop")
117+
118+
User = get_user_model()
119+
user = User.objects.filter(is_superuser=True).first()
120+
if not user:
121+
raise RuntimeError("No superuser found. Create one with manage.py createsuperuser.")
122+
123+
logger.info(f"Starting evaluation: branch={branch}, model={MODEL}, questions={len(QUESTIONS)}")
124+
125+
# ThreadPoolExecutor runs questions concurrently — see run_one docstring
126+
# for trade-off discussion vs asyncio.gather + await run_assistant.
127+
# max_workers=5 stays safely under OpenAI rate limits for gpt-5-nano.
128+
results = []
129+
with ThreadPoolExecutor(max_workers=5) as pool:
130+
futures = {
131+
pool.submit(run_one, question, user, branch): question
132+
for question in QUESTIONS
133+
}
134+
for future in as_completed(futures):
135+
results.append(future.result())
136+
137+
# Import pandas here, not at module top, so that importing this module (e.g.
138+
# run_one from test_eval_assistant.py) does not require pandas. It is only
139+
# needed for the CSV output below, when this script is run directly.
140+
import pandas as pd
141+
142+
df = pd.DataFrame(results)
143+
144+
results_dir = os.path.join(os.path.dirname(__file__), "results")
145+
os.makedirs(results_dir, exist_ok=True)
146+
timestamp = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%S")
147+
output_path = os.path.join(results_dir, f"{branch}-{timestamp}.csv")
148+
df.to_csv(output_path, index=False)
149+
150+
logger.info(f"Results saved to {output_path}")
151+
152+
153+
if __name__ == "__main__":
154+
main()
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# Assistant eval review\n",
8+
"\n",
9+
"Load two result CSVs produced by `eval_assistant.py` (one per branch) and compare the\n",
10+
"assistant's responses side by side.\n",
11+
"\n",
12+
"Set `DEVELOP_CSV` and `FEATURE_CSV` below to the two files in `results/` you want to compare."
13+
]
14+
},
15+
{
16+
"cell_type": "code",
17+
"execution_count": null,
18+
"metadata": {},
19+
"outputs": [],
20+
"source": [
21+
"import pandas as pd\n",
22+
"\n",
23+
"DEVELOP_CSV = \"results/develop-<timestamp>.csv\"\n",
24+
"FEATURE_CSV = \"results/<branch>-<timestamp>.csv\"\n",
25+
"\n",
26+
"develop_df = pd.read_csv(DEVELOP_CSV)\n",
27+
"feature_df = pd.read_csv(FEATURE_CSV)\n",
28+
"\n",
29+
"# Outer join on question so questions missing from either run stay visible.\n",
30+
"comparison = develop_df.merge(\n",
31+
" feature_df, on=\"question\", how=\"outer\", suffixes=(\"_develop\", \"_feature\")\n",
32+
")\n",
33+
"comparison[[\"question\", \"response_output_text_develop\", \"response_output_text_feature\"]]"
34+
]
35+
},
36+
{
37+
"cell_type": "markdown",
38+
"metadata": {},
39+
"source": [
40+
"## TODO (follow-up branch)\n",
41+
"\n",
42+
"Once `eval_assistant.py` records per-row metrics (`tool_calls_made`, token counts,\n",
43+
"`cost_usd`, `duration_s`, `instructions_hash`), extend the comparison to:\n",
44+
"\n",
45+
"- Flag rows where `instructions_hash` differs between branches (prompt changed).\n",
46+
"- Summarize cost/token totals grouped by branch.\n",
47+
"- Highlight rows where `tool_calls_made` differs (retrieval used differently)."
48+
]
49+
}
50+
],
51+
"metadata": {
52+
"language_info": {
53+
"name": "python"
54+
}
55+
},
56+
"nbformat": 4,
57+
"nbformat_minor": 5
58+
}

0 commit comments

Comments
 (0)