app.py

# %%
# imports

import os
import json
from dotenv import load_dotenv
from openai import OpenAI
import gradio as gr
import requests
import urllib.parse  
from datetime import datetime

# %%
# Initialization

load_dotenv(override=True)

openai_api_key = os.getenv('OPENAI_API_KEY')
if openai_api_key:
    print(f"OpenAI API Key exists and begins {openai_api_key[:8]}")
else:
    print("OpenAI API Key not set")
    
MODEL = "gpt-4o-mini"
openai = OpenAI()

# %%
system_message = (
    "You are a clinical trials assistant that uses the ClinicalTrials.gov API to answer questions. "
    "If the user asks for study data, you must always call the `search_studies` tool. "
    "Only use your own knowledge for greetings or general questions."
    "Always return detailed and structured responses."
)

# %%
def search_studies(query):
    import urllib.parse
    import requests

    print("\n==============================")
    print(f"[DEBUG] Received query: {query} (type: {type(query)})")
    print("==============================")

    if isinstance(query, dict):
        filters = []
        query_parts = []

        if "query" in query:
            query_string = urllib.parse.quote_plus(query["query"])
            query_parts.append(f"query.cond={query_string}")

        if "phase" in query:
            filters.append(f"AREA[Phase]{query['phase']}")
        if "country" in query:
            filters.append(f"AREA[LocationCountry]{query['country']}")
        if "study_type" in query:
            filters.append(f"AREA[StudyType]{query['study_type']}")
        if "sex" in query:
            filters.append(f"AREA[Sex]{query['sex']}")
        if "age_group" in query:
            filters.append(f"AREA[StdAge]{query['age_group']}")
        if "status" in query:
            filters.append(f"AREA[OverallStatus]{query['status']}")
        if "sampling_method" in query:
            filters.append(f"AREA[SamplingMethod]{query['sampling_method']}")
        if "ipd_sharing" in query:
            filters.append(f"AREA[IPDSharing]{query['ipd_sharing']}")

        if "start_date_from" in query:
            filters.append(f"AREA[StartDate]RANGE[{query['start_date_from']},MAX]")
        if "start_date_to" in query:
            filters.append(f"AREA[StartDate]RANGE[MIN,{query['start_date_to']}]")
        if "completion_date_from" in query:
            filters.append(f"AREA[CompletionDate]RANGE[{query['completion_date_from']},MAX]")
        if "completion_date_to" in query:
            filters.append(f"AREA[CompletionDate]RANGE[MIN,{query['completion_date_to']}]")

        def normalize_sponsor(s):
            for suffix in ["Inc.", "Inc", "Ltd.", "Ltd", "GmbH", "LLC", "Corp.", "Corp"]:
                s = s.replace(suffix, "")
            return s.strip()

        if "sponsor" in query:
            sponsor_name = normalize_sponsor(query["sponsor"])
            sponsor_string = urllib.parse.quote_plus(sponsor_name)
            query_parts.append(f"query.spons={sponsor_string}")

        page_size = query.get("max_results", 3)
        if ("city" in query or "facility" in query) and "max_results" not in query:
            page_size = 30  # ensure enough studies for post-filtering

        filter_advanced = " AND ".join(filters)
        if filter_advanced:
            filter_advanced = f"({filter_advanced})"
        encoded_filter = urllib.parse.quote(filter_advanced, safe="[]*")

        url = (
            f"https://clinicaltrials.gov/api/v2/studies?"
            f"{'&'.join(query_parts)}"
            f"{f'&filter.advanced={encoded_filter}' if filter_advanced else ''}"
            f"&pageSize={page_size}"
        )

    else:
        encoded_query = urllib.parse.quote_plus(query)
        url = f"https://clinicaltrials.gov/api/v2/studies?query.cond={encoded_query}&pageSize=3"

    print("Requesting:", url)

    try:
        response = requests.get(url)
        print("Status Code:", response.status_code)

        if response.status_code == 400 and any("query.cond=" in part for part in query_parts):
            print("[⚠️ Fallback] Retrying with query.text instead of query.cond...")
            query_parts = [p.replace("query.cond=", "query.text=") for p in query_parts]
            url = (
                f"https://clinicaltrials.gov/api/v2/studies?"
                f"{'&'.join(query_parts)}"
                f"{f'&filter.advanced={encoded_filter}' if filter_advanced else ''}"
                f"&pageSize={page_size}"
            )
            print("Retrying:", url)
            response = requests.get(url)
            print("Retry Status Code:", response.status_code)

        if response.status_code == 200:
            data = response.json()
            trials = data.get("studies", [])
            if not trials:
                return "No studies found."

            def matches(study, key, value):
                section = study.get("protocolSection", {})
                if key == "sponsor":
                    return value.lower() in section.get("sponsorCollaboratorsModule", {}).get("leadSponsor", {}).get("name", "").lower()
                elif key == "intervention":
                    interventions = section.get("armsInterventionsModule", {}).get("interventions", [])
                    return any(value.lower() in i.get("name", "").lower() for i in interventions)
                elif key == "city":
                    locs = section.get("contactsLocationsModule", {}).get("locations", [])
                    return any(value.lower() in loc.get("city", "").lower() for loc in locs)
                elif key == "facility":
                    locs = section.get("contactsLocationsModule", {}).get("locations", [])
                    return any(value.lower() in loc.get("facility", "").lower() for loc in locs)
                return True

            for key in ["sponsor", "intervention", "city", "facility"]:
                if key in query:
                    trials = [s for s in trials if matches(s, key, query[key])]

            if not trials:
                return "No studies found after applying filters."

            result = []
            for study in trials:
                ps = study.get("protocolSection", {})
                id_module = ps.get("identificationModule", {})
                design_module = ps.get("designModule", {})
                status_module = ps.get("statusModule", {})
                elig_module = ps.get("eligibilityModule", {})
                ipd_module = ps.get("ipdSharingStatementModule", {})
                desc_module = ps.get("descriptionModule", {})
                contact_module = ps.get("contactsLocationsModule", {})
                sponsor_module = ps.get("sponsorCollaboratorsModule", {})
                outcomes_module = ps.get("outcomesModule", {})
                arms_module = ps.get("armsInterventionsModule", {})

                nct_id = id_module.get("nctId", "N/A")
                title = id_module.get("briefTitle", "No Title")
                official_title = id_module.get("officialTitle", "N/A")
                phases = design_module.get("phases", [])
                study_type = design_module.get("studyType", "N/A")
                status = status_module.get("overallStatus", "N/A")
                start_date = status_module.get("startDateStruct", {}).get("date", "N/A")
                completion_date = status_module.get("completionDateStruct", {}).get("date", "N/A")
                sex = elig_module.get("sex", "N/A")
                std_ages = elig_module.get("stdAges", [])
                sampling_method = elig_module.get("samplingMethod", "N/A")
                criteria = elig_module.get("eligibilityCriteria", "N/A")
                locations = contact_module.get("locations", [])
                countries = sorted({loc.get("country") for loc in locations if loc.get("country")})

                # Format location info (first 2 entries, truncate the rest)
                location_text_lines = []
                for loc in locations:
                    parts = [loc.get("facility"), loc.get("city"), loc.get("state"), loc.get("country")]
                    clean = [p for p in parts if p]
                    if clean:
                        location_text_lines.append(", ".join(clean))
                if location_text_lines:
                    if len(location_text_lines) > 3:
                        display_lines = location_text_lines[:2]
                        location_text = "\n".join(f"- {line}" for line in display_lines)
                        location_text += f"\n...and {len(location_text_lines)-2} more site(s)"
                    else:
                        location_text = "\n".join(f"- {line}" for line in location_text_lines)
                else:
                    location_text = "N/A"

                description = desc_module.get("detailedDescription", "N/A")
                interventions = arms_module.get("interventions", [])
                intervention_names = [iv.get("name", "") for iv in interventions if iv.get("name")]
                intervention_text = ", ".join(intervention_names) if intervention_names else "N/A"
                sponsor = sponsor_module.get("leadSponsor", {}).get("name", "N/A")
                collaborators = sponsor_module.get("collaborators", [])
                collaborator_names = [c.get("name", "") for c in collaborators]

                result.append(
                    f"### 🧪 {title}\n\n"
                    f"**NCT ID:** `{nct_id}`\n"
                    f"🔗 [View on ClinicalTrials.gov](https://clinicaltrials.gov/study/{nct_id})\n\n"
                    f"**Start Date:** {start_date}\n"
                    f"**Completion Date:** {completion_date}\n\n"
                    f"**Official Title:** {official_title}\n"
                    f"**Type:** {study_type.title()}\n"
                    f"**Phase:** {', '.join(phases) if phases else 'Not reported'}\n"
                    f"**Status:** {status}\n"
                    f"**Countries:** {', '.join(countries) if countries else 'N/A'}\n"
                    f"**Locations:**\n{location_text}\n"
                    f"**Interventions:** {intervention_text}\n"
                    f"**Sponsor:** {sponsor}\n"
                    f"**Collaborators:** {', '.join(collaborator_names) if collaborator_names else 'None'}\n"
                )

            return "\n\n---\n\n".join(result).strip()

        return f"API returned error: {response.status_code}"

    except Exception as e:
        print("Exception occurred:", e)
        return "Error fetching study data."

