Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,5 @@
-# app.py
 import streamlit as st
-from rag import rag_chain  # импортируем цепочку из rag.py
+from rag import rag_chain, retriever  # ← добавьте экспорт retriever из rag.py
 
 st.set_page_config(page_title="Поиск фильмов", layout="centered")
 st.title("🎬 RAG-поиск по фильмам")
@@ -12,15 +11,42 @@ if "messages" not in st.session_state:
 # Отображение истории
 for msg in st.session_state.messages:
     st.chat_message(msg["role"]).write(msg["content"])
+    if "posters" in msg:
+        cols = st.columns(min(5, len(msg["posters"])))
+        for i, url in enumerate(msg["posters"]):
+            if url:
+                cols[i % len(cols)].image(url, width=130)
 
 # Ввод пользователя
 if prompt := st.chat_input("Например: 'Порекомендуй фильм про космос'"):
     st.session_state.messages.append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
 
+    # Получаем документы ДО генерации ответа (чтобы показать те же фильмы, что использовал LLM)
+    docs = retriever.invoke(prompt)
+
+    # Извлекаем poster_url из метаданных
+    posters = [
+        doc.metadata.get("poster_url", "")
+        for doc in docs
+        if doc.metadata.get("poster_url")
+    ]
+
     # Генерация ответа
     with st.spinner("Ищу фильмы..."):
         response = rag_chain.invoke(prompt)
-    
-    st.session_state.messages.append({"role": "assistant", "content": response})
-    st.chat_message("assistant").write(response)
+
+    # Сохраняем в историю с постерами
+    st.session_state.messages.append({
+        "role": "assistant",
+        "content": response,
+        "posters": posters
+    })
+
+    # Вывод ответа и постеров
+    st.chat_message("assistant").write(response)
+    if posters:
+        st.write("##### 🎞️ Найденные фильмы:")
+        cols = st.columns(min(5, len(posters)))
+        for i, url in enumerate(posters):
+            cols[i % len(cols)].image(url, width=130)