Updated the app

app.py

@@ -10,105 +10,96 @@ 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
+# Load environment variables (including OPENAI_API_KEY) from .env if present.
 load_dotenv()
 
-# Define a TypedDict to hold state information.
+# -------------------------------------------------------
+# Helpers to get ChatOpenAI models using the environment key
+# -------------------------------------------------------
+def get_llm():
+    """
+    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.
+    """
+    return ChatOpenAI(
+        model="gpt-4o",   # Adjust the model name as needed
+        temperature=0,
+        openai_api_key=os.getenv("OPENAI_API_KEY")
+    )
+
+
+# -------------------------------------------------------
+# TypedDict for State
+# -------------------------------------------------------
 class State(TypedDict):
     query: str
     category: str
     sentiment: str
     response: str
 
-def init_llm(api_key: str):
-    os.environ["OPENAI_API_KEY"] = api_key
-    return ChatOpenAI(temperature=0)
-
-def process_query(api_key: str, query: str):
-    try:
-        llm = init_llm(api_key)
-        agent = Agent(
-            task=query,
-            llm=llm
-        )
-        result = asyncio.run(agent.run())
-        return result
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Customer Support Agent")
-    with gr.Row():
-        api_key = gr.Textbox(
-            label="OpenAI API Key", 
-            placeholder="Enter your OpenAI API key",
-            type="password"
-        )
-    with gr.Row():
-        query = gr.Textbox(
-            label="Query",
-            placeholder="Enter your query here"
-        )
-    with gr.Row():
-        submit = gr.Button("Submit")
-    with gr.Row():
-        output = gr.Textbox(label="Response")
-    
-    submit.click(
-        fn=process_query,
-        inputs=[api_key, query],
-        outputs=output
-    )
-# Initialize our language models.
-# We use llm_standard for normal tasks and llm_browser for browser-based tasks.
-llm_standard = ChatOpenAI(temperature=0)
-llm_browser = ChatOpenAI(model="gpt-4o", temperature=0)
 
-# Node functions for our workflow.
+# -------------------------------------------------------
+# Node functions for our workflow
+# (using get_llm() or get_llm_browser() on-demand)
+# -------------------------------------------------------
 def categorize(state: State) -> State:
     prompt = ChatPromptTemplate.from_template(
         "Categorize the following customer query into one of these categories: "
         "Technical, Billing, General. Query: {query}"
     )
-    chain = prompt | llm_standard
+    chain = prompt | get_llm()
     category = chain.invoke({"query": state["query"]}).content.strip()
-    return {"category": category}
+    state["category"] = category
+    return 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}"
     )
-    chain = prompt | llm_standard
+    chain = prompt | get_llm()
     sentiment = chain.invoke({"query": state["query"]}).content.strip()
-    return {"sentiment": sentiment}
+    state["sentiment"] = sentiment
+    return state
 
 def handle_technical(state: State) -> State:
     prompt = ChatPromptTemplate.from_template(
         "Provide a technical support response to the following query: {query}"
     )
-    chain = prompt | llm_standard
+    chain = prompt | get_llm()
     response = chain.invoke({"query": state["query"]}).content.strip()
-    return {"response": response}
+    state["response"] = response
+    return state
 
 def handle_billing(state: State) -> State:
     prompt = ChatPromptTemplate.from_template(
         "Provide a billing support response to the following query: {query}"
     )
-    chain = prompt | llm_standard
+    chain = prompt | get_llm()
     response = chain.invoke({"query": state["query"]}).content.strip()
-    return {"response": response}
+    state["response"] = response
+    return state
 
 async def run_browser_agent(task: str) -> str:
-    # Run the browser-use agent asynchronously.
-    agent = Agent(task=task, llm=llm_browser)
+    """
+    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.
-    We call the async function with asyncio.run and then extract only the final answer.
     """
     task = (
         "You are a customer support agent that consults online sources. "
@@ -116,7 +107,7 @@ def handle_general(state: State) -> State:
     )
     result = asyncio.run(run_browser_agent(task))
     final_text = ""
-    
+
     if isinstance(result, str):
         final_text = result.strip()
     elif hasattr(result, "all_results"):
@@ -130,14 +121,18 @@ def handle_general(state: State) -> State:
             final_text = str(result).strip()
     else:
         final_text = str(result).strip()
-    
-    return {"response": final_text}
+
+    state["response"] = final_text
+    return state
 
 def escalate(state: State) -> State:
-    return {"response": "This query has been escalated to a human agent due to negative sentiment."}
+    state["response"] = "This query has been escalated to a human agent due to negative sentiment."
+    return state
 
 def route_query(state: State) -> str:
-    """Determine which node to route to based on sentiment and category."""
+    """
+    Determine which node to route to based on sentiment and category.
+    """
     if state["sentiment"].lower() == "negative":
         return "escalate"
     elif state["category"].lower() == "technical":
@@ -147,7 +142,10 @@ def route_query(state: State) -> str:
     else:
         return "handle_general"
 
-# Create the workflow graph.
+
+# -------------------------------------------------------
+# Create the workflow graph
+# -------------------------------------------------------
 workflow = StateGraph(State)
 workflow.add_node("categorize", categorize)
 workflow.add_node("analyze_sentiment", analyze_sentiment)
@@ -172,31 +170,52 @@ 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 callable function
 app = workflow.compile()
 
+
+# -------------------------------------------------------
+# Gradio callback function
+# -------------------------------------------------------
 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 workflow from 'app' on the user's query.
+    - Return the final response from the workflow.
+    """
+    # If no key is provided by user and none is in environment, show error.
     if not api_key and not os.getenv("OPENAI_API_KEY"):
-        return "Error: Please provide an OpenAI API key"
-    
+        return "Error: Please provide an OpenAI API key."
+
+    # Set user-provided key in environment
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
+
     try:
-        if api_key:
-            os.environ["OPENAI_API_KEY"] = api_key
-            
-        llm = ChatOpenAI(temperature=0)
-        agent = Agent(
-            task=query,
-            llm=llm,
-        )
-        result = await agent.run()
-        return result
+        # Initialize the state
+        state = {
+            "query": query,
+            "category": "",
+            "sentiment": "",
+            "response": ""
+        }
+        final_state = app(state)
+        return final_state["response"]
     except Exception as e:
         return f"Error: {str(e)}"
 
 
-# Build the Gradio UI.
+# -------------------------------------------------------
+# 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.")
+    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():
@@ -218,10 +237,15 @@ with gr.Blocks(title="Customer Support Agent with Browser Use") as demo:
                 interactive=False
             )
     
+    # Note: The order of inputs here matches the function 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()
+    demo.launch()