architect/standards/local-ai/04_CODE.md

04 — КОД: Все файлы с пояснениями

Навигация: ← Установка | README | Тесты →


Содержание

  1. install.sh
  2. docker-compose.yml
  3. config/system_prompt.txt
  4. config/models.yaml
  5. router/router.py
  6. router/tools.py
  7. memory/memory.py

1. install.sh

Скрипт устанавливает весь стек с нуля на чистый Ubuntu.

#!/bin/bash
# install.sh
# Установка локального AI стека на Ubuntu 22.04/24.04
# Запуск: bash install.sh

# Остановить при первой ошибке
set -e

# Цвета для вывода
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m'

ok()   { echo -e "${GREEN}$1${NC}"; }
fail() { echo -e "${RED}$1${NC}"; exit 1; }

echo "=== LOCAL AI STACK INSTALL ==="
echo "Дата: $(date)"
echo ""

# ─── Шаг 1: Проверка системы ───────────────────────────────────

echo "Проверка системы..."

# Проверить Ubuntu
[ -f /etc/os-release ] || fail "Не Ubuntu"
. /etc/os-release
ok "OS: $PRETTY_NAME"

# Проверить AVX2 (нужен для быстрой работы LLM на CPU)
if grep -q avx2 /proc/cpuinfo; then
    ok "CPU: AVX2 поддерживается"
else
    echo "⚠️  AVX2 не найден — модели будут работать медленнее"
fi

# Проверить RAM
RAM_GB=$(free -g | awk '/^Mem:/{print $2}')
if [ "$RAM_GB" -lt 8 ]; then
    fail "Недостаточно RAM: ${RAM_GB}GB, нужно минимум 8GB"
fi
ok "RAM: ${RAM_GB}GB"

# Проверить место на диске
DISK_GB=$(df -BG / | awk 'NR==2{print $4}' | tr -d G)
if [ "$DISK_GB" -lt 20 ]; then
    fail "Недостаточно диска: ${DISK_GB}GB, нужно минимум 20GB"
fi
ok "Диск: ${DISK_GB}GB свободно"

# ─── Шаг 2: Системные пакеты ───────────────────────────────────

echo ""
echo "Установка системных пакетов..."

apt-get update -qq
apt-get install -y -qq \
    curl \
    git \
    docker.io \
    docker-compose \
    python3-pip \
    python3-venv \
    htop

# Запустить Docker
systemctl enable docker
systemctl start docker

ok "Docker установлен и запущен"

# ─── Шаг 3: Ollama ─────────────────────────────────────────────

echo ""
echo "Установка Ollama..."

curl -fsSL https://ollama.com/install.sh | sh

# Настроить Ollama слушать на всех интерфейсах
# (нужно для Docker контейнеров)
mkdir -p /etc/systemd/system/ollama.service.d
cat > /etc/systemd/system/ollama.service.d/override.conf << 'EOF'
[Service]
Environment="OLLAMA_HOST=0.0.0.0:11434"
EOF

systemctl daemon-reload
systemctl enable ollama
systemctl restart ollama

# Подождать пока Ollama запустится
echo "Ожидание запуска Ollama..."
for i in $(seq 1 10); do
    if curl -s http://localhost:11434/api/tags > /dev/null 2>&1; then
        break
    fi
    sleep 2
done

curl -s http://localhost:11434/api/tags > /dev/null || fail "Ollama не запустился"
ok "Ollama запущен на порту 11434"

# ─── Шаг 4: Модели ─────────────────────────────────────────────

echo ""
echo "Скачивание моделей..."
echo "⚠️  Это может занять 20-60 минут зависимости от интернета"

# Основная умная модель
echo "Скачиваю qwen2.5:14b (~8GB)..."
ollama pull qwen2.5:14b
ok "qwen2.5:14b загружен"

# Модель для кода
echo "Скачиваю qwen2.5-coder:7b (~4.5GB)..."
ollama pull qwen2.5-coder:7b
ok "qwen2.5-coder:7b загружен"

