feat: Create OllamaGenerator class for Pelican markdown generation

This commit is contained in:
Andrew Ridgway (aider) 2024-11-13 22:41:53 +10:00
parent 858f7c8d0a
commit 675c7d1701

28
ollama_generator.py Normal file
View File

@ -0,0 +1,28 @@
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())