app.py

import logging
import os
import shutil
import sys
from pathlib import Path

import gradio as gr
from mcp import StdioServerParameters
from smolagents import InferenceClientModel
from smolagents.mcp_client import MCPClient

from src.upgrade_advisor.config import (  # noqa: E402
    AGENT_MODEL,
    CHAT_HISTORY_TURNS_CUTOFF,
    CHAT_HISTORY_WORD_CUTOFF,
    GITHUB_PAT,
    GITHUB_READ_ONLY,
    GITHUB_TOOLSETS,
)

from src.upgrade_advisor.agents.package import PackageDiscoveryAgent  # noqa: E402
from src.upgrade_advisor.chat.chat import (  # noqa: E402
    qn_rewriter,
    run_document_qa,
    summarize_chat_history,
)
from src.upgrade_advisor.misc import (  # noqa: E402
    _monkeypatch_gradio_save_history,
    get_example_questions,
)
from src.upgrade_advisor.theme import Christmas  # noqa: E402

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.addHandler(logging.StreamHandler())

# this is to use the gradio-upload-mcp server for file uploads
uploads_dir = Path("uploads")
uploads_dir.mkdir(exist_ok=True)
uploads_dir = uploads_dir.resolve()

_monkeypatch_gradio_save_history()


def get_agent_model(model_name: str, oauth_token: gr.OAuthToken = None):
    token = os.getenv("HF_TOKEN", None) or oauth_token.token if oauth_token else None
    # provider = os.getenv("HF_INFERENCE_PROVIDER", "together")
    model = InferenceClientModel( 
        model_id=model_name, 
        token=token,
        timeout=1000,
    )
    return model


async def chat_fn(
    message,
    history,
    persisted_attachments=None,
    profile: gr.OAuthProfile = None,
    oauth_token: gr.OAuthToken = None,
):
    # parse incoming history is a list of dicts with 'role' and 'content' keys
    from datetime import datetime

    token = os.getenv("HF_TOKEN", None) or oauth_token.token if oauth_token else None
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    logger.info(f"Received message: {message}")
    logger.info(f"History: {history}")
    if len(history) > 0:
        summarized_history = await summarize_chat_history(
            history,
            turns_cutoff=CHAT_HISTORY_TURNS_CUTOFF,
            word_cutoff=CHAT_HISTORY_WORD_CUTOFF,
            token=token,
        )
    else:
        summarized_history = ""
    incoming_attachments = message.get("files", []) if isinstance(message, dict) else []
    persisted_attachments = persisted_attachments or []
    # If no new attachments are provided, keep using the previously persisted ones.
    attachments = incoming_attachments or persisted_attachments
    latest_attachment = attachments[-1] if attachments else []

    logger.info(f"Summarized chat history:\n{summarized_history}")
    logger.info(f"With attachments: {attachments} (incoming: {incoming_attachments})")
    logger.info(f"Latest attachment: {latest_attachment}")
    logger.info(f"Persisted attachments: {persisted_attachments}")

    # if attachements are present message is a dict with 'text' and 'files' keys
    message = message.get("text", "") if isinstance(message, dict) else message
    # overwrite messages with the text content only
    message = message.strip()
    rewritten_message, is_rewritten_good = await qn_rewriter(
        message,
        summarized_history,
        token=token,
    )

    if is_rewritten_good:
        logger.info(f"Rewritten question: {rewritten_message}")
    else:
        logger.info(f"Using original question: {message}")
        rewritten_message = None
    # Collect events from the agent run
    # add chat summary to message
    message = f"""
        CHAT SUMMARY SO FAR:
        {summarized_history}
        CURRENT QUESTION FROM USER:
        {message}
        """
    if len(attachments) > 0:
        message += """Attached FILE:\n"""
        # use the last file from the list of files only, as
        # the single file is expected to be a pyproject.toml
        # copy to uploads directory
        if latest_attachment:
            # take the last uploaded file
            source_file = latest_attachment
            file_name = f"{timestamp}_{Path(latest_attachment).name}"
        elif len(persisted_attachments) > 0:
            # take the last persisted file if no new uploads
            source_file = persisted_attachments[-1]
            file_name = f"{timestamp}_{Path(persisted_attachments[-1]).name}"
        else:
            source_file = None
            file_name = None

        logger.info(f"Copying uploaded file {source_file} to {uploads_dir}")
        shutil.copy(source_file, uploads_dir / file_name)
        message += f"""
            FILE PATH: {uploads_dir / file_name}\n
            """
    logger.info(f"Final message to agent:\n{message}")
    # Run the package discovery agent to build context
    context = agent.discover_package_info(
        user_input=message,
        reframed_question=rewritten_message,
    )
    # Build a concise context from tool outputs
    logger.info(f"Built context of length {len(context)}")
    logger.info(f"Context content:\n{context}")
    # Run a document QA pass using the user's question
    qa_answer = await run_document_qa(
        question=message,
        context=context,
        rewritten_question=rewritten_message,
        token=token,
    )
    logger.info(f"QA answer: {qa_answer}")
    yield (
        {
            "role": "assistant",
            "content": qa_answer,
        },
        attachments,
    )


