llm_batch_helper 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

llm_batch_helper/__init__.py

@@ -3,7 +3,7 @@ from .config import LLMConfig
 from .input_handlers import get_prompts, read_prompt_files, read_prompt_list
 from .providers import process_prompts_batch
 
-__version__ = "0.1.6"
+__version__ = "0.2.0"
 
 __all__ = [
     "LLMCache",

llm_batch_helper/config.py

@@ -16,6 +16,7 @@ class LLMConfig:
         verification_callback: Optional[Callable[..., bool]] = None,
         verification_callback_args: Optional[Dict] = None,
         max_completion_tokens: Optional[int] = None,
+        **kwargs
     ):
         self.model_name = model_name
         self.temperature = temperature
@@ -30,3 +31,4 @@ class LLMConfig:
         self.verification_callback_args = (
             verification_callback_args if verification_callback_args is not None else {}
         )
+        self.kwargs = kwargs

llm_batch_helper/providers.py

@@ -46,6 +46,7 @@ async def _get_openai_response_direct(
             messages=messages,
             temperature=config.temperature,
             max_completion_tokens=config.max_completion_tokens,
+            **config.kwargs,
         )
         usage_details = {
             "prompt_token_count": response.usage.prompt_tokens,
@@ -94,6 +95,7 @@ async def _get_together_response_direct(
             "messages": messages,
             "temperature": config.temperature,
             "max_tokens": config.max_completion_tokens,
+            **config.kwargs,
         }
 
         response = await client.post(
@@ -116,6 +118,66 @@ async def _get_together_response_direct(
             "usage_details": usage_details,
         }
 
+
+@retry(
+    stop=stop_after_attempt(5),
+    wait=wait_exponential(multiplier=1, min=4, max=60),
+    retry=retry_if_exception_type(
+        (
+            ConnectionError,
+            TimeoutError,
+            httpx.HTTPStatusError,
+            httpx.RequestError,
+        )
+    ),
+    reraise=True,
+)
+async def _get_openrouter_response_direct(
+    prompt: str, config: LLMConfig
+) -> Dict[str, Union[str, Dict]]:
+    api_key = os.environ.get("OPENROUTER_API_KEY")
+    if not api_key:
+        raise ValueError("OPENROUTER_API_KEY environment variable not set")
+
+    async with httpx.AsyncClient(timeout=1000.0) as client:
+        messages = [
+            {"role": "system", "content": config.system_instruction},
+            {"role": "user", "content": prompt},
+        ]
+
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        }
+
+        payload = {
+            "model": config.model_name,
+            "messages": messages,
+            "temperature": config.temperature,
+            "max_tokens": config.max_completion_tokens,
+            **config.kwargs,
+        }
+
+        response = await client.post(
+            "https://openrouter.ai/api/v1/chat/completions",
+            json=payload,
+            headers=headers,
+        )
+        response.raise_for_status()
+        
+        response_data = response.json()
+        usage = response_data.get("usage", {})
+        usage_details = {
+            "prompt_token_count": usage.get("prompt_tokens", 0),
+            "completion_token_count": usage.get("completion_tokens", 0),
+            "total_token_count": usage.get("total_tokens", 0),
+        }
+        
+        return {
+            "response_text": response_data["choices"][0]["message"]["content"],
+            "usage_details": usage_details,
+        }
+
 async def get_llm_response_with_internal_retry(
     prompt_id: str,
     prompt: str,
@@ -135,6 +197,8 @@ async def get_llm_response_with_internal_retry(
             response = await _get_openai_response_direct(prompt, config)
         elif provider.lower() == "together":
             response = await _get_together_response_direct(prompt, config)
+        elif provider.lower() == "openrouter":
+            response = await _get_openrouter_response_direct(prompt, config)
         else:
             raise ValueError(f"Unsupported provider: {provider}")
 
@@ -165,7 +229,7 @@ async def process_prompts_batch(
         prompts: Optional list of prompts in any supported format (string, tuple, or dict)
         input_dir: Optional path to directory containing prompt files
         config: LLM configuration
-        provider: LLM provider to use ("openai", "together", or "gemini")
+        provider: LLM provider to use ("openai", "together", or "openrouter")
         desc: Description for progress bar
         cache_dir: Optional directory for caching responses
         force: If True, force regeneration even if cached response exists

{llm_batch_helper-0.1.6.dist-info → llm_batch_helper-0.2.0.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.3
 Name: llm_batch_helper
-Version: 0.1.6
+Version: 0.2.0
 Summary: A Python package that enables batch submission of prompts to LLM APIs, with built-in async capabilities and response caching.
 License: MIT
-Keywords: llm,openai,together,batch,async,ai,nlp,api
+Keywords: llm,openai,together,openrouter,batch,async,ai,nlp,api
 Author: Tianyi Peng
 Author-email: tianyipeng95@gmail.com
 Requires-Python: >=3.11,<4.0
@@ -245,7 +245,7 @@ Main function for batch processing of prompts.
 ```python
 async def process_prompts_batch(
     config: LLMConfig,
-    provider: str,  # "openai" or "together"
+    provider: str,  # "openai", "together", or "openrouter"
     prompts: Optional[List[str]] = None,
     input_dir: Optional[str] = None,
     cache_dir: str = "llm_cache",

llm_batch_helper-0.2.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+llm_batch_helper/__init__.py,sha256=Cy4p5YE6m3pbpux7GCjpf7GsyUbO_rUtqJsrhTWpPQY,348
+llm_batch_helper/cache.py,sha256=QUODQ1tPCvFThO3yvVOTcorcOrmN2dP5HLF1Y2O1bTQ,1276
+llm_batch_helper/config.py,sha256=RasljP9dzigZpKjm9yW6gU7_e3yjztfokjiDBf77iO4,1372
+llm_batch_helper/exceptions.py,sha256=59_f3jINUhKFble6HTp8pmtLSFE2MYLHWGclwaQKs28,296
+llm_batch_helper/input_handlers.py,sha256=IadA732F1Rw0zcBok5hjZr32RUm8eTUOpvLsRuMvaE4,2877
+llm_batch_helper/providers.py,sha256=lNxAh6hfKjXhHEOMDmoca2dbtcE62UZV-HXlTKemIOE,12230
+llm_batch_helper-0.2.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+llm_batch_helper-0.2.0.dist-info/METADATA,sha256=-x5B1uALZBkM2VLf8tB0GI7JcJkiXu7qgPhkPG427rI,11288
+llm_batch_helper-0.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+llm_batch_helper-0.2.0.dist-info/RECORD,,

llm_batch_helper-0.1.6.dist-info/RECORD

@@ -1,10 +0,0 @@
-llm_batch_helper/__init__.py,sha256=a9jaWNSp0pqLz3gB_mo9r1aYo9f1i-qtuYhJj8AOPHk,348
-llm_batch_helper/cache.py,sha256=QUODQ1tPCvFThO3yvVOTcorcOrmN2dP5HLF1Y2O1bTQ,1276
-llm_batch_helper/config.py,sha256=WcZKTD-Mtocsx1plS9x3hh6MstVmyxD-tyidGUatkPY,1327
-llm_batch_helper/exceptions.py,sha256=59_f3jINUhKFble6HTp8pmtLSFE2MYLHWGclwaQKs28,296
-llm_batch_helper/input_handlers.py,sha256=IadA732F1Rw0zcBok5hjZr32RUm8eTUOpvLsRuMvaE4,2877
-llm_batch_helper/providers.py,sha256=cgccd_4D7J48ClmkigZ7KXOzTnBmaya8soDYF5IlPJs,10212
-llm_batch_helper-0.1.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-llm_batch_helper-0.1.6.dist-info/METADATA,sha256=Msg55neTu6jvxLKW6hicJOp-k7Q6Edp8qdIs_AKfVkM,11262
-llm_batch_helper-0.1.6.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-llm_batch_helper-0.1.6.dist-info/RECORD,,