batch-processing/main.py at main · cho-groq/batch-processing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
"""
Step 1. Set up dependencies
"""

import os
from dotenv import load_dotenv
import requests # pip install requests first!
import time

# Load environment variables from .env file
load_dotenv()

# Access environment variables
api_key = os.getenv("GROQ_API_KEY")


"""
Step 2. Upload the JSONL file to Groq
"""

def upload_file_to_groq(api_key, file_path):
    url = "https://api.groq.com/openai/v1/files"

    headers = {
        "Authorization": f"Bearer {api_key}"
    }

    # Prepare the file and form data
    files = {
        "file": ("batch_file.jsonl", open(file_path, "rb"))
    }

    data = {
        "purpose": "batch"
    }

    # Make the POST request
    response = requests.post(url, headers=headers, files=files, data=data)

    return response.json()

# Usage example
file_path = "batch_input.jsonl"  # Path to your JSONL file
file_id = "" # will be used in the next step

try:
    result = upload_file_to_groq(api_key, file_path)
    file_id = result["id"]
    print("This is the file_id from Step 2: " + file_id)

except Exception as e:
    print(f"Error: {e}")


"""
Step 3. Make a batch object
"""

def create_batch(api_key, input_file_id):
    url = "https://api.groq.com/openai/v1/batches"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    data = {
        "input_file_id": input_file_id,
        "endpoint": "/v1/chat/completions",
        "completion_window": "24h"
    }

    response = requests.post(url, headers=headers, json=data)
    return response.json()

batch_id = "" # will be used in the next step
try:
    result = create_batch(api_key, file_id)
    batch_id = result["id"] # batch result id
    print("This is the Batch object id from Step 3: " + batch_id)

except Exception as e:
    print(f"Error: {e}")


"""
Step 4. Get the batch status
"""

def get_batch_status(api_key, batch_id):
    url = f"https://api.groq.com/openai/v1/batches/{batch_id}"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    response = requests.get(url, headers=headers)
    return response.json()

output_file_id = "" # will be used in the next step
try:
    result = get_batch_status(api_key, batch_id)
    print("\nStep 4 results: ")

    count = 0

    # Corrected condition: use `result["status"]` instead of `result.status`
    while result["status"] != "completed" and count < 100:
        result = get_batch_status(api_key, batch_id)  # Update `result` inside the loop
        time.sleep(3)
        print("Your batch status is: " + result["status"])
        count += 1

    output_file_id = result.get("output_file_id")  # Use .get() to safely access keys
    print("This is your output_file_id from Step 4: " + output_file_id)
except Exception as e:
    print(f"Error: {e}")


"""
Step 5. Retrieve batch results
"""

def download_file_content(api_key, output_file_id, output_file):
    url = f"https://api.groq.com/openai/v1/files/{output_file_id}/content"

    headers = {
        "Authorization": f"Bearer {api_key}"
    }

    response = requests.get(url, headers=headers)

    # Write the content to a file
    with open(output_file, 'wb') as f:
        f.write(response.content)

    return f"\nFile downloaded successfully to {output_file}"

output_file = "batch_output.jsonl" # replace with your own file of choice to download batch job contents to
try:
    result = download_file_content(api_key, output_file_id, output_file)
    print(result)
except Exception as e:
    print(f"Error: {e}")