Update app.py

app.py

@@ -1,102 +1,158 @@
 import os
 import json
 import gradio as gr
-from qwen_agent.llm import get_chat_model
+from huggingface_hub import InferenceClient
+from typing import Dict, Any, List
 
 # Define a static weather tool function
 def get_current_weather(location, unit="fahrenheit"):
     """Get the current weather in a given location"""
     if "tokyo" in location.lower():
-        return json.dumps({"location": "Tokyo", "temperature": "10", "unit": "celsius"})
+        return {"location": "Tokyo", "temperature": "10", "unit": "celsius"}
     elif "san francisco" in location.lower():
-        return json.dumps({"location": "San Francisco", "temperature": "72", "unit": "fahrenheit"})
+        return {"location": "San Francisco", "temperature": "72", "unit": "fahrenheit"}
     elif "paris" in location.lower():
-        return json.dumps({"location": "Paris", "temperature": "22", "unit": "celsius"})
+        return {"location": "Paris", "temperature": "22", "unit": "celsius"}
     else:
-        return json.dumps({"location": location, "temperature": "unknown"})
+        return {"location": location, "temperature": "unknown", "unit": unit}
 
-# Set up the weather function definition
+class HfApiModel:
+    def __init__(self, max_tokens=2096, temperature=0.5, model_id='Qwen/Qwen2.5-Coder-32B-Instruct', token=None):
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.model_id = model_id
+        self.client = InferenceClient(model=model_id, token=token)
+    
+    def generate_with_function_calling(self, messages, functions):
+        try:
+            # Format messages for the model
+            response = self.client.chat_completion(
+                messages=messages,
+                max_tokens=self.max_tokens,
+                temperature=self.temperature,
+                tools=functions,
+                tool_choice="auto"
+            )
+            
+            return response
+        except Exception as e:
+            print(f"Error in generate: {str(e)}")
+            return {"error": str(e)}
+    
+    def call_function(self, function_name, arguments):
+        if function_name == "get_current_weather":
+            location = arguments.get("location", "")
+            unit = arguments.get("unit", "fahrenheit")
+            return get_current_weather(location, unit)
+        return {"error": f"Function {function_name} not found"}
+
+# Initialize the model
+def init_model():
+    token = os.environ.get("HUGGINGFACE_TOKEN")
+    return HfApiModel(
+        max_tokens=2096,
+        temperature=0.5,
+        model_id='Qwen/Qwen2.5-Coder-32B-Instruct'
+    )
+
+# Define the weather function schema
 weather_function = {
-    'name': 'get_current_weather',
-    'description': 'Get the current weather in a given location',
-    'parameters': {
-        'type': 'object',
-        'properties': {
-            'location': {
-                'type': 'string',
-                'description': 'The city and state, e.g. San Francisco, CA',
-            },
-            'unit': {
-                'type': 'string',
-                'enum': ['celsius', 'fahrenheit'],
-                'description': 'The unit of temperature to use. Infer this from the user\'s location.'
+    "type": "function",
+    "function": {
+        "name": "get_current_weather",
+        "description": "Get the current weather in a given location",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "location": {
+                    "type": "string",
+                    "description": "The city and state, e.g. San Francisco, CA"
+                },
+                "unit": {
+                    "type": "string",
+                    "enum": ["celsius", "fahrenheit"],
+                    "description": "The unit of temperature to use. Infer this from the user's location."
+                }
             },
-        },
-        'required': ['location'],
-    },
+            "required": ["location"]
+        }
+    }
 }
 
-# Initialize the Qwen model
-def init_model():
-    llm = get_chat_model({
-        'model': 'Qwen/Qwen2.5-Coder-32B-Instruct',
-        'endpoint_type': 'huggingface_hub',
-        'token': os.environ.get("HUGGINGFACE_TOKEN"),
-    })
-    return llm
-
 # Processing function for Gradio
 def process_message(message, history):
     # Initialize model on first run
-    if not hasattr(process_message, "llm"):
-        process_message.llm = init_model()
+    if not hasattr(process_message, "model"):
+        process_message.model = init_model()
     
-    # Set up the conversation
-    messages = [{'role': 'user', 'content': message}]
-    functions = [weather_function]
+    # Format conversation history for the model
+    formatted_history = []
+    for human, assistant in history:
+        formatted_history.append({"role": "user", "content": human})
+        if assistant:  # Check if assistant response exists
+            formatted_history.append({"role": "assistant", "content": assistant})
+    
+    # Add the current message
+    messages = formatted_history + [{"role": "user", "content": message}]
     
-    # Step 1: Get the initial response
     try:
-        *_, response = process_message.llm.chat(
+        # Get response from the model
+        response = process_message.model.generate_with_function_calling(
             messages=messages,
-            functions=functions,
-            stream=True,
+            functions=[weather_function]
         )
         
-        # Step 2: Check if the model wanted to call a function
-        if response.get('function_call', None):
-            # Step 3: Call the function
-            function_name = response['function_call']['name']
-            function_args = json.loads(response['function_call']['arguments'])
+        # Check if there's a function call
+        if hasattr(response, "choices") and response.choices:
+            message_content = response.choices[0].message
             
-            # Only process weather function calls
-            if function_name == 'get_current_weather':
-                function_response = get_current_weather(
-                    location=function_args.get('location'),
-                    unit=function_args.get('unit'),
-                )
+            # Check if the model wants to call a function
+            if hasattr(message_content, "tool_calls") and message_content.tool_calls:
+                tool_call = message_content.tool_calls[0]
+                function_name = tool_call.function.name
+                function_args = json.loads(tool_call.function.arguments)
+                
+                # Call the function
+                function_result = process_message.model.call_function(function_name, function_args)
+                
+                # Add the function result to messages
+                messages.append({
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [{
+                        "id": tool_call.id,
+                        "type": "function",
+                        "function": {
+                            "name": function_name,
+                            "arguments": tool_call.function.arguments
+                        }
+                    }]
+                })
                 
-                # Step 4: Send the function result back to the model
-                messages.append(response)  # Add the model's response with function call
                 messages.append({
-                    'role': 'function',
-                    'name': function_name,
-                    'content': function_response,
+                    "role": "tool",
+                    "tool_call_id": tool_call.id,
+                    "content": json.dumps(function_result)
                 })
                 
-                # Get final response from the model
-                *_, final_response = process_message.llm.chat(
+                # Get final response
+                final_response = process_message.model.generate_with_function_calling(
                     messages=messages,
-                    functions=functions,
-                    stream=False,
+                    functions=[weather_function]
                 )
-                return final_response['content']
-        
-        # If no function was called, return the initial response
-        return response['content']
+                
+                if hasattr(final_response, "choices") and final_response.choices:
+                    return final_response.choices[0].message.content
+                return "Error processing function result"
+            
+            # If no function call, return the content directly
+            if hasattr(message_content, "content"):
+                return message_content.content
+            
+        return "I couldn't process that request properly. Please try again."
     
     except Exception as e:
-        return f"Error processing your request: {str(e)}"
+        return f"Error: {str(e)}"
 
 # Set up the Gradio interface
 with gr.Blocks(title="Qwen Weather Assistant") as demo: