Update llm.py

llm.py

@@ -1,6 +1,6 @@
-import random
+import os
+import subprocess
 from datetime import datetime
-from transformers import pipeline
 
 # Viral Content Generator Agent
 PREFIX = """You are an Advanced Viral Content Generator with Self-Research and Self-Improvement Capabilities
@@ -43,27 +43,67 @@ def run_gpt(
     purpose,
     **prompt_kwargs,
 ):
-    """Runs the GPT model to generate a response."""
-    
+    """Generate a response using a local Ollama model (gpt-oss by default)."""
+
     content = PREFIX.format(
         date_time_str=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
         purpose=purpose,
     ) + prompt_template.format(**prompt_kwargs)
 
-    generator = pipeline('text-generation', model='gpt2')
+    models = _preferred_models()
+    last_error = None
+    for model in models:
+        try:
+            raw_response = _call_ollama(model, content, max_tokens)
+            response = _strip_prompt(content, raw_response)
+            if stop_tokens:
+                for token in stop_tokens:
+                    if token in response:
+                        response = response.split(token)[0]
+                        break
+            return response.strip()
+        except Exception as exc:  # pylint: disable=broad-except
+            last_error = exc
+            print(f"[LLM]: Failed with model '{model}': {exc}")
+
+    raise RuntimeError(f"All Ollama model attempts failed: {last_error}")
+
+
+def _preferred_models():
+    env_model = os.getenv("OLLAMA_MODEL")
+    candidates = [env_model, "gpt-oss", "qwen3"]
+    return [model for model in candidates if model]
+
+
+def _call_ollama(model: str, prompt: str, max_tokens: int) -> str:
+    """Invoke the local Ollama runtime with the provided prompt."""
+    max_predict = max(64, int(max_tokens or 0))
+    cmd = [
+        "ollama",
+        "run",
+        "--num-predict",
+        str(max_predict),
+        model,
+    ]
+    try:
+        completed = subprocess.run(
+            cmd,
+            input=prompt.encode("utf-8"),
+            capture_output=True,
+            check=False,
+        )
+    except FileNotFoundError as exc:
+        raise RuntimeError("Ollama executable not found. Please install Ollama or adjust PATH.") from exc
 
-    generated_text = generator(
-        content,
-        max_new_tokens=max_tokens,
-        num_return_sequences=1,
-    )[0]['generated_text']
+    if completed.returncode != 0:
+        stderr = completed.stderr.decode("utf-8", errors="ignore")
+        raise RuntimeError(stderr.strip() or f"ollama run {model} failed with code {completed.returncode}")
 
-    response = generated_text[len(content):]
+    return completed.stdout.decode("utf-8", errors="ignore")
 
-    # Manually truncate at the first stop token
-    if stop_tokens:
-        for token in stop_tokens:
-            if token in response:
-                response = response.split(token)[0]
 
-    return response
+def _strip_prompt(prompt: str, raw_response: str) -> str:
+    """Some models echo the prompt; remove it if present."""
+    if raw_response.startswith(prompt):
+        return raw_response[len(prompt):]
+    return raw_response