chat/chatbot.py - master - nofuti/Pichko-System-Module

chat/chatbot.py
78 строк2 KB
Георгий Пичко
В прототип добавлен чат
13 май 2026, 12:38
13 май 2026, 12:380416b5d
#!/usr/bin/env python3
"""
Простой консольный чат-бот: ввод пользователя → ответ модели через API Ollama (HTTP).
cd C:\Users\user\Desktop\AI_Chat\prototype
python -m pip install -r requirements.txt
streamlit run streamlit_app.py
"""

import os
import sys

import requests


SYSTEM_PROMPT = (
    "Ты — эксперт по векторным базам данных. Отвечай кратко и профессионально."
)


def ollama_base_url() -> str:
    return os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip("/")


def send_chat(messages: list[dict[str, str]], model: str, timeout: int = 300) -> str:
    url = f"{ollama_base_url()}/api/chat"
    full_messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
    payload = {"model": model, "messages": full_messages, "stream": False}
    response = requests.post(url, json=payload, timeout=timeout)
    response.raise_for_status()
    data = response.json()
    return data["message"]["content"]


def main() -> None:
    model = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("OLLAMA_MODEL", "llama3.2:1b")
    messages: list[dict[str, str]] = []

    print("Чат с Ollama (через requests).")
    print(f"Сервер: {ollama_base_url()} | Модель: {model}")
    print("Команды: /exit — выход, /reset — очистить историю.\n")

    while True:
        try:
            user_text = input("Вы: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nВыход.")
            break

        if not user_text:
            continue
        lowered = user_text.lower()
        if lowered == "/exit":
            print("Выход.")
            break
        if lowered == "/reset":
            messages.clear()
            print("История диалога очищена.\n")
            continue

        messages.append({"role": "user", "content": user_text})

        try:
            answer = send_chat(messages, model)
        except requests.HTTPError as e:
            messages.pop()
            print(f"Ошибка HTTP: {e.response.status_code} — {e.response.text[:500]}\n")
            continue
        except requests.RequestException as e:
            messages.pop()
            print(f"Ошибка запроса: {e}\n")
            continue

        messages.append({"role": "assistant", "content": answer})
        print(f"Бот: {answer}\n")


if __name__ == "__main__":
    main()