# Модель для embeddings (RAG и память)
echo "Скачиваю nomic-embed-text (~274MB)..."
ollama pull nomic-embed-text
ok "nomic-embed-text загружен"

# ─── Шаг 5: Python зависимости для Router ──────────────────────

echo ""
echo "Установка Python зависимостей..."

pip3 install -q \
    fastapi==0.115.0 \
    uvicorn==0.30.0 \
    requests==2.32.0 \
    pyyaml==6.0.2 \
    chromadb==0.5.0

ok "Python зависимости установлены"

# ─── Шаг 6: Docker сервисы ─────────────────────────────────────

echo ""
echo "Запуск Docker сервисов..."

cd /opt/local-ai
docker-compose up -d --build

# Подождать запуска
sleep 10

ok "Все сервисы запущены"

# ─── Итог ──────────────────────────────────────────────────────

echo ""
echo "=== УСТАНОВКА ЗАВЕРШЕНА ==="
echo ""
echo "Ссылки:"
echo "  Интерфейс:  http://localhost:3000"
echo "  API Router: http://localhost:8000"
echo "  Ollama API: http://localhost:11434"
echo ""
echo "Следующий шаг:"
echo "  bash tests/run_tests.sh"

2. docker-compose.yml

Описывает все сервисы. Docker запускает их как контейнеры.

# docker-compose.yml
version: '3.8'

# Внутренняя сеть — контейнеры видят друг друга по имени
networks:
  ai-network:
    driver: bridge

services:

  # ─── Open WebUI ─────────────────────────────────────────────
  # Веб-интерфейс. Пользователь открывает в браузере.
  open-webui:
    image: ghcr.io/open-webui/open-webui:latest
    container_name: open-webui
    restart: always          # перезапускаться при падении
    ports:
      - "3000:8080"           # порт хоста : порт контейнера
    environment:
      # Куда обращаться за моделями
      # Если используем Router — он проксирует в Ollama
      - OLLAMA_BASE_URL=http://host.docker.internal:11434
      # Секретный ключ для сессий (поменяй на уникальный)
      - WEBUI_SECRET_KEY=local-secret-change-me
    volumes:
      # Данные WebUI (история чатов, настройки)
      - webui-data:/app/backend/data
    networks:
      - ai-network
    extra_hosts:
      # Позволяет обращаться к хостовой машине
      - "host.docker.internal:host-gateway"

  # ─── AI Router ──────────────────────────────────────────────
  # Умный маршрутизатор запросов.
  # Выбирает модель, добавляет промпт, работает с памятью.
  router:
    build: ./router            # собрать из ./router/Dockerfile
    container_name: ai-router
    restart: always
    ports:
      - "8000:8000"
    environment:
      - OLLAMA_URL=http://host.docker.internal:11434
      - CHROMA_URL=http://chromadb:8000
    volumes:
      # Монтируем конфиги (только чтение)
      - ./config:/config:ro
    networks:
      - ai-network
    extra_hosts:
      - "host.docker.internal:host-gateway"
    depends_on:
      - chromadb              # ждать запуска ChromaDB

  # ─── ChromaDB ───────────────────────────────────────────────
  # Векторная база данных.
  # Хранит embeddings для долгой памяти и RAG.
  chromadb:
    image: chromadb/chroma:latest
    container_name: chromadb
    restart: always
    ports:
      - "8001:8000"
    volumes:
      # Данные ChromaDB (векторы, метаданные)
      - chroma-data:/chroma/chroma
    networks:
      - ai-network

# Именованные тома — данные сохраняются между перезапусками
volumes:
  webui-data:
  chroma-data:

3. config/system_prompt.txt

Личность и правила поведения модели.
Это первое что видит модель в каждом разговоре.

# config/system_prompt.txt

Ты — локальный AI-ассистент.
Работаешь офлайн, все данные остаются у пользователя.

## ЛИЧНОСТЬ
- Краткий и точный
- Без похвал: "Отличный вопрос!" — запрещено
- Без воды и вступлений
- Честный: если не знаешь — говори прямо

## АВТОПЕРЕКЛЮЧЕНИЕ РЕЖИМОВ
Определи тип запроса и отвечай соответственно:

