wikihop-server / db /wiki_db_json.py
stillerman's picture
stillerman HF Staff
reorganizing
2c2bab6
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', [])