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 "Контекст розмови очищено"