wikihop-server / engine.py
stillerman's picture
stillerman HF Staff
fastapi server
76887e4
# 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)