Долгосрочная память: векторные хранилища, semantic search, retrieval

# Установка зависимостей
pip install sentence-transformers chromadb numpy

from sentence_transformers import SentenceTransformer
import numpy as np

# Модель скачивается при первом запуске (~90 MB)
# all-MiniLM-L6-v2 — хороший баланс скорости и качества (384 dims)
model = SentenceTransformer("all-MiniLM-L6-v2")


def embed(text: str) -> list[float]:
    """Превращает текст в вектор-эмбеддинг."""
    return model.encode(text, normalize_embeddings=True).tolist()


def embed_batch(texts: list[str]) -> list[list[float]]:
    """
    Векторизует список текстов за один вызов.
    Batch-обработка в 5-10x быстрее, чем вызывать embed() в цикле.
    """
    return model.encode(texts, normalize_embeddings=True, batch_size=64).tolist()


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Косинусное сходство двух нормализованных векторов — просто dot product."""
    return float(np.dot(a, b))   # normalize_embeddings=True делает |a|=|b|=1


# Пример
v1 = embed("Как приготовить борщ?")
v2 = embed("Рецепт традиционного борща с говядиной")
v3 = embed("Квантовая механика в физике")

print(f"Похожие тексты:  {cosine_similarity(v1, v2):.3f}")  # → 0.887
print(f"Разные темы:     {cosine_similarity(v1, v3):.3f}")  # → 0.071

import re


def chunk_by_tokens(
    text: str,
    chunk_size: int = 500,   # символы (не токены — приблизительно)
    overlap: int = 80
) -> list[str]:
    """
    Разбивает текст на чанки фиксированного размера с перекрытием.
    Предпочитает разрыв по пробелу, а не посередине слова.
    """
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))

        # Ищем ближайший пробел перед концом чанка (не режем слово)
        if end < len(text):
            space = text.rfind(' ', start, end)
            if space > start + chunk_size // 2:
                end = space

        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        start = end - overlap

    return chunks


def chunk_by_sentences(
    text: str,
    sentences_per_chunk: int = 5,
    overlap_sentences: int = 1
) -> list[str]:
    """
    Разбивает текст по предложениям — не разрывает смысл в середине фразы.
    """
    # Разбиваем по знакам конца предложения
    sentences = re.split(r'(?<=[.!?])\s+(?=[А-ЯA-Z«"(])', text)
    sentences = [s.strip() for s in sentences if s.strip()]

    chunks = []
    i = 0
    while i < len(sentences):
        chunk_sents = sentences[i : i + sentences_per_chunk]
        chunks.append(' '.join(chunk_sents))
        i += sentences_per_chunk - overlap_sentences

    return chunks


# Пример
text = (
    "Агент — это система, которая воспринимает среду и действует в ней. "
    "В отличие от простого чатбота, агент сохраняет контекст и планирует шаги. "
    "Долгосрочная память позволяет агенту помнить прошлые взаимодействия. "
    "Векторные эмбеддинги — основа семантического поиска по памяти. "
    "RAG объединяет поиск и генерацию в единый пайплайн."
)
chunks = chunk_by_sentences(text, sentences_per_chunk=2, overlap_sentences=1)
for i, c in enumerate(chunks):
    print(f"[{i}] {c[:80]}...")

from dataclasses import dataclass, field


@dataclass
class Document:
    id: str
    content: str
    embedding: list[float]
    metadata: dict = field(default_factory=dict)


class InMemoryVectorStore:
    """
    Простой in-memory vector store с линейным поиском.
    Для продакшна используй ChromaDB, Qdrant или pgvector.
    """
    def __init__(self):
        self._docs: list[Document] = []

    def add(self, doc_id: str, content: str, metadata: dict | None = None) -> None:
        """Добавляет документ: сразу векторизует и сохраняет."""
        self._docs.append(Document(
            id=doc_id,
            content=content,
            embedding=embed(content),
            metadata=metadata or {}
        ))

    def add_texts(self, texts: list[str], metadatas: list[dict] | None = None) -> None:
        """Batch-добавление: векторизует все тексты за один вызов модели."""
        embeddings = embed_batch(texts)
        for i, (text, emb) in enumerate(zip(texts, embeddings)):
            meta = (metadatas[i] if metadatas else {})
            self._docs.append(Document(
                id=f"doc_{len(self._docs)}",
                content=text,
                embedding=emb,
                metadata=meta
            ))

    def search(
        self,
        query: str,
        k: int = 5,
        min_score: float = 0.3
    ) -> list[dict]:
        """
        Семантический поиск Top-K документов.
        Возвращает отсортированный список с полями: id, content, score, metadata.
        """
        if not self._docs:
            return []

        query_emb = embed(query)
        scored = [
            (doc, cosine_similarity(query_emb, doc.embedding))
            for doc in self._docs
        ]
        # Фильтруем ниже порога и сортируем по убыванию
        scored = [(doc, s) for doc, s in scored if s >= min_score]
        scored.sort(key=lambda x: x[1], reverse=True)

        return [
            {
                "id": doc.id,
                "content": doc.content,
                "score": round(score, 4),
                "metadata": doc.metadata
            }
            for doc, score in scored[:k]
        ]

    def __len__(self) -> int:
        return len(self._docs)


# Пример использования
store = InMemoryVectorStore()
store.add_texts([
    "Пользователь предпочитает краткие ответы без лишних слов.",
    "Дедлайн по проекту X — 15 марта 2025 года.",
    "Главная задача на эту неделю: реализовать vector store.",
    "Пользователь работает в стартапе, команда из 4 человек.",
])

results = store.search("Когда дедлайн?", k=2)
for r in results:
    print(f"[{r['score']:.3f}] {r['content']}")

import chromadb
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# PersistentClient сохраняет данные в ./agent_memory_db/
# При следующем запуске — данные на месте
chroma_client = chromadb.PersistentClient(path="./agent_memory_db")

collection = chroma_client.get_or_create_collection(
    name="agent_memory",
    metadata={"hnsw:space": "cosine"}   # метрика сходства
)


def add_to_memory(
    content: str,
    doc_id: str,
    metadata: dict | None = None
) -> None:
    """Добавляет текст в долгосрочную память агента."""
    embedding = model.encode(content, normalize_embeddings=True).tolist()
    collection.add(
        ids=[doc_id],
        embeddings=[embedding],
        documents=[content],
        metadatas=[metadata or {}]
    )


def search_memory(
    query: str,
    k: int = 5,
    where: dict | None = None   # фильтр по метаданным (опционально)
) -> list[dict]:
    """
    Семантический поиск по долгосрочной памяти.
    where: {'type': 'preference'} — ищем только в документах нужного типа.
    """
    query_emb = model.encode(query, normalize_embeddings=True).tolist()

    kwargs: dict = {"query_embeddings": [query_emb], "n_results": k}
    if where:
        kwargs["where"] = where

    results = collection.query(**kwargs)

    if not results["ids"][0]:
        return []

    return [
        {
            "id":       results["ids"][0][i],
            "content":  results["documents"][0][i],
            # ChromaDB возвращает cosine distance (0=одинаковые, 2=противоположные)
            # Переводим в similarity (1=одинаковые, -1=противоположные)
            "score":    round(1 - results["distances"][0][i], 4),
            "metadata": results["metadatas"][0][i]
        }
        for i in range(len(results["ids"][0]))
    ]


# Наполняем память агента разными типами данных
add_to_memory(
    "Пользователь предпочитает краткие ответы без списков",
    doc_id="pref_1",
    metadata={"type": "preference", "user_id": "ivan"}
)
add_to_memory(
    "Дедлайн по проекту Alpha — 15 марта 2025, ответственный Иван",
    doc_id="fact_1",
    metadata={"type": "fact", "date": "2025-01-10"}
)
add_to_memory(
    "В предыдущей сессии обсуждали архитектуру vector store для агента",
    doc_id="ctx_1",
    metadata={"type": "context", "session": "2025-01-09"}
)

# Поиск
results = search_memory("когда дедлайн по проекту?", k=3)
for r in results:
    print(f"[{r['score']:.3f}] ({r['metadata'].get('type', '?')}) {r['content']}")

# Поиск только среди предпочтений пользователя
prefs = search_memory(
    "как отвечать пользователю",
    k=3,
    where={"type": "preference"}
)

# Поиск только контекста из последней сессии
recent_ctx = search_memory(
    "что обсуждали",
    k=5,
    where={"session": "2025-01-09"}
)

# Комбинированный фильтр: тип И пользователь
user_facts = search_memory(
    "дедлайны и задачи",
    k=5,
    where={"$and": [{"type": "fact"}, {"user_id": "ivan"}]}
)

import anthropic

client = anthropic.Anthropic()

SEARCH_MEMORY_TOOL = {
    "name": "search_memory",
    "description": (
        "Поиск релевантной информации в долгосрочной памяти агента. "
        "Используй при вопросах о прошлых разговорах, предпочтениях пользователя, "
        "дедлайнах и фактах из предыдущих сессий."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "Поисковый запрос — что именно нужно найти в памяти"
            },
            "k": {
                "type": "integer",
                "description": "Количество результатов (1-5, по умолчанию 3)"
            },
            "type_filter": {
                "type": "string",
                "enum": ["fact", "preference", "context", "knowledge"],
                "description": "Тип записей для поиска (опционально)"
            }
        },
        "required": ["query"]
    }
}

SAVE_MEMORY_TOOL = {
    "name": "save_to_memory",
    "description": (
        "Сохраняет важную информацию в долгосрочную память. "
        "Используй для фактов, решений, предпочтений пользователя и контекста "
        "который может понадобиться в будущих сессиях."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "content": {
                "type": "string",
                "description": "Текст для сохранения — краткий и конкретный факт"
            },
            "type": {
                "type": "string",
                "enum": ["fact", "preference", "context", "knowledge"],
                "description": "Тип записи"
            }
        },
        "required": ["content", "type"]
    }
}

SYSTEM = """Ты — персональный ассистент с долгосрочной памятью.

