gpt-batch 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

gpt_batch/batcher.py

@@ -20,8 +20,7 @@ class GPTBatcher:
         retry_attempts (int): Number of retries if a request fails. Default is 2.
         miss_index (list): Tracks the indices of requests that failed to process correctly.
     """
-
-    def __init__(self, api_key, model_name="gpt-3.5-turbo-0125", system_prompt="",temperature=1,num_workers=64,timeout_duration=60,retry_attempts=2,api_base_url=None):
+    def __init__(self, api_key, model_name="gpt-3.5-turbo-0125", system_prompt="",temperature=1,num_workers=64,timeout_duration=60,retry_attempts=2,api_base_url=None,**kwargs):
         
         self.is_claude = bool(re.search(r'claude', model_name, re.IGNORECASE))
         
@@ -41,6 +40,7 @@ class GPTBatcher:
         self.timeout_duration = timeout_duration
         self.retry_attempts = retry_attempts
         self.miss_index = []
+        self.extra_params = kwargs 
 
     def get_attitude(self, ask_text):
         index, ask_text = ask_text
@@ -55,6 +55,7 @@ class GPTBatcher:
                     ],
                     system=self.system_prompt if self.system_prompt else None,
                     temperature=self.temperature,
+                    **self.extra_params
                 )
                 return (index, message.content[0].text)
             else:
@@ -66,6 +67,7 @@ class GPTBatcher:
                         {"role": "user", "content": ask_text}
                     ],
                     temperature=self.temperature,
+                    **self.extra_params
                 )
                 return (index, completion.choices[0].message.content)
         except Exception as e:

{gpt_batch-0.1.8.dist-info → gpt_batch-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpt-batch
-Version: 0.1.8
+Version: 0.1.9
 Summary: A package for batch processing with OpenAI API.
 Home-page: https://github.com/fengsxy/gpt_batch
 Author: Ted Yu

gpt_batch-0.1.9.dist-info/RECORD

@@ -0,0 +1,8 @@
+gpt_batch/__init__.py,sha256=zGDItktTxKLSQr44GY78dl5LKsSJig0Q59dzusqhU0U,59
+gpt_batch/batcher.py,sha256=y8B4hIeQJQ16G5PvlNgHE-CtVQzHPhpBssOAg7npQLA,9083
+tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/test_batcher.py,sha256=yRwqe2_VTp4wXXeQRhyDPZ2NZ-H3SSCDAxlNNXh3Aro,5314
+gpt_batch-0.1.9.dist-info/METADATA,sha256=30t3VH_tY1mNWnzBPuQWKmD1o3bcA9yh3htvAWBgyok,3401
+gpt_batch-0.1.9.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+gpt_batch-0.1.9.dist-info/top_level.txt,sha256=FtvJB_L9W_S6jL4G8Em_YWphG1wdKAF20BHUrf4B0yM,16
+gpt_batch-0.1.9.dist-info/RECORD,,

tests/test_batcher.py

@@ -18,6 +18,27 @@ def test_handle_message_list():
     assert len(results) == 2, "There should be two results, one for each message"
     assert all(len(result) >= 2 for result in results), "Each result should be at least two elements"
 
+
+def test_json_format():
+    import json
+    # Initialize the GPTBatcher with hypothetical valid credentials
+    #api_key = #get from system environment
+    api_key = os.getenv('TEST_KEY')
+    if not api_key:
+        raise ValueError("API key must be set in the environment variables")
+    batcher = GPTBatcher(api_key=api_key, model_name='gpt-3.5-turbo-1106', system_prompt="Your task is to discuss privacy questions and provide persuasive answers with supporting reasons.",response_format={ "type": "json_object" })
+    message_list = ["return me a random json object", "return me a random json object"]
+
+    # Call the method under test
+    results = batcher.handle_message_list(message_list)
+    # Assertions to verify the length of the results and the structure of each item
+    assert len(results) == 2, "There should be two results, one for each message"
+    assert all(len(result) >= 2 for result in results), "Each result should be at least two elements"
+    #assert all(isinstance(result, dict) and 'json' in result for result in results), "Each result should be a JSON object with 'json' key"
+    assert all(isinstance(json.loads(result), dict) for result in results), "Each result should be a JSON object with 'json' key"
+
+
+
 def test_handle_embedding_list():
     # Initialize the GPTBatcher with hypothetical valid credentials
     #api_key = #get from system environment

gpt_batch-0.1.8.dist-info/RECORD

@@ -1,8 +0,0 @@
-gpt_batch/__init__.py,sha256=zGDItktTxKLSQr44GY78dl5LKsSJig0Q59dzusqhU0U,59
-gpt_batch/batcher.py,sha256=sruKqK_tY5WNMXvejNz_OAw6AHtJIqcvDKbMbX1rX_M,8959
-tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/test_batcher.py,sha256=FAlV_T5dT2OcM28yWPWWU9yIN0D7SBeYP1oTkvvzXKk,4077
-gpt_batch-0.1.8.dist-info/METADATA,sha256=IIBetdVEVrUx30cxrzB4ZkAy0mGp3PVueJc9tm2WcZA,3401
-gpt_batch-0.1.8.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-gpt_batch-0.1.8.dist-info/top_level.txt,sha256=FtvJB_L9W_S6jL4G8Em_YWphG1wdKAF20BHUrf4B0yM,16
-gpt_batch-0.1.8.dist-info/RECORD,,