add oauth

.env_reference

@@ -1,5 +1,5 @@
 # rename this file to .env and fill in the values below
 GITHUB_PAT = "your_github_personal_access_token_here"
 # see https://github.com/github/github-mcp-server?tab=readme-ov-file#default-toolset
-GITHUB_TOOLSETS = "default"
+GITHUB_TOOLSETS="default,discussions,experiments"
 HF_TOKEN ="hf_"

README.md

@@ -1,4 +1,23 @@
+```
+title: FixMyEnv Agent
+emoji: 🐍
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 5.49.1
+app_file: app.py
+pinned: false
+license: mit
+short_description: MCP for Agents that plan your python package upgrade
+hf_oauth: true 
+tags:
+  - building-mcp-track-enterprise
+  - building-mcp-track-customer
+  - mcp-in-action-track-customer
+  - mcp-in-action-track-enterprise
+---
 # Instructions 
+
 1. Create a scoped Personal Access Token from GitHub from https://github.com/settings/personal-access-tokens/new with the following scopes, that allows access for public repositories. 
 2. Store as GITHUB_PAT in a .env file in the root directory of the project.
 

app.py

@@ -1,8 +1,10 @@
 import logging
+import os
 import shutil
 from pathlib import Path
 
 import gradio as gr
+from huggingface_hub import whoami  # noqa: F401
 from mcp import StdioServerParameters
 from mcpadapt.core import MCPAdapt  # noqa: F401
 from smolagents import InferenceClientModel
@@ -15,7 +17,6 @@ from config import (
     GITHUB_PAT,
     GITHUB_READ_ONLY,
     GITHUB_TOOLSETS,
-    HF_TOKEN,
 )
 from src.upgrade_advisor.agents.package import PackageDiscoveryAgent
 from src.upgrade_advisor.chat.chat import (
@@ -27,7 +28,6 @@ from src.upgrade_advisor.misc import (
     _monkeypatch_gradio_save_history,
     get_example_questions,
 )
-from src.upgrade_advisor.theme import christmas
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -41,10 +41,26 @@ uploads_dir = uploads_dir.resolve()
 _monkeypatch_gradio_save_history()
 
 
-async def chat_fn(message, history, persisted_attachments=None):
+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
+    model = InferenceClientModel(
+        token=token,
+        model_id=model_name,
+    )
+    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}")
@@ -53,6 +69,7 @@ async def chat_fn(message, history, persisted_attachments=None):
             history,
             turns_cutoff=CHAT_HISTORY_TURNS_CUTOFF,
             word_cutoff=CHAT_HISTORY_WORD_CUTOFF,
+            token=token,
         )
     else:
         summarized_history = ""
@@ -72,7 +89,9 @@ async def chat_fn(message, history, persisted_attachments=None):
     # overwrite messages with the text content only
     message = message.strip()
     rewritten_message, is_rewritten_good = await qn_rewriter(
-        message, summarized_history
+        message,
+        summarized_history,
+        token=token,
     )
 
     if is_rewritten_good:
@@ -113,20 +132,27 @@ async def chat_fn(message, history, persisted_attachments=None):
     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
+        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
+        question=message,
+        context=context,
+        rewritten_question=rewritten_message,
+        token=token,
     )
     logger.info(f"QA answer: {qa_answer}")
-    return {
-        "role": "assistant",
-        "content": qa_answer,
-    }, attachments
+    yield (
+        {
+            "role": "assistant",
+            "content": qa_answer,
+        },
+        attachments,
+    )
 
 
 def main():
@@ -162,32 +188,15 @@ def main():
             ],
         )
 
-        # pypi_mcp_client = MCPClient(
-        #     server_parameters=[
-        #         # pypi_mcp_params,
-        #         gh_mcp_params,
-        #         upload_mcp_params,
-        #     ],
-        #     structured_output=True,
-        # )
-
-        model = InferenceClientModel(
-            token=HF_TOKEN,
-            model_id=AGENT_MODEL,
-        )
-
         # Gradio chat interface state to persist uploaded files
         files_state = gr.State([])
-
-        # with MCPAdapt(
-        #     serverparams=[
-        #         gh_mcp_params,
-        #         upload_mcp_params,
-        #     ],
-        #     adapter=SmolAgentsAdapter(structured_output=True),
-        # ) as toolset:
         example_questions = get_example_questions(n=4)
 
+        # TODO: use the gr.Blocks to add login blocks
+        # Add login with huggingface hub to cache token: https://www.gradio.app/guides/sharing-your-app#o-auth-login-via-hugging-face
+        # use a limited token with read-only access to public repoos only
+        # (GITHUB_PAT)
+        # deploy to hf spaces with mcp server enabled
         with MCPClient(
             server_parameters=[
                 gh_mcp_params,
@@ -198,39 +207,63 @@ def main():
             logger.info("MCP clients connected successfully")
 
             global agent
+            model = get_agent_model(model_name=AGENT_MODEL)
             agent = PackageDiscoveryAgent(
                 model=model,
                 tools=toolset,
             )
-            # link package_agent to the chat function
-
-            # attach files from local machine
-            demo = gr.ChatInterface(
-                fn=chat_fn,
-                chatbot=gr.Chatbot(
-                    height=800,
-                    type="messages",
-                ),
-                title="Package Upgrade Advisor",
-                type="messages",
-                additional_inputs_accordion="""
-                You may attach a pyproject.toml or requirements.txt file to get
-                specific upgrade advice for your project.
-                """,
-                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,
-                theme=christmas,
-            )
+            # rewrite with Blocks
+            with gr.Blocks() as demo:
+                gr.LoginButton()
+                gr.Markdown("# 💻 FixMyEnv: Package Upgrade Advisor 🚀🔧🐍📦⚙️")
+                gr.Markdown(
+                    f"""
+                    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.
+
+                    ## How to use:
+                    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,
+                    ),
+                    additional_inputs_accordion="""
+                    You may attach a pyproject.toml or requirements.txt file to get
+                    specific upgrade advice for your project.
+                    """,
+                    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,
+                    # theme=christmas,
+                )
             demo.launch(mcp_server=True, share=False)
 
     finally:

poetry.lock

@@ -129,6 +129,21 @@ files = [
     {file = "audioop_lts-0.2.2.tar.gz", hash = "sha256:64d0c62d88e67b98a1a5e71987b7aa7b5bcffc7dcee65b635823dbdd0a8dbbd0"},
 ]
 
+[[package]]
+name = "authlib"
+version = "1.6.5"
+description = "The ultimate Python library in building OAuth and OpenID Connect servers and clients."
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "authlib-1.6.5-py2.py3-none-any.whl", hash = "sha256:3e0e0507807f842b02175507bdee8957a1d5707fd4afb17c32fb43fee90b6e3a"},
+    {file = "authlib-1.6.5.tar.gz", hash = "sha256:6aaf9c79b7cc96c900f0b284061691c5d4e61221640a948fe690b556a6d6d10b"},
+]
+
+[package.dependencies]
+cryptography = "*"
+
 [[package]]
 name = "backports-asyncio-runner"
 version = "1.2.0"
@@ -799,40 +814,41 @@ tqdm = ["tqdm"]
 
 [[package]]
 name = "gradio"
-version = "5.49.1"
+version = "6.0.1"
 description = "Python library for easily interacting with trained machine learning models"
 optional = false
 python-versions = ">=3.10"
 groups = ["main"]
 files = [
-    {file = "gradio-5.49.1-py3-none-any.whl", hash = "sha256:1b19369387801a26a6ba7fd2f74d46c5b0e2ac9ddef14f24ddc0d11fb19421b7"},
-    {file = "gradio-5.49.1.tar.gz", hash = "sha256:c06faa324ae06c3892c8b4b4e73c706c4520d380f6b9e52a3c02dc53a7627ba9"},
+    {file = "gradio-6.0.1-py3-none-any.whl", hash = "sha256:0f98dc8b414a3f3773cbf3caf5a354507c8ae309ed8266e2f30ca9fa53f379b8"},
+    {file = "gradio-6.0.1.tar.gz", hash = "sha256:5d02e6ac34c67aea26b938b8628c8f9f504871392e71f2db559ab8d6799bdf69"},
 ]
 
 [package.dependencies]
 aiofiles = ">=22.0,<25.0"
 anyio = ">=3.0,<5.0"
 audioop-lts = {version = "<1.0", markers = "python_version >= \"3.13\""}
+authlib = {version = "*", optional = true, markers = "extra == \"oauth\""}
 brotli = ">=1.1.0"
 fastapi = ">=0.115.2,<1.0"
 ffmpy = "*"
-gradio-client = "1.13.3"
+gradio-client = "2.0.0"
 groovy = ">=0.1,<1.0"
 httpx = ">=0.24.1,<1.0"
 huggingface-hub = ">=0.33.5,<2.0"
+itsdangerous = {version = "*", optional = true, markers = "extra == \"oauth\""}
 jinja2 = "<4.0"
 markupsafe = ">=2.0,<4.0"
 numpy = ">=1.0,<3.0"
 orjson = ">=3.0,<4.0"
 packaging = "*"
 pandas = ">=1.0,<3.0"
-pillow = ">=8.0,<12.0"
-pydantic = ">=2.0,<2.12"
+pillow = ">=8.0,<13.0"
+pydantic = ">=2.11.10,<=2.12.4"
 pydub = "*"
 python-multipart = ">=0.0.18"
 pyyaml = ">=5.0,<7.0"
-ruff = ">=0.9.3"
-safehttpx = ">=0.1.6,<0.2.0"
+safehttpx = ">=0.1.7,<0.2.0"
 semantic-version = ">=2.0,<3.0"
 starlette = ">=0.40.0,<1.0"
 tomlkit = ">=0.12.0,<0.14.0"
@@ -841,19 +857,19 @@ typing-extensions = ">=4.0,<5.0"
 uvicorn = ">=0.14.0"
 
 [package.extras]
-mcp = ["mcp (==1.10.1)", "pydantic (>=2.11)"]
+mcp = ["mcp (>=1.21.0,<2.0.0)", "pydantic (>=2.11)"]
 oauth = ["authlib", "itsdangerous"]
 
 [[package]]
 name = "gradio-client"
-version = "1.13.3"
+version = "2.0.0"
 description = "Python library for easily interacting with trained machine learning models"
 optional = false
 python-versions = ">=3.10"
 groups = ["main"]
 files = [
-    {file = "gradio_client-1.13.3-py3-none-any.whl", hash = "sha256:3f63e4d33a2899c1a12b10fe3cf77b82a6919ff1a1fb6391f6aa225811aa390c"},
-    {file = "gradio_client-1.13.3.tar.gz", hash = "sha256:869b3e67e0f7a0f40df8c48c94de99183265cf4b7b1d9bd4623e336d219ffbe7"},
+    {file = "gradio_client-2.0.0-py3-none-any.whl", hash = "sha256:77bedf20edcc232d8e7986c1a22165b2bbca1c7c7df10ba808a093d5180dae18"},
+    {file = "gradio_client-2.0.0.tar.gz", hash = "sha256:56b462183cb8741bd3e69b21db7d3b62c5abb03c2c2bb925223f1eb18f950e89"},
 ]
 
 [package.dependencies]
@@ -862,7 +878,6 @@ httpx = ">=0.24.1"
 huggingface-hub = ">=0.19.3,<2.0"
 packaging = "*"
 typing-extensions = ">=4.0,<5.0"
-websockets = ">=13.0,<16.0"
 
 [[package]]
 name = "groovy"
@@ -1096,6 +1111,18 @@ files = [
     {file = "iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730"},
 ]
 
+[[package]]
+name = "itsdangerous"
+version = "2.2.0"
+description = "Safely pass data to untrusted environments and back."
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "itsdangerous-2.2.0-py3-none-any.whl", hash = "sha256:c6242fc49e35958c8b15141343aa660db5fc54d4f13a1db01a3f5891b98700ef"},
+    {file = "itsdangerous-2.2.0.tar.gz", hash = "sha256:e0050c0b7da1eea53ffaf149c0cfbb5c6e2e2b69c4bef22c81fa6eb73e5f6173"},
+]
+
 [[package]]
 name = "jinja2"
 version = "3.1.6"
@@ -2708,35 +2735,6 @@ files = [
     {file = "rpds_py-0.28.0.tar.gz", hash = "sha256:abd4df20485a0983e2ca334a216249b6186d6e3c1627e106651943dbdb791aea"},
 ]
 
-[[package]]
-name = "ruff"
-version = "0.14.5"
-description = "An extremely fast Python linter and code formatter, written in Rust."
-optional = false
-python-versions = ">=3.7"
-groups = ["main"]
-files = [
-    {file = "ruff-0.14.5-py3-none-linux_armv6l.whl", hash = "sha256:f3b8248123b586de44a8018bcc9fefe31d23dda57a34e6f0e1e53bd51fd63594"},
-    {file = "ruff-0.14.5-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:f7a75236570318c7a30edd7f5491945f0169de738d945ca8784500b517163a72"},
-    {file = "ruff-0.14.5-py3-none-macosx_11_0_arm64.whl", hash = "sha256:6d146132d1ee115f8802356a2dc9a634dbf58184c51bff21f313e8cd1c74899a"},
-    {file = "ruff-0.14.5-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e2380596653dcd20b057794d55681571a257a42327da8894b93bbd6111aa801f"},
-    {file = "ruff-0.14.5-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:2d1fa985a42b1f075a098fa1ab9d472b712bdb17ad87a8ec86e45e7fa6273e68"},
-    {file = "ruff-0.14.5-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:88f0770d42b7fa02bbefddde15d235ca3aa24e2f0137388cc15b2dcbb1f7c7a7"},
-    {file = "ruff-0.14.5-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:3676cb02b9061fee7294661071c4709fa21419ea9176087cb77e64410926eb78"},
-    {file = "ruff-0.14.5-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b595bedf6bc9cab647c4a173a61acf4f1ac5f2b545203ba82f30fcb10b0318fb"},
-    {file = "ruff-0.14.5-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f55382725ad0bdb2e8ee2babcbbfb16f124f5a59496a2f6a46f1d9d99d93e6e2"},
-    {file = "ruff-0.14.5-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7497d19dce23976bdaca24345ae131a1d38dcfe1b0850ad8e9e6e4fa321a6e19"},
-    {file = "ruff-0.14.5-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:410e781f1122d6be4f446981dd479470af86537fb0b8857f27a6e872f65a38e4"},
-    {file = "ruff-0.14.5-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:c01be527ef4c91a6d55e53b337bfe2c0f82af024cc1a33c44792d6844e2331e1"},
-    {file = "ruff-0.14.5-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:f66e9bb762e68d66e48550b59c74314168ebb46199886c5c5aa0b0fbcc81b151"},
-    {file = "ruff-0.14.5-py3-none-musllinux_1_2_i686.whl", hash = "sha256:d93be8f1fa01022337f1f8f3bcaa7ffee2d0b03f00922c45c2207954f351f465"},
-    {file = "ruff-0.14.5-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:c135d4b681f7401fe0e7312017e41aba9b3160861105726b76cfa14bc25aa367"},
-    {file = "ruff-0.14.5-py3-none-win32.whl", hash = "sha256:c83642e6fccfb6dea8b785eb9f456800dcd6a63f362238af5fc0c83d027dd08b"},
-    {file = "ruff-0.14.5-py3-none-win_amd64.whl", hash = "sha256:9d55d7af7166f143c94eae1db3312f9ea8f95a4defef1979ed516dbb38c27621"},
-    {file = "ruff-0.14.5-py3-none-win_arm64.whl", hash = "sha256:4b700459d4649e2594b31f20a9de33bc7c19976d4746d8d0798ad959621d64a4"},
-    {file = "ruff-0.14.5.tar.gz", hash = "sha256:8d3b48d7d8aad423d3137af7ab6c8b1e38e4de104800f0d596990f6ada1a9fc1"},
-]
-
 [[package]]
 name = "safehttpx"
 version = "0.1.7"
@@ -3356,4 +3354,4 @@ files = [
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.10,<3.15"
-content-hash = "8c1bb4d6fd49cbec6bfde5966b6f2490e520425a5c5e491617163b523ec7f85f"
+content-hash = "ea00e08a220f0b1891b043f1a06bcd40dd78bf93323c52b20ed3775d4ac97ccb"

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "anyio (>=4.11.0,<5.0.0)",
     "transformers (>=4.57.1,<5.0.0)",
     "typer[all] (>=0.20.0,<0.21.0)",
-    "gradio (>=5.49.1,<6.0.0)",
+    "gradio[oauth] (>=6.0.1,<7.0.0)",
     "ddgs (>=9.9.1,<10.0.0)",
     "uv (>=0.9.11,<0.10.0)",
     "markdownify (>=1.2.2,<2.0.0)",

src/upgrade_advisor/agents/package.py

@@ -28,8 +28,6 @@ logger.setLevel(logging.INFO)
 logger.addHandler(logging.StreamHandler())
 logger.addHandler(logging.FileHandler("package_agent.log"))
 
-# TODO: Fix the async issues with smolagents using mcpadapt
-
 
 class PackageDiscoveryAgent:
     """Agent that discovers metadata about Python packages using MCP tools."""

src/upgrade_advisor/agents/prompts.py

@@ -13,9 +13,12 @@ def get_package_discovery_prompt(
         user_input += f"\nREFRAMED QUESTION (LLM-generated):\n{reframed_question}\n"
 
     return f"""
-    You are a package discovery agent and an upgrade advisor for Python
-    packages.
-    Your goal is to find relevant metadata about Python packages using the
+    You are a package discovery and an upgrade advisor agent for Python
+    packages. You are called "FixMyEnv" and specialize in helping developers
+    identify package upgrade issues, compatibility problems, known bugs, and
+    best practices for managing Python package dependencies.
+    You find relevant metadata about Python packages and look up their github
+    repositories and relevant discussions, issues, releases etc using the
     available tools and use that context to provide helpful answers
     based on the user's question. If the user asks about upgrade recommendations,
     compatibility issues, known bugs, or best practices, you should gather
@@ -38,26 +41,30 @@ def get_package_discovery_prompt(
     discovery or upgrade advice. Politely inform the user that you can only help
     with Python package-related queries.
 
+    The first step in answering such questions is to gather relevant information about
+    the packages involved using the available tools.
+    Some tools like the `resolve_pyproject_toml` can directly analyze a
+    pyproject.toml content to find compatibility issues and upgrade suggestions.
+    Always prepare a plan before executing any tools. Iterate over the plan
+    step-by-step, using the tools as needed
+    to gather more information. When evidence is sufficient, provide a final answer
+    that directly addresses the user's question with clear recommendations or
+    insights. The recommendations should be in a structured markdown format with
+    bullet points or numbered lists or code blocks where appropriate.
+
+    IMPORTANT CONTEXT AND GUIDELINES:
+    Any issues in converting requirements.txt to pyproject.toml or parsing
+    pyproject.toml is your responsibility to handle using the available tools.
+    Make fixes as needed but report such things later in your final answer.
+
     Your knowledge cutoff may prevent you from knowing what's recent.
-    NO MATTER WHAT, always use the current date (ISO format YYYY-MM-DD): {today_date}
+    NO MATTER WHAT, always use the todays date: {today_date}
     when reasoning about dates and
-    releases. Some tools also provide you the release date information, which
+    releases (dates are in ISO format YYYY-MM-DD).
+    Some tools also provide you the release date information of packages and
+    repo releases, which
     you can transform to ISO format and make comparisons.
 
-    The first step to tackle such questions is to gather relevant information about the
-    packages involved using the available tools. Some tools like the
-    `resolve_pyproject_toml`
-    can directly analyze a pyproject.toml content to find
-    compatibility issues and upgrade suggestions.
-    Use the tools to fetch
-    package metadata, version history, release notes, compatibility info, and
-    known issues. Then, analyze the collected data to identify any potential
-    issues, improvements, or recommendations related to the user's question.
-
-    Always prepare a plan before executing any tools. Iterate over the plan
-    step-by-step, using the tools as needed to gather more information. When evidence
-    is sufficient, provide a final answer that directly addresses the user's
-    question with clear recommendations or insights.
 
     IMPORTANT EXECUTION GUIDELINES:
     - Do NOT print intermediate tool outputs. Do not wrap results in strings
@@ -66,13 +73,16 @@ def get_package_discovery_prompt(
     - The `output_schema` of each tool describes the expected output structure.
     - Make sure your final answer contains a "reasoning" field that explains
     how you arrived at your final answer. Do not omit this field.
-    - Do not mention the tool names, rather mention what the tool helped you discover.
     - Return the final structured object using: final_answer(<python_dict>)
     - Ensure the returned object STRICTLY matches the expected schema for that tool.
       Do not add or rename keys. Keep value types correct.
     - To read the contents of any uploaded files, call the `read_upload_file` tool
     with the path you received (direct file IO like `open()` is blocked).
 
+    In your final answer, mention what steps you followed (enumerated) and
+    the findings from each of those steps. Based on these findings, state your
+    final recommendations.
+
     {user_input}
 
     HINTS:

src/upgrade_advisor/chat/chat.py

@@ -21,10 +21,14 @@ logger.setLevel(logging.INFO)
 logger.addHandler(logging.StreamHandler())
 
 
-async def query(payload):
+async def query(payload, token=None):
     API_URL = "https://router.huggingface.co/v1/chat/completions"
+    token = token or os.environ.get("HF_TOKEN", "")  # use passed token or env var
+    if token is None or token == "":
+        logger.error("No Huggingface token provided for API request.")
+
     headers = {
-        "Authorization": f"Bearer {os.environ['HF_TOKEN']}",
+        "Authorization": f"Bearer {token}",
     }
     try:
         response = requests.post(API_URL, headers=headers, json=payload, timeout=300)
@@ -78,7 +82,7 @@ def extract_answer_content(answer_content: str) -> str:
 
 
 async def run_document_qa(
-    question: str, context: str, rewritten_question: str = None
+    question: str, context: str, rewritten_question: str = None, token: str = None
 ) -> str:
     response = await query(
         {
@@ -93,14 +97,17 @@ async def run_document_qa(
                 },
             ],
             "model": CHAT_MODEL,
-        }
+        },
+        token=token,
     )
 
     answer = parse_response(response)
     return extract_answer_content(answer["content"])
 
 
-async def qn_rewriter_judge(original_question: str, rewritten_question: str) -> str:
+async def qn_rewriter_judge(
+    original_question: str, rewritten_question: str, token: str = None
+) -> str:
     response = await query(
         {
             "messages": [
@@ -112,7 +119,8 @@ async def qn_rewriter_judge(original_question: str, rewritten_question: str) ->
                 },
             ],
             "model": CHAT_MODEL,
-        }
+        },
+        token=token,
     )
     answer = parse_response(response)
     answer_text = extract_answer_content(answer["content"])
@@ -123,7 +131,9 @@ async def qn_rewriter_judge(original_question: str, rewritten_question: str) ->
         return False
 
 
-async def qn_rewriter(original_question: str, summarized_history: str = "") -> str:
+async def qn_rewriter(
+    original_question: str, summarized_history: str = "", token: str = None
+) -> str:
     response = await query(
         {
             "messages": [
@@ -135,12 +145,15 @@ async def qn_rewriter(original_question: str, summarized_history: str = "") -> s
                 },
             ],
             "model": CHAT_MODEL,
-        }
+        },
+        token=token,
     )
 
     answer = parse_response(response)
     rewritten_question = extract_answer_content(answer["content"])
-    is_good = await qn_rewriter_judge(original_question, rewritten_question)
+    is_good = await qn_rewriter_judge(
+        original_question, rewritten_question, token=token
+    )
 
     return rewritten_question, is_good
 
@@ -149,6 +162,7 @@ async def summarize_chat_history(
     history: list[dict],
     turns_cutoff=10,
     word_cutoff=100,
+    token: str = None,
 ) -> str:
     # history is a list of dicts with 'role' and 'content' keys
     # [{"role": "user", "content": "..."}, {"role": "assistant", "content":
@@ -176,7 +190,8 @@ async def summarize_chat_history(
                 },
             ],
             "model": CHAT_MODEL,
-        }
+        },
+        token=token,
     )
 
     answer = parse_response(response)

src/upgrade_advisor/misc.py

@@ -147,11 +147,28 @@ def get_example_questions(n: int = 3) -> str:
     ]
     choices = random.sample(all_questions, k=n)
     # format as list of lists for gradio examples
-    choices = [[choice] for choice in choices]
+    choices = [[q] for q in choices]
 
     return choices
 
 
+def to_openai_message_format(role: str, content: str, append_to: list = None) -> dict:
+    message = {
+        "role": role,
+        "content": [
+            {
+                "type": "text",
+                "text": content,
+            }
+        ],
+    }
+    if append_to is not None:
+        # if its supposed to be appended to a list
+        append_to.append(message)
+        return append_to
+    return message
+
+
 def _monkeypatch_gradio_save_history():
     """Guard against non-int indices in Gradio's chat history saver.