agent/tools.py

import asyncio
import json
from datetime import datetime

import blood_rag
import reports
from rag import search_supplements as _search_supplements

# ── Notepad state ──────────────────────────────────────────────────────────────
_notes: list    = []   # [{"id": int, "text": str, "time": str}]
_listeners: list = []  # one asyncio.Queue per connected SSE client
_counter: int   = 0


def _broadcast(payload: dict):
    data = json.dumps(payload)
    for q in list(_listeners):
        try:
            q.put_nowait(data)
        except asyncio.QueueFull:
            pass


def add_listener(q: asyncio.Queue):
    _listeners.append(q)


def remove_listener(q: asyncio.Queue):
    if q in _listeners:
        _listeners.remove(q)


def get_notes_snapshot() -> dict:
    return {"notes": list(_notes)}


# ── Tool implementations ───────────────────────────────────────────────────────
def write_to_notepad(text: str) -> str:
    global _counter
    _counter += 1
    entry = {
        "id":   _counter,
        "text": text.strip(),
        "time": datetime.now().strftime("%I:%M:%S %p"),
    }
    _notes.append(entry)
    _broadcast(get_notes_snapshot())
    return f'Written to notepad: "{text.strip()}"'


def read_notepad() -> str:
    if not _notes:
        return "The notepad is empty."
    lines = [f"{n['time']} — {n['text']}" for n in _notes]
    return "Current notes:\n" + "\n".join(lines)


def clear_notepad() -> str:
    _notes.clear()
    _broadcast(get_notes_snapshot())
    return "Notepad cleared."


def get_datetime() -> str:
    return datetime.now().strftime("%A, %B %d %Y — %I:%M %p")


# ── Clinical tools ────────────────────────────────────────────────────────────
async def get_blood_report_summary() -> str:
    report = reports.get_current_report()
    if not report.get("has_report"):
        return (
            "No blood test report has been uploaded yet. Ask the user to upload their "
            "blood test report (PDF) using the Blood Test Report uploader."
        )
    return report["summary"]


async def search_supplements(query: str) -> str:
    return await _search_supplements(query)


async def search_blood_knowledge(query: str) -> str:
    """Search the locally indexed blood test documents."""
    return await blood_rag.search(query)


async def run_tool(name: str, args: dict) -> str:
    if name == "write_to_notepad":
        return write_to_notepad(args.get("text", ""))
    if name == "read_notepad":
        return read_notepad()
    if name == "clear_notepad":
        return clear_notepad()
    if name == "get_datetime":
        return get_datetime()
    if name == "get_blood_report_summary":
        return await get_blood_report_summary()
    if name == "search_supplements":
        return await search_supplements(args.get("query", ""))
    if name == "search_blood_knowledge":
        return await search_blood_knowledge(args.get("query", ""))
    return f"Unknown tool: {name}"


# ── OpenAI tool schemas ────────────────────────────────────────────────────────
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "write_to_notepad",
            "description": (
                "Write a note to the live notepad. Use this whenever the user says "
                "'note that', 'write down', 'remember', 'save this', 'keep track of', "
                "or asks you to record anything."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": (
                            "The note text. Be concise — capture the key information "
                            "in plain language, no bullet points or markdown."
                        ),
                    }
                },
                "required": ["text"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read_notepad",
            "description": (
                "Read all notes currently on the notepad. Use when the user asks "
                "'what did I note', 'read my notes', or 'what have you written down'."
            ),
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "clear_notepad",
            "description": "Erase all notes from the notepad. Use only when the user explicitly asks to clear or delete all notes.",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_datetime",
            "description": "Get the current date and time.",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_blood_report_summary",
            "description": (
                "Get a summary of the user's most recently uploaded blood test report, "
                "including which markers are low, high, or normal. Call this first "
                "whenever the user asks about their blood test, lab results, deficiencies, "
                "or what supplements they should take."
            ),
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search_blood_knowledge",
            "description": (
                "Search the locally indexed blood test RAG for relevant passages about "
                "a patient's blood test results, markers, deficiencies, or lab values. "
                "Use this when the user asks about their blood work, specific markers, "
                "or what their results mean."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The question or marker to search for, e.g. 'iron deficiency anemia' or 'vitamin D levels'.",
                    }
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search_supplements",
            "description": (
                "Search the supplement knowledge base for recommendations related to a "
                "specific marker or deficiency (e.g. 'low vitamin D', 'low iron / "
                "ferritin'). Call this once per abnormal marker found in the user's blood "
                "report before recommending supplements."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The marker or deficiency to look up, e.g. 'low vitamin D' or 'high triglycerides'.",
                    }
                },
                "required": ["query"],
            },
        },
    },
]