blog_creator/ollama_generator.py

29 lines
976 B
Python

import os
import requests
class OllamaGenerator:
def __init__(self, title: str, content: str):
self.title = title
self.content = content
self.ollama_url = f"{os.getenv('OLLAMA_PROTOCOL')}://{os.getenv('OLLAMA_HOST')}:{os.getenv('OLLAMA_PORT')}/api/generate"
def generate_markdown(self) -> str:
data = {
"prompt": f"Title: {self.title}\n\nContent:\n{self.content}",
"max_tokens": 1000,
"temperature": 0.7,
"top_p": 0.9,
"presence_penalty": 0,
"frequency_penalty": 0
}
response = requests.post(self.ollama_url, json=data)
if response.status_code == 200:
return response.json()["choices"][0]["text"]
else:
raise Exception(f"Failed to generate markdown: {response.text}")
def save_to_file(self, filename: str) -> None:
with open(filename, "w") as f:
f.write(self.generate_markdown())