diff chatgpt.py @ 3:7770a4bd42e2 draft default tip

planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools/chatgpt commit c21d9a2cb410ee30dc47f4a13247862481816266
author bgruening
date Wed, 11 Sep 2024 16:36:21 +0000
parents dab494dce303
children
line wrap: on
line diff
--- a/chatgpt.py	Fri Aug 23 10:21:21 2024 +0000
+++ b/chatgpt.py	Wed Sep 11 16:36:21 2024 +0000
@@ -1,9 +1,10 @@
+import json
 import os
 import sys
 
-from openai import OpenAI
+from openai import AuthenticationError, OpenAI
 
-context_files = sys.argv[1].split(",")
+context_files = json.loads(sys.argv[1])
 question = sys.argv[2]
 model = sys.argv[3]
 with open(sys.argv[4], "r") as f:
@@ -14,51 +15,39 @@
 
 client = OpenAI(api_key=openai_api_key)
 
-file_search_sup_ext = [
-    "c",
-    "cs",
-    "cpp",
-    "doc",
-    "docx",
-    "html",
-    "java",
-    "json",
-    "md",
-    "pdf",
-    "php",
-    "pptx",
-    "py",
-    "rb",
-    "tex",
-    "txt",
-    "css",
-    "js",
-    "sh",
-    "ts",
-]
-
-vision_sup_ext = ["jpg", "jpeg", "png", "webp", "gif"]
-
 file_search_file_streams = []
 image_files = []
 
-for path in context_files:
-    ext = path.split(".")[-1].lower()
-    if ext in vision_sup_ext:
+for path, type in context_files:
+    if type == "image":
         if os.path.getsize(path) > 20 * 1024 * 1024:
             print(f"File {path} exceeds the 20MB limit and will not be processed.")
             sys.exit(1)
         file = client.files.create(file=open(path, "rb"), purpose="vision")
         promt = {"type": "image_file", "image_file": {"file_id": file.id}}
         image_files.append(promt)
-    elif ext in file_search_sup_ext:
+    else:
         file_search_file_streams.append(open(path, "rb"))
 
-assistant = client.beta.assistants.create(
-    instructions="You will receive questions about files from file searches and image files. For file search queries, identify and retrieve the relevant files based on the question. For image file queries, analyze the image content and provide relevant information or insights based on the image data.",
-    model=model,
-    tools=[{"type": "file_search"}] if file_search_file_streams else [],
-)
+try:
+    assistant = client.beta.assistants.create(
+        instructions=(
+            "You will receive questions about files from file searches "
+            "and image files. For file search queries, identify and "
+            "retrieve the relevant files based on the question. "
+            "For image file queries, analyze the image content and "
+            "provide relevant information or insights based on the image data."
+        ),
+        model=model,
+        tools=[{"type": "file_search"}] if file_search_file_streams else [],
+    )
+except AuthenticationError as e:
+    print(f"Authentication error: {e.message}")
+    sys.exit(1)
+except Exception as e:
+    print(f"An error occurred: {str(e)}")
+    sys.exit(1)
+
 if file_search_file_streams:
     vector_store = client.beta.vector_stores.create()
     file_batch = client.beta.vector_stores.file_batches.upload_and_poll(
@@ -88,7 +77,12 @@
 assistant_messages = list(
     client.beta.threads.messages.list(thread_id=thread.id, run_id=run.id)
 )
-
+if not assistant_messages:
+    print(
+        "No output was generated!\nPlease ensure that your OpenAI account has sufficient credits.\n"
+        "You can check your balance here: https://platform.openai.com/settings/organization/billing"
+    )
+    sys.exit(1)
 message_content = assistant_messages[0].content[0].text.value
 print("Output has been saved!")
 with open("output.txt", "w") as f: