Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,312 @@
+import os
+from typing import Dict, Callable, List, Union, TypedDict
+
+import gradio as gr
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
+from langchain_core.pydantic_v1 import BaseModel
+
+from langgraph.checkpoint.sqlite import SqliteSaver
+
+class Node:
+    def __init__(self, id: str, function: Callable):
+        """
+        Initialize a Node with an ID and a function to execute.
+
+        Args:
+            id (str): The unique identifier for the node.
+            function (Callable): The function to execute for this node.
+        """
+        self.id = id
+        self.function = function
+
+    def execute(self, state: Dict) -> Dict:
+        """
+        Execute the node's function with the given state.
+
+        Args:
+            state (Dict): The current state of the chatbot.
+
+        Returns:
+            Dict: The updated state after executing the node's function.
+        """
+        return self.function(state)
+
+class Edge:
+    def __init__(self, source: str, target: str, condition: Callable[[Dict], bool] = None):
+        """
+        Initialize an Edge with a source node, target node, and an optional condition.
+
+        Args:
+            source (str): The ID of the source node.
+            target (str): The ID of the target node.
+            condition (Callable[[Dict], bool], optional): A condition function that determines if the edge should be traversed.
+        """
+        self.source = source
+        self.target = target
+        self.condition = condition
+
+    def is_active(self, state: Dict) -> bool:
+        """
+        Check if the edge is active based on the given state.
+
+        Args:
+            state (Dict): The current state of the chatbot.
+
+        Returns:
+            bool: True if the edge is active, False otherwise.
+        """
+        if self.condition:
+            return self.condition(state)
+        return True
+
+class Graph:
+    def __init__(self):
+        """
+        Initialize an empty Graph with dictionaries to hold nodes and edges.
+        """
+        self.nodes = {}
+        self.edges = {}
+
+    def add_node(self, node: Node):
+        """
+        Add a node to the graph.
+
+        Args:
+            node (Node): The node to add.
+        """
+        self.nodes[node.id] = node
+
+    def add_edge(self, edge: Edge):
+        """
+        Add an edge to the graph.
+
+        Args:
+            edge (Edge): The edge to add.
+        """
+        if edge.source not in self.edges:
+            self.edges[edge.source] = []
+        self.edges[edge.source].append(edge)
+
+    def get_next_node(self, current_node_id: str, state: Dict) -> Union[Node, None]:
+        """
+        Get the next node to traverse to based on the current state.
+
+        Args:
+            current_node_id (str): The ID of the current node.
+            state (Dict): The current state of the chatbot.
+
+        Returns:
+            Union[Node, None]: The next node to traverse to, or None if no valid edge is found.
+        """
+        if current_node_id in self.edges:
+            for edge in self.edges[current_node_id]:
+                if edge.is_active(state):
+                    return self.nodes[edge.target]
+        return None
+
+    def execute(self, start_node_id: str, state: Dict) -> Dict:
+        """
+        Execute the graph starting from the specified node.
+
+        Args:
+            start_node_id (str): The ID of the starting node.
+            state (Dict): The initial state of the chatbot.
+
+        Returns:
+            Dict: The final state after traversing the graph.
+        """
+        current_node = self.nodes.get(start_node_id)
+        while current_node:
+            state = current_node.execute(state)
+            next_node = self.get_next_node(current_node.id, state)
+            if next_node is None:
+                break
+            current_node = next_node
+        return state
+
+class State(TypedDict):
+    """
+    Define the State type using TypedDict to specify the structure of the state dictionary.
+    """
+    messages: List[Union[Dict, BaseMessage, ToolMessage]]
+    ask_human: bool
+
+class RequestAssistance(BaseModel):
+    """
+    Define RequestAssistance model inheriting from BaseModel for schema validation.
+    """
+    request: str
+
+def chatbot_function(state: State) -> State:
+    """
+    Chatbot function definition which processes the current state and generates a response.
+
+    Args:
+        state (State): The current state of the chatbot including messages and ask_human flag.
+
+    Returns:
+        State: The updated state after processing the response.
+    """
+    response = llm_with_tools.invoke(state["messages"])
+    ask_human = False
+    
+    if response.tool_calls:
+        tool_name = response.tool_calls[0].get("name")
+        if tool_name == "RequestAssistance":
+            ask_human = True
+        else:
+            tool_response = DuckDuck_tool.run(response.tool_calls[0]["args"]["query"])
+            response.content = tool_response  # Update AI message content with tool response
+    
+    new_state = {"messages": state["messages"] + [response], "ask_human": ask_human}
+    return new_state
+
+def create_response(response: str, ai_message: AIMessage) -> ToolMessage:
+    """
+    Create a ToolMessage from a given response and AI message.
+
+    Args:
+        response (str): The response content to be included in the ToolMessage.
+        ai_message (AIMessage): The original AI message containing tool call information.
+
+    Returns:
+        ToolMessage: The created ToolMessage.
+    """
+    return ToolMessage(content=response, tool_call_id=ai_message.tool_calls[0].get("id"))
+
+def human_node_function(state: State) -> State:
+    """
+    Process the state if human assistance is required.
+
+    Args:
+        state (State): The current state of the chatbot including messages and ask_human flag.
+
+    Returns:
+        State: The updated state after processing human assistance.
+    """
+    new_messages = state["messages"]
+    
+    if state["messages"] and not isinstance(state["messages"][-1], ToolMessage):
+        new_response = create_response("No response from human.", state["messages"][-1])
+        new_messages.append(new_response)
+    
+    new_state = {"messages": new_messages, "ask_human": False}
+    return new_state
+
+def tools_condition(state: State) -> str:
+    """
+    Determine the next node in the state graph based on the current state.
+
+    Args:
+        state (State): The current state of the chatbot including messages and ask_human flag.
+
+    Returns:
+        str: The identifier of the next node to process.
+    """
+    # Define your condition to choose the next node here
+    # Example: Check if the state contains a specific tool call
+    for message in state["messages"]:
+        if isinstance(message, AIMessage) and message.tool_calls:
+            return "tools"
+    return "chatbot"
+
+def tool_node_function(state: State) -> State:
+    """
+    Process the state by executing the appropriate tool function.
+
+    Args:
+        state (State): The current state of the chatbot including messages and ask_human flag.
+
+    Returns:
+        State: The updated state after processing the tool function.
+    """
+    new_messages = state["messages"]
+    
+    for message in state["messages"]:
+        if isinstance(message, AIMessage) and message.tool_calls:
+            tool_response = DuckDuck_tool.run(message.tool_calls[0]["args"]["query"])
+            new_response = create_response(tool_response, message)
+            new_messages.append(new_response)
+    
+    new_state = {"messages": new_messages, "ask_human": False}
+    return new_state
+
+def format_message(msg: Union[Dict, BaseMessage, ToolMessage]) -> Dict[str, str]:
+    """
+    Format a message for display in the chat.
+
+    Args:
+        msg (Union[Dict, BaseMessage, ToolMessage]): The message to be formatted.
+
+    Returns:
+        Dict[str, str]: The formatted message as a dictionary with role and content.
+    """
+    if isinstance(msg, dict):
+        formatted_msg = {"role": msg["role"], "content": msg["content"]}
+    else:
+        role = "assistant" if isinstance(msg, AIMessage) else "user"
+        formatted_msg = {"role": role, "content": msg.content}
+    return formatted_msg
+
+def update_chat(message: str, chatbot_state: Dict) -> List[List[str]]:
+    """
+    Update the chat with a new user message and process it through the chatbot.
+
+    Args:
+        message (str): The user's message to be added to the chat.
+        chatbot_state (Dict): The current state of the chatbot.
+
+    Returns:
+        List[List[str]]: The formatted messages for display in the chat.
+    """
+    state = {"messages": [{"role": "user", "content": message}], "ask_human": False}
+    chatbot_state["messages"].append(state["messages"][0])
+
+    new_state = graph.execute("chatbot", chatbot_state)
+    chatbot_state["messages"] = new_state["messages"]
+    chatbot_state["ask_human"] = new_state["ask_human"]
+
+    formatted_messages = [format_message(msg) for msg in chatbot_state["messages"]]
+    return [[msg["role"], msg["content"]] for msg in formatted_messages]
+
+def init_chatbot() -> Dict:
+    """
+    Initialize the chatbot with an empty state.
+
+    Returns:
+        Dict: The initial state of the chatbot.
+    """
+    initial_state = {"messages": [], "ask_human": False}
+    return initial_state
+
+# Initialize the tools and chatbot
+llm_with_tools = ChatGoogleGenerativeAI(model="gemini-1.5-flash", google_api_key=os.getenv("GOOGLE_API_KEY"))
+DuckDuck_tool = DuckDuckGoSearchRun()
+toolset = [DuckDuck_tool]
+
+# Initialize the graph and add nodes and edges
+graph = Graph()
+graph.add_node(Node("chatbot", chatbot_function))
+graph.add_node(Node("toolset", tool_node_function))
+graph.add_node(Node("human", human_node_function))
+
+graph.add_edge(Edge("chatbot", "toolset", lambda state: not state.get("ask_human", False)))
+graph.add_edge(Edge("toolset", "chatbot"))
+graph.add_edge(Edge("human", "chatbot"))
+
+# Initialize Gradio interface
+with gr.Blocks() as iface:
+    chatbot_state = gr.State(init_chatbot())
+
+    with gr.Row():
+        with gr.Column():
+            user_input = gr.Textbox(label="Your message")
+            send_button = gr.Button("Send")
+            chat_output = gr.Chatbot(label="Chatbot conversation")
+
+    send_button.click(update_chat, inputs=[user_input, chatbot_state], outputs=[chat_output])
+
+iface.launch()
+