db/wiki_db_api.py

import requests

class WikiDBAPI:
    def __init__(self, api_endpoint):
        """Initialize the database with API endpoint URL"""
        self.api_endpoint = api_endpoint.rstrip('/')
        # Verify connection and get initial article count
        self._article_count = self._get_article_count()
        print(f"Connected to Wiki API at {api_endpoint} with {self._article_count} articles")
    
    def __del__(self):
        """Clean up resources when object is destroyed"""
        # No persistent connection to close in the API version
        pass
    
    def _get_article_count(self):
        """Get the number of articles via API"""
        response = requests.get(f"{self.api_endpoint}/article_count")
        response.raise_for_status()
        return response.json()["count"]
    
    def get_article_count(self):
        """Return the number of articles"""
        return self._article_count
        
    def get_all_article_titles(self):
        """Return a list of all article titles"""
        response = requests.get(f"{self.api_endpoint}/article_titles")
        response.raise_for_status()
        return response.json()["titles"]
        
    def get_article(self, title):
        """Get article data by title"""
        response = requests.get(f"{self.api_endpoint}/article", params={"title": title})
        if response.status_code == 404:
            return {}
        response.raise_for_status()
        return response.json()
        
    def article_exists(self, title):
        """Check if an article exists"""
        response = requests.get(f"{self.api_endpoint}/article_exists", params={"title": title})
        response.raise_for_status()
        return response.json()["exists"]
        
    def get_article_text(self, title):
        """Get the text of an article"""
        response = requests.get(f"{self.api_endpoint}/article_text", params={"title": title})
        if response.status_code == 404:
            return ''
        response.raise_for_status()
        return response.json()["text"]
        
    def get_article_links(self, title):
        """Get the links of an article"""
        response = requests.get(f"{self.api_endpoint}/article_links", params={"title": title})
        if response.status_code == 404:
            return []
        response.raise_for_status()
        return response.json()["links"]