# Установка
pip install chromadb openai

# Базовое использование — 5 строк кода
import chromadb
from chromadb.config import Settings

# Персистентный клиент — данные хранятся на диске
client = chromadb.PersistentClient(
    path="/opt/data/chromadb_store",
    settings=Settings(anonymized_telemetry=False)
)

# Создаём коллекцию (аналог таблицы)
collection = client.get_or_create_collection(
    name="agent_memory",
    metadata={"hnsw:space": "cosine"}  # cosine / l2 / ip
)

print(f"✅ ChromaDB: {client.count_collections()} коллекций")
# Вывод: ✅ ChromaDB: 3 коллекций

from openai import OpenAI
import uuid

client_openai = OpenAI()

def add_documents(collection, documents: list, metadatas: list = None):
    """Добавляет документы с авто-эмбеддингом через OpenAI"""
    if metadatas is None:
        metadatas = [{}] * len(documents)
    
    # 1. Получаем эмбеддинги батчем (дешевле и быстрее)
    resp = client_openai.embeddings.create(
        model="text-embedding-3-small",
        input=documents
    )
    embeddings = [d.embedding for d in resp.data]
    
    # 2. Генерируем уникальные ID
    ids = [str(uuid.uuid4()) for _ in documents]
    
    # 3. Добавляем в коллекцию
    collection.add(
        ids=ids,
        embeddings=embeddings,
        documents=documents,
        metadatas=metadatas
    )
    return {"added": len(documents)}

# Пример: индексация базы знаний магазина
docs = [
    "Возврат товара возможен в течение 14 дней после получения.",
    "Доставка по Москве — бесплатно при заказе от 3000₽.",
    "Гарантия на электронику — 1 год с момента покупки.",
]
meta = [
    {"topic": "returns", "source": "policy"},
    {"topic": "delivery", "source": "policy"},
    {"topic": "warranty", "source": "policy"},
]
add_documents(collection, docs, meta)
# Вывод: {'added': 3}

def semantic_search(collection, query: str, n_results=5, where_filter=None):
    """Семантический поиск с опциональной фильтрацией по метаданным"""
    # 1. Embedding запроса
    query_emb = get_embedding(query)
    
    # 2. Поиск по коллекции
    kwargs = {
        "query_embeddings": [query_emb],
        "n_results": n_results,
        "include": ["documents", "metadatas", "distances"]
    }
    if where_filter:
        kwargs["where"] = where_filter
    
    results = collection.query(**kwargs)
    
    # 3. Форматируем результат
    formatted = []
    for i in range(len(results["documents"][0])):
        score = round(1 - results["distances"][0][i], 4)  # Косинусное сходство
        formatted.append({
            "content": results["documents"][0][i],
            "metadata": results["metadatas"][0][i],
            "score": score
        })
    return formatted

def get_embedding(text: str) -> list:
    resp = client_openai.embeddings.create(model="text-embedding-3-small", input=text)
    return resp.data[0].embedding

# 1. Обычный семантический поиск
results = semantic_search(collection, "хочу вернуть бракованный товар")
# [{'content': 'Возврат товара возможен в течение 14 дней...', 'score': 0.92}, ...]

# 2. Поиск с фильтром по метаданным
results = semantic_search(collection, "сколько стоит доставка", 
    where_filter={"topic": "delivery"})
# Только документы с topic='delivery'

# 3. Сложный фильтр: $and, $or, $gte, $in
results = semantic_search(collection, "правила магазина",
    where_filter={
        "$and": [
            {"source": "policy"},
            {"topic": {"$in": ["returns", "warranty"]}}
        ]
    })

def rag_query(collection, question: str, system_prompt: str = "") -> dict:
    """Полный RAG-пайплайн: Retrieve → Augment → Generate"""
    
    # Шаг 1: Retrieve — ищем релевантные документы
    retrieved = semantic_search(collection, question, n_results=3)
    
    # Шаг 2: Augment — формируем контекст для LLM
    context_parts = []
    for i, doc in enumerate(retrieved):
        context_parts.append(
            f"[Источник {i+1}, релевантность: {doc['score']:.2f}] {doc['content']}"
        )
    context = "\n\n".join(context_parts)
    
    # Шаг 3: Generate — генерируем ответ с контекстом
    if not system_prompt:
        system_prompt = """Ты — ассистент, который отвечает СТРОГО на основе контекста.
Если ответа нет в контексте, скажи 'У меня недостаточно информации'.
Всегда указывай источник информации в ответе."""
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"""Контекст:
{context}

Вопрос: {question}

Ответ (на русском):"""}
    ]
    
    resp = client_openai.chat.completions.create(
        model="gpt-4o-mini", temperature=0.3, messages=messages
    )
    
    return {
        "answer": resp.choices[0].message.content,
        "sources": [{"content": d["content"][:100], "score": d["score"]} for d in retrieved]
    }

