untested ollama

This commit is contained in:
Andrew Ridgway 2024-11-14 17:48:11 +10:00
parent 675c7d1701
commit 1bd0ae9b1b
3 changed files with 50 additions and 11 deletions

View File

View File

@ -0,0 +1,39 @@
import os
from ollama import Client
class OllamaGenerator:
def __init__(self, title: str, content: str):
self.title = title
self.content = content
self.ollama_client = Client(host=f"""
{os.getenv('OLLAMA_PROTOCOL')}://{os.getenv('OLLAMA_HOST')}:{os.getenv('OLLAMA_PORT')}
"""
)
def generate_markdown(self) -> str:
prompt = f"""
Generate a 1000 word blog as a markdown document.
The title for the blog is {self.title}.
Do not output the title in the markdown.
The content for the blog is:
{self.content}
Only output markdown DO NOT GENERATE AN EXPLANATION
"""
try:
self.response = self.ollama_client.chat(model=self.ollama_model,
messages=[
{
'role': 'user',
'content': f'{prompt}',
},
])
except Exception as e:
raise Exception(f"Failed to generate markdown: {e}")
def save_to_file(self, filename: str) -> None:
with open(filename, "w") as f:
f.write(self.generate_markdown())

View File

@ -1,5 +1,4 @@
from ollama import Client
import ai_generators.ollama_md_generator as omg
import trilium.notes as tn
tril = tn.TrilumNotes()
@ -7,17 +6,18 @@ tril = tn.TrilumNotes()
tril.get_new_notes()
tril_notes = tril.get_notes_content()
def convert_to_lowercase_with_underscores(string):
return string.lower().replace(" ", "_")
for note in tril_notes:
print(tril_notes[note]['title'])
print(tril_notes[note]['content'])
client = Client(host='http://192.168.178.45:11434')
print("Generating Document")
ai_gen = omg.OllamaGenerator(tril_notes[note]['title'],
tril_notes[note]['content'])
os_friendly_title = convert_to_lowercase_with_underscores(tril_notes[note]['title'])
omg.save_to_file(f"{os_friendly_title}.md")
user_prompt = input("Ask mistral-nemo a question: ")
response = client.chat(model='mistral-nemo', messages=[
{
'role': 'user',
'content': f'{user_prompt}',
},
])
print(response['message']['content'])