Create weather_with_chat.py

example/weather_with_chat.py

@@ -0,0 +1,246 @@
+from dataclasses import dataclass
+import json
+from typing import List, Dict, Any, Optional
+from openai import OpenAI
+"""
+EXAMPLE OUTPUT:
+
+****************************************
+RUNNING QUERY: What's the weather for Paris, TX in fahrenheit?
+
+Agent Issued Step 1
+----------------------------------------
+
+Agent Issued Step 2
+----------------------------------------
+
+Agent Issued Step 3
+----------------------------------------
+AGENT MESSAGE:  The current weather in Paris, TX is 85 degrees fahrenheit. It is partly cloudy, with highs in the 90s.
+Conversation Complete
+
+
+****************************************
+RUNNING QUERY: Who won the most recent PGA?
+
+Agent Issued Step 1
+----------------------------------------
+
+Agent Issued Step 2
+----------------------------------------
+AGENT MESSAGE:  I'm sorry, but I don't have the ability to provide sports information. I can help you with weather and location data. Is there anything else I can assist you with?
+Conversation Complete
+"""
+
+@dataclass
+class WeatherConfig:
+    """Configuration for OpenAI and API settings"""
+    api_key: str = "" # The VLLM api_key
+    api_base: str = "" # The VLLM api_base URL
+    model: Optional[str] = None
+    max_steps: int = 5
+
+class WeatherTools:
+    """Collection of available tools/functions for the weather agent"""
+
+    @staticmethod
+    def get_current_weather(latitude: List[float], longitude: List[float], unit: str) -> str:
+        """Get weather for given coordinates"""
+        # We are mocking the weather here, but in the real world, you will submit a request here.
+        return f"The weather is 85 degrees {unit}. It is partly cloudy, with highs in the 90's."
+
+    @staticmethod
+    def get_geo_coordinates(city: str, state: str) -> str:
+        """Get coordinates for a given city"""
+        coordinates = {
+            "Dallas": {"TX": (32.7767, -96.7970)},
+            "San Francisco": {"CA": (37.7749, -122.4194)},
+            "Paris": {"TX": (33.6609, 95.5555)}
+        }
+        lat, lon = coordinates.get(city, {}).get(state, (0, 0))
+        # We are mocking the weather here, but in the real world, you will submit a request here.
+        return f"The coordinates for {city}, {state} are: latitude {lat}, longitude {lon}"
+
+    @staticmethod
+    def no_relevant_function(user_query_span : str) -> str:
+        return "No relevant function for your request was found. We will stop here."
+
+    @staticmethod
+    def chat(chat_string : str):
+        print ("AGENT MESSAGE: ", chat_string)
+
+class ToolRegistry:
+    """Registry of available tools and their schemas"""
+
+    @property
+    def available_functions(self) -> Dict[str, callable]:
+        return {
+            "get_current_weather": WeatherTools.get_current_weather,
+            "get_geo_coordinates": WeatherTools.get_geo_coordinates,
+            "no_relevant_function" : WeatherTools.no_relevant_function,
+            "chat" : WeatherTools.chat
+        }
+
+    @property
+    def tool_schemas(self) -> List[Dict[str, Any]]:
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_current_weather",
+                    "description": "Get the current weather in a given location. Use exact coordinates.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "latitude": {"type": "array", "description": "The latitude for the city."},
+                            "longitude": {"type": "array", "description": "The longitude for the city."},
+                            "unit": {
+                                "type": "string",
+                                "description": "The unit to fetch the temperature in",
+                                "enum": ["celsius", "fahrenheit"]
+                            }
+                        },
+                        "required": ["latitude", "longitude", "unit"]
+                    }
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_geo_coordinates",
+                    "description": "Get the latitude and longitude for a given city",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "city": {"type": "string", "description": "The city to find coordinates for"},
+                            "state": {"type": "string", "description": "The two-letter state abbreviation"}
+                        },
+                        "required": ["city", "state"]
+                    }
+                }
+            },
+            {
+                "type": "function",
+                "function" : {
+                  "name": "no_relevant_function",
+                  "description": "Call this when no other provided function can be called to answer the user query.",
+                  "parameters": {
+                    "type": "object",
+                    "properties": {
+                      "user_query_span": {
+                        "type": "string",
+                        "description": "The part of the user_query that cannot be answered by any other function calls."
+                      }
+                    },
+                    "required": ["user_query_span"]
+                  }
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "chat",
+                    "description": "Call this tool when you want to chat with the user. The user won't see anything except for whatever you pass into this function.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "chat_string": {
+                                "type": "string",
+                                "description": "The string to send to the user to chat back to them.",
+                            }
+                        },
+                        "required": ["chat_string"],
+                    },
+                },
+            },
+        ]
+
+class WeatherAgent:
+    """Main agent class that handles the conversation and tool execution"""
+
+    def __init__(self, config: WeatherConfig):
+        self.config = config
+        self.client = OpenAI(api_key=config.api_key, base_url=config.api_base)
+        self.tools = ToolRegistry()
+        self.messages = []
+
+        if not config.model:
+            models = self.client.models.list()
+            self.config.model = models.data[0].id
+
+    def _serialize_tool_call(self, tool_call) -> Dict[str, Any]:
+        """Convert tool call to serializable format"""
+        return {
+            "id": tool_call.id,
+            "type": tool_call.type,
+            "function": {
+                "name": tool_call.function.name,
+                "arguments": tool_call.function.arguments
+            }
+        }
+
+    def process_tool_calls(self, message) -> None:
+        """Process and execute tool calls from assistant"""
+        for tool_call in message.tool_calls:
+            function_name = tool_call.function.name
+            function_args = json.loads(tool_call.function.arguments)
+
+            function_response = self.tools.available_functions[function_name](**function_args)
+
+            self.messages.append({
+                "role": "tool",
+                "content": json.dumps(function_response),
+                "tool_call_id": tool_call.id,
+                "name": function_name
+            })
+
+    def run_conversation(self, initial_query: str) -> None:
+        """Run the main conversation loop"""
+        self.messages = [
+                {"role" : "system", "content" : "Make sure to use the chat() function to provide the final answer to the user."},
+                {"role": "user", "content": initial_query}]
+
+        print ("\n" * 5)
+        print ("*" * 40)
+        print (f"RUNNING QUERY: {initial_query}")
+
+        for step in range(self.config.max_steps):
+
+            response = self.client.chat.completions.create(
+                messages=self.messages,
+                model=self.config.model,
+                tools=self.tools.tool_schemas,
+                temperature=0.0,
+            )
+
+            message = response.choices[0].message
+
+            if not message.tool_calls:
+                print("Conversation Complete")
+                break
+
+            print(f"\nAgent Issued Step {step + 1}")
+            print("-" * 40)
+
+            self.messages.append({
+                "role": "assistant",
+                "content": json.dumps(message.content),
+                "tool_calls": [self._serialize_tool_call(tc) for tc in message.tool_calls]
+            })
+
+            self.process_tool_calls(message)
+
+        if step >= self.config.max_steps - 1:
+            print("Maximum steps reached")
+
+def main():
+    # Example usage
+    config = WeatherConfig()
+    agent = WeatherAgent(config)
+    agent.run_conversation("What's the weather for Paris, TX in fahrenheit?")
+
+    # Example OOD usage
+    agent.run_conversation("Who won the most recent PGA?")
+
+if __name__ == "__main__":
+    main()