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

gpt_batch/batcher.py

@@ -1,57 +1,73 @@
 from openai import OpenAI
+import anthropic
 from concurrent.futures import ThreadPoolExecutor, wait
 from functools import partial
 from tqdm import tqdm
+import re
 
 class GPTBatcher:
     """
-    A class to handle batching and sending requests to the OpenAI GPT model efficiently.
+    A class to handle batching and sending requests to the OpenAI GPT model and Anthropic Claude models efficiently.
 
     Attributes:
-        client (OpenAI): The client instance to communicate with the OpenAI API using the provided API key.
-        model_name (str): The name of the GPT model to be used. Default is 'gpt-3.5-turbo-0125'.
+        client: The client instance to communicate with the API (OpenAI or Anthropic).
+        is_claude (bool): Flag to indicate if using a Claude model.
+        model_name (str): The name of the model to be used. Default is 'gpt-3.5-turbo-0125'.
         system_prompt (str): Initial prompt or context to be used with the model. Default is an empty string.
         temperature (float): Controls the randomness of the model's responses. Higher values lead to more diverse outputs. Default is 1.
         num_workers (int): Number of worker threads used for handling concurrent requests. Default is 64.
         timeout_duration (int): Maximum time (in seconds) to wait for a response from the API before timing out. Default is 60 seconds.
         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.
-
-    Parameters:
-        api_key (str): API key for authenticating requests to the OpenAI API.
-        model_name (str, optional): Specifies the GPT model version. Default is 'gpt-3.5-turbo-0125'.
-        system_prompt (str, optional): Initial text or question to seed the model with. Default is empty.
-        temperature (float, optional): Sets the creativity of the responses. Default is 1.
-        num_workers (int, optional): Number of parallel workers for request handling. Default is 64.
-        timeout_duration (int, optional): Timeout for API responses in seconds. Default is 60.
-        retry_attempts (int, optional): How many times to retry a failed request. Default is 2.
     """
 
     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):
         
-        self.client = OpenAI(api_key=api_key)
+        self.is_claude = bool(re.search(r'claude', model_name, re.IGNORECASE))
+        
+        if self.is_claude:
+            self.client = anthropic.Anthropic(api_key=api_key)
+            # Anthropic doesn't support custom base URL the same way
+            # If needed, this could be implemented differently
+        else:
+            self.client = OpenAI(api_key=api_key)
+            if api_base_url:
+                self.client.base_url = api_base_url
+        
         self.model_name = model_name
         self.system_prompt = system_prompt
         self.temperature = temperature
         self.num_workers = num_workers
         self.timeout_duration = timeout_duration
         self.retry_attempts = retry_attempts
-        self.miss_index =[]
-        if api_base_url:
-            self.client.base_url = api_base_url
+        self.miss_index = []
 
     def get_attitude(self, ask_text):
         index, ask_text = ask_text
         try:
-            completion = self.client.chat.completions.create(
-                model=self.model_name,
-                messages=[
-                    {"role": "system", "content": self.system_prompt},
-                    {"role": "user", "content": ask_text}
-                ],
-                temperature=self.temperature,
-            )
-            return (index, completion.choices[0].message.content)
+            if self.is_claude:
+                # Use the Anthropic Claude API
+                message = self.client.messages.create(
+                    model=self.model_name,
+                    max_tokens=1024,  # You can make this configurable if needed
+                    messages=[
+                        {"role": "user", "content": ask_text}
+                    ],
+                    system=self.system_prompt if self.system_prompt else None,
+                    temperature=self.temperature,
+                )
+                return (index, message.content[0].text)
+            else:
+                # Use the OpenAI API as before
+                completion = self.client.chat.completions.create(
+                    model=self.model_name,
+                    messages=[
+                        {"role": "system", "content": self.system_prompt},
+                        {"role": "user", "content": ask_text}
+                    ],
+                    temperature=self.temperature,
+                )
+                return (index, completion.choices[0].message.content)
         except Exception as e:
             print(f"Error occurred: {e}")
             self.miss_index.append(index)
@@ -61,7 +77,7 @@ class GPTBatcher:
         new_list = []
         num_workers = self.num_workers
         timeout_duration = self.timeout_duration
-        retry_attempts = 2
+        retry_attempts = self.retry_attempts
     
         executor = ThreadPoolExecutor(max_workers=num_workers)
         message_chunks = list(self.chunk_list(message_list, num_workers))
