context_adding.py

import os
from openai import OpenAI
from loguru import logger

from config import OPENAI_KEY


client = OpenAI(api_key=OPENAI_KEY)
OPENAI_MODEL = "gpt-4o-mini"

# Function to extract context and prepend it to the text
def generate_context_and_prepend(input_text):
    prompt = f"Provide a short context summary for the following dialogue:\n\n{input_text}\n\nContext:"

    # Call OpenAI API
    response = client.chat.completions.create(
        model=OPENAI_MODEL,
        messages=[
            {"role": "system", "content": "You are a helpful assistant that summarizes dialogue."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.5
    )

    # Extract the generated context
    context_summary = response.choices[0].message.content.strip()

    # Prepend the context to the original text
    output_text = f"Context: {context_summary}\n\n{input_text}"
    return output_text


def add_file_context(input_dir, output_dir, file_name):
    if file_name.endswith(".txt"):
        input_file_path = os.path.join(input_dir, file_name)
        output_file_path = os.path.join(output_dir, file_name)

        with open(input_file_path, "r", encoding="utf-8") as f:
            original_text = f.read()

        # Skip files with context already present
        if original_text.lstrip().startswith("Context:"):
            logger.info(f"Skipped file (already has context): {file_name}")
            return

        # Generate context and prepend it
        updated_text = generate_context_and_prepend(original_text)

        with open(output_file_path, "w", encoding="utf-8") as f:
            f.write(updated_text)

        logger.info(f"Added context to: {file_name}")


# Process all text files in the directory
def process_directory(input_dir, output_dir):
    # Ensure output directory exists
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for filename in os.listdir(input_dir):
        add_file_context(input_dir, output_dir, filename)


# Example usage
if __name__ == "__main__":
    input_directory = os.path.join(os.getcwd(), "output_texts/done")
    output_directory = os.path.join(os.getcwd(), "output_texts/done")
    process_directory(input_directory, output_directory)
    print("All files have been processed and saved in the 'processed files' folder.")