main.py

######## Imports ########
#########################

from fasthtml.common import (
    Script,
    fast_app,
    Div,
    H1,
    P,
    Meta,
    Form,
    Input,
    serve,
    EventStream,
    Dialog,
    Button,
    Hidden,
    Title,
    Img,
    Span,
    A,
    Br,
    Link,
    H2,
    HttpHeader,
)

from dotenv import load_dotenv
from fh_heroicons import Heroicon
import os
import mlflow

from src.lib import setup_db, ConversationManager

from src.components import (
    QuestionComponent,
    DocumentationComponent,
    AnswerComponent,
    Hero,
    Navbar,
    Features,
    ExampleQuestions,
    CTA,
    Footer,
    bg_style,
)

from src.utils import answer_query_sse

load_dotenv()

######## DB STUFF ########
##########################

c = setup_db(os.getenv("DB_PATH", "data/cache.db"))

######## Conversation Manager ########
######################################
conversation_manager = ConversationManager()

######## MLFlow Setup ########
##################################
mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI", "http://127.0.0.1:5000"))
mlflow.set_experiment(os.getenv("MLFLOW_EXPERIMENT", "wilhelmai-dev"))
mlflow.dspy.autolog()


######## FastHTML App Init ########
###################################


fonts = (
    Link(rel="preconnect", href="https://fonts.googleapis.com"),
    Link(rel="preconnect", href="https://fonts.gstatic.com", crossorigin=True),
    Link(
        href="https://fonts.googleapis.com/css2?family=Geist:wght@100..900&display=swap",
        rel="stylesheet",
    ),
)

description = "wilhelm.ai - Revolutionizing Radiology Education"

meta_tags = (
    Meta(name="description", content=description),
    Meta(property="og:title", content=description),
    Meta(property="og:description", content=description),
    Meta(property="og:url", content="https://radrag-production.up.railway.app/"),
    Meta(property="og:type", content="website"),
    Meta(name="twitter:card", content="summary"),
    Meta(name="twitter:title", content=description),
    Meta(name="twitter:description", content=description),
)

twcss = Script(src="https://cdn.tailwindcss.com/?plugins=typography")
sse = Script(src="https://unpkg.com/htmx-ext-sse@2.2.2/sse.js")
app, rt = fast_app(
    pico=False,
    hdrs=(twcss, sse, *fonts, *meta_tags),
    live=os.getenv("DEVELOPMENT", False),
    bodykw={"style": bg_style + "font-family: 'Geist', sans-serif;"},
)


######## Routes ########
########################


