working agent

.gitignore

@@ -1,3 +1,4 @@
 *.bz2
 data
-__pycache__
+__pycache__
+.env

agent.py

@@ -1,136 +1,186 @@
-# LLM agent player using HuggingFace's smolagent
-from wiki_run_engine import WikiRunEnvironment
+import litellm
 from rich.console import Console
+from rich.panel import Panel
+from rich.markdown import Markdown
+from rich.table import Table
+from rich.box import SIMPLE
+from wiki_run_engine import WikiRunEnvironment
+from langsmith import traceable
+import os
+from openai import OpenAI
+from langsmith.wrappers import wrap_openai
 
-console = Console()
-
-try:
-    from smolagent import Agent, AgentConfig
-except ImportError:
-    console.print("[red]smolagent package not found. Please install with 'uv pip install smolagent'[/red]")
-    raise
+openai_client = wrap_openai(OpenAI())
 
-class AgentPlayer:
-    def __init__(self, wiki_data_path, model_name="HuggingFaceH4/zephyr-7b-beta"):
-        """Initialize agent player"""
+class WikiRunAgent:
+    def __init__(self, wiki_data_path, model="gemini/gemini-2.5-pro-exp-03-25"):
         self.env = WikiRunEnvironment(wiki_data_path)
+        self.model = model
+        self.console = Console()
         
-        # Initialize LLM agent
-        config = AgentConfig(
-            model=model_name,
-            output_parser="json"
-        )
-        self.agent = Agent(config)
-        
-    def play(self, start_article=None, target_article=None, max_steps=20):
-        """Play a game of Wiki Run using the LLM agent"""
-        # Reset environment
+    @traceable(name="WikiRun Game")
+    def run_game(self, start_article=None, target_article=None):
+        """Play the WikiRun game with LLM agent"""
         state = self.env.reset(start_article, target_article)
         
-        console.print("[bold]Agent Wiki Run[/bold]")
-        console.print(f"Starting article: [cyan]{state['current_article']}[/cyan]")
-        console.print(f"Target article: [red]{state['target_article']}[/red]")
-        console.print()
+        self.console.print(Panel(f"[bold cyan]Starting WikiRun![/bold cyan]"))
+        self.console.print(f"[bold green]Starting at:[/bold green] {state['current_article']}")
+        self.console.print(f"[bold red]Target:[/bold red] {state['target_article']}\n")
         
-        steps = 0
-        while not state['is_complete'] and steps < max_steps:
-            console.print(f"[bold]Step {steps + 1}:[/bold]")
-            console.print(f"Current article: [cyan]{state['current_article']}[/cyan]")
-            
-            # Create prompt for agent
-            prompt = self._create_agent_prompt(state)
-            
-            # Get agent's decision
-            tool_result = self.agent.run(
-                prompt,
-                tools=[
-                    {
-                        "name": "choose_next_article",
-                        "description": "Choose the next Wikipedia article to navigate to",
-                        "parameters": {
-                            "type": "object",
-                            "properties": {
-                                "article": {
-                                    "type": "string",
-                                    "description": "The title of the next article to navigate to"
-                                },
-                                "reasoning": {
-                                    "type": "string",
-                                    "description": "Explanation of why this article was chosen"
-                                }
-                            },
-                            "required": ["article", "reasoning"]
-                        }
-                    }
-                ]
-            )
+        while not state['is_complete']:
+            # Display current game status
+            self._display_game_status(state)
             
-            # Extract agent's choice
-            choice = tool_result.get("choose_next_article", {})
-            next_article = choice.get("article", "")
-            reasoning = choice.get("reasoning", "")
+            # Get LLM's choice
+            choice = self._get_llm_choice(state)
+            self.console.print(f"\n[bold yellow]Agent chooses:[/bold yellow] {choice}")
             
-            console.print(f"[yellow]Agent reasoning: {reasoning}[/yellow]")
-            console.print(f"Agent chooses: [green]{next_article}[/green]")
-            
-            # Verify the choice is valid
-            if next_article in state['available_links']:
-                state, message = self.env.step(next_article)
-                if message:
-                    console.print(f"[bold]{message}[/bold]")
-            else:
-                console.print("[red]Invalid choice! Agent selected an article that's not in the available links.[/red]")
-                # Choose a random valid link as fallback
-                import random
-                next_article = random.choice(state['available_links'])
-                console.print(f"[yellow]Falling back to random choice: {next_article}[/yellow]")
-                state, _ = self.env.step(next_article)
+            # Process the choice
+            available_links = self._get_available_links(state['available_links'])
+            if not available_links:
+                self.console.print("[bold red]No available links to choose from![/bold red]")
+                break
                 
-            steps += 1
-            console.print()
-            
-        # Game complete
-        if state['is_complete']:
-            console.print("[bold green]Success! Agent reached the target article![/bold green]")
-        else:
-            console.print("[bold red]Failed to reach target within step limit.[/bold red]")
-            
-        console.print(f"Steps taken: [bold]{state['steps_taken']}[/bold]")
-        console.print(f"Path: [italic]{' → '.join(state['path_taken'])}[/italic]")
-        
+            try:
+                # If choice is a number
+                idx = int(choice) - 1
+                if 0 <= idx < len(available_links):
+                    next_article = available_links[idx]
+                    self.console.print(f"[bold cyan]Moving to:[/bold cyan] {next_article}\n")
+                    state, message = self.env.step(next_article)
+                    if message:
+                        self.console.print(f"[bold]{message}[/bold]")
+                else:
+                    self.console.print("[bold red]Invalid choice. Trying again.[/bold red]\n")
+            except ValueError:
+                self.console.print("[bold red]Invalid choice format. Trying again.[/bold red]\n")
+                
+        self.console.print(Panel(f"[bold green]Game completed in {state['steps_taken']} steps[/bold green]"))
+        self.console.print(f"[bold]Path:[/bold] {' → '.join(state['path_taken'])}")
         return state
     
