31 lines
884 B
Python
31 lines
884 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
|
||
|
|
return self.ollama.invoke(messanges).content
|