Update example/vllm_v2_extraction_agent.py

example/vllm_v2_extraction_agent.py

@@ -10,53 +10,53 @@ EXAMPLE OUTPUT:
 
 What is the current population for the city  where Einstein was born?
 
-Turn 1
+Step 1
 ----------------------------------------
 
 Executing: fetch_wiki_content
 Arguments: {'title': 'Albert Einstein'}
 
-Turn 2
+Step 2
 ----------------------------------------
 
 Executing: deliver_answer
 Arguments: {'fields': ['Ulm, German Empire']}
 ANSWER FROM THE ASSISTANT: ['Ulm, German Empire']
 
-Turn 3
+Step 3
 ----------------------------------------
 
 Executing: fetch_wiki_content
 Arguments: {'title': 'Ulm'}
 
-Turn 4
+Step 4
 ----------------------------------------
 
 Executing: deliver_answer
 Arguments: {'fields': ['128,928']}
 ANSWER FROM THE ASSISTANT: ['128,928']
 
-Turn 5
+Step 5
 ----------------------------------------
 Extraction Complete
 
 
 Why was Einstein famous?
 
-Turn 1
+Step 1
 ----------------------------------------
 
 Executing: fetch_wiki_content
 Arguments: {'title': 'Albert Einstein'}
 
-Turn 2
+Step 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
+Step 3
 ----------------------------------------
 Extraction Complete
 """
@@ -67,7 +67,7 @@ class WikiConfig:
     api_key: str = "sk-123"
     api_base: str = "{info}/v1"
     model: Optional[str] = None
-    max_turns: int = 5
+    max_steps: int = 5
     wikipedia_base_url: str = "https://en.wikipedia.org/wiki/"
 
 class WikiTools:
@@ -239,8 +239,8 @@ class WikiExtractionAgent:
 
         all_results = []
 
-        for turn in range(self.config.max_turns):
-            print(f"\nTurn {turn + 1}")
+        for step in range(self.config.max_steps):
+            print(f"\nStep {step + 1}")
             print("-" * 40)
 
             response = self.client.chat.completions.create(
@@ -273,11 +273,17 @@ def main():
     agent = WikiExtractionAgent(config)
 
     # Multi-step query example
+    # The model should first issue a call to wikipedia for Einstein, extract the part from the document about where he was born
+    # and use the value from that extraction (which could contain the city name) to call another wikipedia article for the city
+    # and pull the population from it.
+    # See lines 11 to 41 for the full trace of this actual query that Athene-V2-Agent issues.
     results = agent.extract_information(
         query="""What is the current population for the city  where Einstein was born?"""
     )
-
+    
     # Single query example
+    # Here, the model should just issue a call to Einstein's wikipedia page, and extract the parts regarding his
+    # accomplishment. 
     results = agent.extract_information(
         query="Why was Einstein famous?"
     )