# %%
# There's a particular dictionary structure that's required to describe our function:

search_function = {
    "name": "search_studies",
    "description": "Search for clinical trials with strict filtering on all key metadata fields such as condition, country, phase, study type, sex, age group, sampling method, sponsor, intervention, locations, start dates, completion dates, etc.",
    "parameters": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "Condition or keyword to search for. (e.g., 'lung cancer', 'IBD')"
            },
            "phase": {
                "type": "string",
                "description": "Clinical trial phase. (e.g., 'Phase 1', 'Phase 2', 'Phase 3')"
            },
            "status": {
                "type": "string",
                "description": "Recruitment status. (e.g., 'RECRUITING', 'COMPLETED')"
            },
            "country": {
                "type": "string",
                "description": "Country where the trial is conducted. (e.g., 'Italy')"
            },
            "study_type": {
                "type": "string",
                "description": "Type of study. (e.g., 'INTERVENTIONAL', 'OBSERVATIONAL')"
            },
            "sex": {
                "type": "string",
                "description": "Sex eligibility. (e.g., 'Male', 'Female', 'All')"
            },
            "age_group": {
                "type": "string",
                "description": "Standard age group. (e.g., 'CHILD', 'ADULT', 'OLDER_ADULT')"
            },
            "sampling_method": {
                "type": "string",
                "description": "Participant sampling method. (e.g., 'PROBABILITY_SAMPLE', 'NON_PROBABILITY_SAMPLE')"
            },
            "intervention": {
                "type": "string",
                "description": "Intervention or treatment keyword. (e.g., 'aspirin', 'TAE')"
            },
            "sponsor": {
                "type": "string",
                "description": "Name of the lead sponsor or organization. (e.g., 'Pfizer', 'NIH')"
            },
            "ipd_sharing": {
                "type": "string",
                "description": "Will individual participant data (IPD) be shared? (e.g., 'YES', 'NO', 'UND')"
            },
            "city": {
                "type": "string",
                "description": "City where the trial site is located. (e.g., 'Chicago')"
            },
            "facility": {
                "type": "string",
                "description": "Facility or hospital name where trial is conducted. (e.g., 'Mayo Clinic')"
            },
            "start_date_from": {
                "type": "string",
                "description": "Earliest start date allowed (format: YYYY-MM or YYYY-MM-DD)"
            },
            "start_date_to": {
                "type": "string",
                "description": "Latest start date allowed"
            },
            "completion_date_from": {
                "type": "string",
                "description": "Earliest completion date allowed"
            },
            "completion_date_to": {
                "type": "string",
                "description": "Latest completion date allowed"
            },
            "max_results": {
                "type": "integer",
                "description": "Maximum number of studies to return"
            }
        },
        "required": ["query"],
        "additionalProperties": False
    }
}

