respect ratelimit and change output extension to .gpt

1 files changed, 38 insertions, 40 deletions

diff --git a/main.py b/main.py
index e344a59..af76aca 100644
--- a/main.py
+++ b/main.py
@@ -1,10 +1,16 @@
 import os
 import re
 import sys
-from openai import OpenAI
+import backoff
+import openai
 from datetime import datetime, timedelta
 from tqdm import tqdm
 
+client = openai.OpenAI(
+    #base_url="http://localhost:8080/v1",
+    base_url="https://api.openai.com/v1",
+)
+
 SYSTEM_PROMPT_TEMPLATE = """
 You are an expert podcast content analyzer. Your task is to identify and extract pre-recorded dynamic advertising segments from podcast transcripts. These ads typically have a distinct tone shift, often using more direct, persuasive language, and frequently include calls to action, product mentions, or specific brand names.
 
@@ -202,28 +208,29 @@ def chunk_transcript(
         
     return chunks
 
-def call_openai_api(client: OpenAI, chunk_text: str, system_prompt: str) -> str:
+@backoff.on_exception(backoff.expo, openai.RateLimitError)
+def completions_with_backoff(**kwargs):
+    return client.chat.completions.create(**kwargs)
+
+def call_openai_api(chunk_text: str, system_prompt: str) -> str:
     """
     Makes an API call to OpenAI with the given system prompt and transcript chunk.
     """
-    try:
-        response = client.chat.completions.create(
-            #model="Gemma-3-4B_32K",
-            #model="google/gemini-2.0-flash-exp:free",
-            model="gpt-4.1",
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": f"Input Transcript:\n{chunk_text}"}
-            ],
-            temperature=0.0,
-        )
-        #print(response)
-        return response.choices[0].message.content.strip()
-    except Exception as e:
-        print(f"Error calling OpenAI API: {e}")
-        return ""
-
-def process_file(input_file: str, output_file: str, client: OpenAI):
+    response = completions_with_backoff(
+        #model="Gemma-3-4B_32K",
+        #model="google/gemini-2.0-flash-exp:free",
+        model="gpt-4.1",
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"Input Transcript:\n{chunk_text}"}
+        ],
+        temperature=0.0,
+        max_tokens=1024,    # realistically it'd be under like 512
+    )
+    #print(response)
+    return response.choices[0].message.content.strip()
+
+def process_file(input_file: str, output_file: str):
     with open(input_file, "r", encoding="utf-8") as f:
         full_transcript_content = f.read()
 
@@ -247,7 +254,11 @@ def process_file(input_file: str, output_file: str, client: OpenAI):
         # Add the "Your Turn" header before the chunk content
         user_prompt_content = f"**Your Turn - Input Transcript:**\n{chunk_text_for_api}"
 
-        ad_timestamps_str = call_openai_api(client, user_prompt_content, SYSTEM_PROMPT_TEMPLATE)
+        try:
+            ad_timestamps_str = call_openai_api(user_prompt_content, SYSTEM_PROMPT_TEMPLATE)
+        except Exception as e:
+            print(f"Error calling OpenAI API: {e}")
+            continue
         
         if ad_timestamps_str:
             # Split by newlines in case the model returns multiple segments
@@ -258,29 +269,16 @@ def process_file(input_file: str, output_file: str, client: OpenAI):
         f.write("\n".join(all_ad_timestamps))
 
 def main():
-    # --- Configuration ---
-    api_key = os.getenv("OPENAI_KEY")
-    if not api_key:
-        print("OPENAI_KEY empty")
-        sys.exit(1)
-
-    client = OpenAI(
-        #api_key="none",
-        api_key=api_key,
-        #base_url="http://localhost:8080/v1",
-        base_url="https://api.openai.com/v1",
-    )
-
     paths = sys.argv[1:]
-    for i, p in enumerate(paths):
-        out = "".join(p.split(".")[:-1]) + "_gpt.txt"
-        if os.path.isfile(out)
-            paths.remove(i)
+    for p in paths:
+        out = "".join(p.split(".")[:-1]) + ".gpt"
+        if os.path.isfile(out):
+            paths.remove(p)
     print(paths)
 
     for p in tqdm(paths):
-        out = "".join(p.split(".")[:-1]) + "_gpt.txt"
-        process_file(p, out, client)
+        out = "".join(p.split(".")[:-1]) + ".gpt"
+        process_file(p, out)
 
 if __name__ == "__main__":
     main()