StockDocs/ai_processor/ai_processor.py

import os
import requests
import json
import datetime
import time

LOCAL_AI_SERVICE_URL = os.getenv("AI_SERVICE_URL") 

class ArticleFile:
    def __init__(self, filename, content, source="Unfiltered"):
        self.filename = filename
        self.content = content
        self.source = source

    def __str__(self):
        return f"ArticleFile(filename={self.filename}, source={self.source}, content_length={len(self.content)})"

def load_articles_from_folder(folder_path):
    """
    Loads all articles from text files in the specified folder and returns them as a list of ArticleFile objects.
    """
    articles = []
    # Skip non-directory files in the root folder
    for newspaper in os.listdir(folder_path):
        newspaper_path = os.path.join(folder_path, newspaper)
        if not os.path.isdir(newspaper_path):
            continue
            
        for filename in os.listdir(newspaper_path):
            file_path = os.path.join(newspaper_path, filename)
            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):
                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()
                    articles.append(ArticleFile(filename, content, source))
    
    print(f"Loaded {len(articles)} new articles from folder {folder_path}.")
    return articles

def call_local_ai_service(article):
    """
    Sends the article to the local Ollama AI server and returns the parsed response.
    Trims any 'thinking' section from the DeepSeek model response.
    """
    url = LOCAL_AI_SERVICE_URL + "/api/generate"
    prompt = (
        "Give a detailed summary of the article and then list out the tickers and their movement in the article. Do not say anything at the end of your response."
        "Format the response as JSON with keys 'summary' and 'tickersAndMovements', where 'tickersAndMovements' is a list of "
        "['Ticker', 'Movement direction', 'Percent change']. Here's an example of the expected output:\n"
        "{'summary': '...', 'tickersAndMovements': [['AAPL', '+', '2.5%'], ['GOOGL', '-', '1.2%']]}\n\n" \
        "Validate the JSON format and ensure it is parsable before giving the answer. For every entry in "
        "tickersAndMovements, it must have 3 entries. If you don't know the ticker symbol, don't make one up, just give "
        "the full name, you must have either +, - or = for the second entry, and you must have a decimal with percent (like 2.4%) or (?%) if you're unsure" \
        "Do not say anything at the end of your response."
    )
    payload = {
        "model": "llama3.2:latest",  
        "prompt": f"{prompt}\n\nArticle:\n{article}",
        "stream": False,
        "format": "json",
        "options": {
            "temperature": 0.1
        }
    }

    print(f"Sending request to local AI service ...")

    # Send the request to the local AI service
    response = requests.post(url, json=payload, timeout=180)  # Increased timeout to 300 seconds
    response.raise_for_status()
    # The AI should return a JSON string, so parse it
    raw_response = response.json()["response"]
    # Remove the 'thinking' section if present
    think_marker = "</think>"
    if think_marker in raw_response:
        raw_response = raw_response.split(think_marker, 1)[1]  # Take the second half after the marker

    # Remove any leading code block markers or whitespace (e.g., ```json or ``` or whitespace)
    raw_response = raw_response.lstrip("` \njson")
    # Remove any trailing code block marker ```
    if raw_response.endswith("```"):
        raw_response = raw_response[:-3].rstrip()

    #print(f"Received response (trimmed): {raw_response}")
    print(f"Received response from local AI service, processing...")
    ai_processed_results = json.loads(raw_response)
    return ai_processed_results

def process_article(article_file):
    processed_article = None

    # Check if we already processed this article
    # TODO : potentially expensive, consider using a database or cache
    if os.path.exists(f"output/{article_file.filename}.json"):
        print(f"Article {article_file.filename} already processed, skipping.")
        
        # open file and return the content
        with open(f"output/{article_file.filename}.json", 'r', encoding='utf-8') as f:
            processed_article = json.load(f)
        return processed_article

    for attempt in range(3):
        try:
            if not article_file.content.strip() or len(article_file.content) < 30:
                print(f"Skipping empty article or article with insufficient content: {article_file.filename}")
                return None

            # Pass the article content to the AI service
            processed_article = call_local_ai_service(article_file.content)
            print(f"Processed article: {processed_article['summary'][:100]}...")

            # Add metadata to the processed article
            processed_article['filename'] = article_file.filename
            processed_article['original_content'] = article_file.content
            processed_article['source'] = article_file.source  # <-- Add source here
            break  # Success, exit retry loop
        except Exception as e:
            print(f"\tError processing article (attempt {attempt+1}/3): {e}")
        if processed_article is not None:
            return processed_article
        else:
            print(f"\tSkipping article after 3 failed attempts.")

    return processed_article

while True:
    print("=========================================")

    # Retrieve the current archive of pulled articles
    articles_folder = os.path.join("/app/articles")
    if not os.path.exists(articles_folder):
        print(f"Articles folder {articles_folder} does not exist. Please check the path.")
        exit(1)
    print("Loading articles from folder " + articles_folder + " ...")
    articles = load_articles_from_folder(articles_folder)

    # Process the articles retrieved
    results = []
    failedProcessCount = 0
    for article in articles: 
        try:
            result = process_article(article)

            if result is None:
                failedProcessCount += 1
                #print(f"Skipping article {article.filename} due to processing error.")
                continue
                
            # Save the processed result to a JSON file
            with open("output/" + result['filename'] + ".json", "w", encoding="utf-8") as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"Error processing or saving article {article.filename}: {e}")

    print(f"Processed {len(articles) - failedProcessCount} articles successfully, {failedProcessCount} failed to process.") 

    time.delay(60 * 5)  # Wait for 5 minutes before the next iteration