Update app.py

app.py

@@ -1,117 +1,93 @@
-import streamlit as st
+import chainlit as cl
 import os
 from huggingface_hub import InferenceClient
 from qdrant_client import QdrantClient
 from sentence_transformers import SentenceTransformer
 
 # --- НАСТРОЙКИ ---
-st.set_page_config(page_title="Sales AI + RAG", page_icon="💼", layout="centered")
-st.title("💼 Виртуальный Отдел Продаж")
-st.caption("Чат с базой знаний (RAG) на Qwen 2.5")
-
-# --- КЛЮЧИ ---
 HF_TOKEN = os.getenv("HF_TOKEN")
 QDRANT_URL = os.getenv("QDRANT_URL")
 QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
+MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
 
-if not HF_TOKEN:
-    st.error("⚠️ Нет HF_TOKEN!")
-    st.stop()
+# --- ИНИЦИАЛИЗАЦИЯ ПРИ СТАРТЕ ЧАТА ---
+@cl.on_chat_start
+async def start():
+    # 1. Показываем загрузку
+    msg = cl.Message(content="Загружаю базы знаний и подключаюсь к нейросетям...")
+    await msg.send()
 
-# --- ЗАГРУЗКА РЕСУРСОВ (ОДИН РАЗ) ---
-@st.cache_resource
-def load_resources():
-    # 1. Чат-модель
-    hf_client = InferenceClient("Qwen/Qwen2.5-7B-Instruct", token=HF_TOKEN)
+    # 2. Инициализируем клиентов (сохраняем в сессию пользователя)
+    hf_client = InferenceClient(MODEL_ID, token=HF_TOKEN)
     
-    # 2. База знаний (Qdrant)
     q_client = None
     if QDRANT_URL and QDRANT_API_KEY:
         try:
             q_client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
-            print("✅ Qdrant подключен")
-        except Exception as e:
-            print(f"❌ Ошибка Qdrant: {e}")
+            print("✅ Qdrant OK")
+        except:
+            print("❌ Qdrant Error")
             
-    # 3. Модель для поиска (Embeddings)
     encoder = SentenceTransformer('all-MiniLM-L6-v2')
     
-    return hf_client, q_client, encoder
+    # Сохраняем в сессию, чтобы использовать при каждом сообщении
+    cl.user_session.set("hf_client", hf_client)
+    cl.user_session.set("q_client", q_client)
+    cl.user_session.set("encoder", encoder)
 
-client, qdrant, encoder = load_resources()
+    # 3. Обновляем сообщение на приветствие
+    msg.content = "👋 Привет! Я готов к работе. Задай вопрос по базе знаний."
+    await msg.update()
 
-# --- ФУНКЦИЯ ПОИСКА В БАЗЕ (RAG) ---
-def get_context(query):
-    if not qdrant: 
-        return ""
+# --- ФУНКЦИЯ RAG (ПОИСК) ---
+def get_context(query, q_client, encoder):
+    if not q_client: return ""
     try:
-        # Превращаем вопрос в вектор
         vector = encoder.encode(query).tolist()
-        
-        # Ищем в коллекции "sales_knowledge" (или создадим её позже)
-        search_result = qdrant.search(
+        hits = q_client.search(
             collection_name="sales_knowledge", 
             query_vector=vector, 
             limit=3
         )
-        # Собираем текст из найденного
-        return "\n\n".join([hit.payload.get("text", "") for hit in search_result])
-    except Exception as e:
-        print(f"Ошибка поиска: {e}")
-        return ""
-
-# --- ЧАТ ---
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+        return "\n".join([hit.payload.get("text", "") for hit in hits])
+    except: return ""
 
-# Показываем историю
-for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
+# --- ОБРАБОТКА СООБЩЕНИЯ ---
+@cl.on_message
+async def main(message: cl.Message):
+    # Достаем инструменты из сессии
+    hf_client = cl.user_session.get("hf_client")
+    q_client = cl.user_session.get("q_client")
+    encoder = cl.user_session.get("encoder")
 
-# Обработка ввода
-if prompt := st.chat_input("Ваш вопрос..."):
-    # 1. Сохраняем вопрос юзера
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.markdown(prompt)
-
-    # 2. ИЩЕМ КОНТЕКСТ В QDRANT
-    context = get_context(prompt)
+    # 1. Ищем контекст в RAG
+    context = get_context(message.content, q_client, encoder)
     
-    # 3. Формируем системный промпт
-    system_instruction = "Ты — менеджер по продажам. Отвечай коротко и по делу."
+    # 2. Формируем промпт
+    system_prompt = "Ты полезный ассистент. Отвечай на русском языке."
     if context:
-        system_instruction += f"\n\nИспользуй эту информацию из базы знаний для ответа:\n{context}"
-        print(f"Нашел в базе: {context[:100]}...") # Для отладки в логах
+        system_prompt += f"\n\nИспользуй эту информацию для ответа:\n{context}"
+        # Можно показать найденный контекст в UI (опционально)
+        # await cl.Message(content=f"📚 Нашел в базе:\n{context[:100]}...").send()
 
-    # Собираем сообщения для API
-    api_messages = [{"role": "system", "content": system_instruction}]
-    for m in st.session_state.messages:
-        api_messages.append({"role": m["role"], "content": m["content"]})
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": message.content}
+    ]
 
-    # 4. Генерируем ответ
-    with st.chat_message("assistant"):
-        message_placeholder = st.empty()
-        full_response = ""
-        
-        try:
-            stream = client.chat_completion(
-                messages=api_messages, 
-                max_tokens=512, 
-                stream=True, 
-                temperature=0.7
-            )
-            
-            for chunk in stream:
-                content = chunk.choices[0].delta.content
-                if content:
-                    full_response += content
-                    message_placeholder.markdown(full_response + "▌")
-            
-            message_placeholder.markdown(full_response)
-            
-        except Exception as e:
-            st.error(f"Ошибка API: {e}")
+    # 3. Отправляем в Qwen и стримим ответ
+    msg = cl.Message(content="")
+    await msg.send() # Создаем пустой пузрь
 
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
+    try:
+        stream = hf_client.chat_completion(messages, max_tokens=1024, stream=True, temperature=0.7)
+        
+        for chunk in stream:
+            token = chunk.choices[0].delta.content
+            if token:
+                await msg.stream_token(token)
+                
+        await msg.update() # Финализируем сообщение
+        
+    except Exception as e:
+        await cl.Message(content=f"Ошибка: {str(e)}").send()