asd

app.py

@@ -1,120 +1,558 @@
 import os
+import logging
+import logging.config
+from typing import Any
+from uuid import uuid4, UUID
+import json
+
 import gradio as gr
-from huggingface_hub import InferenceClient
+from dotenv import load_dotenv
 from langchain_core.messages import HumanMessage
-from langchain.agents import AgentExecutor
-from agent import build_graph
-from PIL import Image
+from langgraph.types import RunnableConfig
+from pydantic import BaseModel
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+load_dotenv()
 
+# There are tools set here dependent on environment variables
+from graph import graph, weak_model, search_enabled # noqa
 
-DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+FOLLOWUP_QUESTION_NUMBER = 3
+TRIM_MESSAGE_LENGTH = 16  # Includes tool messages
+USER_INPUT_MAX_LENGTH = 10000  # Characters
 
-class BasicAgent:
-    """A langgraph agent."""
-    def __init__(self):
-        print("BasicAgent initialized.")
-        self.graph = build_graph()
+# We need the same secret for data persistance
+# If you store sensitive data, you should store your secret in .env
+BROWSER_STORAGE_SECRET = "itsnosecret"
 
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        # Wrap the question in a HumanMessage from langchain_core
-        messages = [HumanMessage(content=question)]
-        config = {"recursion_limit": 27}
+with open('logging-config.json', 'r') as fh:
+    config = json.load(fh)
+logging.config.dictConfig(config)
+logger = logging.getLogger(__name__)
 
-        messages = self.graph.invoke({"messages": messages}, config=config)
+async def chat_fn(user_input: str, history: dict, input_graph_state: dict, uuid: UUID, prompt: str, search_enabled: bool, download_website_text_enabled: bool):
+    """
+    Args:
+        user_input (str): The user's input message
+        history (dict): The history of the conversation in gradio
+        input_graph_state (dict): The current state of the graph. This includes tool call history
+        uuid (UUID): The unique identifier for the current conversation. This can be used in conjunction with langgraph or for memory
+        prompt (str): The system prompt
+    Yields:
+        str: The output message
+        dict|Any: The final state of the graph
+        bool|Any: Whether to trigger follow up questions
 
-        answer = messages['messages'][-1].content
-        return answer[14:]
+        We do not use gradio history in the graph since we want the ToolMessage in the history
+        ordered properly. GraphProcessingState.messages is used as history instead
+    """
+    try:
+        logger.info(f"Prompt: {prompt}")
+        input_graph_state["tools_enabled"] = {
+            "download_website_text": download_website_text_enabled,
+            "tavily_search_results_json": search_enabled,
+        }
+        if prompt:
+            input_graph_state["prompt"] = prompt
+        if "messages" not in input_graph_state:
+            input_graph_state["messages"] = []
+        input_graph_state["messages"].append(
+            HumanMessage(user_input[:USER_INPUT_MAX_LENGTH])
+        )
+        input_graph_state["messages"] = input_graph_state["messages"][-TRIM_MESSAGE_LENGTH:]
+        config = RunnableConfig(
+            recursion_limit=20,
+            run_name="user_chat",
+            configurable={"thread_id": uuid}
+        )
 
-try:
-    agent = BasicAgent()
-except Exception as e:
-        print(f"Error instantiating agent: {e}")
+        output: str = ""
+        final_state: dict | Any = {}
+        waiting_output_seq: list[str] = []
 
