index.py

import os
import json
import requests
import logging
from datetime import datetime
from typing import List, Dict
from firecrawl import FirecrawlApp
from openai import OpenAI
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Add logging configuration at the top of the file
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class PredictionResearcher:
    def __init__(self):
        self.perplexity_api_key = os.getenv('PERPLEXITY_API_KEY')
        self.firecrawl_api_key = os.getenv('FIRECRAWL_API_KEY')
        self.openai_api_key = os.getenv('OPENAI_API_KEY')
        
        if not all([self.perplexity_api_key, self.firecrawl_api_key, self.openai_api_key]):
            raise ValueError("Please set PERPLEXITY_API_KEY, FIRECRAWL_API_KEY and OPENAI_API_KEY in your .env file")

        logger.info("Initializing PredictionResearcher")

    def get_research_report(self, question: str) -> Dict:
        """Get research report from Perplexity API"""
        logger.info(f"Getting research report for question: {question}")
        
        headers = {
            "Authorization": f"Bearer {self.perplexity_api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": "sonar-pro",
            "messages": [
                {
                    "role": "system",
                    "content": "You are an expert researcher helping write research reports for predictions for 2025. Provide detailed reports with citations."
                },
                {
                    "role": "user", 
                    "content": f"Research question: {question}\nProvide a detailed research report with specific data points, trends, and citations."
                }
            ],
            "max_tokens": 2048,
            "return_citations": True
        }
        
        try:
            response = requests.post(
                "https://api.perplexity.ai/chat/completions",
                headers=headers,
                json=data
            )
            response.raise_for_status()
            logger.info("Successfully received research report from Perplexity API")
            return response.json()
        except requests.exceptions.RequestException as e:
            logger.error(f"Error getting research report: {str(e)}")
            raise

    def scrape_citations(self, urls: List[str]) -> List[Dict]:
        """Scrape cited URLs using Firecrawl API"""
        logger.info(f"Starting to scrape {len(urls)} citations")
        logger.info(f"URLs to scrape: {urls}")
        
        app = FirecrawlApp(api_key=self.firecrawl_api_key)
        scraped_data = []
        
        try:
            # First attempt batch scraping
            batch_result = app.batch_scrape_urls(
                urls,
                {'formats': ['markdown']}
            )
            
            # Process batch results
            if isinstance(batch_result, dict) and 'data' in batch_result:
                for item in batch_result['data']:
                    if item.get('markdown'):
                        url = item.get('metadata', {}).get('sourceURL', '')
                        scraped_data.append({
                            'url': url,
                            'data': {
                                'markdown': item['markdown']
                            }
                        })
                        logger.info(f"Successfully scraped URL in batch: {url}")
        
        except Exception as e:
            logger.warning(f"Batch scraping failed: {str(e)}. Falling back to individual scraping.")
            
            # Fall back to scraping URLs individually
            for url in urls:
                try:
                    result = app.scrape_url(url, {'formats': ['markdown']})
                    if isinstance(result, dict) and result.get('markdown'):
                        scraped_data.append({
                            'url': url,
                            'data': {
                                'markdown': result['markdown']
                            }
                        })
                        logger.info(f"Successfully scraped individual URL: {url}")
                    else:
                        logger.warning(f"No markdown content in response for URL: {url}")
                except Exception as individual_error:
                    logger.error(f"Error scraping individual URL {url}: {str(individual_error)}")
                    continue
        
        logger.info(f"Completed scraping {len(scraped_data)} URLs successfully")
        return scraped_data

    def generate_gpt_prompt(self, question: str, research_data: Dict, citations: List[Dict]) -> str:
        """Generate detailed GPT prompt with research and citations"""
        logger.info("Generating GPT prompt")
        
        # Extract content and citations from Perplexity response
        content = research_data["choices"][0]["message"]["content"]
        
        prompt = f"""You are a master analyst who can process large amounts of information to make accurate predictions about the future. Analyze the following prediction question, research report, and citations to provide a specific probability estimate (0-100%) for its occurrence.

<task>
- Carefully analyze the provided research and citations
- Highlight key factors that could change your estimate
- Provide a specific probability estimate between 0 and 100%. The estimate will be graded using the Brier score
</task>

<format>
Please provide your detailed analysis following this format, ending with a specific probability estimate.

1. Initial Analysis:
   - Key trends
   - Supporting factors
   - Opposing factors

2. Scenario Analysis:
   - Best case scenario
   - Base case scenario  
   - Worst case scenario

3. Risk Factors:
   - Upside catalysts
   - Downside risks
   - Key metrics to monitor

4. Probability Estimate:
   - Specific percentage between 0 and 100%
</format>

<question>
{question}
</question>

<research>
{content}
</research>

<citations>
"""
        
        # Add scraped citation content with token limit
        for i, citation in enumerate(citations):
            citation_content = citation.get("data", {}).get("markdown", "")
            if citation_content:
                # Roughly estimate tokens (1 token ≈ 4 chars)
                max_chars = 4 * 2048
                if len(citation_content) > max_chars:
                    citation_content = citation_content[:max_chars] + "\n[Content truncated to 2048 tokens...]"
                prompt += f"\n<citation{i+1}>\n{citation_content}\n</citation{i+1}>"
        
        prompt += "\n</citations>"
        
        logger.info(f"Generated prompt with {len(citations)} citations")
        return prompt

    def get_o3_mini_analysis(self, prompt: str) -> str:
        """Get analysis from OpenAI's o3-mini model"""
        logger.info("Getting analysis from o3-mini model")
        
        try:
            client = OpenAI(api_key=self.openai_api_key)
            
            response = client.chat.completions.create(
                model="o3-mini",
                reasoning_effort="high",
                messages=[
                    {
                        "role": "user",
                        "content": prompt
                    }
                ]
            )
            
            logger.info("Successfully received o3-mini analysis")
            return response.choices[0].message.content
            
        except Exception as e:
            logger.error(f"Error getting o3-mini analysis: {str(e)}")
            raise

    def analyze_prediction(self, question: str) -> None:
        """Analyze a single prediction question"""
        logger.info(f"Starting prediction analysis for question: {question}")
        
        try:
            # Get research report
            research_data = self.get_research_report(question)
            
            # Extract URLs from citations field in response
            urls = research_data.get("citations", [])
            logger.info(f"Extracted {len(urls)} URLs from research report citations")
            
            # Scrape citations
            citations = self.scrape_citations(urls)
            
            # Generate GPT prompt
            prompt = self.generate_gpt_prompt(question, research_data, citations)
            
            # Create a filename-safe version of the question in snake_case
            safe_question = "".join(c.lower() for c in question if c.isalnum() or c in (' ', '-', '_')).strip()
            safe_question = safe_question.replace(' ', '_').replace('-', '_')
            safe_question = safe_question[:50]  # Limit length to avoid extremely long filenames
            
            # Generate timestamp
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            
            # Save prompt to markdown file
            prompt_filename = f"{timestamp}_{safe_question}_prompt.md"
            with open(prompt_filename, "w") as f:
                f.write(prompt)
            logger.info(f"Prompt saved to {prompt_filename}")

            # Get o3-mini analysis
            o3_analysis = self.get_o3_mini_analysis(prompt)
            
            # Save o3 analysis to text file
            analysis_filename = f"{timestamp}_{safe_question}_o3_prediction.txt"
            with open(analysis_filename, "w") as f:
                f.write(o3_analysis)
            logger.info(f"O3 analysis saved to {analysis_filename}")
            
        except Exception as e:
            logger.error(f"Error analyzing question '{question}': {str(e)}")
            
    def analyze_multiple_predictions(self, questions: List[str]) -> None:
        """Analyze multiple prediction questions"""
        logger.info(f"Starting analysis of {len(questions)} predictions")
        
        for i, question in enumerate(questions, 1):
            logger.info(f"Processing question {i} of {len(questions)}")
            self.analyze_prediction(question)
            
        logger.info("Completed analysis of all questions")

if __name__ == "__main__":
    logger.info("Starting prediction research program")
    researcher = PredictionResearcher()
    
    questions = [
        """Will Maersk resume shipping in the Red Sea in 2025?""",
        """Will the poverty rate in Argentina be lower in the first half of 2025 compared to the second half of 2023?""",
        """Will an application to ban AfD (Germany) be filed at the Federal Constitutional Court before 2026?"""
    ]
    
    researcher.analyze_multiple_predictions(questions)