# %%
# And this is included in a list of tools:

tools = [{"type": "function", "function": search_function}]

# %%
def chat(message, history):
    messages = [{"role": "system", "content": system_message}] + history + [{"role": "user", "content": message}]

    # 🔄 First attempt: try to stream the LLM output
    response_stream = openai.chat.completions.create(
        model=MODEL,
        messages=messages,
        tools=tools,
        tool_choice="auto",
        stream=True
    )

    full_response = ""
    tool_call_detected = False

    for chunk in response_stream:
        choice = chunk.choices[0]
        delta = choice.delta

        # 🧠 Detect tool call request during stream
        if hasattr(delta, "tool_calls") and delta.tool_calls:
            tool_call_detected = True
            break  # Exit streaming — can't continue past tool call

        if delta.content:
            full_response += delta.content
            yield full_response  # Live stream to user

    # 🧰 Tool call fallback (non-streamed)
    if tool_call_detected:
        fallback = openai.chat.completions.create(
            model=MODEL,
            messages=messages,
            tools=tools,
            tool_choice="auto"  # No stream here, required to get tool_calls
        )

        message = fallback.choices[0].message
        print("Finish reason:", fallback.choices[0].finish_reason)
        print("Tool calls:", message.tool_calls if hasattr(message, 'tool_calls') else None)

        # 🔧 Call the tool(s)
        tool_responses = handle_tool_call(message)

        # Add the assistant tool call message and all corresponding tool responses
        messages.append(message)
        messages.extend(tool_responses)

        # 🧠 Now ask GPT to summarize the tool result(s)
        final_response_stream = openai.chat.completions.create(
            model=MODEL,
            messages=messages,
            stream=True
        )

        final_output = ""
        for chunk in final_response_stream:
            delta = chunk.choices[0].delta
            if delta.content:
                final_output += delta.content
                yield final_output  # Stream final GPT summary

    # 🧯 Final fallback if nothing streamed
    elif not full_response:
        fallback = openai.chat.completions.create(
            model=MODEL,
            messages=messages,
            tools=tools,
            tool_choice="auto"
        )
        yield fallback.choices[0].message.content

