Update app.py

app.py

@@ -1,12 +1,15 @@
 import os
+import gradio as gr
+import json
 from typing import Annotated
 from typing_extensions import TypedDict
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_community.tools.tavily_search import TavilySearchResults
 from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
-from langchain_huggingface import HuggingFaceEndpoint
+from langchain_core.messages import ToolMessage
 from dotenv import load_dotenv
 import logging
-import gradio as gr
 
 # Initialize logging
 logging.basicConfig(level=logging.INFO)
@@ -15,7 +18,7 @@ logging.basicConfig(level=logging.INFO)
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Initialize Hugging Face endpoint
+# Initialize the HuggingFace model
 llm = HuggingFaceEndpoint(
     repo_id="mistralai/Mistral-7B-Instruct-v0.3",
     huggingfacehub_api_token=HF_TOKEN.strip(),
@@ -23,6 +26,10 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=200
 )
 
+# Initialize Tavily Search tool
+tool = TavilySearchResults(max_results=2)
+tools = [tool]
+
 # Define the state structure
 class State(TypedDict):
     messages: Annotated[list, add_messages]
@@ -33,80 +40,129 @@ graph_builder = StateGraph(State)
 # Define the chatbot function
 def chatbot(state: State):
     try:
-        logging.info(f"Input Messages: {state['messages']}")
-        response = llm.invoke(state["messages"])
+        # Get the last message and ensure it's a string
+        input_message = state["messages"][-1] if state["messages"] else ""
+        
+        # Ensure that input_message is a string (check the type)
+        if isinstance(input_message, str):
+            query = input_message  # If it's already a string, use it directly
+        elif hasattr(input_message, 'content') and isinstance(input_message.content, str):
+            query = input_message.content  # Extract the content if it's a HumanMessage object
+        else:
+            raise ValueError("Input message is not in the correct format")
+
+        logging.info(f"Input Message: {query}")
+
+        # Invoke the LLM for a response
+        response = llm.invoke([query])
         logging.info(f"LLM Response: {response}")
-        return {"messages": [response]}
+
+        # Now, invoke Tavily Search and get the results
+        search_results = tool.invoke({"query": query})
+        
+        # Extract URLs from search results
+        urls = [result.get("url", "No URL found") for result in search_results]
+        
+        # Prepare the result to include URL information
+        result_with_url = {
+            "role": "assistant",  # Set the role to 'assistant'
+            "content": response,  # Set the response as content
+            "urls": urls  # Include the URLs of the search results
+        }
+
+        return {"messages": state["messages"] + [result_with_url]}
+    
     except Exception as e:
         logging.error(f"Error: {str(e)}")
-        return {"messages": [f"Error: {str(e)}"]}
+        return {"messages": state["messages"] + [f"Error: {str(e)}"]}
 
-# Add nodes and edges to the state graph
-graph_builder.add_node("chatbot", chatbot)
-graph_builder.add_edge(START, "chatbot")
-graph_builder.add_edge("chatbot", END)
+# Add tool node to the graph
+class BasicToolNode:
+    """A node that runs the tools requested in the last AIMessage."""
+    def __init__(self, tools: list) -> None:
+        self.tools_by_name = {tool.name: tool for tool in tools}
 
-# Compile the state graph
-graph = graph_builder.compile()
+    def __call__(self, inputs: dict):
+        if messages := inputs.get("messages", []):
+            message = messages[-1]
+        else:
+            raise ValueError("No message found in input")
+        
+        outputs = []
+        for tool_call in message.tool_calls:
+            tool_result = self.tools_by_name[tool_call["name"]].invoke(
+                tool_call["args"]
+            )
+            outputs.append(
+                ToolMessage(
+                    content=json.dumps(tool_result),
+                    name=tool_call["name"],
+                    tool_call_id=tool_call["id"],
+                )
+            )
+        return {"messages": outputs}
 
