Update app.py

app.py

@@ -1,105 +1,112 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from smolagents import CodeAgent, DuckDuckGoSearchTool
+from smolagents.models import TransformersModel
 
-class FixedAIAgent:
+class WorkingAIAgent:
     def __init__(self):
-        print("Loading DialoGPT-small...")
+        print("Starting AI agent with web search...")
         
-        # Загружаем модель и токенизатор
-        self.tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
-        self.model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
+        # Используем маленькую модель
+        self.model = TransformersModel("microsoft/DialoGPT-small")
         
-        # Добавляем pad token если его нет
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-            
-        print("DialoGPT-small loaded successfully!")
+        # Инструмент для поиска в интернете
+        self.search_tool = DuckDuckGoSearchTool()
+        
+        # Создаем агента с поиском
+        self.agent = CodeAgent(
+            tools=[self.search_tool],
+            model=self.model,
+            max_steps=3
+        )
+        
+        print("AI agent with web search is ready!")
     
     def chat(self, message, history):
+        print(f"User: {message}")
+        
         try:
-            # Форматируем историю для DialoGPT
-            input_text = self.format_conversation(message, history)
+            # Простой промпт для агента
+            prompt = f"""
+            Please answer this question: {message}
             
-            # Токенизируем входной текст
-            inputs = self.tokenizer.encode(input_text + self.tokenizer.eos_token, return_tensors='pt')
+            If you need to find current information, use the search tool.
+            Give a clear and helpful response.
+            """
             
-            # Генерируем ответ
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    inputs,
-                    max_length=1000,
-                    pad_token_id=self.tokenizer.eos_token_id,
-                    do_sample=True,
-                    temperature=0.7,
-                    top_k=50,
-                    top_p=0.95,
-                    repetition_penalty=1.2
-                )
+            # Получаем ответ от агента
+            response = self.agent.run(prompt)
             
-            # Декодируем ответ
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Очищаем ответ
+            clean_response = self.clean_answer(response)
+            print(f"AI: {clean_response[:100]}...")
             
-            # Извлекаем только новый ответ (убираем историю)
-            bot_response = self.extract_new_response(input_text, response)
-            
-            return bot_response
+            return clean_response
             
         except Exception as e:
-            return f"Error: {str(e)}"
+            error_msg = f"I encountered an error. Please try again with a different question."
+            print(f"Error: {e}")
+            return error_msg
     
-    def format_conversation(self, message, history):
-        """Форматирует историю чата для DialoGPT"""
-        # Начинаем с нового сообщения
-        conversation = f"User: {message}"
+    def clean_answer(self, text):
+        """Убираем технические детали из ответа"""
+        if not text:
+            return "I couldn't find an answer to that question."
         
-        # Добавляем историю (последние 2-3 сообщения)
-        if history:
-            # Берем последние 2 обмена
-            recent_history = history[-2:] if len(history) > 2 else history
-            
-            for user_msg, bot_msg in recent_history:
-                conversation = f"User: {user_msg}\nBot: {bot_msg}\n" + conversation
+        # Разбиваем на строки и фильтруем
+        lines = text.split('\n')
+        clean_lines = []
         
-        return conversation
-    
-    def extract_new_response(self, input_text, full_response):
-        """Извлекает только новый ответ из полного ответа модели"""
-        if input_text in full_response:
-            # Убираем входной текст чтобы оставить только новый ответ
-            new_response = full_response[len(input_text):].strip()
-            # Убираем возможные префиксы
-            if new_response.startswith("Bot:"):
-                new_response = new_response[4:].strip()
-            return new_response
-        else:
-            # Если не нашли входной текст, возвращаем как есть
-            return full_response
+        for line in lines:
+            line_lower = line.lower()
+            # Пропускаем технические строки
+            if any(word in line_lower for word in ['tool:', 'searching', 'step', 'using tool']):
+                continue
+            if line.strip():  # Добавляем только непустые строки
+                clean_lines.append(line.strip())
+        
+        # Собираем обратно
+        result = ' '.join(clean_lines)
+        
+        # Ограничиваем длину
+        if len(result) > 800:
+            result = result[:797] + "..."
+        
+        return result if result else "I found some information but couldn't format it properly."
 
 # Создаем агента
-agent = FixedAIAgent()
+ai_agent = WorkingAIAgent()
 
-# Создаем интерфейс
-with gr.Blocks() as app:
-    gr.Markdown("# AI Chat Assistant")
-    gr.Markdown("Powered by DialoGPT-small")
+# Создаем интерфейс чата
+with gr.Blocks(title="AI Assistant with Web Search") as app:
+    gr.Markdown("# AI Assistant with Web Search")
+    gr.Markdown("Ask me anything and I'll search the internet for answers!")
+    
+    # Чат интерфейс
+    chatbot = gr.Chatbot(
+        height=400,
+        label="Conversation"
+    )
     
-    chatbot = gr.Chatbot(height=400, label="Chat History")
     msg = gr.Textbox(
-        label="Your message", 
-        placeholder="Type your message here...",
+        label="Your question",
+        placeholder="Ask me anything... (I can search the web!)",
         lines=2
     )
-    clear = gr.Button("Clear Chat")
+    
+    clear_btn = gr.Button("Clear Chat")
     
     def respond(message, chat_history):
-        bot_message = agent.chat(message, chat_history)
-        chat_history.append((message, bot_message))
+        # Получаем ответ от агента
+        bot_response = ai_agent.chat(message, chat_history)
+        # Добавляем в историю чата
+        chat_history.append((message, bot_response))
         return "", chat_history
     
+    # Настраиваем обработчики
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: None, None, chatbot)
+    clear_btn.click(lambda: None, None, chatbot, queue=False)
 
+# Запускаем приложение
 if __name__ == "__main__":
-    print("Starting AI Chat...")
-    app.lalunch()
+    print("AI Assistant starting...")
+    app.launch(share=True)