# %%
def handle_tool_call(message):
    import json

    tool_responses = []

    for tool_call in message.tool_calls:
        arguments = json.loads(tool_call.function.arguments)
        result = search_studies(arguments)

        tool_responses.append({
            "role": "tool",
            "tool_call_id": tool_call.id,
            "content": result if isinstance(result, str) else json.dumps(result)
        })

    return tool_responses

# %%
example_prompts = [
    "Show me trials in Taiwan studying Vedolizumab.",
    "List studies for Crohn's disease that started after 2015.",
    "Give me 5 completed trials on lung cancer in Japan.",
    "Provide latest Perennial Allergic Rhinitis study from Eli Lily",
    "Find interventional Phase 3 studies for breast cancer in France.",
    "List observational studies in Thailand with female participants over 65.",
    "Show me trials in United States, Houston, at MD Anderson.",
    "Show studies that started after 2022 for Asthma and are still ongoing.",
    "Show me studies that use HS135 in Canada.",
    "Show me trials that share individual participant data (IPD) in the USA.",
    "Get trial details for NCT06083857.",
    "Find completed prostate cancer trials for adults in Germany.",
    "List Phase 4 trials with probability sampling in South Korea."
]

description = """\
<div style="font-family: sans-serif; line-height: 1.8;">
  <strong>Ask about any of the following criteria to find clinical trials:</strong><br><br>

  <style>
    table.ctg-guide-table, table.ctg-guide-table td {
      border: none !important;
      padding: 6px;
      vertical-align: top;
    }
  </style>

  <table class="ctg-guide-table" style="width: 100%; table-layout: fixed; border-collapse: collapse;">
    <tr>
      <td>🩺 <strong>Medical Conditions</strong><br><small>e.g., 'lung cancer', 'prostate cancer'</small></td>
      <td>🔢 <strong>NCT ID</strong><br><small>e.g., 'NCT01234567'</small></td>
      <td>🧪 <strong>Trial Phase</strong><br><small>e.g., 'Phase 1', 'Phase 2', 'Phase 3', 'Phase 4'</small></td>
    </tr>
    <tr>
      <td>🧫 <strong>Study Type</strong><br><small>e.g., 'INTERVENTIONAL', 'OBSERVATIONAL'</small></td>
      <td>📋 <strong>Status</strong><br><small>e.g., 'RECRUITING', 'COMPLETED', 'TERMINATED'</small></td>
      <td>💊 <strong>Interventions</strong><br><small>e.g., 'aspirin', 'HS135', 'Vedolizumab'</small></td>
    </tr>
    <tr>
      <td>🏢 <strong>Sponsor</strong><br><small>e.g., 'Pfizer', 'NIH', 'Amgen'</small></td>
      <td>🚻 <strong>Sex</strong><br><small>e.g., 'Male', 'Female', 'All'</small></td>
      <td>🧍 <strong>Age Group</strong><br><small>e.g., 'CHILD', 'ADULT', 'OLDER_ADULT'</small></td>
    </tr>
    <tr>
      <td>🎯 <strong>Sampling Method</strong><br><small>e.g., 'PROBABILITY_SAMPLE', 'NON_PROBABILITY_SAMPLE'</small></td>
      <td>📤 <strong>IPD Sharing</strong><br><small>e.g., 'YES', 'NO', 'UNDECIDED' (Individual Participant Data)</small></td>
      <td>🌍 <strong>Location</strong><br><small>e.g., country, city, or facility</small></td>
    </tr>
    <tr>
      <td colspan="3">📅 <strong>Start/Completion Date</strong><br><small>e.g., '2020', '2020-05', or '2020-05-20'</small></td>
    </tr>
  </table>

  <br>
  💬 Ask your question naturally, but include keywords like condition, location, phase, sex, or sponsor for best results.<br>
  You can combine filters (e.g., <em>'recruiting lung cancer trials in Canada by Pfizer'</em>) for more precise answers.<br><br>
  💡 <em>Try one of the examples below to get started!</em>
</div>
"""

gr.ChatInterface(
    fn=chat,  # Your actual function
    type="messages",
    title="ClinicalTrials.gov Agent",
    description=description,
    chatbot=gr.Chatbot(label="CTGagent", type="messages", height=1000),
    examples=example_prompts
).launch(app_kwargs={"title": "CTGagent"})

# %%