File size: 3,352 Bytes
fa70ae5
 
 
 
 
 
e91ced9
2c2bab6
76887e4
fa70ae5
 
 
e91ced9
 
76887e4
e91ced9
76887e4
 
 
 
e91ced9
fa70ae5
 
 
 
 
 
 
 
 
4432909
fa70ae5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4432909
 
fa70ae5
 
 
 
 
 
 
 
 
 
4432909
fa70ae5
 
 
 
 
 
 
 
 
 
 
 
 
4432909
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
# Core game engine that manages:
# 1. Loading/parsing Wikipedia data
# 2. Game state management
# 3. Navigation between articles

import random
from db.wiki_db_sqlite import WikiDBSqlite
from db.wiki_db_json import WikiDBJson
from db.wiki_db_api import WikiDBAPI
class WikiRunEnvironment:
    def __init__(self, wiki_data_path):
        """Initialize with path to Wikipedia data"""
        if wiki_data_path.endswith('.json'):
            self.db = WikiDBJson(wiki_data_path)
        elif wiki_data_path.endswith('.db'):
            self.db = WikiDBSqlite(wiki_data_path)
        elif wiki_data_path.startswith('http'):
            self.db = WikiDBAPI(wiki_data_path)
        else:
            raise ValueError("Invalid file extension. Must be .json or .db")

        self.current_article = None
        self.target_article = None
        self.path_taken = []
        self.steps = 0
        
    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 = self.db.get_all_article_titles()
            
            if start_article is None:
                start_article = random.choice(available_articles)
                
            if target_article is None:
                # Ensure target is different from start
                remaining = [a for a in available_articles if a != start_article]
                target_article = random.choice(remaining)
        
        self.current_article = start_article
        self.target_article = target_article
        self.path_taken = [start_article]
        self.steps = 0
        
        return self.get_current_state()
    
    def get_current_state(self):
        """Get the current state of the environment"""
        if self.current_article is None:
            return None
            
        return {
            'current_article': self.current_article,
            'target_article': self.target_article,
            '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
        }
    
    def step(self, next_article):
        """Take a step to the next article"""
        if self.current_article is None:
            return None, "Game not initialized. Call reset() first."
            
        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"
            
        # Update state
        self.current_article = next_article
        self.path_taken.append(next_article)
        self.steps += 1
        
        # Check if we've reached the target
        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)