diff --git a/generator/src/llm_utils.py b/generator/src/llm_utils.py index 7ed1d29..7e54c9e 100644 --- a/generator/src/llm_utils.py +++ b/generator/src/llm_utils.py @@ -1,14 +1,19 @@ import requests import logging +import asyncio +from functools import partial from generator.src.config import OLLAMA_MODEL logger = logging.getLogger(__name__) +# 🔥 Limit concurrency +semaphore = asyncio.Semaphore(5) + + def generate_text(prompt, model=None, url_override=None): """Generate text from prompt using local Ollama.""" target_model = model or OLLAMA_MODEL - - # Use centralized config unless overridden + if url_override: url = url_override else: @@ -24,27 +29,28 @@ def generate_text(prompt, model=None, url_override=None): "num_ctx": 8192 } } - - logger.info(f"Sending request to Ollama (model: {target_model}, prompt length: {len(prompt)} chars)") - + try: resp = requests.post(url, json=payload, timeout=600) resp.raise_for_status() - except Exception as e: - logger.error(f"Error calling Ollama at {url}: {e}", exc_info=True) - return "" - - try: data = resp.json() - response_text = data.get("response", "") - - logger.info(f"Received response from Ollama (length: {len(response_text)} chars)") - logger.info(f"Response done: {data.get('done', False)}") - - if not response_text: - logger.warning(f"Empty response from Ollama. Full data: {data}") - - return response_text + return data.get("response", "") except Exception as e: - logger.error(f"Error parsing Ollama response: {e}", exc_info=True) + logger.error(f"LLM error: {e}", exc_info=True) return "" + + +async def generate_text_async(prompt, model=None, url_override=None): + loop = asyncio.get_event_loop() + func = partial(generate_text, prompt, model, url_override) + return await loop.run_in_executor(None, func) + + +async def controlled_generate_text(prompt): + async with semaphore: + return await generate_text_async(prompt) + + +async def process_batch(prompts): + tasks = [controlled_generate_text(p) for p in prompts] + return await asyncio.gather(*tasks, return_exceptions=True)