Upload 3 files

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 tokenizer.json filter=lfs diff=lfs merge=lfs -text
+agent.png filter=lfs diff=lfs merge=lfs -text

example/vllm_v2_extraction_agent.py

@@ -0,0 +1,287 @@
+from dataclasses import dataclass
+from typing import List, Dict, Any, Optional
+import json
+import requests
+from bs4 import BeautifulSoup
+from openai import OpenAI
+
+"""
+EXAMPLE OUTPUT:
+
+What is the current population for the city  where Einstein was born?
+
+Turn 1
+----------------------------------------
+
+Executing: fetch_wiki_content
+Arguments: {'title': 'Albert Einstein'}
+
+Turn 2
+----------------------------------------
+
+Executing: deliver_answer
+Arguments: {'fields': ['Ulm, German Empire']}
+ANSWER FROM THE ASSISTANT: ['Ulm, German Empire']
+
+Turn 3
+----------------------------------------
+
+Executing: fetch_wiki_content
+Arguments: {'title': 'Ulm'}
+
+Turn 4
+----------------------------------------
+
+Executing: deliver_answer
+Arguments: {'fields': ['128,928']}
+ANSWER FROM THE ASSISTANT: ['128,928']
+
+Turn 5
+----------------------------------------
+Extraction Complete
+
+
+Why was Einstein famous?
+
+Turn 1
+----------------------------------------
+
+Executing: fetch_wiki_content
+Arguments: {'title': 'Albert Einstein'}
+
+Turn 2
+----------------------------------------
+
+Executing: deliver_answer
+Arguments: {'fields': ['Best known for developing the theory of relativity, Einstein also made important contributions to quantum mechanics.', 'His mass–energy equivalence formula E = mc2, which arises from special relativity, has been called "the world\'s most famous equation."', 'He received the 1921 Nobel Prize in Physics.']}
+ANSWER FROM THE ASSISTANT: ['Best known for developing the theory of relativity, Einstein also made important contributions to quantum mechanics.', 'His mass–energy equivalence formula E = mc2, which arises from special relativity, has been called "the world\'s most famous equation."', 'He received the 1921 Nobel Prize in Physics.']
+
+Turn 3
+----------------------------------------
+Extraction Complete
+"""
+
+@dataclass
+class WikiConfig:
+    """Configuration for OpenAI and Wikipedia settings"""
+    api_key: str = "sk-123"
+    api_base: str = "{info}/v1"
+    model: Optional[str] = None
+    max_turns: int = 5
+    wikipedia_base_url: str = "https://en.wikipedia.org/wiki/"
+
+class WikiTools:
+    """Collection of Wikipedia and extraction tools"""
+
+    def __init__(self, base_url: str):
+        self.base_url = base_url
+
+    def fetch_wiki_content(self, title: str, section: Optional[str] = None) -> str:
+        """Fetch and clean Wikipedia article content, optionally from a specific section"""
+        url = f"{self.base_url}{title.replace(' ', '_')}"
+        response = requests.get(url)
+        soup = BeautifulSoup(response.content, 'html.parser')
+
+        # Remove unwanted sections
+        for unwanted in soup.find_all(['script', 'style', 'footer', 'header']):
+            unwanted.decompose()
+
+        if section:
+            # Find specific section if requested
+            section_tag = soup.find('span', {'id': section})
+            if section_tag:
+                content = section_tag.parent.find_next_siblings()
+                text = ' '.join(tag.get_text() for tag in content)
+            else:
+                return "Section not found"
+        else:
+            # Get main content
+            content = soup.find(id='mw-content-text')
+            if content:
+                text = content.get_text()
+            else:
+                return "Content not found"
+
+        # Clean and normalize text
+        text = ' '.join(text.split())
+        return text[:8000]  # Truncate to avoid token limits
+
+    @staticmethod
+    def deliver_answer(fields: List[str]) -> Dict[str, Any]:
+        """Extract specific information from text spans"""
+        print (f"ANSWER FROM THE ASSISTANT: {fields}")
+        return {
+            "extracted_fields": "Provided fields was delivered to the user successfully."
+        }
+
+class ToolRegistry:
+    """Registry of available tools and their schemas"""
+
+    def __init__(self, wiki_tools: WikiTools):
+        self.wiki_tools = wiki_tools
+
+    @property
+    def available_functions(self) -> Dict[str, callable]:
+        return {
+            "fetch_wiki_content": self.wiki_tools.fetch_wiki_content,
+            "deliver_answer": self.wiki_tools.deliver_answer
+        }
+
+    @property
+    def tool_schemas(self) -> List[Dict[str, Any]]:
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "fetch_wiki_content",
+                    "description": "Fetch content from a Wikipedia article",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "title": {
+                                "type": "string",
+                                "description": "The title of the Wikipedia article"
+                            },
+                            "section": {
+                                "type": "string",
+                                "description": "Optional: Specific section ID to fetch",
+                                "optional": True
+                            }
+                        },
+                        "required": ["title"]
+                    }
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "deliver_answer",
+                    "description": "Extract specific information from the fetched text",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "fields": {
+                                "type": "array",
+                                "items": {"type": "string"},
+                                "description": "List of text spans from the article that are relevant to the query"
+                            }
+                        },
+                        "required": ["fields"]
+                    }
+                }
+            }
+        ]
+
+class WikiExtractionAgent:
+    """Main agent class that handles the extraction process"""
+
+    def __init__(self, config: WikiConfig):
+        self.config = config
+        self.client = OpenAI(api_key=config.api_key, base_url=config.api_base)
+        self.wiki_tools = WikiTools(config.wikipedia_base_url)
+        self.tools = ToolRegistry(self.wiki_tools)
+        self.messages = [{"system" : "1. First fetch any wikipedia pages you might need to answer the user query. Do not answer from parametric knowledge.\n\n2.Then, provide the answer to the user using the deliver_answer from the retrieved wikipedia page.\n\n3. You may need to issue multiple calls to wikipedia after extracting answers if there are nested dependencies for information."}]
+
+        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) -> List[Dict[str, Any]]:
+        """Process and execute tool calls from assistant"""
+        results = []
+
+        for tool_call in message.tool_calls:
+            function_name = tool_call.function.name
+            function_args = json.loads(tool_call.function.arguments)
+
+            print(f"\nExecuting: {function_name}")
+            print(f"Arguments: {function_args}")
+
+            function_response = self.tools.available_functions[function_name](**function_args)
+            results.append({
+                "tool": function_name,
+                "args": function_args,
+                "response": function_response
+            })
+
+            self.messages.append({
+                "role": "tool",
+                "content": json.dumps(function_response),
+                "tool_call_id": tool_call.id,
+                "name": function_name
+            })
+
+        return results
+
+    def extract_information(self, query: str) -> List[Dict[str, Any]]:
+        """Main method to handle the extraction process"""
+        self.messages = [{
+            "role": "user",
+            "content": f"""Extract information from Wikipedia to answer this query: {query}
+
+            You can use these tools:
+            1. fetch_wiki_content: Get article content
+            2. deliver_answer: deliver relevant information
+
+            Please fetch content first, and iterate as needed to get to the webpage with the correct answer and then deliver the relevant information."""
+        }]
+
+        all_results = []
+
+        for turn in range(self.config.max_turns):
+            print(f"\nTurn {turn + 1}")
+            print("-" * 40)
+
+            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("Extraction Complete")
+                break
+
+            self.messages.append({
+                "role": "assistant",
+                "content": json.dumps(message.content),
+                "tool_calls": [self._serialize_tool_call(tc) for tc in message.tool_calls]
+            })
+
+            results = self.process_tool_calls(message)
+            all_results.extend(results)
+
+        return all_results
+
+def main():
+    # Example usage
+    config = WikiConfig()
+    agent = WikiExtractionAgent(config)
+
+    # Multi-step query example
+    results = agent.extract_information(
+        query="""What is the current population for the city  where Einstein was born?"""
+    )
+
+    # Single query example
+    results = agent.extract_information(
+        query="Why was Einstein famous?"
+    )
+
+
+if __name__ == "__main__":
+    main()