@@ -75,14 +91,14 @@ class GPTBatcher:
                     new_list.extend(future.result() for future in done if future.done())
                     if len(not_done) == 0:
                         break
-                    future_to_message = {executor.submit(self.get_attitude, future_to_message[future]): future for future in not_done}
+                    future_to_message = {executor.submit(self.get_attitude, future_to_message[future]): future_to_message[future] for future in not_done}
         except Exception as e:
             print(f"Error occurred: {e}")
         finally:
             executor.shutdown(wait=False)
             return new_list
 
-    def complete_attitude_list(self,attitude_list, max_length):
+    def complete_attitude_list(self, attitude_list, max_length):
         completed_list = []
         current_index = 0
         for item in attitude_list:
@@ -106,7 +122,7 @@ class GPTBatcher:
         for i in range(0, len(lst), n):
             yield lst[i:i + n]
 
-    def handle_message_list(self,message_list):
+    def handle_message_list(self, message_list):
         indexed_list = [(index, data) for index, data in enumerate(message_list)]
         max_length = len(indexed_list)
         attitude_list = self.process_attitude(indexed_list)
@@ -115,32 +131,50 @@ class GPTBatcher:
         attitude_list = [x[1] for x in attitude_list]
         return attitude_list
     
-    def process_embedding(self,message_list):
-            new_list = []
-            executor = ThreadPoolExecutor(max_workers=self.num_workers)
-            # Split message_list into chunks
-            message_chunks = list(self.chunk_list(message_list, self.num_workers))
-            fixed_get_embedding = partial(self.get_embedding)
-            for chunk in tqdm(message_chunks, desc="Processing messages"):
-                future_to_message = {executor.submit(fixed_get_embedding, message): message for message in chunk}
-                for i in range(self.retry_attempts):
-                    done, not_done = wait(future_to_message.keys(), timeout=self.timeout_duration)
-                    for future in not_done:
-                        future.cancel()
-                    new_list.extend(future.result() for future in done if future.done())
-                    if len(not_done) == 0:
-                        break
-                    future_to_message = {executor.submit(fixed_get_embedding, future_to_message[future]): future_to_message[future] for future in not_done}
-            executor.shutdown(wait=False)
-            return new_list
-    def get_embedding(self,text):
-        index,text = text
-        response = self.client.embeddings.create(
-        input=text,
-        model=self.model_name)
-        return (index,response.data[0].embedding)
+    def get_embedding(self, text):
+        index, text = text
+        try:
+            if self.is_claude:
+                # Use Anthropic's embedding API if available
+                # Note: As of March 2025, make sure to check Anthropic's latest API
+                # for embeddings, as the format might have changed
+                response = self.client.embeddings.create(
+                    model=self.model_name,
+                    input=text
+                )
+                return (index, response.embedding)
+            else:
+                # Use OpenAI's embedding API
+                response = self.client.embeddings.create(
+                    input=text,
+                    model=self.model_name
+                )
+                return (index, response.data[0].embedding)
+        except Exception as e:
+            print(f"Error getting embedding: {e}")
+            self.miss_index.append(index)
+            return (index, None)
+    
+    def process_embedding(self, message_list):
+        new_list = []
+        executor = ThreadPoolExecutor(max_workers=self.num_workers)
+        # Split message_list into chunks
+        message_chunks = list(self.chunk_list(message_list, self.num_workers))
+        fixed_get_embedding = partial(self.get_embedding)
+        for chunk in tqdm(message_chunks, desc="Processing messages"):
+            future_to_message = {executor.submit(fixed_get_embedding, message): message for message in chunk}
+            for i in range(self.retry_attempts):
+                done, not_done = wait(future_to_message.keys(), timeout=self.timeout_duration)
+                for future in not_done:
+                    future.cancel()
+                new_list.extend(future.result() for future in done if future.done())
+                if len(not_done) == 0:
+                    break
+                future_to_message = {executor.submit(fixed_get_embedding, future_to_message[future]): future_to_message[future] for future in not_done}
+        executor.shutdown(wait=False)
+        return new_list
 
-    def handle_embedding_list(self,message_list):
+    def handle_embedding_list(self, message_list):
         indexed_list = [(index, data) for index, data in enumerate(message_list)]
         max_length = len(indexed_list)
         attitude_list = self.process_embedding(indexed_list)