@rt("/app")
def get():
    @rt
    def ask(query: str, conv_id: str = None):
        if not query or not query.strip():
            return Div(cls="m-4 text-zinc-800")("The query cannot be empty.")

        is_followup = bool(conv_id)

        conv_id, _ = conversation_manager.get_or_create_conversation(conv_id)

        ac = AnswerComponent(
            hx_ext="sse",
            sse_connect=receive_answer.to(query=query, conv_id=conv_id),
            sse_swap="message",
            sse_close="close",
            hx_swap="innerHTML show:bottom",
        )

        if is_followup:
            response = (
                QuestionComponent(query),
                ac,
                HttpHeader("Hx-Reswap", "beforeend"),
            )

            return response

        response = (
            QuestionComponent(query),
            ac,
            SubmitForm(conv_id=conv_id, hx_swap_oob="true"),
        )

        return response

    @rt
    def receive_answer(
        query: str,
        conv_id: str,
    ):
        return EventStream(answer_query_sse(query, conv_id, conversation_manager))

    def ExampleQuestion(qtext=""):
        return Div(
            Form(hx_post=ask.to(), hx_target="#result", hx_trigger="click")(
                Hidden(name="query", value=qtext),
                Hidden(name="conv_id", value=""),
                qtext,
            ),
            cls="p-4 bg-white rounded border border-zinc-200 cursor-pointer hover:opacity-70",
        )

    example_queries = [
        "How do I differentiate between type 1 and type 2 endoleaks on CTA?",
        "what critical shoulder angle is normal?",
        "how can I differentiate hepatic adenoma from FNH?",
        "how do I classify hemorrhagic transformation of ischemic stroke?",
    ]

    def SubmitForm(conv_id: str = None, **kwargs):
        return Form(
            hx_post=ask.to(),
            hx_target="#result",
            hx_swap="innerHTML",
            id="submitform",
            cls="flex gap-2 mt-4 border-t border-t-zinc-200 pt-4",
            **{"hx-on::after-request": "this.reset()"},
            **kwargs,
        )(
            Input(
                type="text",
                name="query",
                placeholder="Ask here ...",
                required=True,
                cls="border rounded border-zinc-200 p-2 bg-white flex-1",
            ),
            Hidden(name="conv_id", value=conv_id if conv_id else ""),
        )

    return Title(description), Div(
        cls="max-w-[680px] mx-auto p-4 max-h-screen flex flex-col"
    )(
        Div(cls="flex justify-between items-center mb-4")(
            H1(cls="text-3xl font-bold text-zinc-800 flex gap-2 items-center")(
                Heroicon("cube-transparent", cls="text-green-500 w-6 h-6"),
                A(href="/")("wilhelm", Span(".ai", cls="text-green-500")),
            ),
            Button(
                Heroicon("information-circle", cls="w-6 h-6 text-zinc-800"),
                Script("me().on('click', () => me('#info-modal').showModal())"),
                cls="hover:opacity-70",
            ),
            A(
                href="/app",
                cls="flex gap-2 border border-zinc-200 bg-zinc-100 text-zinc-800 p-2 rounded font-semibold hover:opacity-70",
            )(Heroicon("plus", cls="w-6 h-6"), "New Chat"),
        ),
        Dialog(
            id="info-modal",
            cls="p-2 max-w-sm bg-zinc-100 border rounded border-zinc-200",
        )(
            Button(
                Heroicon("x-mark"),
                Script("me().on('click', () => me('#info-modal').close())"),
                cls="absolute top-2 right-2 cursor-pointer",
            ),
            DocumentationComponent(),
        ),
        Div(
            cls="space-y-2 p-2 text-zinc-800 bg-zinc-100 border border-zinc-200 rounded flex-1 min-h-0 flex flex-col"
        )(
            Div(id="result", cls="rounded overflow-auto flex-1 min-h-0")(
                Div(cls="flex py-4 gap-2")(
                    Img(
                        src="assets/roentgen_sketch.svg",
                        cls="hidden md:block w-32 mx-auto",
                        id="wilhelm",
                    ),
                    Div(
                        H1(
                            "Reliable Radiology Knowledge, ",
                            Br(),
                            "One Question Away.",
                            cls="text-zinc-800 font-semibold text-2xl mb-2",
                        ),
                        P(
                            "Wilhelm.ai finds and summarizes authoritative Radiopaedia content to answer your questions — fast, transparent, and educational.",
                            cls="text-zinc-400 text-sm",
                        ),
                    ),
                ),
                Div(cls="w-full border-1 border-t border-zinc-200 my-2")(),
                Div(cls="my-4")(
                    H2(
                        "Explore Example Questions",
                        cls="text-lg text-zinc-800 font-semibold",
                    ),
                    P(cls="text-sm text-zinc-400")(
                        "Click on a question to get a trusted, AI-powered answer."
                    ),
                ),
                Div(cls="grid grid-cols-2 gap-2 text-zinc-800")(
                    *[ExampleQuestion(q) for q in example_queries],
                ),
            ),
            Div(cls="shrink-0")(SubmitForm()),
        ),
        P(cls="text-xs text-zinc-400 text-center mt-4")(
            "wilhelm.ai is an educational tool and does not replace clinical judgment. Always verify findings with official guidelines and medical literature."
        ),
    )


@rt
def index():
    return (
        Title(description),
        Div(style=bg_style)(
            Navbar(), Hero(), Features(), ExampleQuestions(), CTA(), Footer()
        ),
    )


serve()