-def show_graph():
-    if not os.path.exists("graph.png"):
-        return None
-    return Image.open("graph.png")
+        async for stream_mode, chunk in graph.astream(
+                    input_graph_state,
+                    config=config,
+                    stream_mode=["values", "messages"],
+                ):
+            if stream_mode == "values":
+                final_state = chunk
+                last_message = chunk["messages"][-1]
+                if hasattr(last_message, "tool_calls"):
+                    for msg_tool_call in last_message.tool_calls:
+                        tool_name: str = msg_tool_call['name']
+                        if tool_name == "tavily_search_results_json":
+                            query = msg_tool_call['args']['query']
+                            waiting_output_seq.append(f"Searching for '{query}'...")
+                            yield "\n".join(waiting_output_seq), gr.skip(), gr.skip()
+                        # download_website_text is the name of the function defined in graph.py
+                        elif tool_name == "download_website_text":
+                            url = msg_tool_call['args']['url']
+                            waiting_output_seq.append(f"Downloading text from '{url}'...")
+                            yield "\n".join(waiting_output_seq), gr.skip(), gr.skip()
+                        else:
+                            waiting_output_seq.append(f"Running {tool_name}...")
+                            yield "\n".join(waiting_output_seq), gr.skip(), gr.skip()
+            elif stream_mode == "messages":
+                msg, metadata = chunk
+                # print("output: ", msg, metadata)
+                # assistant_node is the name we defined in the langgraph graph
+                if metadata['langgraph_node'] == "assistant_node" and msg.content:
+                    output += msg.content
+                    yield output, gr.skip(), gr.skip()
+        # Trigger for asking follow up questions
+        # + store the graph state for next iteration
+        # yield output, dict(final_state), gr.skip()
+        yield output + " ", dict(final_state), True
+    except Exception:
+        logger.exception("Exception occurred")
+        user_error_message = "There was an error processing your request. Please try again."
+        yield user_error_message, gr.skip(), False
 
-config = {"configurable": {"thread_id": "1"}}
+def clear():
+    return dict(), uuid4()
 
+class FollowupQuestions(BaseModel):
+    """Model for langchain to use for structured output for followup questions"""
+    questions: list[str]
 