example/vllm_v2_weather_agent.py

@@ -0,0 +1,234 @@
+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?
+Turn 1
+----------------------------------------
+
+Executing: get_geo_coordinates
+Arguments: {'city': 'Paris', 'state': 'TX'}
+Response: The coordinates for Paris, TX are: latitude 33.6609, longitude 95.5555
+
+Turn 2
+----------------------------------------
+
+Executing: get_current_weather
+Arguments: {'latitude': [33.6609], 'longitude': [95.5555], 'unit': 'fahrenheit'}
+Response: The weather is 85 degrees fahrenheit. It is partly cloudy, with highs in the 90's.
+
+Turn 3
+----------------------------------------
+Conversation Complete
+
+
+****************************************
+RUNNING QUERY: Who won the most recent PGA?
+Turn 1
+----------------------------------------
+
+Executing: no_relevant_function
+Arguments: {'user_query_span': 'Who won the most recent PGA?'}
+Response: No relevant function for your request was found. We will stop here.
+
+Turn 2
+----------------------------------------
+Conversation Complete
+"""
+
+@dataclass
+class WeatherConfig:
+    """Configuration for OpenAI and API settings"""
+    api_key: str = "" # FILL IN WITH YOUR VLLM_ENDPOINT_KEY
+    api_base: str = "" # FILL IN WITH YOUR VLLM_ENDPOINT
+    model: Optional[str] = None
+    max_turns: 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."
+
+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,
+        }
+
+    @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"]
+                  }
+                }
+            }
+        ]
+
+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)
+
+            print(f"\nExecuting: {function_name}")
+            print(f"Arguments: {function_args}")
+
+            function_response = self.tools.available_functions[function_name](**function_args)
+            print(f"Response: {function_response}")
+
+            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": "user", "content": initial_query}]
+
+        print ("\n" * 5)
+        print ("*" * 40)
+        print (f"RUNNING QUERY: {initial_query}")
+
+        for turn in range(self.config.max_turns):
+            print(f"\nTurn {turn + 1}")
+            print("-" * 40)
+
+            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
+
+            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 turn >= self.config.max_turns - 1:
+            print("Maximum turns 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()