@@ -152,5 +186,16 @@ class GPTBatcher:
     def get_miss_index(self):
         return self.miss_index
 
-    # Add other necessary methods similar to the above, refactored to fit within this class structure.
-
+# Example usage:
+if __name__ == "__main__":
+    # For OpenAI
+    openai_batcher = GPTBatcher(
+        api_key="your_openai_api_key",
+        model_name="gpt-4-turbo"
+    )
+    
+    # For Claude
+    claude_batcher = GPTBatcher(
+        api_key="your_anthropic_api_key",
+        model_name="claude-3-7-sonnet-20250219"
+    )

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpt-batch
-Version: 0.1.6
+Version: 0.1.8
 Summary: A package for batch processing with OpenAI API.
 Home-page: https://github.com/fengsxy/gpt_batch
 Author: Ted Yu
@@ -10,8 +10,8 @@ Platform: UNKNOWN
 Description-Content-Type: text/markdown
 Requires-Dist: openai
 Requires-Dist: tqdm
+Requires-Dist: anthropic
 
-Certainly! Here's a clean and comprehensive README for your `GPTBatcher` tool, formatted in Markdown:
 
 ```markdown
 # GPT Batcher
@@ -62,6 +62,22 @@ print(result)
 # Expected output: ["embedding_1", "embedding_2", "embedding_3", "embedding_4"]
 ```
 
+### Handling Message Lists with different API
+
+This example demonstrates how to send a list of questions and receive answers with different api:
+
+```python
+from gpt_batch.batcher import GPTBatcher
+
+# Initialize the batcher
+batcher = GPTBatcher(api_key='sk-', model_name='deepseek-chat',api_base_url="https://api.deepseek.com/v1")
+
+
+# Send a list of messages and receive answers
+result = batcher.handle_message_list(['question_1', 'question_2', 'question_3', 'question_4'])
+
+# Expected output: ["answer_1", "answer_2", "answer_3", "answer_4"]
+```
 ## Configuration
 
 The `GPTBatcher` class can be customized with several parameters to adjust its performance and behavior:

gpt_batch-0.1.8.dist-info/RECORD

@@ -0,0 +1,8 @@
+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,,

{gpt_batch-0.1.6.dist-info → gpt_batch-0.1.8.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: bdist_wheel (0.45.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tests/test_batcher.py

@@ -51,5 +51,27 @@ def test_get_miss_index():
     miss_index = batcher.get_miss_index()
     assert miss_index == [], "The miss index should be empty"
 # Optionally, you can add a test configuration if you have specific needs
+
+
+def test_claude_handle_message_list():
+    # Initialize the GPTBatcher with Claude model
+    api_key = os.getenv('ANTHROPIC_API_KEY')
+    if not api_key:
+        raise ValueError("Anthropic API key must be set in the environment variables as ANTHROPIC_API_KEY")
+    
+    batcher = GPTBatcher(
+        api_key=api_key, 
+        model_name='claude-3-7-sonnet-20250219', 
+        system_prompt="Your task is to discuss privacy questions and provide persuasive answers with supporting reasons."
+    )
+    message_list = ["I think privacy is important", "I don't think privacy is important"]
+
+    # 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(isinstance(result, str) and len(result) > 0 for result in results if result is not None), "Each result should be a non-empty string if not None"
+    assert batcher.is_claude, "Should recognize model as Claude"
 if __name__ == "__main__":
     pytest.main()

gpt_batch-0.1.6.dist-info/RECORD

@@ -1,8 +0,0 @@
-gpt_batch/__init__.py,sha256=zGDItktTxKLSQr44GY78dl5LKsSJig0Q59dzusqhU0U,59
-gpt_batch/batcher.py,sha256=jKLK-iuByg3Mc2ZungT5aZYzO60c5yO-YXCOf_70O6w,7591
-tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/test_batcher.py,sha256=N88RZrSuBaDti6Lry7xipyGXHn3jKg85O12mjcHHZA0,3006
-gpt_batch-0.1.6.dist-info/METADATA,sha256=Q0EhkVe8YbKac3JjhASu3_wY3y9hV_YJqqwVEzlf9wc,2932
-gpt_batch-0.1.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-gpt_batch-0.1.6.dist-info/top_level.txt,sha256=FtvJB_L9W_S6jL4G8Em_YWphG1wdKAF20BHUrf4B0yM,16
-gpt_batch-0.1.6.dist-info/RECORD,,