agent.py

from smolagents import CodeAgent
import logging
from typing import Optional, List, Dict, Any
from datetime import datetime
import os

logger = logging.getLogger(__name__)

class ResearchAgent(CodeAgent):
    """
    Дослідницький агент для наукового пошуку та аналізу з підтримкою контексту розмови.
    """
    
    def __init__(self, model, tools, **kwargs):
        """
        Ініціалізація дослідницького агента.
        
        Args:
            model: Модель для обробки запитів
            tools: Інструменти для пошуку та обробки інформації
            **kwargs: Додаткові параметри
        """
        super().__init__(model=model, tools=tools, **kwargs)
        self.available_tools = {tool.name: tool for tool in tools}
        self.model_initializer = None
        self.step_callback = None
        # Ініціалізація контексту розмови
        self.conversation_context = {
            'previous_queries': [],      # Історія запитів
            'search_results': {},        # Результати пошуку для кожного запиту
            'generated_content': {},     # Згенерований контент для кожного запиту
            'last_topic': None          # Остання обговорювана тема
        }
        logger.info(f"ResearchAgent ініціалізовано з інструментами: {list(self.available_tools.keys())}")

    def _update_steps(self, step_info: str, step_result: Any = None):
        """
        Оновлення інформації про кроки з результатами.
        
        Args:
            step_info: Інформація про поточний крок
            step_result: Результат виконання кроку
        """
        if self.step_callback:
            step_data = {
                'info': step_info,
                'result': self._format_step_result(step_result)
            }
            self.step_callback(step_data)
            logger.info(f"Step: {step_info}")
            if step_result:
                logger.info(f"Result: {str(step_result)[:200]}...")

    def _format_step_result(self, result: Any) -> str:
        """
        Форматування результату кроку для відображення.
        
        Args:
            result: Результат для форматування
        Returns:
            str: Відформатований результат
        """
        if result is None:
            return ""
        if isinstance(result, dict):
            return "\n".join([f"{k}: {str(v)}" for k, v in result.items()])
        elif isinstance(result, list):
            return "\n".join([str(item) for item in result])
        return str(result)

    def _is_related_query(self, current_query: str, previous_context: Dict) -> bool:
        """
        Визначає, чи є поточний запит пов'язаним з попереднім контекстом.
        
        Args:
            current_query: Поточний запит
            previous_context: Попередній контекст
        Returns:
            bool: True якщо запит пов'язаний з попереднім контекстом
        """
        if not previous_context['last_topic']:
            return False

        prompt = f"""Визначте, чи є це уточнюючим запитом до попередньої теми.
        Попередня тема: {previous_context['last_topic']}
        Новий запит: {current_query}
        Відповідь має бути лише True або False."""

        try:
            response = self.model(prompt)
            return 'true' in response.lower()
        except Exception as e:
            logger.error(f"Помилка при визначенні зв'язку запитів: {e}")
            return False

    def _build_context_prompt(self, query: str, context: Dict) -> str:
        """
        Формує промпт з урахуванням попереднього контексту.
        
        Args:
            query: Поточний запит
            context: Контекст розмови
        Returns:
            str: Сформований промпт
        """
        prompt_parts = []
        
        if context['previous_queries']:
            prompt_parts.append("Попередні запити та результати:")
            for prev_query in context['previous_queries'][-3:]:
                prompt_parts.append(f"Запит: {prev_query}")
                if prev_query in context['search_results']:
                    prompt_parts.append("Знайдена інформація:")
                    prompt_parts.append(context['search_results'][prev_query])
                if prev_query in context['generated_content']:
                    prompt_parts.append("Згенерований контент:")
                    prompt_parts.append(context['generated_content'][prev_query])
                prompt_parts.append("---")

        prompt_parts.append(f"Поточний запит: {query}")
        prompt_parts.append("На основі всієї наведеної інформації створіть детальну відповідь.")
        
        return "\n".join(prompt_parts)

    def _update_context(self, query: str, search_results: str, generated_content: str):
        """
        Оновлює контекст розмови.
        
        Args:
            query: Поточний запит
            search_results: Результати пошуку
            generated_content: Згенерований контент
        """
        self.conversation_context['previous_queries'].append(query)
        self.conversation_context['search_results'][query] = search_results
        self.conversation_context['generated_content'][query] = generated_content
        self.conversation_context['last_topic'] = query

    def process_query(self, query: str, available_files: Optional[List[str]] = None) -> str:
        """
        Обробка дослідницького запиту з підтримкою контексту.
        
        Args:
            query: Запит для обробки
            available_files: Список доступних файлів
        Returns:
            str: Результат обробки запиту
        """
        try:
            logger.info(f"Обробка дослідницького запиту: {query}")
            self._update_steps("Початок обробки запиту")
            
            # Перевіряємо зв'язок з попереднім контекстом
            is_related = self._is_related_query(query, self.conversation_context)
            self._update_steps("Аналіз зв'язку з попереднім контекстом", 
                             {'is_related': is_related})

            # Виконуємо пошук
            self._update_steps("Пошук інформації в інтернеті")
            search_results = self.available_tools['web_search'](query=query)
            self._update_steps("Отримано результати пошуку", search_results)

            # Формуємо промпт
            if is_related:
                prompt = self._build_context_prompt(query, self.conversation_context)
                self._update_steps("Формування запиту з урахуванням контексту", 
                                 {'context_length': len(self.conversation_context['previous_queries'])})
            else:
                prompt = f"Task: {query}\n\nSearch Results:\n{search_results}"
                self._update_steps("Формування нового запиту")

            # Генеруємо відповідь
            self._update_steps("Генерація відповіді")
            result = self.model(prompt)
            
            # Оновлюємо контекст
            self._update_steps("Оновлення контексту розмови")
            self._update_context(query, search_results, result)
            
            self._update_steps("Запит успішно оброблено", {
                'context_size': len(self.conversation_context['previous_queries']),
                'is_related_query': is_related
            })
            
            return result
            
        except Exception as e:
            error_msg = f"Помилка при обробці запиту: {str(e)}"
            logger.error(error_msg)
            self._update_steps(f"Помилка", {'error': error_msg})
            return error_msg

    def clear_context(self):
        """
        Очищення контексту розмови.
        
        Returns:
            str: Повідомлення про очищення контексту
        """
        self.conversation_context = {
            'previous_queries': [],
            'search_results': {},
            'generated_content': {},
            'last_topic': None
        }
        logger.info("Контекст розмови очищено")
        return "Контекст розмови очищено"