app.py

import os
import re
from datetime import datetime
from typing import Any

import gradio as gr
from openai import OpenAI


NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY", "").strip()
NVIDIA_BASE_URL = os.getenv(
    "NVIDIA_BASE_URL",
    "https://integrate.api.nvidia.com/v1",
).strip()
MODEL_NAME = os.getenv(
    "MODEL_NAME",
    "nvidia/nemotron-3-ultra-550b-a55b",
).strip()

MAX_ITERATIONS = int(os.getenv("MAX_ITERATIONS", "10"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "16384"))
TEMPERATURE = float(os.getenv("TEMPERATURE", "1"))
TOP_P = float(os.getenv("TOP_P", "0.95"))
REASONING_BUDGET = int(os.getenv("REASONING_BUDGET", "16384"))


PERSONALITY = """
You are WVY, a curious deep reasoning agent designed to think & update your
memory consistently to evolve your understanding by using mental notes.

When giving a response use your tools to make a decision. Use mental notes to
prompt yourself & adapt to complex situations. Ask for help when needed, and
do not rush into action. Use your mental notes to plan executions step by step.

Do not try to answer questions immediately. Ask questions before giving your
final answer when the request is complex.

Do not rely only on pretrained patterns. Reason using first-principles thinking
based on your observations.
""".strip()


def new_session_state() -> dict[str, Any]:
    return {
        "chat_thread": [],
        "mental_notes": {},
    }


def parse_mental_notes(
    response: str,
    mental_notes: dict[int, str],
) -> tuple[str, dict[int, str], list[str]]:
    mental_notes = dict(mental_notes)

    mental_note_pattern = re.compile(
        r"<mentalnote:(?P<number>\d+)>\s*"
        r"(?P<context>.*?)\s*"
        r"</mentalnote:(?P=number)>",
        re.DOTALL | re.IGNORECASE,
    )

    delete_note_pattern = re.compile(
        r"<deletenote:(?P<number>\d+)\s*>",
        re.IGNORECASE,
    )

    actions = []

    for match in mental_note_pattern.finditer(response):
        note_number = int(match.group("number"))
        note_context = match.group("context").strip()

        action = "Edited" if note_number in mental_notes else "Created"

        mental_notes[note_number] = note_context
        actions.append(f"{action} mental note {note_number}")

    for match in delete_note_pattern.finditer(response):
        note_number = int(match.group("number"))

        if note_number in mental_notes:
            del mental_notes[note_number]
            actions.append(f"Deleted mental note {note_number}")

    clean_response = mental_note_pattern.sub("", response)
    clean_response = delete_note_pattern.sub("", clean_response).strip()

    return clean_response, mental_notes, actions


def chat_with_user(chat_thread: list[dict[str, str]]) -> str:
    if not chat_thread:
        return "No conversation yet."

    return "\n\n".join(
        f"{message['role'].capitalize()}: {message['content']}"
        for message in chat_thread
    )


def mental_notes_text(mental_notes: dict[int, str]) -> str:
    if not mental_notes:
        return "No mental notes have been created."

    return "\n".join(
        f"{number}. {mental_notes[number]}"
        for number in sorted(mental_notes)
    )


def tools_text() -> str:
    return f"""
When you generate a response, any message outside of a tool may be sent to the
user.

To make a new mental note or edit an existing note, use:

<mentalnote:insert-number-here>
insert-context-of-memory-here
</mentalnote:insert-number-here>

To delete a mental note, use:

<deletenote:insert-number-here>

Update your memory consistently to evolve your understanding.

To end your response, use:

<endresponse>

You can generate up to {MAX_ITERATIONS} outputs before your final response.
""".strip()


def perspective(
    chat_thread: list[dict[str, str]],
    mental_notes: dict[int, str],
) -> str:
    current_time = datetime.now().strftime("%Y-%m-%d %I:%M:%S %p")

    return f"""
This is the current time:
{current_time}

This is the chat with the user:
{chat_with_user(chat_thread)}

These are your mental notes:
{mental_notes_text(mental_notes)}

These are your tools:
{tools_text()}
""".strip()


def notes_for_display(mental_notes: dict[int, str]) -> str:
    if not mental_notes:
        return "_No mental notes have been created._"

    return "\n\n".join(
        f"### Mental Note {number}\n{mental_notes[number]}"
        for number in sorted(mental_notes)
    )


def get_client() -> OpenAI:
    if not NVIDIA_API_KEY:
        raise RuntimeError(
            "NVIDIA_API_KEY is missing. Add it as a Secret in the "
            "Hugging Face Space settings."
        )

    return OpenAI(
        base_url=NVIDIA_BASE_URL,
        api_key=NVIDIA_API_KEY,
    )


def call_model(messages: list[dict[str, str]]) -> str:
    client = get_client()

    completion = client.chat.completions.create(
        model=MODEL_NAME,
        messages=messages,
        temperature=TEMPERATURE,
        top_p=TOP_P,
        max_tokens=MAX_TOKENS,
        extra_body={
            "chat_template_kwargs": {
                "enable_thinking": True,
            },
            "reasoning_budget": REASONING_BUDGET,
        },
        stream=True,
    )

    response_parts = []

    for chunk in completion:
        if not chunk.choices:
            continue

        content = chunk.choices[0].delta.content

        if content is not None:
            response_parts.append(content)

    return "".join(response_parts)


