77 lines
3.7 KiB
Python
77 lines
3.7 KiB
Python
import ollama
|
|
import os
|
|
import requests
|
|
import json
|
|
from .logger import get_logger
|
|
logger = get_logger()
|
|
|
|
class OllamaClient():
|
|
def __init__(self):
|
|
self.host = os.getenv("ENDPOINT_OLLAMA", "https://ollamamodel.matitos.org")
|
|
self.client = ollama.Client(host=self.host)
|
|
self.options = {"temperature": 0, "seed": 13579}
|
|
|
|
def _get_default_model(self):
|
|
return os.getenv("OLLAMA_MODEL_DEFAULT", "llama3.2:3b")
|
|
|
|
def get_models(self):
|
|
try:
|
|
# Get models
|
|
models = sorted([m.model for m in self.client.list().models])
|
|
# r = requests.get( os.path.join(endpoint, "models") )
|
|
# r.json().get("models")
|
|
|
|
# Default within it?
|
|
if (self._get_default_model() in models):
|
|
return [self._get_default_model()] + [m for m in models if m != self._get_default_model()]
|
|
else:
|
|
return models
|
|
except Exception as e:
|
|
return []
|
|
|
|
def get_prompt(self, content):
|
|
return "Provide, in one sentence each, the what, why, who, when, where, and a detailed summary of the content below:\n\n{}".format(content)
|
|
return "First, provide a detailed summary of the content below in one paragraph. Second, specify in one sentence each the who, what, when, where and why of the story. Do not mention or reference the original text, its source, or any phrases like 'According to' or 'The text states':\n\n{}".format(content)
|
|
return "First, provide a summary of the content below in one paragraph. Second, specify the Who, What, When, Where and Why of the story:\n\n{}".format(content)
|
|
# First, provide a summary of the content below in one paragraph. Second, specify the who, what, when, where and why of the story in one sentence each. Do not mention or reference the original text, its source, or any phrases like 'According to' or 'The text states':
|
|
'''
|
|
return ("Rewrite the content below into a clear and concise summary of one paragraph maximum, presenting the key points as if they are newly written insights. "
|
|
"Do not mention or reference the original text, its source, or any phrases like 'According to' or 'The text states'. "
|
|
"Write in a natural, standalone format that feels like an original explanation. "
|
|
"Keep it brief, engaging, informative, in the style of a news article:\n\n{}".format(content)
|
|
)
|
|
'''
|
|
|
|
def generate(self, model, prompt, format=None):
|
|
try:
|
|
# Generate response
|
|
response = self.client.generate(model=model, prompt=prompt, format=format, options=self.options)
|
|
# Extract response
|
|
response = response.response
|
|
# Json? -> Dict
|
|
if (format == "json"):
|
|
# Dict
|
|
response = json.loads(response)
|
|
# Force unload
|
|
r = requests.post( os.path.join(self.host, "unload_model") )
|
|
except Exception as e:
|
|
logger.warning("Exception while generating LLM response: {}".format(str(e)))
|
|
if (format == "json"):
|
|
response = {}
|
|
else:
|
|
response = None
|
|
# Text
|
|
return response
|
|
|
|
def generate_stream(self, model, prompt):
|
|
try:
|
|
# Generate response
|
|
response = self.client.generate(model=model, prompt=prompt, format="json", stream=True, options=self.options)
|
|
# Streamed chunks
|
|
for chunk in response:
|
|
yield chunk.response
|
|
# Force unload
|
|
r = requests.post( os.path.join(self.host, "unload_model") )
|
|
except Exception as e:
|
|
logger.warning("Exception while generating LLM response: {}".format(str(e)))
|