Updated Script

app.py

@@ -5,50 +5,46 @@ from typing import Dict, TypedDict
 
 import gradio as gr
 
-from langgraph.graph import StateGraph, END
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_openai import ChatOpenAI
 from browser_use import Agent
 
-# Load environment variables (including OPENAI_API_KEY) from .env if present.
+# ──────────────────────────────────────────────────────────────────────────
+# 1) Load environment
+# ──────────────────────────────────────────────────────────────────────────
 load_dotenv()
 
-# -------------------------------------------------------
-# Helpers to get ChatOpenAI models using the environment key
-# -------------------------------------------------------
+# ──────────────────────────────────────────────────────────────────────────
+# 2) Helper to get ChatOpenAI from environment
+# ──────────────────────────────────────────────────────────────────────────
 def get_llm():
-    """
-    Returns a ChatOpenAI instance using the OPENAI_API_KEY from environment.
-    """
+    """Returns a ChatOpenAI instance using the OPENAI_API_KEY from environment."""
     return ChatOpenAI(
         temperature=0,
         openai_api_key=os.getenv("OPENAI_API_KEY")
     )
 
 def get_llm_browser():
-    """
-    Returns a ChatOpenAI instance (GPT-4 or your custom model)
-    using the OPENAI_API_KEY from environment.
-    """
+    """Returns a ChatOpenAI instance for the browser agent (e.g., GPT-4) from environment."""
     return ChatOpenAI(
-        model="gpt-4o",   # Adjust the model name as needed
+        model="gpt-4o",  # Adjust if needed
         temperature=0,
         openai_api_key=os.getenv("OPENAI_API_KEY")
     )
 
-# -------------------------------------------------------
-# TypedDict for State
-# -------------------------------------------------------
+# ──────────────────────────────────────────────────────────────────────────
+# 3) TypedDict for internal state
+# ──────────────────────────────────────────────────────────────────────────
 class State(TypedDict):
     query: str
     category: str
     sentiment: str
     response: str
 