def generate_response(
    user_message: str,
    chatbot_history: list[dict[str, str]] | None,
    session_state: dict[str, Any] | None,
):
    user_message = (user_message or "").strip()
    chatbot_history = list(chatbot_history or [])
    session_state = dict(session_state or new_session_state())

    if not user_message:
        return (
            chatbot_history,
            "",
            notes_for_display(
                session_state.get("mental_notes", {})
            ),
            session_state,
            "",
        )

    chat_thread = list(
        session_state.get("chat_thread", [])
    )

    mental_notes = dict(
        session_state.get("mental_notes", {})
    )

    chat_thread.append({
        "role": "user",
        "content": user_message,
    })

    chatbot_history.append({
        "role": "user",
        "content": user_message,
    })

    iteration_responses = []
    visible_responses = []
    status_lines = []

    try:
        for iteration_number in range(
            1,
            MAX_ITERATIONS + 1,
        ):
            status_lines.append(
                f"Iteration {iteration_number}/{MAX_ITERATIONS}"
            )

            if iteration_responses:
                previous_iterations = "\n\n".join(
                    f"Iteration {index + 1}:\n{response}"
                    for index, response in enumerate(
                        iteration_responses
                    )
                )
            else:
                previous_iterations = (
                    "No previous iterations exist "
                    "for this response."
                )

            iteration_context = f"""
This is iteration {iteration_number} of a maximum of {MAX_ITERATIONS}.

These are your previous iterations for this response:

{previous_iterations}

Continue processing the user's request.

You may create, edit, or delete mental notes when needed.

When the complete response is finished, include:

<endresponse>
""".strip()

            response = call_model([
                {
                    "role": "system",
                    "content": PERSONALITY,
                },
                {
                    "role": "user",
                    "content": perspective(
                        chat_thread,
                        mental_notes,
                    ),
                },
                {
                    "role": "user",
                    "content": iteration_context,
                },
            ])

            end_response_found = bool(
                re.search(
                    r"<endresponse\s*>",
                    response,
                    flags=re.IGNORECASE,
                )
            )

            response_without_end = re.sub(
                r"<endresponse\s*>",
                "",
                response,
                flags=re.IGNORECASE,
            ).strip()

            clean_response, mental_notes, _ = (
                parse_mental_notes(
                    response_without_end,
                    mental_notes,
                )
            )

            iteration_responses.append(clean_response)

            if clean_response:
                visible_responses.append(clean_response)

            if end_response_found:
                break

        final_response = "\n\n".join(
            visible_responses
        ).strip()

        if not final_response:
            final_response = (
                "The response ended without visible content."
            )

    except Exception as error:
        final_response = f"Error: {error}"

    chat_thread.append({
        "role": "assistant",
        "content": final_response,
    })

    chatbot_history.append({
        "role": "assistant",
        "content": final_response,
    })

    session_state = {
        "chat_thread": chat_thread,
        "mental_notes": mental_notes,
    }

    return (
        chatbot_history,
        "",
        notes_for_display(mental_notes),
        session_state,
        " · ".join(status_lines),
    )


def reset_session():
    state = new_session_state()

    return (
        [],
        "",
        notes_for_display({}),
        state,
        "",
    )


CSS = """
.gradio-container {
    max-width: 1500px !important;
}

#notes-panel {
    height: 620px;
    overflow-y: auto;
}

#status-line {
    min-height: 24px;
    opacity: 0.7;
}
"""


with gr.Blocks(
    title="Starpower SDK — Perspective Feedback Loop",
) as demo:
    session_state = gr.State(new_session_state())

    gr.Markdown(
        """
# Starpower SDK
### Conceptual Introspection — Direct Feedback Loop

WVY uses an iterative perspective prompt and editable mental notes to process
each request.
"""
    )

    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="Chat",
                height=520,
                placeholder="Send a message to WVY.",
            )

            user_input = gr.Textbox(
                label="Message",
                placeholder="Send a message to WVY...",
                lines=3,
            )

            with gr.Row():
                send_button = gr.Button(
                    "Send",
                    variant="primary",
                )

                reset_button = gr.Button(
                    "Reset session"
                )

            status = gr.Markdown(
                "",
                elem_id="status-line",
            )

        with gr.Column(scale=1):
            mental_notes_display = gr.Markdown(
                notes_for_display({}),
                label="Mental Notes",
                elem_id="notes-panel",
            )

    submit_event = send_button.click(
        fn=generate_response,
        inputs=[
            user_input,
            chatbot,
            session_state,
        ],
        outputs=[
            chatbot,
            user_input,
            mental_notes_display,
            session_state,
            status,
        ],
        show_progress="full",
        concurrency_limit=1,
    )

    user_input.submit(
        fn=generate_response,
        inputs=[
            user_input,
            chatbot,
            session_state,
        ],
        outputs=[
            chatbot,
            user_input,
            mental_notes_display,
            session_state,
            status,
        ],
        show_progress="full",
        concurrency_limit=1,
    )

    reset_button.click(
        fn=reset_session,
        inputs=[],
        outputs=[
            chatbot,
            user_input,
            mental_notes_display,
            session_state,
            status,
        ],
        cancels=[submit_event],
    )


if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=1
    ).launch(
        css=CSS
    )