# Пример RAG-запроса
result = rag_query(collection, "Можно ли вернуть телефон, если он не понравился?")
print(result["answer"])
# "Да, возврат товара возможен в течение 14 дней после получения [Источник 1].
#  Для электроники также действует гарантия 1 год [Источник 2]."

from datetime import datetime

# Коллекция для хранения истории диалогов
memory = client.get_or_create_collection(
    name="conversation_memory",
    metadata={"hnsw:space": "cosine"}
)

def store_interaction(user_id: str, session_id: str, role: str, content: str):
    """Сохраняет сообщение в память агента"""
    embedding = get_embedding(content)
    memory.add(
        ids=[str(uuid.uuid4())],
        embeddings=[embedding],
        documents=[content],
        metadatas=[{
            "user_id": user_id,
            "session_id": session_id,
            "role": role,
            "timestamp": datetime.now().isoformat()
        }]
    )

def recall_memories(user_id: str, query: str, n=5) -> list:
    """Извлекает релевантные прошлые взаимодействия пользователя"""
    query_emb = get_embedding(query)
    results = memory.query(
        query_embeddings=[query_emb],
        n_results=n,
        where={"user_id": user_id},  # Только диалоги этого пользователя
        include=["documents", "metadatas", "distances"]
    )
    memories = []
    for i in range(len(results["documents"][0])):
        memories.append({
            "role": results["metadatas"][0][i]["role"],
            "content": results["documents"][0][i],
            "score": round(1 - results["distances"][0][i], 3)
        })
    return sorted(memories, key=lambda x: x["score"], reverse=True)

# Пример: сохраняем диалог и извлекаем контекст
store_interaction("user_42", "sess_001", "user", "Какой у вас самый быстрый ноутбук?")
store_interaction("user_42", "sess_001", "assistant", "MacBook Pro M4 — от 189 990₽")
past = recall_memories("user_42", "мощный компьютер для работы")
# Агент «вспомнит» предыдущий диалог про ноутбуки

pip install sentence-transformers

from sentence_transformers import SentenceTransformer
import numpy as np

# Локальная embedding-модель (русский язык)
# Варианты: 'intfloat/multilingual-e5-large', 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'
model = SentenceTransformer('intfloat/multilingual-e5-large')

def local_embed(texts: list) -> list:
    """Бесплатные эмбеддинги через SentenceTransformer"""
    # Важно: для e5 моделей добавляем префикс 'query:' или 'passage:'
    prefixed = [f"passage: {t}" for t in texts]
    embeddings = model.encode(prefixed, normalize_embeddings=True)
    return embeddings.tolist()

# Использование с ChromaDB
# Создаём embedding function для ChromaDB
from chromadb.utils import embedding_functions

class LocalEmbeddingFunction:
    def __call__(self, texts: list) -> list:
        return local_embed(texts)

# ChromaDB может автоматически эмбеддить при добавлении
collection_local = client.get_or_create_collection(
    name="local_embeddings",
    embedding_function=LocalEmbeddingFunction(),
    metadata={"hnsw:space": "cosine"}
)

# Теперь add() автоматически создаст эмбеддинги локально
collection_local.add(documents=["Бесплатная доставка от 5000₽"], ids=["doc_1"])
print(f"✅ Локальные эмбеддинги: {collection_local.count()} документов")

def hybrid_search(collection, query: str, n_results=5, alpha=0.5):
    """Гибридный поиск: векторы (семантика) + BM25 (ключевые слова)
    alpha=0.0 → только BM25, alpha=1.0 → только векторный"""
    
    # 1. Векторный поиск (семантический)
    query_emb = get_embedding(query)
    vector_results = collection.query(
        query_embeddings=[query_emb],
        n_results=n_results * 2,  # Берём с запасом для слияния
        include=["documents", "distances"]
    )
    
    # 2. Полнотекстовый поиск (BM25)
    fulltext_results = collection.query(
        query_texts=[query],
        n_results=n_results * 2,
        include=["documents", "distances"]
    )
    
    # 3. Reciprocal Rank Fusion (RRF) для объединения
    def rrf_score(rank, k=60):
        return 1.0 / (k + rank)
    
    scores = {}
    for rank, doc in enumerate(vector_results["documents"][0]):
        scores[doc] = scores.get(doc, 0) + alpha * rrf_score(rank)
    
    for rank, doc in enumerate(fulltext_results["documents"][0]):
        scores[doc] = scores.get(doc, 0) + (1 - alpha) * rrf_score(rank)
    
    # 4. Сортируем и возвращаем top-N
    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    return [{"content": doc, "hybrid_score": round(score, 4)} 
            for doc, score in ranked[:n_results]]

# Пример: гибридный поиск
results = hybrid_search(collection, "ORD-12345 статус заказа", alpha=0.4)
# alpha=0.4: больше веса точным совпадениям (ORD-12345),
# но учитываем и семантику ("статус заказа")