-# -------------------------------------------------------
-# Node functions for our workflow
-# (using get_llm() or get_llm_browser() on-demand)
-# -------------------------------------------------------
+# ──────────────────────────────────────────────────────────────────────────
+# 4) Individual node-like functions
+# ──────────────────────────────────────────────────────────────────────────
+
 def categorize(state: State) -> State:
     prompt = ChatPromptTemplate.from_template(
         "Categorize the following customer query into one of these categories: "
@@ -62,7 +58,8 @@ def categorize(state: State) -> State:
 def analyze_sentiment(state: State) -> State:
     prompt = ChatPromptTemplate.from_template(
         "Analyze the sentiment of the following customer query. "
-        "Respond with either 'Positive', 'Neutral', or 'Negative'. Query: {query}"
+        "Respond with either 'Positive', 'Neutral', or 'Negative'. "
+        "Query: {query}"
     )
     chain = prompt | get_llm()
     sentiment = chain.invoke({"query": state["query"]}).content.strip()
@@ -88,17 +85,13 @@ def handle_billing(state: State) -> State:
     return state
 
 async def run_browser_agent(task: str) -> str:
-    """
-    Helper to run the browser-use Agent asynchronously.
-    """
+    """Helper to run the browser-use Agent asynchronously."""
     agent = Agent(task=task, llm=get_llm_browser())
     result = await agent.run()
     return result
 
 def handle_general(state: State) -> State:
-    """
-    For general queries, we use the browser agent to consult online resources.
-    """
+    """For general queries, we use the browser agent to consult online resources."""
     task = (
         "You are a customer support agent that consults online sources. "
         f"Provide a detailed, informed response to this customer query: {state['query']}"
@@ -109,12 +102,10 @@ def handle_general(state: State) -> State:
     if isinstance(result, str):
         final_text = result.strip()
     elif hasattr(result, "all_results"):
-        # Iterate over the list of ActionResults to extract the final done answer
+        # Check if any ActionResults are "done" with extracted content
         for action in result.all_results:
-            # Check if the action is marked as done and has extracted content
             if action.get("is_done") and action.get("extracted_content"):
-                final_text = action.get("extracted_content").strip()
-        # Fallback in case no done action is found
+                final_text = action["extracted_content"].strip()
         if not final_text:
             final_text = str(result).strip()
     else:
@@ -129,7 +120,7 @@ def escalate(state: State) -> State:
 
 def route_query(state: State) -> str:
     """
-    Determine which node to route to based on sentiment and category.
+    Determine which function to use based on sentiment and category.
     """
     if state["sentiment"].lower() == "negative":
         return "escalate"
@@ -140,78 +131,78 @@ def route_query(state: State) -> str:
     else:
         return "handle_general"
 
-# -------------------------------------------------------
-# Create the workflow graph
-# -------------------------------------------------------
-workflow = StateGraph(State)
-workflow.add_node("categorize", categorize)
-workflow.add_node("analyze_sentiment", analyze_sentiment)
-workflow.add_node("handle_technical", handle_technical)
-workflow.add_node("handle_billing", handle_billing)
-workflow.add_node("handle_general", handle_general)
-workflow.add_node("escalate", escalate)
-
-workflow.add_edge("categorize", "analyze_sentiment")
-workflow.add_conditional_edges(
-    "analyze_sentiment",
-    route_query,
-    {
-        "handle_technical": "handle_technical",
-        "handle_billing": "handle_billing",
-        "handle_general": "handle_general",
-        "escalate": "escalate"
-    }
-)
-workflow.add_edge("handle_technical", END)
-workflow.add_edge("handle_billing", END)
-workflow.add_edge("handle_general", END)
-workflow.add_edge("escalate", END)
-workflow.set_entry_point("categorize")
-
-# Compile the workflow into a "CompiledStateGraph" object.
-app = workflow.compile()
-
-
-# -------------------------------------------------------
-# Gradio callback function
-# -------------------------------------------------------
+# ──────────────────────────────────────────────────────────────────────────
+# 5) A simple "workflow" function (manual, no langgraph)
+# ──────────────────────────────────────────────────────────────────────────
+
+def run_workflow(state: State) -> State:
+    """
+    Manually steps through:
+      1) categorize
+      2) sentiment
+      3) route
+      4) handle_x or escalate
+    """
+    # Step 1
+    state = categorize(state)
+    # Step 2
+    state = analyze_sentiment(state)
+    # Step 3 - route
+    next_step = route_query(state)
+    
+    if next_step == "handle_technical":
+        state = handle_technical(state)
+    elif next_step == "handle_billing":
+        state = handle_billing(state)
+    elif next_step == "handle_general":
+        state = handle_general(state)
+    else:
+        state = escalate(state)
+    
+    return state
+
+# ──────────────────────────────────────────────────────────────────────────
+# 6) Gradio callback
+# ──────────────────────────────────────────────────────────────────────────
+
 async def run_customer_support(query: str, api_key: str = "") -> str:
     """
     Main function called by Gradio upon submit.
     - If user provided an API key, set it in the environment.
-    - Then run the compiled workflow on the user's query.
-    - Return the final response from the workflow.
+    - Then run the manual "workflow" on the user's query.
+    - Return the final response from the final state.
     """
+    # Check key
     if not api_key and not os.getenv("OPENAI_API_KEY"):
         return "Error: Please provide an OpenAI API key."
     
+    # Set user-provided key
     if api_key:
         os.environ["OPENAI_API_KEY"] = api_key
 
     try:
         # Initialize the state
-        state = {
+        state: State = {
             "query": query,
             "category": "",
             "sentiment": "",
             "response": ""
         }
-        # Instead of app(state) or app.run(state), call .execute(...)
-        final_state = app.execute(state)
+        final_state = run_workflow(state)  # Manually run the chain of steps
         return final_state["response"]
     except Exception as e:
         return f"Error: {str(e)}"
 
-# -------------------------------------------------------
-# Build the Gradio UI
-# -------------------------------------------------------
+# ──────────────────────────────────────────────────────────────────────────
+# 7) Build the Gradio UI
+# ──────────────────────────────────────────────────────────────────────────
 with gr.Blocks(title="Customer Support Agent with Browser Use") as demo:
     gr.Markdown("# Customer Support Agent with Browser Use")
     gr.Markdown(
         "This agent categorizes customer queries and uses a browser-based agent "
         "to provide informed answers (when the query is general)."
     )
-    
+
     with gr.Row():
         with gr.Column():
             api_key_input = gr.Textbox(
@@ -231,16 +222,13 @@ with gr.Blocks(title="Customer Support Agent with Browser Use") as demo:
                 lines=10,
                 interactive=False
             )
-    
-    # Note: The order of inputs here matches the function signature.
+
+    # The order of inputs in submit_btn.click must match run_customer_support signature
     submit_btn.click(
         fn=run_customer_support,
         inputs=[query_input, api_key_input],
         outputs=output_box
     )
 
-# -------------------------------------------------------
-# Launch in local or Hugging Face Spaces
-# -------------------------------------------------------
 if __name__ == "__main__":
     demo.launch()