You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
26 lines
1003 B
26 lines
1003 B
from ollama import Client |
|
import os |
|
import env_manager |
|
env_manager.set_env() |
|
|
|
class LLM: |
|
def __init__(self, system_message=None, num_ctx=2000, temperature=0, chat=True) -> None: |
|
self.llm_model = os.getenv("LLM_MODEL") |
|
self.system_message = system_message |
|
self.options = {"temperature": temperature, "num_ctx": num_ctx} |
|
self.messages = [{'role': 'system', 'content': self.system_message}] |
|
self.chat = chat |
|
self.ollama = Client(host=f'{os.getenv("LLM_URL")}:{os.getenv("LLM_PORT")}') |
|
|
|
def generate(self, prompt: str) -> str: |
|
self.messages.append({"role": "user", "content": prompt}) |
|
|
|
result = self.ollama.chat(model=self.llm_model, messages=self.messages, options=self.options, ) |
|
|
|
answer = result['message']['content'] |
|
self.messages.append({"role": "assistant", "content": answer}) |
|
if not self.chat: |
|
self.messages = [{'role': 'system', 'content': self.system_message}] |
|
|
|
return answer |
|
|
|
|