Update app.py

app.py

@@ -1,15 +1,18 @@
 import gradio as gr
-from transformers.agents import Tool, HfAgent
 import requests
 import os
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+# --- Load an open LLM (here Mistral 7B) ---
+model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 # --- SAP Sales Order Header tool ---
-def query_sales_order_header(query: str) -> str:
-    """
-    Query SAP API for Sales Order Header information (top 5 records).
-    """
+def query_sales_order_header():
     api_url = "https://sandbox.api.sap.com/s4hanacloud/sap/opu/odata/sap/API_SALES_ORDER_SRV/A_SalesOrder?$top=5&$inlinecount=allpages"
-    api_key = os.getenv("SAP_SANDBOX_API_KEY", "YOUR_API_KEY")  # Set in Space secrets!
+    api_key = os.getenv("SAP_SANDBOX_API_KEY", "YOUR_API_KEY")
     headers = {
         "APIKey": api_key,
         "Accept": "application/json"
@@ -36,48 +39,36 @@ def query_sales_order_header(query: str) -> str:
     except Exception as e:
         return f"Error fetching Sales Orders: {e}"
 
-# --- Register the tool for the agent ---
-sap_so_tool = Tool(
-    name="SalesOrderHeader",
-    description="Fetch SAP Sales Order Header data.",
-    function=query_sales_order_header
-)
-
-# --- Agent setup (add more tools if needed) ---
-agent = HfAgent(
-    repo_id="mistralai/Mistral-7B-Instruct-v0.1",
-    tools=[sap_so_tool]
-)
-
-# --- Chat with agent & simple memory ---
-def run_agent(input_text, chat_history):
-    # Prepend chat memory for better context
-    full_input = ""
-    if chat_history:
-        for turn in chat_history:
-            full_input += f"User: {turn[0]}\nAgent: {turn[1]}\n"
-    full_input += f"User: {input_text}\nAgent:"
-    output = agent.run(full_input)
-    chat_history = chat_history or []
-    chat_history.append((input_text, output))
-    return chat_history, chat_history
+# --- Main chat function ---
+def chat_agent(message, history):
+    # Rudimentary tool use: check for sales order question
+    tool_trigger = ["sales order", "sap order", "sap sales"]
+    if any(kw in message.lower() for kw in tool_trigger):
+        tool_response = query_sales_order_header()
+        response = f"(SAP Sales Orders)\n{tool_response}"
+    else:
+        # Pure LLM response
+        response = llm(message, max_new_tokens=256)[0]["generated_text"]
+    history = history or []
+    history.append((message, response))
+    return history, history
 
 # --- Gradio UI ---
 with gr.Blocks() as demo:
     gr.Markdown(
         """
         # SAP Sales Order Chat Agent  
-        Ask about sales orders!  
+        Ask about SAP sales orders or general questions.  
         Example:  
         - `Show me the latest SAP sales orders.`  
-        - `Give me the net value of top sales orders.`  
+        - `Tell me about SAP order values.`  
         """
     )
     chatbot = gr.Chatbot()
     txt = gr.Textbox(label="Your question")
     clear = gr.Button("Clear chat")
 
-    txt.submit(run_agent, [txt, chatbot], [chatbot, chatbot])
+    txt.submit(chat_agent, [txt, chatbot], [chatbot, chatbot])
     clear.click(lambda: ([], []), None, [chatbot, chatbot])
 
 demo.launch()