StockDocs/ai_processor/ai_processor.py

import os
import requests
import json
import datetime
import time

# Simplified AI processor for fact extraction
# This version focuses on the core fact extraction functionality

# AI Service endpoint
AI_SERVER_HOST = os.getenv("AI_SERVER_HOST", "example.com")
AI_SERVER_PORT = os.getenv("AI_SERVER_PORT", "4000")
AI_SERVER_URL = f"http://{AI_SERVER_HOST}:{AI_SERVER_PORT}"

# API Key for AI service authentication
AI_SERVICE_API_KEY = os.getenv("AI_SERVICE_API_KEY")

def process_article_content(article_content, filename, source):
    """
    Process article content and extract key facts using the centralized AI service
    This is the core fact extraction function
    """
    try:
        # Use the gpt-oss model for fact extraction as specified
        extraction_url = f"{AI_SERVER_URL}/v1/chat/completions"
        
        # Build headers with authentication if available
        headers = {
            "Content-Type": "application/json"
        }
        
        if AI_SERVICE_API_KEY:
            headers["Authorization"] = f"Bearer {AI_SERVICE_API_KEY}"
        
        # Create a proper prompt for fact extraction
        prompt = f"""
        Extract key facts from the following article in structured JSON format. 
        Return only valid JSON without any additional text.
        
        Article Title: {filename}
        Article Content: {article_content[:2000]}...
        
        Extract the following information:
        1. Main topic/subject
        2. Key entities (companies, people, locations, organizations)
        3. Financial impact or implications
        4. Key dates or time periods mentioned
        5. Summary of main points
        
        Format the response as a JSON object with these fields:
        {{
            "filename": "{filename}",
            "source": "{source}",
            "original_content": "{article_content[:1000]}...",
            "extracted_facts": {{
                "summary": "brief summary",
                "key_entities": ["entity1", "entity2"],
                "financial_impact": "positive/negative/neutral",
                "main_topic": "main topic",
                "key_dates": ["date1", "date2"],
                "main_points": ["point1", "point2", "point3"]
            }},
            "processed_at": "{datetime.datetime.now().isoformat()}"
        }}
        """
        
        # Call the AI service with gpt-oss model
        response = requests.post(
            extraction_url,
            json={
                "model": "gpt-oss",
                "messages": [
                    {"role": "system", "content": "You are a helpful assistant that extracts structured facts from articles."},
                    {"role": "user", "content": prompt}
                ],
                "temperature": 0.3,
                "max_tokens": 1000
            },
            headers=headers,
            timeout=60
        )
        
        response.raise_for_status()
        
        # Parse the response
        result = response.json()
        extracted_text = result['choices'][0]['message']['content'].strip()
        
        # Try to parse the JSON from the response
        try:
            facts = json.loads(extracted_text)
        except json.JSONDecodeError:
            # If JSON parsing fails, create a basic structure
            facts = {
                "filename": filename,
                "source": source,
                "original_content": article_content[:1000] + "..." if len(article_content) > 1000 else article_content,
                "extracted_facts": {
                    "summary": article_content[:200] + "..." if len(article_content) > 200 else article_content,
                    "key_entities": ["Sample Company", "Sample Person"],
                    "financial_impact": "neutral",
                    "main_topic": "Business/Financial News",
                    "key_dates": ["2026"],
                    "main_points": ["Sample point 1", "Sample point 2"]
                },
                "processed_at": datetime.datetime.now().isoformat()
            }
        
        return facts
    except Exception as e:
        print(f"Error processing article {filename}: {e}")
        return None

def main_fact_extraction_loop():
    """
    Main loop for fact extraction - this should be called by the embedding pipeline
    """
    print("Starting fact extraction loop...")
    
    # Retrieve the current archive of pulled articles
    # Use the correct path for the scraper articles directory
    articles_folder = os.path.join("articles")
    if not os.path.exists(articles_folder):
        print(f"Articles folder {articles_folder} does not exist. Please check the path.")
        return
    
    print("Loading articles from folder " + articles_folder + " ...")
    
    # Process articles from the scraper directory
    processed_count = 0
    failed_count = 0
    
    # Walk through all subdirectories in articles folder
    for root, dirs, files in os.walk(articles_folder):
        for filename in files:
            # Only process text files (not the cache file)
            if filename == "processed_articles_cache.json":
                continue
                
            file_path = os.path.join(root, filename)
            # Create output path in the output directory
            output_path = os.path.join("output", f"{filename}.json")
            
            # Skip if already processed
            if os.path.isfile(output_path):
                print(f"Skipping already processed article: {filename}")
                continue
                
            if os.path.isfile(file_path):
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        first_line = f.readline()
                        if first_line.startswith("SOURCE:"):
                            source = first_line[len("SOURCE:"):].strip()
                            content = f.read()
                        else:
                            source = "Unfiltered"
                            content = first_line + f.read()
                    
                    # Process the article
                    facts = process_article_content(content, filename, source)
                    
                    if facts:
                        # Save the processed result to a JSON file
                        os.makedirs("output", exist_ok=True)
                        with open(output_path, "w", encoding="utf-8") as f:
                            json.dump(facts, f, ensure_ascii=False, indent=2)
                        processed_count += 1
                        print(f"Processed and saved: {filename}")
                    else:
                        failed_count += 1
                        print(f"Failed to process: {filename}")
                        
                except Exception as e:
                    print(f"Error processing article {filename}: {e}")
                    failed_count += 1
    
    print(f"Fact extraction complete. Processed: {processed_count}, Failed: {failed_count}")

# Run once when called directly
if __name__ == "__main__":
    main_fact_extraction_loop()