Entity Memory: отслеживание сущностей в разговоре

from dataclasses import dataclass, field
from typing import Any
import json
import time


@dataclass
class Entity:
    """Одна именованная сущность в памяти агента."""
    name: str
    entity_type: str          # "person", "project", "org", "concept", "location"
    description: str = ""     # краткое описание одной строкой
    attributes: dict[str, Any] = field(default_factory=dict)
    first_seen: float = field(default_factory=time.time)
    last_seen: float = field(default_factory=time.time)
    mention_count: int = 1

    def update(self, new_attrs: dict[str, Any], new_description: str = "") -> None:
        """Мержим новые атрибуты с существующими."""
        self.attributes.update(new_attrs)
        if new_description:
            self.description = new_description
        self.last_seen = time.time()
        self.mention_count += 1

    def to_text(self) -> str:
        """Текстовое представление для вставки в промпт."""
        lines = [f"[{self.entity_type.upper()}] {self.name}"]
        if self.description:
            lines.append(f"  Описание: {self.description}")
        for k, v in self.attributes.items():
            lines.append(f"  {k}: {v}")
        return "\n".join(lines)


class EntityStore:
    """Хранилище сущностей для агента."""

    def __init__(self):
        self._entities: dict[str, Entity] = {}

    def upsert(self, name: str, entity_type: str,
               description: str = "",
               attributes: dict[str, Any] | None = None) -> Entity:
        """Создаёт или обновляет сущность."""
        key = name.lower().strip()  # нормализация ключа
        if key not in self._entities:
            self._entities[key] = Entity(
                name=name,
                entity_type=entity_type,
                description=description,
                attributes=attributes or {},
            )
        else:
            self._entities[key].update(attributes or {}, description)
        return self._entities[key]

    def get(self, name: str) -> Entity | None:
        return self._entities.get(name.lower().strip())

    def all(self) -> list[Entity]:
        return sorted(self._entities.values(), key=lambda e: -e.mention_count)

    def to_prompt_section(self) -> str:
        """Форматирует все сущности для вставки в системный промпт."""
        if not self._entities:
            return ""
        parts = ["## Известные сущности из разговора\n"]
        for entity in self.all():
            parts.append(entity.to_text())
        return "\n".join(parts)

    def save(self, path: str) -> None:
        """Сохраняем в JSON для персистентности между сессиями."""
        data = {}
        for key, e in self._entities.items():
            data[key] = {
                "name": e.name, "entity_type": e.entity_type,
                "description": e.description, "attributes": e.attributes,
                "first_seen": e.first_seen, "last_seen": e.last_seen,
                "mention_count": e.mention_count,
            }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    @classmethod
    def load(cls, path: str) -> "EntityStore":
        """Загружаем из JSON."""
        store = cls()
        try:
            with open(path, encoding="utf-8") as f:
                data = json.load(f)
            for key, d in data.items():
                e = Entity(**d)
                store._entities[key] = e
        except FileNotFoundError:
            pass
        return store

import anthropic
import json
import re

client = anthropic.Anthropic()

EXTRACTION_SYSTEM = """Ты — экстрактор сущностей. Анализируй сообщение пользователя
и возвращай JSON-список сущностей с их атрибутами.

Типы сущностей: person, project, org, concept, location, product

Формат ответа — только JSON, без пояснений:
[
  {
    "name": "Иван Петров",
    "type": "person",
    "description": "Python-разработчик",
    "attributes": {
      "роль": "Python-разработчик",
      "стек": "FastAPI"
    }
  }
]

Правила:
- Извлекай только явно упомянутые факты, не додумывай
- Если новых сущностей нет — верни пустой список []
- Один человек = одна запись, даже если у него несколько атрибутов
- "Я" / "меня" относится к пользователю (используй имя если оно известно)"""


def extract_entities(message: str, known_entities: list[str] = None) -> list[dict]:
    """
    Извлекает сущности из сообщения с помощью LLM.
    known_entities — список имён уже известных сущностей (для помощи с кореференцией).
    """
    context = ""
    if known_entities:
        context = f"\nУже известные сущности: {', '.join(known_entities)}\n"

    response = client.messages.create(
        model="claude-haiku-4-5-20251001",   # быстрая модель для extraction
        max_tokens=1024,
        system=EXTRACTION_SYSTEM,
        messages=[{
            "role": "user",
            "content": f"{context}Сообщение: {message}"
        }]
    )

    text = response.content[0].text.strip()

    # Вырезаем JSON даже если модель добавила пояснения
    match = re.search(r'\[.*\]', text, re.DOTALL)
    if not match:
        return []

    try:
        return json.loads(match.group())
    except json.JSONDecodeError:
        return []


# Пример использования
message = "Я работаю в компании Иннотех, веду проект 'Трекер задач' на FastAPI + PostgreSQL"
entities = extract_entities(message)

for e in entities:
    print(f"{e['type']}: {e['name']} — {e.get('description', '')}")
    for k, v in e.get('attributes', {}).items():
        print(f"  {k}: {v}")

