2024-05-18 19:04:39 +00:00
|
|
|
from langchain_community.chat_models import ChatOllama
|
|
|
|
|
|
|
|
class OllamaChatBot:
|
|
|
|
def __init__(self, base_url, model, headers):
|
|
|
|
self.base_url = base_url
|
|
|
|
self.model = model
|
2024-05-19 12:03:47 +00:00
|
|
|
if self.is_empty(headers):
|
|
|
|
self.headers = ""
|
|
|
|
else:
|
|
|
|
self.headers = headers
|
2024-05-18 19:04:39 +00:00
|
|
|
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
|
|
|
|
)
|
2024-05-19 12:03:47 +00:00
|
|
|
|
|
|
|
def is_empty(self, dictionary):
|
|
|
|
return len(dictionary) == 1 and list(dictionary.keys())[0] == '' and list(dictionary.values())[0] == ''
|
|
|
|
|
2024-05-18 19:04:39 +00:00
|
|
|
|
|
|
|
|
|
|
|
def get_request(self, prompt):
|
|
|
|
messanges = []
|
|
|
|
self.messanges.append(prompt)
|
|
|
|
if len(self.messanges) > 5:
|
|
|
|
messanges = messanges[:5]
|
|
|
|
else:
|
|
|
|
messanges = self.messanges
|
2024-05-19 12:03:20 +00:00
|
|
|
try:
|
|
|
|
return self.ollama.invoke(messanges).content
|
|
|
|
except ValueError:
|
|
|
|
return "An unexpected Error occuried"
|