db refactor

agent.py

@@ -84,7 +84,7 @@ class WikiRunAgent:
         
         for i, link in enumerate(links):
             # Check if link is available
-            is_available = link in self.env.wiki_data
+            is_available = self.env.article_exists(link)
             status = "[green]✓[/green]" if is_available else "[red]✗[/red]"
             color = "green" if is_available else "red"
             table.add_row(
@@ -97,7 +97,7 @@ class WikiRunAgent:
     
     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]
+        return [link for link in links if self.env.article_exists(link)]
     
     @traceable(name="LLM Decision")
     def _get_llm_choice(self, state):

cli.py

@@ -47,7 +47,7 @@ class WikiRunCLI:
         row = []
         for i, link in enumerate(links):
             # Check if link is available in the current article
-            is_available = link in self.env.wiki_data
+            is_available = self.env.article_exists(link)
             color = "green" if is_available else "red"
             row.append(f"[{color}]{i+1}. {link}[/{color}]")
             if len(row) == 3:
@@ -80,7 +80,7 @@ class WikiRunCLI:
         
         if choice == "p":
             # Solver mode
-            available_articles = list(self.env.wiki_data.keys())
+            available_articles = self.env.db.get_all_article_titles()
             console.print(f"Available articles: {len(available_articles)}")
             
             # Show a sample of available articles
@@ -94,7 +94,7 @@ class WikiRunCLI:
             target = Prompt.ask("Target article")
             
             # Initialize solver and find path
-            solver = WikiSolver(self.env.wiki_data)
+            solver = WikiSolver(self.env)
             path, error = solver.find_path(start, target)
             
             if error:
@@ -111,7 +111,7 @@ class WikiRunCLI:
             state = self.env.reset()
         else:
             # Get start and target articles
-            available_articles = list(self.env.wiki_data.keys())
+            available_articles = self.env.db.get_all_article_titles()
             console.print(f"Available articles: {len(available_articles)}")
             
             # Show a sample of available articles

wiki_db_json.py

@@ -0,0 +1,40 @@
+import json
+
+class WikiDBJson:
+    def __init__(self, wiki_data_path):
+        """Initialize the database with path to Wikipedia data"""
+        self.wiki_data = self._load_wiki_data(wiki_data_path)
+        
+    def _load_wiki_data(self, path):
+        """Load Wikipedia data from JSON file"""
+        print(f"Loading wiki data from {path}...")
+        with open(path, 'r', encoding='utf-8') as f:
+            wiki_data = json.load(f)
+        print(f"Loaded {len(wiki_data)} articles")
+        return wiki_data
+    
+    def get_article_count(self):
+        """Return the number of articles in the database"""
+        return len(self.wiki_data)
+        
+    def get_all_article_titles(self):
+        """Return a list of all article titles"""
+        return list(self.wiki_data.keys())
+        
+    def get_article(self, title):
+        """Get article data by title"""
+        return self.wiki_data.get(title, {})
+        
+    def article_exists(self, title):
+        """Check if an article exists in the database"""
+        return title in self.wiki_data
+        
+    def get_article_text(self, title):
+        """Get the text of an article"""
+        article = self.get_article(title)
+        return article.get('text', '')
+        
+    def get_article_links(self, title):
+        """Get the links of an article"""
+        article = self.get_article(title)
+        return article.get('links', []) 

wiki_run_engine.py

@@ -3,31 +3,23 @@
 # 2. Game state management
 # 3. Navigation between articles
 
-import json
 import random
+from wiki_db_json import WikiDBJson
 
 class WikiRunEnvironment:
     def __init__(self, wiki_data_path):
         """Initialize with path to Wikipedia data"""
-        self.wiki_data = self._load_wiki_data(wiki_data_path)
+        self.db = WikiDBJson(wiki_data_path)
         self.current_article = None
         self.target_article = None
         self.path_taken = []
         self.steps = 0
         
-    def _load_wiki_data(self, path):
-        """Load Wikipedia data from JSON file"""
-        print(f"Loading wiki data from {path}...")
-        with open(path, 'r', encoding='utf-8') as f:
-            wiki_data = json.load(f)
-        print(f"Loaded {len(wiki_data)} articles")
-        return wiki_data
-    
     def reset(self, start_article=None, target_article=None):
         """Reset the environment with new start/target articles"""
         if start_article is None or target_article is None:
             # Choose random articles if not specified
-            available_articles = list(self.wiki_data.keys())
+            available_articles = self.db.get_all_article_titles()
             
             if start_article is None:
                 start_article = random.choice(available_articles)
@@ -49,12 +41,11 @@ class WikiRunEnvironment:
         if self.current_article is None:
             return None
             
-        current = self.wiki_data.get(self.current_article, {})
         return {
             'current_article': self.current_article,
             'target_article': self.target_article,
-            'article_text': current.get('text', ''),
-            'available_links': current.get('links', []),
+            'article_text': self.db.get_article_text(self.current_article),
+            'available_links': self.db.get_article_links(self.current_article),
             'steps_taken': self.steps,
             'path_taken': self.path_taken,
             'is_complete': self.current_article == self.target_article
@@ -65,8 +56,7 @@ class WikiRunEnvironment:
         if self.current_article is None:
             return None, "Game not initialized. Call reset() first."
             
-        current = self.wiki_data.get(self.current_article, {})
-        available_links = current.get('links', [])
+        available_links = self.db.get_article_links(self.current_article)
         
         if next_article not in available_links:
             return self.get_current_state(), f"Invalid link: {next_article} not in available links"
@@ -80,3 +70,7 @@ class WikiRunEnvironment:
         is_complete = self.current_article == self.target_article
         
         return self.get_current_state(), "Target reached!" if is_complete else ""
+    
+    def article_exists(self, article_title):
+        """Check if an article exists in the database"""
+        return self.db.article_exists(article_title)

wiki_solver.py

@@ -7,15 +7,15 @@ from rich.progress import Progress, SpinnerColumn, TextColumn
 console = Console()
 
 class WikiSolver:
-    def __init__(self, wiki_data):
-        """Initialize the solver with wiki data"""
-        self.wiki_data = wiki_data
+    def __init__(self, wiki_run_engine):
+        """Initialize the solver with a WikiRunEnvironment instance"""
+        self.engine = wiki_run_engine
         
     def find_path(self, start_article, target_article):
         """Find the shortest path using BFS"""
-        if start_article not in self.wiki_data:
+        if not self.engine.article_exists(start_article):
             return None, "Start article not found in wiki data"
-        if target_article not in self.wiki_data:
+        if not self.engine.article_exists(target_article):
             return None, "Target article not found in wiki data"
             
         # Initialize BFS
@@ -36,11 +36,12 @@ class WikiSolver:
                     return path, None
                     
                 # Get all links from current article
-                if current in self.wiki_data:
-                    for next_article in self.wiki_data[current].get('links', []):
-                        if next_article not in visited and next_article in self.wiki_data:
-                            visited.add(next_article)
-                            queue.append((next_article, path + [next_article]))
+                # Need to set current article to get links
+                state = self.engine.reset(current, target_article)
+                for next_article in state['available_links']:
+                    if next_article not in visited and self.engine.article_exists(next_article):
+                        visited.add(next_article)
+                        queue.append((next_article, path + [next_article]))
                             
         return None, "No path found"