File size: 2,100 Bytes
11b0dac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
"""
Test script for Socratic mode functionality
"""
import os
os.environ["OMP_NUM_THREADS"] = "1"
import warnings
warnings.filterwarnings("ignore")
from dotenv import load_dotenv
load_dotenv()
from langgraph.checkpoint.memory import MemorySaver
from medrax.models import ModelFactory
from medrax.agent import Agent
from medrax.utils import load_prompts_from_file
# Initialize checkpointer
checkpointer = MemorySaver()
# Load LLM
llm = ModelFactory.create_model(
model_name="gemini-2.0-flash",
temperature=0.7,
max_tokens=5000
)
# Load prompts
prompts = load_prompts_from_file("medrax/docs/system_prompts.txt")
# Test Assistant Mode
print("\n" + "="*60)
print("Testing ASSISTANT MODE")
print("="*60)
assistant_prompt = prompts.get("MEDICAL_ASSISTANT", "You are a helpful medical imaging assistant.")
assistant_agent = Agent(
llm,
tools=[], # Using empty tools for quick testing
system_prompt=assistant_prompt,
checkpointer=checkpointer,
mode="assistant"
)
config = {"configurable": {"thread_id": "test_assistant"}}
response = assistant_agent.workflow.invoke(
{"messages": [("user", "What is pneumonia?")]},
config=config
)
print("\nAssistant Response:")
print(response["messages"][-1].content)
# Test Socratic Mode
print("\n" + "="*60)
print("Testing SOCRATIC TUTOR MODE")
print("="*60)
socratic_prompt = prompts.get("SOCRATIC_TUTOR", "You are a Socratic medical educator.")
socratic_agent = Agent(
llm,
tools=[], # Using empty tools for quick testing
system_prompt=socratic_prompt,
checkpointer=checkpointer,
mode="socratic"
)
config = {"configurable": {"thread_id": "test_socratic"}}
response = socratic_agent.workflow.invoke(
{"messages": [("user", "What is pneumonia?")]},
config=config
)
print("\nSocratic Tutor Response:")
print(response["messages"][-1].content)
print("\n" + "="*60)
print("Test Complete!")
print("="*60)
print("\nThe key difference:")
print("- Assistant Mode: Provides direct answers and explanations")
print("- Socratic Mode: Guides learning through questions without giving direct answers") |