src/rag.py

from qdrant_client import QdrantClient
from langchain_huggingface import HuggingFaceEmbeddings
# from langchain_community.chat_models import GigaChat
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_qdrant import QdrantVectorStore
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain.retrievers import ContextualCompressionRetriever
from sentence_transformers import CrossEncoder
from langchain_groq import ChatGroq







client = QdrantClient(
    url="https://7acfa434-9e7d-4ff3-bc16-98679211cca6.europe-west3-0.gcp.cloud.qdrant.io",
    api_key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.Fb6pTw9_wQqEPMe3kcAW0o-VobmCUOjlpIMHJep5UpU",  # ← обязательно!
    https=True
)


model_name = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': True, 'batch_size':1024}

embeddings_model = HuggingFaceEmbeddings(
    model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': True, 'batch_size':1024},
)


vector_store = QdrantVectorStore(
    client=client,
    collection_name='film_col_1',
    embedding=embeddings_model
)

llm = ChatGroq(
    model="llama-3.1-8b-instant",        # или "llama3-70b-8192", "mixtral-8x7b-32768"
    api_key="your_key",  # ← замените на свой!
    temperature=0.7,
    max_tokens=512,
)



base_retriever = vector_store.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 20}
)

model_r = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
compressor = CrossEncoderReranker(model=model_r, top_n=5)

retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=base_retriever
)

def format_docs(docs):
    """Форматирует документы для передачи в промпт"""
    formatted = []

    for i, doc in enumerate(docs, 1):
        metadata = doc.metadata

        vacancy_info = f"""
        === ФИЛЬМ {i} ===
        Название: {metadata.get('title', 'Не указано')}
        Год: {metadata.get('year', 'Не указано')}
        Жанр: {metadata.get('genres', 'Не указано')}
        Директор: {metadata.get('directors', 'Не указано')}
        Оценка: {metadata.get('vote_average', 'Не указано')}
        Количество оценок: {metadata.get('vote_count', 'Не указано')}
        Ссылка: {metadata.get('tmdb_url', 'Не указано')}
        Описание: {doc.page_content[:300]}...
        """

        formatted.append(vacancy_info)

    return "\n".join(formatted)


rag_prompt = ChatPromptTemplate.from_messages([
    (
    "system",
    """Ты — кинокритик с 20-летним стажем, бывший сценарист, а ныне саркастичный эксперт по мировому кинематографу! 🎬
    Твоя задача — проанализировать предоставленные фильмы и дать остроумную, но профессиональную оценку с лёгкой долей цинизма и любви к кино.

    Стиль анализа:
    - Глубоко разбирай контекст: жанры, режиссёры, годы, рейтинги — но без занудства
    - Используй кинематографические мемы и отсылки («Оскар убежал», «это не фильм — это терапия», «как в советском»)
    - Подмечай абсурдные или трогательные детали: «режиссёр снял 3 фильма, а актёр — 200», «рейтинг 2.0, но постер шедевр»
    - Давай рекомендации с иронией: «Смотреть только если вы фанат страданий» или «Идеально для просмотра после третьего кофе»
    - Структурируй ответ с эмодзи, краткими заголовками и живыми комментариями
    - Отвечай на русском языке — умно, ярко и с характером!
    

    Помни: сарказм — да, злоба — нет. Цель — чтобы читатель улыбнулся и захотел посмотреть фильм (или хотя бы посмеялся над ним). 😏

    Если среди фильмов есть что-то эпически странное, гениальное или настолько плохое, что становится хорошим — обязательно выдели это! 🍿"""
),

(
    "human",
    """📽️ КОНТЕКСТ (или как говорят в индустрии — «наши материалы»):
{context}

🎯 ЗАПРОС ОТ РЕДАКЦИИ: {question}"""
)
])



# Создаем RAG цепочку
rag_chain = (
    {
        "context": retriever | format_docs,
        "question": RunnablePassthrough()
    } # словарь, в котором ключи - это переменные, которые будут переданы в промпт
    | rag_prompt # промпт для RAG
    | llm # тут можно поставить любую llm-модель
    | StrOutputParser() # для вывода ответа в читаемом виде
)