StockDocs/ai_processor/ai_processor.py

import os
import requests
import json

LOCAL_AI_SERVICE_URL = "http://192.168.8.124:11434"

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 = []
    for newspaper in os.listdir(folder_path):
        for filename in os.listdir(os.path.join(folder_path, newspaper)):
            file_path = os.path.join(folder_path, newspaper, filename)
            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)} 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
        # TODO: In the future, we should force the response to be json and adjust the temperature
    }

    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 : protentially 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:
            # 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

# Retrieve the current archive of pulled articles
articles_folder = os.path.join(os.path.dirname(__file__), "../scraper/articles")
print("Loading articles from folder " + articles_folder + " ...")
articles = load_articles_from_folder(articles_folder)

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

        # 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}")