57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
import ollama
|
|
|
|
class OllamaProvider:
|
|
|
|
def __init__(self, model):
|
|
self.model = model
|
|
self.sys_prompt = "You are a helpful assistant."
|
|
|
|
def generate_stream(self, messages, temperature=0.8, max_tokens=300):
|
|
"""Генерирует ответ от Ollama и возвращает его"""
|
|
try:
|
|
stream = ollama.chat(
|
|
model=self.model,
|
|
messages=messages,
|
|
stream=True,
|
|
options={
|
|
"temperature": temperature,
|
|
"num_predict": max_tokens,
|
|
},
|
|
)
|
|
|
|
# Собираем ответ в строку
|
|
response_text = ""
|
|
for chunk in stream:
|
|
content = chunk['message']['content']
|
|
print(content, end='', flush=True)
|
|
response_text += content
|
|
return response_text
|
|
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
if "502" in error_msg:
|
|
return "Ошибка: Проверьте наличие VPN или Proxy."
|
|
return f"Ошибка: {error_msg}"
|
|
|
|
def generate_stream_tokens(self, messages, temperature=0.8, max_tokens=300):
|
|
try:
|
|
stream = ollama.chat(
|
|
model=self.model,
|
|
messages=messages,
|
|
stream=True,
|
|
options={
|
|
"temperature": temperature,
|
|
"num_predict": max_tokens,
|
|
},
|
|
)
|
|
|
|
for chunk in stream:
|
|
content = chunk['message']['content']
|
|
|
|
if content:
|
|
yield content # 🔥 ВАЖНО: внутри цикла
|
|
|
|
except Exception as e:
|
|
yield f"Error: {str(e)}"
|
|
if "502" in e:
|
|
yield "Ошибка: Проверьте наличие VPN или Proxy." |