app.py

from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Custom Tool to fetch datasets related to body parts or imaging types

@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """
    Search and retrieve publicly available medical datasets from Hugging Face based on any medical-related keyword.

    Args:
        arg1: A keyword related to medical data (e.g., 'cancer', 'diabetes', 'CT scan', 'radiology', 'dermoscopy').
        arg2: The maximum number of datasets to retrieve.

    Returns:
        A list of dataset names matching the search query, or a message stating that no datasets were found.
    """
    try:
        keyword = arg1.strip().lower()
        limit = int(arg2)

        # Define a basic list of medically relevant terms
        medical_terms = [
            # Anatomy / Body Parts
            "skin", "brain", "lung", "chest", "abdomen", "spine", "bone", "heart", "liver", "kidney",
            "bladder", "stomach", "colon", "rectum", "esophagus", "pancreas", "breast", "ear", "eye", 
            "retina", "tooth", "teeth", "tongue", "jaw", "neck", "wrist", "hand", "leg", "arm", "shoulder", "pelvis",
        
            # Diseases / Conditions
            "cancer", "tumor", "stroke", "diabetes", "pneumonia", "covid", "asthma", "eczema", "melanoma",
            "hypertension", "alzheimer", "parkinson", "arthritis", "scoliosis", "epilepsy", "glaucoma",
            "ulcer", "hepatitis", "leukemia", "lymphoma", "tuberculosis", "anemia", "obesity", "depression",
            "anxiety", "bipolar", "autism", "adhd", "ptsd", "psychosis", "schizophrenia",
        
            # Imaging Modalities
            "mri", "ct", "xray", "x-ray", "ultrasound", "pet", "fmri", "mammo", "angiography", "radiography",
            "echocardiogram", "spect", "dermoscopy", "colonoscopy", "endoscopy", "biopsy", "histopathology",
        
            # Medical Specialties
            "radiology", "pathology", "oncology", "cardiology", "neurology", "dermatology", "dentistry",
            "ophthalmology", "urology", "orthopedics", "gastroenterology", "pulmonology", "nephrology",
            "psychiatry", "pediatrics", "geriatrics", "infectious disease",
        
            # Symptoms / Signs
            "lesion", "infection", "fever", "pain", "inflammation", "rash", "headache", "swelling", 
            "cough", "seizure", "dizziness", "vomiting", "diarrhea", "nausea", "fatigue", "itching",
        
            # Common Specific Diseases
            "breast cancer", "prostate cancer", "lung cancer", "skin cancer", "colon cancer", 
            "brain tumor", "liver cancer", "cervical cancer", "bladder cancer", "thyroid cancer",
        
            # Procedures / Interventions
            "surgery", "chemotherapy", "radiation", "transplant", "dialysis", "intubation", "stenting",
            "ventilation", "vaccination", "anesthesia", "rehabilitation", "prosthetics", "orthotics",
        
            # Lab Tests / Biomarkers
            "blood test", "cbc", "glucose", "hemoglobin", "cholesterol", "biomarker", "urinalysis",
            "pcr", "serology", "antibody", "antigen",
        
            # Clinical Settings / Roles
            "icu", "hospital", "emergency", "clinical notes", "nursing", "physician", "patient",
            "medical record", "electronic health record", "ehr", "vitals",
        
            # Age-based Terms
            "pediatric", "neonatal", "infant", "child", "adolescent", "geriatrics", "elderly",
        
            # Epidemiology / Public Health
            "epidemiology", "prevalence", "incidence", "mortality", "public health", "health disparity",
            "risk factor", "social determinant",
        
            # Pharmacology / Medications
            "drug", "medication", "pharmacology", "side effect", "adverse event", "dose", "tablet",
            "vaccine", "clinical trial", "placebo"
        ]


        # Check if keyword is in known medical terms
        if not any(term in keyword for term in medical_terms):
            return f"No medical datasets found for '{arg1}'."

        # Fetch datasets from Hugging Face
        response = requests.get(
            f"https://huggingface.co/api/datasets?search={keyword}&limit={limit}"
        )
        response.raise_for_status()
        datasets = response.json()

        # Return message if no datasets found
        if not datasets:
            return f"No medical datasets found for '{arg1}'."

        # Collect and return dataset names
        results = [f"- {ds.get('id', 'Unknown')}" for ds in datasets[:limit]]
        return f"Medical datasets related to '{arg1}':\n" + "\n".join(results)

    except Exception as e:
        return f"Error searching medical datasets for '{arg1}': {str(e)}"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    A tool that fetches the current local time in a specified timezone.

    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').

    Returns:
        A string showing the current local time in the specified timezone.
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

final_answer = FinalAnswerTool()

# Model setup
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',  # this model may be overloaded
    custom_role_conversions=None,
)

# Load tool from hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Load prompt templates
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# Create the agent
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone, my_custom_tool],  # add your tools here
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# Launch the UI
GradioUI(agent).launch()