-def run_langgraph_agent(user_input: str):
-    graph = build_graph()
-    result = graph.stream(
-        {"messages": [HumanMessage(content=user_input)]},
-        config,
-        stream_mode="values",
+async def populate_followup_questions(end_of_chat_response: bool, messages: dict[str, str], uuid: UUID):
+    """
+    This function gets called a lot due to the asynchronous nature of streaming
+
+    Only populate followup questions if streaming has completed and the message is coming from the assistant
+    """
+    if not end_of_chat_response or not messages or messages[-1]["role"] != "assistant":
+        return *[gr.skip() for _ in range(FOLLOWUP_QUESTION_NUMBER)], False
+    config = RunnableConfig(
+        run_name="populate_followup_questions",
+        configurable={"thread_id": uuid}
     )
-    return result["messages"][-1].content if "messages" in result else result
+    weak_model_with_config = weak_model.with_config(config)
+    follow_up_questions = await weak_model_with_config.with_structured_output(FollowupQuestions).ainvoke([
+        ("system", f"suggest {FOLLOWUP_QUESTION_NUMBER} followup questions for the user to ask the assistant. Refrain from asking personal questions."),
+        *messages,
+    ])
+    if len(follow_up_questions.questions) != FOLLOWUP_QUESTION_NUMBER:
+        raise ValueError("Invalid value of followup questions")
+    buttons = []
+    for i in range(FOLLOWUP_QUESTION_NUMBER):
+        buttons.append(
+            gr.Button(follow_up_questions.questions[i], visible=True, elem_classes="chat-tab"),
+        )
+    return *buttons, False
 
+async def summarize_chat(end_of_chat_response: bool, messages: dict, sidebar_summaries: dict, uuid: UUID):
+    """Summarize chat for tab names"""
+    # print("\n------------------------")
+    # print("not end_of_chat_response", not end_of_chat_response)
+    # print("not messages", not messages)
+    # if messages:
+    #     print("messages[-1][role] != assistant", messages[-1]["role"] != "assistant")
+    # print("isinstance(sidebar_summaries, type(lambda x: x))", isinstance(sidebar_summaries, type(lambda x: x)))
+    # print("uuid in sidebar_summaries", uuid in sidebar_summaries)
+    should_return = (
+        not end_of_chat_response or
+        not messages or
+        messages[-1]["role"] != "assistant" or
+        # This is a bug with gradio
+        isinstance(sidebar_summaries, type(lambda x: x)) or
+        # Already created summary
+        uuid in sidebar_summaries
+    )
+    if should_return:
+        return gr.skip(), gr.skip()
+    config = RunnableConfig(
+        run_name="summarize_chat",
+        configurable={"thread_id": uuid}
+    )
+    weak_model_with_config = weak_model.with_config(config)
+    summary_response = await weak_model_with_config.ainvoke([
+        ("system", "summarize this chat in 7 tokens or less. Refrain from using periods"),
+        *messages,
+    ])
+    if uuid not in sidebar_summaries:
+        sidebar_summaries[uuid] = summary_response.content
+    return sidebar_summaries, False
 
-demo = gr.Interface(
-    fn=run_langgraph_agent,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your message..."),
-    outputs="text",
-    title="LangGraph Agent Chat",
-)
+async def new_tab(uuid, gradio_graph, messages, tabs, prompt, sidebar_summaries):
+    new_uuid = uuid4()
+    new_graph = {}
+    if uuid not in sidebar_summaries:
+        sidebar_summaries, _ = await summarize_chat(True, messages, sidebar_summaries, uuid)
+    tabs[uuid] = {
+        "graph": gradio_graph,
+        "messages": messages,
+        "prompt": prompt,
+    }
+    suggestion_buttons = []
+    for _ in range(FOLLOWUP_QUESTION_NUMBER):
+        suggestion_buttons.append(gr.Button(visible=False))
+    new_messages = {}
+    new_prompt = "You are a helpful assistant."
+    return new_uuid, new_graph, new_messages, tabs, new_prompt, sidebar_summaries, *suggestion_buttons
 
-if __name__ == "__main__":
-    demo.launch()
+def switch_tab(selected_uuid, tabs, gradio_graph, uuid, messages, prompt):
+    # I don't know of another way to lookup uuid other than
+    # by the button value
 
+    # Save current state
+    if messages:
+        tabs[uuid] = {
+            "graph": gradio_graph,
+            "messages": messages,
+            "prompt": prompt
+        }
 
+    if selected_uuid not in tabs:
+        logger.error(f"Could not find the selected tab in offloaded_tabs_data_storage {selected_uuid}")
+        return gr.skip(), gr.skip(), gr.skip(), gr.skip()
+    selected_tab_state = tabs[selected_uuid]
+    selected_graph = selected_tab_state["graph"]
+    selected_messages = selected_tab_state["messages"]
+    selected_prompt = selected_tab_state.get("prompt", "")
+    suggestion_buttons = []
+    for _ in range(FOLLOWUP_QUESTION_NUMBER):
+        suggestion_buttons.append(gr.Button(visible=False))
+    return selected_graph, selected_uuid, selected_messages, tabs, selected_prompt, *suggestion_buttons
 
+def delete_tab(current_chat_uuid, selected_uuid, sidebar_summaries, tabs):
+    output_messages = gr.skip()
+    if current_chat_uuid == selected_uuid:
+        output_messages = dict()
+    if selected_uuid in tabs:
+        del tabs[selected_uuid]
+    if selected_uuid in sidebar_summaries:
+        del sidebar_summaries[selected_uuid]
+    return sidebar_summaries, tabs, output_messages
 
-# def respond(
-#     message,
-#     history: list[tuple[str, str]],
-#     system_message,
-#     max_tokens,
-#     temperature,
-#     top_p,
-# ):
-#     messages = [{"role": "system", "content": system_message}]
+def submit_edit_tab(selected_uuid, sidebar_summaries, text):
+    sidebar_summaries[selected_uuid] = text
+    return sidebar_summaries, ""
 
-#     for val in history:
-#         if val[0]:
-#             messages.append({"role": "user", "content": val[0]})
-#         if val[1]:
-#             messages.append({"role": "assistant", "content": val[1]})
+CSS = """
+footer {visibility: hidden}
+.followup-question-button {font-size: 12px }
+.chat-tab {
+    font-size: 12px;
+    padding-inline: 0;
+}
+.chat-tab.active {
+    background-color: #654343;
+}
+#new-chat-button { background-color: #0f0f11; color: white; }
 
-#     messages.append({"role": "user", "content": message})
+.tab-button-control {
+    min-width: 0;
+    padding-left: 0;
+    padding-right: 0;
+}
+"""
 
-#     response = ""
+# We set the ChatInterface textbox id to chat-textbox for this to work
+TRIGGER_CHATINTERFACE_BUTTON = """
+function triggerChatButtonClick() {
 
-#     for message in client.chat_completion(
-#         messages,
-#         max_tokens=max_tokens,
-#         stream=True,
-#         temperature=temperature,
-#         top_p=top_p,
-#     ):
-#         token = message.choices[0].delta.content
+  // Find the div with id "chat-textbox"
+  const chatTextbox = document.getElementById("chat-textbox");
 
-#         response += token
-#         yield response
+  if (!chatTextbox) {
+    console.error("Error: Could not find element with id 'chat-textbox'");
+    return;
+  }
 
+  // Find the button that is a descendant of the div
+  const button = chatTextbox.querySelector("button");
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-# demo = gr.ChatInterface(
-#     respond,
-#     additional_inputs=[
-#         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-#         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-#         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-#         gr.Slider(
-#             minimum=0.1,
-#             maximum=1.0,
-#             value=0.95,
-#             step=0.05,
-#             label="Top-p (nucleus sampling)",
-#         ),
-#     ],
-# )
+  if (!button) {
+    console.error("Error: No button found inside the chat-textbox element");
+    return;
+  }
+
+  // Trigger the click event
+  button.click();
+}"""
+
+if __name__ == "__main__":
+    logger.info("Starting the interface")
+    with gr.Blocks(title="Langgraph Template", fill_height=True, css=CSS) as app:
+        current_prompt_state = gr.BrowserState(
+            storage_key="current_prompt_state",
+            secret=BROWSER_STORAGE_SECRET,
+        )
+        current_uuid_state = gr.BrowserState(
+            uuid4,
+            storage_key="current_uuid_state",
+            secret=BROWSER_STORAGE_SECRET,
+        )
+        current_langgraph_state = gr.BrowserState(
+            dict(),
+            storage_key="current_langgraph_state",
+            secret=BROWSER_STORAGE_SECRET,
+        )
+        end_of_assistant_response_state = gr.State(
+            bool(),
+        )
+        # [uuid] -> summary of chat
+        sidebar_names_state = gr.BrowserState(
+            dict(),
+            storage_key="sidebar_names_state",
+            secret=BROWSER_STORAGE_SECRET,
+        )
+        # [uuid] -> {"graph": gradio_graph, "messages": messages}
+        offloaded_tabs_data_storage = gr.BrowserState(
+            dict(),
+            storage_key="offloaded_tabs_data_storage",
+            secret=BROWSER_STORAGE_SECRET,
+        )
+
+        chatbot_message_storage = gr.BrowserState(
+            [],
+            storage_key="chatbot_message_storage",
+            secret=BROWSER_STORAGE_SECRET,
+        )
+        with gr.Column():
+            prompt_textbox = gr.Textbox(show_label=False, interactive=True)
+        with gr.Row():
+            checkbox_search_enabled = gr.Checkbox(
+                value=True,
+                label="Enable search",
+                show_label=True,
+                visible=search_enabled,
+                scale=1,
+            )
+            checkbox_download_website_text = gr.Checkbox(
+                value=True,
+                show_label=True,
+                label="Enable downloading text from urls",
+                scale=1,
+            )
+        chatbot = gr.Chatbot(
+            type="messages",
+            scale=1,
+            show_copy_button=True,
+            height=600,
+            editable="all",
+        )
+        tab_edit_uuid_state = gr.State(
+            str()
+        )
+        prompt_textbox.change(lambda prompt: prompt, inputs=[prompt_textbox], outputs=[current_prompt_state])
+        with gr.Sidebar() as sidebar:
+            @gr.render(inputs=[tab_edit_uuid_state, end_of_assistant_response_state, sidebar_names_state, current_uuid_state, chatbot, offloaded_tabs_data_storage])
+            def render_chats(tab_uuid_edit, end_of_chat_response, sidebar_summaries, active_uuid, messages, tabs):
+                current_tab_button_text = ""
+                if active_uuid not in sidebar_summaries:
+                    current_tab_button_text = "Current Chat"
+                elif active_uuid not in tabs:
+                    current_tab_button_text = sidebar_summaries[active_uuid]
+                if current_tab_button_text:
+                    gr.Button(current_tab_button_text, elem_classes=["chat-tab", "active"])
+                for chat_uuid, tab in reversed(tabs.items()):
+                    elem_classes = ["chat-tab"]
+                    if chat_uuid == active_uuid:
+                        elem_classes.append("active")
+                    button_uuid_state = gr.State(chat_uuid)
+                    with gr.Row():
+                        clear_tab_button = gr.Button(
+                            "🗑",
+                            scale=0,
+                            elem_classes=["tab-button-control"]
+                        )
+                        clear_tab_button.click(
+                            fn=delete_tab,
+                            inputs=[
+                                current_uuid_state,
+                                button_uuid_state,
+                                sidebar_names_state,
+                                offloaded_tabs_data_storage
+                            ],
+                            outputs=[
+                                sidebar_names_state,
+                                offloaded_tabs_data_storage,
+                                chat_interface.chatbot_value
+                            ]
+                        )
+                        chat_button_text = sidebar_summaries.get(chat_uuid)
+                        if not chat_button_text:
+                            chat_button_text = str(chat_uuid)
+                        if chat_uuid != tab_uuid_edit:
+                            set_edit_tab_button = gr.Button(
+                                "✎",
+                                scale=0,
+                                elem_classes=["tab-button-control"]
+                            )
+                            set_edit_tab_button.click(
+                                fn=lambda x: x,
+                                inputs=[button_uuid_state],
+                                outputs=[tab_edit_uuid_state]
+                            )
+                            chat_tab_button = gr.Button(
+                                chat_button_text,
+                                elem_id=f"chat-{chat_uuid}-button",
+                                elem_classes=elem_classes,
+                                scale=2
+                            )
+                            chat_tab_button.click(
+                                fn=switch_tab,
+                                inputs=[
+                                    button_uuid_state,
+                                    offloaded_tabs_data_storage,
+                                    current_langgraph_state,
+                                    current_uuid_state,
+                                    chatbot,
+                                    prompt_textbox
+                                ],
+                                outputs=[
+                                    current_langgraph_state,
+                                    current_uuid_state,
+                                    chat_interface.chatbot_value,
+                                    offloaded_tabs_data_storage,
+                                    prompt_textbox,
+                                    *followup_question_buttons
+                                ]
+                            )
+                        else:
+                            chat_tab_text = gr.Textbox(
+                                chat_button_text,
+                                scale=2,
+                                interactive=True,
+                                show_label=False
+                            )
+                            chat_tab_text.submit(
+                                fn=submit_edit_tab,
+                                inputs=[
+                                    button_uuid_state,
+                                    sidebar_names_state,
+                                    chat_tab_text
+                                ],
+                                outputs=[
+                                    sidebar_names_state,
+                                    tab_edit_uuid_state
+                                ]
+                            )
+                    # )
+                # return chat_tabs, sidebar_summaries
+            new_chat_button = gr.Button("New Chat", elem_id="new-chat-button")
+        chatbot.clear(fn=clear, outputs=[current_langgraph_state, current_uuid_state])
+        with gr.Row():
+            followup_question_buttons = []
+            for i in range(FOLLOWUP_QUESTION_NUMBER):
+                btn = gr.Button(f"Button {i+1}", visible=False)
+                followup_question_buttons.append(btn)
+
+        multimodal = False
+        textbox_component = (
+            gr.MultimodalTextbox if multimodal else gr.Textbox
+        )
+        with gr.Column():
+            textbox = textbox_component(
+                show_label=False,
+                label="Message",
+                placeholder="Type a message...",
+                scale=7,
+                autofocus=True,
+                submit_btn=True,
+                stop_btn=True,
+                elem_id="chat-textbox",
+                lines=1,
+            )
+        chat_interface = gr.ChatInterface(
+            chatbot=chatbot,
+            fn=chat_fn,
+            additional_inputs=[
+                current_langgraph_state,
+                current_uuid_state,
+                prompt_textbox,
+                checkbox_search_enabled,
+                checkbox_download_website_text,
+            ],
+            additional_outputs=[
+                current_langgraph_state,
+                end_of_assistant_response_state
+            ],
+            type="messages",
+            multimodal=multimodal,
+            textbox=textbox,
+        )
+
+        new_chat_button.click(
+            new_tab,
+            inputs=[
+                current_uuid_state,
+                current_langgraph_state,
+                chatbot,
+                offloaded_tabs_data_storage,
+                prompt_textbox,
+                sidebar_names_state,
+            ],
+            outputs=[
+                current_uuid_state,
+                current_langgraph_state,
+                chat_interface.chatbot_value,
+                offloaded_tabs_data_storage,
+                prompt_textbox,
+                sidebar_names_state,
+                *followup_question_buttons,
+            ]
+        )
+
+        def click_followup_button(btn):
+            buttons = [gr.Button(visible=False) for _ in range(len(followup_question_buttons))]
+            return btn, *buttons
+        for btn in followup_question_buttons:
+            btn.click(
+                fn=click_followup_button,
+                inputs=[btn],
+                outputs=[
+                    chat_interface.textbox,
+                    *followup_question_buttons
+                ]
+            ).success(lambda: None, js=TRIGGER_CHATINTERFACE_BUTTON)
+
+        chatbot.change(
+            fn=populate_followup_questions,
+            inputs=[
+                end_of_assistant_response_state,
+                chatbot,
+                current_uuid_state
+            ],
+            outputs=[
+                *followup_question_buttons,
+                end_of_assistant_response_state
+            ],
+            trigger_mode="multiple"
+        )
+        chatbot.change(
+            fn=summarize_chat,
+            inputs=[
+                end_of_assistant_response_state,
+                chatbot,
+                sidebar_names_state,
+                current_uuid_state
+            ],
+            outputs=[
+                sidebar_names_state,
+                end_of_assistant_response_state
+            ],
+            trigger_mode="multiple"
+        )
+        chatbot.change(
+            fn=lambda x: x,
+            inputs=[chatbot],
+            outputs=[chatbot_message_storage],
+            trigger_mode="always_last"
+        )
+
+        @app.load(inputs=[chatbot_message_storage], outputs=[chat_interface.chatbot_value])
+        def load_messages(messages):
+            return messages
+
+        @app.load(inputs=[current_prompt_state], outputs=[prompt_textbox])
+        def load_prompt(current_prompt):
+            return current_prompt
+
+    app.launch(
+        server_name="127.0.0.1",
+        server_port=int(os.getenv("GRADIO_SERVER_PORT", 7860)),
+        # favicon_path="assets/favicon.ico"
+    )

graph.py

@@ -0,0 +1,116 @@
+import logging
+import os
+from typing import Annotated
+
+import aiohttp
+from langchain_core.messages import AnyMessage
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.prebuilt import ToolNode
+from langgraph.graph import StateGraph, END, add_messages
+from langchain_community.tools import TavilySearchResults
+from pydantic import BaseModel, Field
+from trafilatura import extract
+
+logger = logging.getLogger(__name__)
+ASSISTANT_SYSTEM_PROMPT_BASE = """"""
+search_enabled = bool(os.environ.get("TAVILY_API_KEY"))
+
+@tool
+async def download_website_text(url: str) -> str:
+    """Download the text from a website"""
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url) as response:
+                response.raise_for_status()
+                downloaded = await response.text()
+        result = extract(downloaded, include_formatting=True, include_links=True, output_format='json', with_metadata=True)
+        return result or "No text found on the website"
+    except Exception as e:
+        logger.error(f"Failed to download {url}: {str(e)}")
+        return f"Error retrieving website content: {str(e)}"
+
+tools = [download_website_text]
+
+if search_enabled:
+    tavily_search_tool = TavilySearchResults(
+        max_results=5,
+        search_depth="advanced",
+        include_answer=True,
+        include_raw_content=True,
+    )
+    tools.append(tavily_search_tool)
+else:
+    print("TAVILY_API_KEY environment variable not found. Websearch disabled")
+
+weak_model = ChatOpenAI(model="gpt-4o-mini", tags=["assistant"])
+model = weak_model
+assistant_model = weak_model
+
+class GraphProcessingState(BaseModel):
+    # user_input: str = Field(default_factory=str, description="The original user input")
+    messages: Annotated[list[AnyMessage], add_messages] = Field(default_factory=list)
+    prompt: str = Field(default_factory=str, description="The prompt to be used for the model")
+    tools_enabled: dict = Field(default_factory=dict, description="The tools enabled for the assistant")
+    search_enabled: bool = Field(default=True, description="Whether to enable search tools")
+
+async def assistant_node(state: GraphProcessingState, config=None):
+    assistant_tools = []
+    if state.tools_enabled.get("download_website_text", True):
+        assistant_tools.append(download_website_text)
+    if search_enabled and state.tools_enabled.get("tavily_search_results_json", True):
+        assistant_tools.append(tavily_search_tool)
+    assistant_model = model.bind_tools(assistant_tools)
+    if state.prompt:
+        final_prompt = "\n".join([state.prompt, ASSISTANT_SYSTEM_PROMPT_BASE])
+    else:
+        final_prompt = ASSISTANT_SYSTEM_PROMPT_BASE
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", final_prompt),
+            MessagesPlaceholder(variable_name="messages"),
+        ]
+    )
+    chain = prompt | assistant_model
+    response = await chain.ainvoke({"messages": state.messages}, config=config)
+
+    return {
+        "messages": response
+    }
+
+def assistant_cond_edge(state: GraphProcessingState):
+    last_message = state.messages[-1]
+    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
+        logger.info(f"Tool call detected: {last_message.tool_calls}")
+        return "tools"
+    return END
+
+def define_workflow() -> CompiledStateGraph:
+    """Defines the workflow graph"""
+    # Initialize the graph
+    workflow = StateGraph(GraphProcessingState)
+
+    # Add nodes
+    workflow.add_node("assistant_node", assistant_node)
+    workflow.add_node("tools", ToolNode(tools))
+
+    # Edges
+    workflow.add_edge("tools", "assistant_node")
+
+    # Conditional routing
+    workflow.add_conditional_edges(
+        "assistant_node",
+        # If the latest message (result) from assistant is a tool call -> assistant_cond_edge routes to tools
+        # If the latest message (result) from assistant is a not a tool call -> assistant_cond_edge routes to END
+        assistant_cond_edge,
+    )
+    # Set end nodes
+    workflow.set_entry_point("assistant_node")
+    # workflow.set_finish_point("assistant_node")
+
+    return workflow.compile()
+
+graph = define_workflow()