def process_message_with_coref(
    message: str,
    store: EntityStore
) -> list[dict]:
    """
    Извлекает сущности с учётом уже известных — для разрешения кореференции.
    """
    # Передаём имена известных сущностей как подсказку
    known = [e.name for e in store.all()]
    entities = extract_entities(message, known_entities=known)

    for e in entities:
        store.upsert(
            name=e["name"],
            entity_type=e["type"],
            description=e.get("description", ""),
            attributes=e.get("attributes", {}),
        )

    return entities

def build_system_prompt(base_prompt: str, store: EntityStore) -> str:
    """
    Собирает системный промпт: базовые инструкции + секция с сущностями.
    """
    entity_section = store.to_prompt_section()
    if entity_section:
        return f"{base_prompt}\n\n{entity_section}"
    return base_prompt

import anthropic
from pathlib import Path


BASE_SYSTEM = """Ты — ассистент-разработчик. Помогаешь с архитектурой, кодом, отладкой.
Отвечай точно, конкретно, со ссылкой на технологии пользователя если они известны."""

ENTITY_STORE_PATH = "entity_store.json"


class EntityMemoryAgent:
    def __init__(self):
        self.client = anthropic.Anthropic()
        self.store = EntityStore.load(ENTITY_STORE_PATH)
        self.history: list[dict] = []

    def chat(self, user_message: str) -> str:
        # Шаг 1: Извлечь и обновить сущности
        process_message_with_coref(user_message, self.store)

        # Шаг 2: Сохранить обновлённый store (персистентность)
        self.store.save(ENTITY_STORE_PATH)

        # Шаг 3: Добавить сообщение в историю
        self.history.append({"role": "user", "content": user_message})

        # Шаг 4: Собрать контекст = base prompt + entities
        system = build_system_prompt(BASE_SYSTEM, self.store)

        # Шаг 5: Вызвать LLM
        response = self.client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=2048,
            system=system,
            messages=self.history,
        )

        answer = response.content[0].text

        # Шаг 6: Добавить ответ в историю
        self.history.append({"role": "assistant", "content": answer})

        return answer

    def show_entities(self) -> None:
        """Отладка: показать текущее состояние хранилища."""
        for entity in self.store.all():
            print(entity.to_text())
            print()


# Использование
agent = EntityMemoryAgent()

# Первое сообщение — агент запоминает факты
r1 = agent.chat("Привет! Я Иван, пишу FastAPI-сервис для Иннотех.")
print(r1)

# Второе сообщение — новые факты
r2 = agent.chat("Проект называется 'Трекер задач', дедлайн 30 марта.")
print(r2)

# 20 сообщений спустя — агент всё ещё знает контекст
r3 = agent.chat("Какую базу данных посоветуешь для моего проекта?")
print(r3)
# → Ответ будет с учётом: "Трекер задач", FastAPI, Иннотех

# Смотрим что накопилось
agent.show_entities()

import sqlite3
import json


class SQLiteEntityStore:
    """Entity store на базе SQLite с поддержкой нескольких пользователей."""

    def __init__(self, db_path: str = "entities.db"):
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self._init_db()

    def _init_db(self):
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS entities (
                user_id     TEXT NOT NULL,
                name_key    TEXT NOT NULL,
                name        TEXT NOT NULL,
                entity_type TEXT NOT NULL,
                description TEXT DEFAULT '',
                attributes  TEXT DEFAULT '{}',
                mention_count INTEGER DEFAULT 1,
                last_seen   REAL,
                PRIMARY KEY (user_id, name_key)
            )
        """)
        self.conn.commit()

    def upsert(self, user_id: str, name: str, entity_type: str,
               description: str = "", attributes: dict | None = None) -> None:
        key = name.lower().strip()
        attrs_json = json.dumps(attributes or {}, ensure_ascii=False)

        self.conn.execute("""
            INSERT INTO entities (user_id, name_key, name, entity_type, description, attributes, last_seen)
            VALUES (?, ?, ?, ?, ?, ?, unixepoch())
            ON CONFLICT(user_id, name_key) DO UPDATE SET
                attributes    = json_patch(attributes, excluded.attributes),
                description   = CASE WHEN excluded.description != '' THEN excluded.description ELSE description END,
                mention_count = mention_count + 1,
                last_seen     = unixepoch()
        """, (user_id, key, name, entity_type, description, attrs_json))
        self.conn.commit()

    def load_for_user(self, user_id: str) -> list[Entity]:
        cursor = self.conn.execute(
            "SELECT name, entity_type, description, attributes, mention_count, last_seen "
            "FROM entities WHERE user_id = ? ORDER BY mention_count DESC",
            (user_id,)
        )
        result = []
        for row in cursor:
            name, etype, desc, attrs_json, count, last_seen = row
            e = Entity(
                name=name, entity_type=etype, description=desc,
                attributes=json.loads(attrs_json),
                mention_count=count, last_seen=last_seen,
            )
            result.append(e)
        return result

# Минимальный рабочий паттерн entity memory

store = EntityStore.load("store.json")           # загружаем между сессиями

def agent_turn(user_msg: str) -> str:
    # 1. Extract + Update
    process_message_with_coref(user_msg, store)
    store.save("store.json")                     # персистентность

    # 2. Build context
    system = build_system_prompt(BASE_SYSTEM, store)

    # 3. Generate
    return llm_call(system=system, messages=history)