-    def _create_agent_prompt(self, state):
-        """Create prompt for the agent"""
+    def _display_game_status(self, state):
+        """Display current game status with rich formatting"""
+        # Display current article
+        self.console.print(Panel(f"[bold cyan]{state['current_article']}[/bold cyan]", 
+                           expand=False, 
+                           border_style="cyan"))
+        
+        # Display article links
+        self.console.print("[bold green]Available Links:[/bold green]")
+        self._display_links(state['available_links'])
+        
+        # Display path so far
+        self.console.print(f"\n[bold yellow]Steps taken:[/bold yellow] {state['steps_taken']}")
+        if state['path_taken']:
+            self.console.print(f"[bold yellow]Path so far:[/bold yellow] {' → '.join(state['path_taken'])}")
+    
+    def _display_links(self, links):
+        """Display links in a nicely formatted table"""
+        table = Table(show_header=False, box=SIMPLE)
+        table.add_column("Number", style="dim")
+        table.add_column("Link", style="green")
+        table.add_column("Available", style="bold")
+        
+        for i, link in enumerate(links):
+            # Check if link is available
+            is_available = link in self.env.wiki_data
+            status = "[green]✓[/green]" if is_available else "[red]✗[/red]"
+            color = "green" if is_available else "red"
+            table.add_row(
+                f"{i+1}", 
+                f"[{color}]{link}[/{color}]", 
+                status
+            )
+            
+        self.console.print(table)
+    
+    def _get_available_links(self, links):
+        """Filter links to only those available in the wiki data"""
+        return [link for link in links if link in self.env.wiki_data]
+    
+    @traceable(name="LLM Decision")
+    def _get_llm_choice(self, state):
+        """Ask LLM for next move"""
         current = state['current_article']
         target = state['target_article']
-        links = state['available_links']
+        all_links = state['available_links']
+        available_links = self._get_available_links(all_links)
+        path_so_far = state['path_taken']
         
-        prompt = f"""You are playing the Wiki Run game. Your goal is to navigate from the current Wikipedia article to the target article using only the available links.
-
+        # Create prompt with relevant context (not the full article)
+        prompt = f"""You are playing WikiRun, trying to navigate from one Wikipedia article to another using only links.
 Current article: {current}
 Target article: {target}
+Available links (numbered):
+{self._format_links(available_links)}
 
-Available links (choose one):
-{', '.join(links)}
-
-Choose the link that you think will get you closest to the target article. Consider:
-1. Direct connections to the target
-2. Articles that might be in the same category as the target
-3. General articles that might have many links to other topics
+Your path so far: {' -> '.join(path_so_far)}
 
-Use the choose_next_article tool to make your selection."""
+Think about which link is most likely to lead you toward the target article.
+First, think step by step about your strategy.
+Then output your choice as a number in this format: <choice>N</choice> where N is the link number.
+"""
         
-        return prompt
+        # Call LLM via litellm with langsmith tracing
+        response = litellm.completion(
+            model=self.model,
+            messages=[{"role": "user", "content": prompt}],
+            # metadata={
+            #     "current_article": current,
+            #     "target_article": target,
+            #     "available_links": available_links,
+            #     "steps_taken": state['steps_taken'],
+            #     "path_so_far": path_so_far
+            # }
+        )
+        
+        # Extract the choice from response
+        content = response.choices[0].message.content
+        self.console.print(Panel(Markdown(content), title="[bold]Agent Thinking[/bold]", border_style="yellow"))
+        
+        # Extract choice using format <choice>N</choice>
+        import re
+        choice_match = re.search(r'<choice>(\d+)</choice>', content)
+        if choice_match:
+            return choice_match.group(1)
+        else:
+            # Fallback: try to find any number in the response
+            numbers = re.findall(r'\d+', content)
+            if numbers:
+                for num in numbers:
+                    if 1 <= int(num) <= len(available_links):
+                        return num
+            # Default to first link if no valid choice found
+            return "1" if available_links else "0"
+    
+    def _format_links(self, links):
+        """Format the list of links for the prompt"""
+        return "\n".join([f"{i+1}. {link}" for i, link in enumerate(links)])
+
+def setup_langsmith():
+    """Print instructions for setting up LangSmith tracing"""
+    console = Console()
+    console.print(Panel("[bold yellow]LangSmith Setup Instructions[/bold yellow]"))
+    console.print("To enable LangSmith tracing, set the following environment variables:")
+    console.print("[bold]export LANGSMITH_API_KEY='your-api-key'[/bold]")
+    console.print("Get your API key from: https://smith.langchain.com/settings")
+    console.print("Once set, your WikiRun agent will log traces to your LangSmith dashboard")
 
 if __name__ == "__main__":
     import sys
     
     if len(sys.argv) < 2:
-        console.print("[red]Please provide the path to Wikipedia data[/red]")
+        console = Console()
+        console.print("[bold red]Please provide the path to Wikipedia data[/bold red]")
         console.print("Usage: python agent.py <wiki_data_path>")
         sys.exit(1)
+    
+    # Remind about LangSmith setup
+    if not os.environ.get("LANGSMITH_API_KEY"):
+        setup_langsmith()
         
     wiki_data_path = sys.argv[1]
-    agent = AgentPlayer(wiki_data_path)
-    agent.play()
+    agent = WikiRunAgent(wiki_data_path)
+    agent.run_game(start_article="Peanut", target_article="Silicon Valley")
+    # agent.run_game(start_article="Silicon Valley", target_article="Peanut")