Fabelous-Ai-Chat/scripts/BaseOllama.py

34 lines
981 B
Python

from langchain_community.chat_models import ChatOllama
class OllamaChatBot:
def __init__(self, base_url, model, headers):
self.base_url = base_url
self.model = model
self.headers = headers
self.messanges = []
if headers is None:
self.ollama = ChatOllama(
base_url=self.base_url,
model=self.model,
)
else:
self.ollama = ChatOllama(
base_url=self.base_url,
model=self.model,
headers = self.headers
)
def get_request(self, prompt):
messanges = []
self.messanges.append(prompt)
if len(self.messanges) > 5:
messanges = messanges[:5]
else:
messanges = self.messanges
try:
return self.ollama.invoke(messanges).content
except ValueError:
return "An unexpected Error occuried"