-# Function to stream updates from the graph
-def stream_graph_updates(user_input: str):
-    """
-    Stream updates from the graph based on user input and return the assistant's reply.
-    """
-    assistant_reply = ""
-    for event in graph.stream({"messages": [("user", user_input)]}):
-        for value in event.values():
-            if isinstance(value["messages"][-1], dict):
-                # If it's a dict, extract 'content'
-                assistant_reply = value["messages"][-1].get("content", "")
-            elif isinstance(value["messages"][-1], str):
-                # If it's a string, use it directly
-                assistant_reply = value["messages"][-1]
-    return assistant_reply
-
-# Gradio chatbot function using the streaming updates
-def gradio_chatbot(user_message: str):
-    """
-    Handle Gradio user input, process through the graph, and return only the assistant's reply.
-    """
-    try:
-        return stream_graph_updates(user_message)
-    except Exception as e:
-        logging.error(f"Error in Gradio chatbot: {str(e)}")
-        return f"Error: {str(e)}"
+# Add tool node to the graph
+tool_node = BasicToolNode(tools=tools)
+graph_builder.add_node("tools", tool_node)
 
-# Terminal-based chat loop
-def terminal_chat():
+# Define the conditional routing function
+def route_tools(state: State):
     """
-    Chat loop for terminal interaction.
+    Route to the ToolNode if the last message has tool calls.
+    Otherwise, route to the end.
     """
-    while True:
-        try:
-            user_input = input("User: ").strip()
-            # Process input through the chatbot
-            response = stream_graph_updates(user_input)
-            print(f"Assistant: {response}")
-        except Exception as e:
-            print(f"Error: {str(e)}")
-            break
-
-# Create Gradio interface
-interface = gr.Interface(
-    fn=gradio_chatbot,
-    inputs=gr.Textbox(placeholder="Enter your message", label="Your Message"),
-    outputs=gr.Textbox(label="Assistant's Reply"),
-    title="Chatbot",
-    description="Interactive chatbot using a state graph and Hugging Face Endpoint."
+    if isinstance(state, list):
+        ai_message = state[-1]
+    elif messages := state.get("messages", []):
+        ai_message = messages[-1]
+    else:
+        raise ValueError(f"No messages found in input state to tool_edge: {state}")
+    
+    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
+        return "tools"
+    
+    return END
+
+# Add nodes and conditional edges to the state graph
+graph_builder.add_node("chatbot", chatbot)
+graph_builder.add_conditional_edges(
+    "chatbot",
+    route_tools,
+    {"tools": "tools", END: END}
 )
+graph_builder.add_edge("tools", "chatbot")
+graph_builder.add_edge(START, "chatbot")
+graph = graph_builder.compile()
 
-if __name__ == "__main__":
-    # Launch terminal chat and Gradio interface in separate threads
-    import threading
+# Gradio interface
+def chat_interface(input_text, state):
+    # Prepare state if not provided
+    if state is None:
+        state = {"messages": []}
+    
+    # Append user input to state
+    state["messages"].append(input_text)
+    
+    # Process state through the graph
+    updated_state = graph.invoke(state)
+    return updated_state["messages"][-1], updated_state
 
-    # Terminal chat thread
-    terminal_thread = threading.Thread(target=terminal_chat, daemon=True)
-    terminal_thread.start()
+# Create Gradio app
+with gr.Blocks() as demo:
+    gr.Markdown("### Chatbot with Tavily Search Integration")
+    chat_state = gr.State({"messages": []})
+    
+    with gr.Row():
+        with gr.Column():
+            user_input = gr.Textbox(label="Your Message", placeholder="Type your message here...", lines=2)
+            submit_button = gr.Button("Submit")
+        
+        with gr.Column():
+            chatbot_output = gr.Textbox(label="Chatbot Response", interactive=False, lines=4)
+    
+    submit_button.click(chat_interface, inputs=[user_input, chat_state], outputs=[chatbot_output, chat_state])
 
-    # Launch Gradio interface
-    interface.launch(share=True)
+# Launch the Gradio app
+if __name__ == "__main__":
+    demo.launch()