notapp.py

@@ -0,0 +1,120 @@
+import os
+import gradio as gr
+from huggingface_hub import InferenceClient
+from langchain_core.messages import HumanMessage
+from langchain.agents import AgentExecutor
+from agent import build_graph
+from PIL import Image
+
+"""
+For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+"""
+client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+
+
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+
+class BasicAgent:
+    """A langgraph agent."""
+    def __init__(self):
+        print("BasicAgent initialized.")
+        self.graph = build_graph()
+
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        # Wrap the question in a HumanMessage from langchain_core
+        messages = [HumanMessage(content=question)]
+        config = {"recursion_limit": 27}
+
+        messages = self.graph.invoke({"messages": messages}, config=config)
+
+        answer = messages['messages'][-1].content
+        return answer[14:]
+
+try:
+    agent = BasicAgent()
+except Exception as e:
+        print(f"Error instantiating agent: {e}")
+
+def show_graph():
+    if not os.path.exists("graph.png"):
+        return None
+    return Image.open("graph.png")
+
+config = {"configurable": {"thread_id": "1"}}
+
+
+def run_langgraph_agent(user_input: str):
+    graph = build_graph()
+    result = graph.stream(
+        {"messages": [HumanMessage(content=user_input)]},
+        config,
+        stream_mode="values",
+    )
+    return result["messages"][-1].content if "messages" in result else result
+
+
+demo = gr.Interface(
+    fn=run_langgraph_agent,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your message..."),
+    outputs="text",
+    title="LangGraph Agent Chat",
+)
+
+if __name__ == "__main__":
+    demo.launch()
+
+
+
+
+# def respond(
+#     message,
+#     history: list[tuple[str, str]],
+#     system_message,
+#     max_tokens,
+#     temperature,
+#     top_p,
+# ):
+#     messages = [{"role": "system", "content": system_message}]
+
+#     for val in history:
+#         if val[0]:
+#             messages.append({"role": "user", "content": val[0]})
+#         if val[1]:
+#             messages.append({"role": "assistant", "content": val[1]})
+
+#     messages.append({"role": "user", "content": message})
+
+#     response = ""
+
+#     for message in client.chat_completion(
+#         messages,
+#         max_tokens=max_tokens,
+#         stream=True,
+#         temperature=temperature,
+#         top_p=top_p,
+#     ):
+#         token = message.choices[0].delta.content
+
+#         response += token
+#         yield response
+
+
+"""
+For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+"""
+# demo = gr.ChatInterface(
+#     respond,
+#     additional_inputs=[
+#         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+#         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+#         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+#         gr.Slider(
+#             minimum=0.1,
+#             maximum=1.0,
+#             value=0.95,
+#             step=0.05,
+#             label="Top-p (nucleus sampling)",
+#         ),
+#     ],
+# )