llm_batch_helper 0.1.6__tar.gz → 0.2.0__tar.gz

{llm_batch_helper-0.1.6 → llm_batch_helper-0.2.0}/PKG-INFO

@@ -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.1.6 → llm_batch_helper-0.2.0}/README.md

@@ -218,7 +218,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.1.6 → llm_batch_helper-0.2.0}/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-0.1.6 → llm_batch_helper-0.2.0}/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-0.1.6 → llm_batch_helper-0.2.0}/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 → llm_batch_helper-0.2.0}/pyproject.toml

@@ -1,13 +1,13 @@
 [tool.poetry]
 name = "llm_batch_helper"
-version = "0.1.6"
+version = "0.2.0"
 description = "A Python package that enables batch submission of prompts to LLM APIs, with built-in async capabilities and response caching."
 authors = ["Tianyi Peng <tianyipeng95@gmail.com>"]
 readme = "README.md"
 license = "MIT"
 homepage = "https://github.com/TianyiPeng/LLM_batch_helper"
 repository = "https://github.com/TianyiPeng/LLM_batch_helper"
-keywords = ["llm", "openai", "together", "batch", "async", "ai", "nlp", "api"]
+keywords = ["llm", "openai", "together", "openrouter", "batch", "async", "ai", "nlp", "api"]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Developers",