КОД (триггеры: код, функция, баг, python, js, sql, class, def):
→ Пиши рабочий код сразу, без псевдокода
→ Язык всегда указывай (```python)
→ Добавляй пример использования
→ Обрабатывай edge cases

ОБЪЯСНЕНИЕ (триггеры: объясни, почему, что такое, как работает):
→ От простого к сложному
→ Аналогии из жизни
→ Один концепт — один абзац
→ Резюме в конце

ТЕКСТ (триггеры: напиши, составь, письмо, документ):
→ Чёткая структура
→ Конкретика вместо общих слов
→ Стиль под задачу

АНАЛИЗ (триггеры: найди, сравни, таблица, список):
→ Таблица или нумерованный список
→ Факты отдельно от выводов
→ Итог одной строкой

ПЛАН (триггеры: план, шаги, как сделать, последовательность):
→ Нумерованные шаги
→ Каждый шаг — конкретное действие
→ Результат каждого шага

## ПРАВИЛА
- Сначала ответ, потом объяснение
- Максимум 50 строк если не просят больше
- Отвечай на языке вопроса (русский по умолчанию)
- Если неясно — задай 1 уточняющий вопрос
- Не придумывай факты и данные
- Код в блоках с указанием языка

4. config/models.yaml

Настройки моделей и правила роутинга.

# config/models.yaml

# Список доступных моделей
models:
  # Основная умная модель — для большинства задач
  default: qwen2.5:14b

  # Специализированная на коде
  # Обучена на GitHub, лучше понимает паттерны
  code: qwen2.5-coder:7b

  # Большая модель для сложных задач
  # Требует 40GB+ RAM
  heavy: qwen2.5:72b

  # Только для создания векторов (embeddings)
  # Не для чата! Маленькая и быстрая
  embed: nomic-embed-text

# Правила автоматического выбора модели
routing:

  # Если любое из этих слов в запросе → используем code модель
  code_keywords:
    - "код"
    - "функция"
    - "python"
    - "javascript"
    - "typescript"
    - "sql"
    - "баг"
    - "ошибка в коде"
    - "class"
    - "def "
    - "dockerfile"
    - "docker-compose"
    - "скрипт"
    - "алгоритм"
    - "рефактор"
    - "nginx"
    - "linux команда"

  # Если любое из этих слов → используем heavy модель
  heavy_keywords:
    - "проанализируй детально"
    - "сравни подробно"
    - "архитектура системы"
    - "стратегия развития"
    - "напиши большой документ"
    - "полный анализ"

5. router/router.py

Главный файл. Принимает запросы, выбирает модель,
добавляет память и системный промпт.

# router/router.py
"""
AI Router — умный маршрутизатор запросов.

Принимает: HTTP POST /chat с промптом и историей
Делает:    выбирает модель, добавляет контекст, вызывает Ollama
Возвращает: ответ + имя использованной модели
"""
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import requests
import yaml
import os
import json
import logging

# Настройка логирования
logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

app = FastAPI(title="AI Router", version="1.0.0")

# ─── Конфигурация ────────────────────────────────────────────────────────────

# URL сервисов (берём из переменных окружения, с дефолтами)
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
CHROMA_URL = os.getenv("CHROMA_URL", "http://localhost:8001")

# Загружаем конфиг моделей
def load_config():
    config_path = os.getenv("CONFIG_PATH", "/config/models.yaml")
    with open(config_path) as f:
        return yaml.safe_load(f)

# Загружаем системный промпт (личность модели)
def load_system_prompt():
    prompt_path = os.getenv("PROMPT_PATH", "/config/system_prompt.txt")
    with open(prompt_path) as f:
        return f.read()

CFG = load_config()
SYSTEM_PROMPT = load_system_prompt()

# ─── Схемы данных (Pydantic валидация) ───────────────────────────────────────

class Message(BaseModel):
    role: str    # "user", "assistant", "system"
    content: str

class ChatRequest(BaseModel):
    prompt: str
    history: list[Message] = []   # история разговора
    model: str | None = None      # принудительно задать модель
    stream: bool = False          # стриминг токенов

class ChatResponse(BaseModel):
    response: str
    model_used: str

# ─── Логика выбора модели ─────────────────────────────────────────────────────

def pick_model(prompt: str) -> str:
    """
    Выбрать оптимальную модель по содержимому запроса.

    Алгоритм:
    1. Если есть ключевые слова тяжёлых задач → heavy модель
    2. Если есть ключевые слова кода → code модель
    3. Иначе → default модель
    """
    p = prompt.lower()

    # Сначала проверяем "тяжёлые" задачи (они важнее)
    for kw in CFG["routing"].get("heavy_keywords", []):
        if kw in p:
            log.info(f"Выбрана модель heavy по ключевому слову: '{kw}'")
            return CFG["models"]["heavy"]

    # Затем задачи с кодом
    for kw in CFG["routing"].get("code_keywords", []):
        if kw in p:
            log.info(f"Выбрана модель code по ключевому слову: '{kw}'")
            return CFG["models"]["code"]

    # По умолчанию — основная модель
    log.info("Выбрана дефолтная модель")
    return CFG["models"]["default"]


# ─── Работа с памятью (ChromaDB) ─────────────────────────────────────────────

def get_embedding(text: str) -> list[float]:
    """
    Превратить текст в вектор чисел через nomic-embed-text.

    Вектор — это 768 чисел которые кодируют смысл текста.
    Похожие по смыслу тексты → близкие векторы.
    """
    r = requests.post(f"{OLLAMA_URL}/api/embed", json={
        "model": CFG["models"]["embed"],
        "input": text
    }, timeout=30)
    return r.json()["embeddings"][0]


def search_memory(query: str, top_k: int = 3) -> list[str]:
    """
    Найти в долгой памяти тексты похожие на запрос.

    Использует косинусное расстояние между векторами.
    """
    try:
        # Превратить запрос в вектор
        query_vector = get_embedding(query)

        # Поиск в ChromaDB
        r = requests.post(f"{CHROMA_URL}/api/v1/collections/memory/query", json={
            "query_embeddings": [query_vector],
            "n_results": top_k
        }, timeout=10)

        if r.status_code == 200:
            results = r.json()
            docs = results.get("documents", [[]])[0]
            return docs
    except Exception as e:
        log.warning(f"Ошибка поиска в памяти: {e}")

    return []


def save_to_memory(text: str, metadata: dict = None):
    """
    Сохранить важный факт в долгую память.

    Текст → вектор → ChromaDB.
    При следующих запросах будет найден если тема совпадёт.
    """
    try:
        # Убедиться что коллекция существует
        requests.post(f"{CHROMA_URL}/api/v1/collections", json={
            "name": "memory",
            "metadata": {}
        })

        vector = get_embedding(text)
        import uuid
        doc_id = str(uuid.uuid4())

        requests.post(f"{CHROMA_URL}/api/v1/collections/memory/add", json={
            "ids": [doc_id],
            "embeddings": [vector],
            "documents": [text],
            "metadatas": [metadata or {}]
        })
        log.info(f"Сохранено в память: {text[:50]}...")
    except Exception as e:
        log.warning(f"Ошибка сохранения в память: {e}")


# ─── Формирование промпта ─────────────────────────────────────────────────────

def build_messages(history: list, prompt: str, memory_context: str = "") -> list:
    """
    Собрать список сообщений для Ollama.

    Структура:
    [system prompt] + [память из ChromaDB] + [история чата] + [новый вопрос]
    """
    messages = []

    # 1. Системный промпт — всегда первый
    system_content = SYSTEM_PROMPT

    # 2. Добавить релевантную память если есть
    if memory_context:
        system_content += f"\n\n## Из памяти предыдущих разговоров:\n{memory_context}"

    messages.append({"role": "system", "content": system_content})

    # 3. История разговора (не более последних 20 сообщений)
    recent_history = history[-20:] if len(history) > 20 else history
    for msg in recent_history:
        messages.append({"role": msg.role, "content": msg.content})

    # 4. Текущий вопрос
    messages.append({"role": "user", "content": prompt})

    return messages


# ─── Основные эндпоинты ───────────────────────────────────────────────────────

@app.get("/health")
def health():
    """
    Проверка здоровья всего стека.
    Возвращает статус каждого сервиса.
    """
    status = {"status": "ok", "services": {}}

    # Проверить Ollama
    try:
        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=3)
        models = [m["name"] for m in r.json().get("models", [])]
        status["services"]["ollama"] = {"status": "ok", "models": models}
    except Exception as e:
        status["services"]["ollama"] = {"status": "error", "detail": str(e)}
        status["status"] = "degraded"

    # Проверить ChromaDB
    try:
        r = requests.get(f"{CHROMA_URL}/api/v1/heartbeat", timeout=3)
        status["services"]["chromadb"] = {"status": "ok"}
    except Exception as e:
        status["services"]["chromadb"] = {"status": "error", "detail": str(e)}

    return status


@app.post("/chat", response_model=ChatResponse)
def chat(body: ChatRequest):
    """
    Основной эндпоинт чата.

    Полный цикл:
    1. Выбрать модель
    2. Найти релевантную память
    3. Собрать промпт
    4. Вызвать Ollama
    5. Вернуть ответ
    """
    # 1. Выбор модели (или использовать указанную)
    model = body.model or pick_model(body.prompt)

    # 2. Поиск в долгой памяти
    memory_items = search_memory(body.prompt)
    memory_context = "\n".join(memory_items) if memory_items else ""

    # 3. Формирование промпта
    messages = build_messages(body.history, body.prompt, memory_context)

    log.info(f"Запрос: модель={model}, слов={len(body.prompt.split())}, "
             f"история={len(body.history)}, память={len(memory_items)}")

    # 4. Вызов Ollama
    try:
        r = requests.post(f"{OLLAMA_URL}/api/chat", json={
            "model": model,
            "messages": messages,
            "stream": False,
            "options": {
                "temperature": 0.7,    # 0=детерминированный, 1=творческий
                "num_ctx": 8192,       # размер контекстного окна
            }
        }, timeout=120)

        r.raise_for_status()
        response_text = r.json()["message"]["content"]

    except requests.Timeout:
        raise HTTPException(503, "Ollama не ответил вовремя")
    except Exception as e:
        raise HTTPException(500, f"Ошибка Ollama: {e}")

    # 5. Опционально сохранить ответ в память
    # (сохраняем только длинные важные ответы)
    if len(response_text) > 200:
        save_to_memory(
            f"Вопрос: {body.prompt[:100]}\nОтвет: {response_text[:300]}",
            metadata={"type": "conversation"}
        )

    return ChatResponse(response=response_text, model_used=model)


@app.post("/chat/stream")
def chat_stream(body: ChatRequest):
    """
    Стриминг токенов — ответ появляется по мере генерации.

    Для интерактивных интерфейсов: пользователь видит текст
    сразу, не ждёт пока модель сгенерирует всё целиком.
    """
    model = body.model or pick_model(body.prompt)
    memory_items = search_memory(body.prompt)
    memory_context = "\n".join(memory_items) if memory_items else ""
    messages = build_messages(body.history, body.prompt, memory_context)

    def generate():
        # Запрос с stream=True — Ollama отдаёт токены по одному
        r = requests.post(f"{OLLAMA_URL}/api/chat", json={
            "model": model,
            "messages": messages,
            "stream": True
        }, stream=True, timeout=120)

        for line in r.iter_lines():
            if line:
                chunk = json.loads(line)
                token = chunk.get("message", {}).get("content", "")
                if token:
                    yield token

    return StreamingResponse(generate(), media_type="text/plain")


# ─── OpenAI-совместимый API ───────────────────────────────────────────────────
# Open WebUI ожидает этот формат

@app.post("/v1/chat/completions")
def openai_compat(body: dict):
    """
    OpenAI-совместимый эндпоинт.
    Open WebUI и большинство клиентов используют этот формат.
    """
    messages = body.get("messages", [])
    model = body.get("model", CFG["models"]["default"])

    # Найти последнее сообщение пользователя
    user_messages = [m for m in messages if m["role"] == "user"]
    last_prompt = user_messages[-1]["content"] if user_messages else ""

    # Выбрать модель умно
    if model in ["gpt-4", "gpt-3.5-turbo", "default"]:
        model = pick_model(last_prompt)

    # Вызвать Ollama
    r = requests.post(f"{OLLAMA_URL}/api/chat", json={
        "model": model,
        "messages": messages,
        "stream": False
    }, timeout=120)

    content = r.json()["message"]["content"]

    # Вернуть в формате OpenAI
    return {
        "id": "chatcmpl-local",
        "object": "chat.completion",
        "model": model,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": content},
            "finish_reason": "stop"
        }]
    }

6. router/tools.py

Инструменты которые может использовать модель-агент.

# router/tools.py
"""
Инструменты для агентского режима.

Модель может вызывать эти функции через Function Calling
(Qwen2.5 поддерживает tool_calls в API).
"""
import subprocess
import requests
import json


def web_search(query: str) -> str:
    """
    Поиск в интернете через DuckDuckGo (без API ключа).

    Возвращает краткие описания первых результатов.
    Не передаёт данные никуда кроме DuckDuckGo.
    """
    try:
        # DuckDuckGo Instant Answer API — бесплатный
        r = requests.get(
            "https://api.duckduckgo.com/",
            params={"q": query, "format": "json", "no_html": 1},
            timeout=10
        )
        data = r.json()

        results = []

        # Моментальный ответ (для известных фактов)
        if data.get("AbstractText"):
            results.append(data["AbstractText"][:300])

        # Связанные темы
        for topic in data.get("RelatedTopics", [])[:3]:
            if isinstance(topic, dict) and topic.get("Text"):
                results.append(topic["Text"][:200])

        return "\n".join(results) if results else "Результаты не найдены"

    except Exception as e:
        return f"Ошибка поиска: {e}"


def run_python(code: str) -> str:
    """
    Выполнить Python код в изолированном процессе.

    БЕЗОПАСНОСТЬ:
    - Таймаут 10 секунд
    - Ограничение вывода 2KB
    - Нет доступа к сети (можно усилить с помощью nsjail)
    """
    try:
        result = subprocess.run(
            ["python3", "-c", code],
            capture_output=True,
            text=True,
            timeout=10  # убить если зависло
        )

        output = result.stdout[:2000]  # ограничить вывод
        errors = result.stderr[:500]

        if errors:
            return f"Вывод:\n{output}\n\nОшибки:\n{errors}"
        return output or "(нет вывода)"

    except subprocess.TimeoutExpired:
        return "Ошибка: превышен таймаут 10 секунд"
    except Exception as e:
        return f"Ошибка выполнения: {e}"


def read_file(path: str) -> str:
    """
    Прочитать файл с диска.

    Ограничения:
    - Только файлы в /data/ (настраивается)
    - Максимум 50KB
    """
    ALLOWED_BASE = "/data"  # можно читать только отсюда

    # Защита от path traversal: /data/../etc/passwd
    import os
    abs_path = os.path.realpath(path)
    if not abs_path.startswith(ALLOWED_BASE):
        return f"Ошибка: нет доступа к {path}"

    try:
        with open(abs_path, 'r', encoding='utf-8') as f:
            content = f.read(50000)  # максимум 50KB
        return content
    except FileNotFoundError:
        return f"Файл не найден: {path}"
    except Exception as e:
        return f"Ошибка чтения: {e}"


# Список инструментов в формате для Ollama tool_calls
TOOLS_SCHEMA = [
    {
        "type": "function",
        "function": {
            "name": "web_search",
            "description": "Поиск информации в интернете",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Поисковый запрос"
                    }
                },
                "required": ["query"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "run_python",
            "description": "Выполнить Python код и получить результат",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "Python код для выполнения"
                    }
                },
                "required": ["code"]
            }
        }
    }
]


# Маппинг имя → функция
TOOLS_MAP = {
    "web_search": web_search,
    "run_python": run_python,
    "read_file": read_file,
}


def call_tool(name: str, args: dict) -> str:
    """Вызвать инструмент по имени."""
    if name not in TOOLS_MAP:
        return f"Неизвестный инструмент: {name}"
    return TOOLS_MAP[name](**args)

7. memory/memory.py

Управление долгой памятью через ChromaDB.

# memory/memory.py
"""
Долгая память для AI-ассистента.

Принцип:
1. Важные факты из разговоров → превращаем в векторы
2. Сохраняем в ChromaDB
3. При новых вопросах → ищем похожее
4. Найденное → добавляем в контекст

Это позволяет AI "помнить" факты между сессиями
без дообучения модели.
"""
import chromadb
import requests
import os
from datetime import datetime

CHROMA_HOST = os.getenv("CHROMA_HOST", "localhost")
CHROMA_PORT = int(os.getenv("CHROMA_PORT", "8001"))
OLLAMA_URL  = os.getenv("OLLAMA_URL", "http://localhost:11434")
EMBED_MODEL = "nomic-embed-text"


class Memory:
    """Работа с долгой памятью."""

    def __init__(self, collection_name: str = "main"):
        # Подключиться к ChromaDB
        self.client = chromadb.HttpClient(
            host=CHROMA_HOST,
            port=CHROMA_PORT
        )

        # Получить или создать коллекцию
        # Коллекция = таблица в векторной БД
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}  # метрика сходства
        )

    def embed(self, text: str) -> list[float]:
        """Текст → вектор через nomic-embed-text."""
        r = requests.post(f"{OLLAMA_URL}/api/embed", json={
            "model": EMBED_MODEL,
            "input": text
        }, timeout=30)
        return r.json()["embeddings"][0]

    def save(self, text: str, metadata: dict = None) -> str:
        """
        Сохранить текст в память.

        Возвращает ID сохранённого документа.
        """
        import uuid
        doc_id = str(uuid.uuid4())

        vector = self.embed(text)

        meta = {
            "timestamp": datetime.now().isoformat(),
            "text_preview": text[:100]
        }
        if metadata:
            meta.update(metadata)

        self.collection.add(
            ids=[doc_id],
            embeddings=[vector],
            documents=[text],
            metadatas=[meta]
        )

        return doc_id

    def search(self, query: str, top_k: int = 3) -> list[dict]:
        """
        Найти в памяти тексты похожие на запрос.

        Возвращает список {text, distance, metadata}.
        distance: 0 = идентично, 1 = совсем разное
        """
        query_vector = self.embed(query)

        results = self.collection.query(
            query_embeddings=[query_vector],
            n_results=min(top_k, self.collection.count() or 1)
        )

        items = []
        for i, doc in enumerate(results["documents"][0]):
            items.append({
                "text": doc,
                "distance": results["distances"][0][i],
                "metadata": results["metadatas"][0][i]
            })

        # Фильтровать слишком далёкие результаты
        # distance > 0.5 = скорее всего нерелевантно
        return [i for i in items if i["distance"] < 0.5]

    def forget(self, doc_id: str):
        """Удалить конкретный факт из памяти."""
        self.collection.delete(ids=[doc_id])

    def count(self) -> int:
        """Сколько фактов в памяти."""
        return self.collection.count()

    def clear(self):
        """Очистить всю память."""
        self.client.delete_collection(self.collection.name)
        self.collection = self.client.get_or_create_collection(
            name=self.collection.name
        )


# ─── Пример использования ────────────────────────────────────────────────────

if __name__ == "__main__":
    mem = Memory()

    # Сохранить факты
    mem.save("Пользователь работает с Drupal 11 и Bootstrap")
    mem.save("Главная БД проекта — PostgreSQL 16")
    mem.save("Сервер Ubuntu 24.04, IP 91.218.142.168")

    print(f"В памяти: {mem.count()} фактов")

    # Поиск
    results = mem.search("какая база данных используется?")
    for r in results:
        print(f"[{r['distance']:.2f}] {r['text']}")

Следующий документ: 05_TESTS.md