def main():
    logger.info("Starting MCP client...")

    try:
        gh_mcp_params = dict(
            url="https://api.githubcopilot.com/mcp/",
            transport="streamable-http",
            headers={
                "Authorization": f"Bearer {GITHUB_PAT}",
                "X-MCP-Toolsets": GITHUB_TOOLSETS,
                "X-MCP-Readonly": GITHUB_READ_ONLY,
            },
        )
        # pypi_mcp_params = dict(
        #     # url="https://mcp-1st-birthday-pypi-mcp.hf.space/gradio_api/mcp/",
        #     url="https://mcp-1st-birthday-uv-pypi-mcp.hf.space/gradio_api/mcp/",
        #     transport="streamable-http",
        # )
        upload_mcp_params = StdioServerParameters(
            command="uvx",
            args=[
                "--from",
                "gradio[mcp]",
                "gradio",
                "upload-mcp",
                # Base must be the Gradio root; upload-mcp adds
                # /gradio_api/upload.
                # The docs are misleading here, it has gradio_api/upload as the base.
                "https://mcp-1st-birthday-uv-pypi-mcp.hf.space/",
                uploads_dir.as_posix(),
            ],
        )

        # Gradio chat interface state to persist uploaded files
        files_state = gr.State([])
        example_questions = get_example_questions(n=4)
        christmas = Christmas()

        with MCPClient(
            server_parameters=[
                gh_mcp_params,
                upload_mcp_params,
            ],
            structured_output=True,
        ) as toolset:
            logger.info("MCP clients connected successfully")

            global agent
            model = get_agent_model(model_name=AGENT_MODEL)
            agent = PackageDiscoveryAgent(
                model=model,
                tools=toolset,
            )
            # rewrite with Blocks
            with gr.Blocks() as demo:
                gr.LoginButton()
                gr.Markdown(f"""
                    # 💻 FixMyEnv: Package Upgrade Advisor 🚀🔧🐍📦⚙️

                    Welcome to the Package Upgrade Advisor!
                    This AI-powered assistant helps you identify and resolve
                    outdated or vulnerable packages in your Python projects.
                    Simply ask a question about package upgrades, and if you
                    have a `pyproject.toml` or `requirements.txt` file, feel free
                    to attach it for more tailored advice.
                    """, 
                    container=True,
                ) 
                with gr.Accordion("How to Use", open=False):
                    gr.Markdown(f""" 
                        1. Type your question in the chat box below.
                        2. (Optional) Attach your `pyproject.toml` or `requirements.txt`
                        file using the upload button. Uploaded files are
                        immediately removed after the session ends.
                        3. Click "Submit" and wait for the AI to analyze your query
                        and provide recommendations.

                        Note: The assistant uses Huggingface Inference API for
                        [{AGENT_MODEL}](https://huggingface.co/{AGENT_MODEL}) LLM
                        capabilities with Smolagents Tool calling and GitHub MCP
                        for package data retrieval. Huggingface login is therefore
                        required to use the app. This gradio app serves as an MCP Server
                        as well!
                        """, 
                    )
                gr.ChatInterface(
                    fn=chat_fn,
                    chatbot=gr.Chatbot(
                        height=600,
                    ),
                    textbox=gr.MultimodalTextbox(
                        label="pyproject.toml or requirements.txt file can be attached",
                        file_types=[".toml", ".txt"],
                        file_count="single",
                        min_width=100,
                        sources="upload",
                    ),
                    additional_inputs=[files_state],
                    additional_outputs=[files_state],
                    save_history=True,
                    examples=example_questions,
                    stop_btn=True,
                    cache_examples=False,
                )
            demo.launch(mcp_server=True, share=False, theme=christmas)

    finally:
        logger.info("Cleaning up MCP client resources")
        # remove contents of uploads_dir
        for f in uploads_dir.iterdir():
            if f.is_dir():
                try:
                    shutil.rmtree(f)
                except Exception:
                    logger.exception(f"Failed to delete uploaded directory: {f}")
            else:
                try:
                    f.unlink()
                except Exception:
                    logger.exception(f"Failed to delete uploaded file: {f}")
        logger.info("Shutdown complete.")


if __name__ == "__main__":
    main()