Используй инструменты памяти:
- search_memory: ищи контекст прошлых разговоров перед ответом на вопросы о них
- save_to_memory: сохраняй важные факты, предпочтения и решения из текущего диалога

При поиске — формулируй запрос чётко. При сохранении — пиши кратко и конкретно."""


def run_memory_agent(user_message: str, conversation: list[dict]) -> str:
    """
    Агент с долгосрочной памятью.
    conversation — краткосрочная история текущей сессии (список messages).
    """
    import uuid

    messages = list(conversation)
    messages.append({"role": "user", "content": user_message})

    while True:
        response = client.messages.create(
            model="claude-opus-4-6",
            max_tokens=1024,
            system=SYSTEM,
            tools=[SEARCH_MEMORY_TOOL, SAVE_MEMORY_TOOL],
            messages=messages
        )

        if response.stop_reason == "end_turn":
            answer = response.content[0].text
            messages.append({"role": "assistant", "content": answer})
            return answer

        # Обрабатываем tool calls
        messages.append({"role": "assistant", "content": response.content})
        tool_results = []

        for block in response.content:
            if block.type != "tool_use":
                continue

            if block.name == "search_memory":
                where = None
                if tf := block.input.get("type_filter"):
                    where = {"type": tf}

                results = search_memory(
                    query=block.input["query"],
                    k=block.input.get("k", 3),
                    where=where
                )
                result_text = "\n\n".join(
                    f"[score={r['score']:.2f}, type={r['metadata'].get('type','?')}]\n{r['content']}"
                    for r in results
                ) or "Ничего не найдено в памяти."

                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": result_text
                })

            elif block.name == "save_to_memory":
                doc_id = f"{block.input['type']}_{uuid.uuid4().hex[:8]}"
                add_to_memory(
                    content=block.input["content"],
                    doc_id=doc_id,
                    metadata={"type": block.input["type"]}
                )
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": f"Сохранено: {block.input['content'][:60]}..."
                })

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

import uuid
from pathlib import Path


def index_markdown_file(filepath: str | Path) -> int:
    """
    Индексирует один Markdown-файл: разбивает на чанки и сохраняет в ChromaDB.
    Возвращает число созданных чанков.
    """
    path = Path(filepath)
    text = path.read_text(encoding="utf-8")

    # Разбиваем по параграфам (двойной перенос строки)
    paragraphs = [p.strip() for p in text.split("\n\n") if len(p.strip()) > 50]

    # Дополнительно режем слишком длинные параграфы
    chunks = []
    for para in paragraphs:
        if len(para) > 800:
            chunks.extend(chunk_by_tokens(para, chunk_size=600, overlap=80))
        else:
            chunks.append(para)

    # Batch-индексирование
    embeddings = embed_batch(chunks)
    ids = [f"{path.stem}_{i}" for i in range(len(chunks))]
    metadatas = [
        {"source": path.name, "type": "knowledge", "chunk_idx": i}
        for i in range(len(chunks))
    ]

    collection.add(
        ids=ids,
        embeddings=embeddings,
        documents=chunks,
        metadatas=metadatas
    )

    return len(chunks)


def index_conversation_history(messages: list[dict], session_id: str) -> None:
    """
    Сохраняет финальную историю диалога как контекстную память.
    Вызывается в конце каждой сессии.
    """
    # Группируем по парам user+assistant
    pairs = []
    for i in range(0, len(messages) - 1, 2):
        if messages[i]["role"] == "user" and messages[i+1]["role"] == "assistant":
            pairs.append(
                f"User: {messages[i]['content']}\nAssistant: {messages[i+1]['content']}"
            )

    if not pairs:
        return

    embeddings = embed_batch(pairs)
    ids = [f"conv_{session_id}_{i}" for i in range(len(pairs))]
    metadatas = [{"type": "context", "session": session_id} for _ in pairs]

    collection.add(ids=ids, embeddings=embeddings, documents=pairs, metadatas=metadatas)


# Использование
n = index_markdown_file("docs/architecture.md")
print(f"Проиндексировано {n} чанков из architecture.md")