diff --git a/README.md b/README.md index 8cbba60..ca814b1 100644 --- a/README.md +++ b/README.md @@ -1,128 +1,92 @@ -# 🤖 PyAgentX: Autonomous Multi-Agent System for Software Development +# PyAgentX: Продвинутая многоагентная AI-система -**PyAgentX** — это продвинутая система на базе LLM, предназначенная для автоматизации полного цикла разработки программного обеспечения. Она использует команду специализированных ИИ-агентов, которые совместно работают над решением задач: от декомпозиции высокоуровневых целей до написания, проверки, тестирования и оценки кода. +**PyAgentX** — это фреймворк для создания автономных AI-агентов, способных решать сложные, многошаговые задачи. Система использует архитектуру многоагентной команды, где центральный Оркестратор управляет командой узкоспециализированных агентов для достижения глобальной цели. -Ключевой особенностью системы является **гибридная RAG-система**, которая совмещает семантический и ключевой поиск по базе знаний. Это позволяет агентам генерировать код, соответствующий внутренним стандартам, архитектурным паттернам и лучшим практикам вашей команды. +Проект основан на самых современных концепциях в области Agentic AI, включая иерархическое планирование, самокоррекцию, продвинутый RAG с переранжированием и долгосрочную память. -## 🚀 Ключевые возможности +--- -- **Динамическая команда агентов**: Система гибко настраивается через YAML-конфиги. Вы можете легко добавлять, удалять или изменять роли агентов. -- **Продвинутый RAG с гибридным поиском**: Использует комбинацию векторного поиска (семантика) и BM25 (ключевые слова) с Reciprocal Rank Fusion для максимально релевантных результатов из базы знаний. -- **Динамическая конфигурация**: Все параметры системы, от моделей LLM до настроек RAG для каждого агента, задаются в YAML-файлах в директории `configs/`. -- **Интерфейс командной строки (CLI)**: Удобный запуск и управление через `main.py` с использованием `typer`. -- **Самостоятельная регистрация инструментов**: Агенты сами определяют и регистрируют свой инструментарий, что делает систему более модульной и инкапсулированной. -- **Работа с файловой системой**: Агенты могут читать, создавать и редактировать файлы прямо в вашем проекте. -- **Надежное тестирование**: Встроенный набор модульных тестов (`pytest`) с использованием моков позволяет проверять логику, не затрагивая реальные API. +## 🏛️ Ключевые концепции -## 🏛️ Архитектура +Мы разделили документацию на несколько разделов для удобства: -Система построена на модульных принципах, где каждый компонент имеет четкую зону ответственности. +1. [**Архитектура многоагентной системы**](./docs/01_multi_agent_architecture.md): Обзор Оркестратора, Планировщика, Исполнителей и принципов их взаимодействия. +2. [**Ядро агента: Цикл ReAct и Рефлексия**](./docs/02_agent_core_loop.md): Описание внутреннего цикла работы каждого агента и механизма самокоррекции при ошибках. +3. [**Продвинутый RAG-пайплайн**](./docs/03_advanced_rag.md): Детальное описание нашего двухэтапного RAG с семантическим чанкированием и переранжированием. +4. [**Долгосрочная память**](./docs/04_long_term_memory.md): Как агенты запоминают информацию между сессиями для повышения эффективности. +5. [**Оценка и Безопасность**](./docs/04_evaluation_and_safety.md): Описание наших собственных легковесных систем для оценки качества и модерации ответов. -``` -PyAgentX/ -├── app/ -│ ├── agents/ # Логика и роли специализированных агентов -│ │ ├── roles/ # Классы для конкретных ролей (Coding, Testing, etc.) -│ │ ├── agent.py # Базовый класс Agent с основной логикой -│ │ └── tools.py # Определения инструментов (read_file, edit_file) -│ ├── factory/ # Фабрика для создания агентов -│ │ └── agent_factory.py -│ └── rag/ # Логика для Retrieval-Augmented Generation -│ └── retriever.py # Гибридный ретривер (Vector + BM25) -├── configs/ # Конфигурационные файлы YAML -│ ├── agents/ # Конфиги для каждого агента -│ └── config.yaml # Главный конфигурационный файл -├── db/ # Локальная база данных (создается автоматически) -│ ├── chunks.json # Текстовые чанки -│ ├── embeddings.npy # Векторные эмбеддинги -│ └── bm25_index.pkl # Индекс BM25 для ключевого поиска -├── knowledge/ # Исходники для базы знаний (документация в .md) -├── scripts/ # Вспомогательные скрипты -│ └── build_knowledge_base.py -├── tests/ # Модульные тесты pytest -├── .env # Файл для секретных ключей -├── main.py # Точка входа в приложение (CLI на Typer) -├── poetry.lock # Файл зависимостей Poetry -├── pyproject.toml # Конфигурация проекта и зависимостей -└── README.md -``` +--- -**Логика работы:** -1. Пользователь запускает `main.py`, передавая задачу через CLI. -2. `agent_factory` читает `configs/config.yaml` и конфигурации агентов, создавая команду. -3. `TaskDecomposer` получает задачу и разбивает ее на последовательность шагов. -4. `Orchestrator` (в `main.py`) последовательно передает каждый шаг соответствующему агенту. -5. Каждый агент перед выполнением задачи может обратиться к `KnowledgeRetriever` для получения релевантного контекста из базы знаний. -6. Агент выполняет шаг, используя свои инструменты (`tools.py`). -7. Цикл повторяется, пока все шаги не будут выполнены. +## 🛠️ Технологический стек -## 🛠️ Быстрый старт +- **Язык**: Python 3.10+ +- **LLM API**: OpenAI +- **Оркестрация и ядро агентов**: Кастомная реализация +- **Семантическое чанкирование**: `semchunk` +- **RAG**: + - `unstructured` для парсинга документов + - `rank-bm25` для sparse-поиска + - `numpy` для векторных операций + - `sentence-transformers` для Cross-Encoder Re-ranking +- **Долгосрочная память**: `sqlite3` +- **Безопасность**: `guardrails-ai` +- **Оценка**: `deepeval`, `pytest` +- **Зависимости**: `python-dotenv`, `tiktoken` -### 1. Клонирование репозитория +--- +## 🚀 Как запустить + +### 1. Установка зависимостей ```bash -git clone https://github.com/your-username/PyAgentX.git -cd PyAgentX +pip install -r requirements.txt ``` ### 2. Настройка окружения - -Проект использует [Poetry](https://python-poetry.org/) для управления зависимостями. - -- Установите Poetry (если не установлен): - ```bash - pip install poetry - ``` -- Создайте виртуальное окружение и установите зависимости: - ```bash - poetry install - ``` - -### 3. Конфигурация - -- Создайте файл `.env` в корне проекта, скопировав `.env.example` (если он есть) или создав новый. -- Добавьте ваш ключ OpenAI: - ``` - OPENAI_API_KEY="sk-..." - ``` - -### 4. Создание базы знаний - -Перед первым запуском необходимо проиндексировать вашу документацию из папки `knowledge/`. - +Скопируйте `.env.example` в `.env` и укажите ваш `OPENAI_API_KEY`. ```bash -poetry run python -m scripts.build_knowledge_base +cp .env.example .env ``` -Этот шаг нужно повторять только при изменении файлов в папке `knowledge/`. - -### 5. Запуск - -Запустите главный скрипт, передав ему задачу. +### 3. Создание базы знаний +Для работы RAG-пайплайна необходимо один раз создать векторную базу знаний из документов в директории `knowledge_base/`. ```bash -poetry run python -m main --briefing "Создай новую функцию в 'app/utils.py' для сложения двух чисел и напиши на нее тест в 'tests/test_utils.py'." +python -m scripts.build_knowledge_base ``` -Система начнет выполнение, и вы увидите логи работы агентов в консоли. - -## 🧩 Расширение системы - -### Добавление нового агента - -1. Создайте новый класс агента в `app/agents/roles/`, унаследовав его от `Agent`. -2. В `__init__` нового агента зарегистрируйте необходимые ему инструменты с помощью `self.add_tool()`. -3. Создайте для него YAML-конфиг в `configs/agents/`. -4. Добавьте нового агента в главный `configs/config.yaml`. - -### Добавление нового инструмента - -1. Определите функцию-инструмент и ее JSON-схему в `app/agents/tools.py`. -2. Добавьте вызов `self.add_tool()` в `__init__` того агента, который должен использовать этот инструмент. - -### Добавление знаний +### 4. Запуск тестов (Опционально) +Чтобы убедиться, что все работает корректно: +```bash +pytest +``` -Просто добавьте новый `.md` файл в папку `knowledge/` и перезапустите скрипт `scripts/build_knowledge_base.py`. +### 5. Запуск веб-сервера +```bash +uvicorn app.main:app --reload +``` +После запуска API будет доступно по адресу `http://127.0.0.1:8000/docs`. -## 💡 Следующий шаг: `PROJECT_CONTEXT.md` +--- -Для дальнейшего улучшения согласованности действий агентов планируется внедрение файла `PROJECT_CONTEXT.md`. Это будет "конституция" проекта, содержащая глобальные правила и стандарты, которая будет автоматически добавляться в системный промпт каждого агента. +## 🏗️ Структура проекта +``` +PyAgentX/ +├── app/ +│ ├── agents/ # Логика агентов, их роли и инструменты +│ │ ├── roles/ +│ │ └── prompts/ +│ ├── evaluation/ # Собственный фреймворк оценки +│ ├── memory/ # Менеджер долгосрочной памяти (SQLite) +│ ├── rag/ # RAG-пайплайн (ретривер) +│ ├── safety/ # Собственные "ограждения" (Guardrails) +│ ├── main.py # FastAPI приложение +│ └── orchestrator.py # Оркестратор +├── docs/ # Детальная документация +├── knowledge_base/ # Исходные документы для базы знаний +├── scripts/ # Скрипты (например, для создания RAG-базы) +├── tests/ # Тесты +├── .env.example +├── README.md +└── requirements.txt +``` diff --git a/app/agents/agent.py b/app/agents/agent.py index fddd0af..81f5f9f 100644 --- a/app/agents/agent.py +++ b/app/agents/agent.py @@ -16,32 +16,61 @@ read_file_tool, read_file_tool_def, list_files_tool, list_files_tool_def, ) +from app.agents.web_search_tool import web_search_tool, web_search_tool_def +from app.agents.memory_tool import save_memory_tool, save_memory_tool_def +from app.memory.memory_manager import MemoryManager +from app.safety.custom_guardrails import CustomGuardrailManager -# Новый, улучшенный системный промпт, превращающий агента в программиста. -SYSTEM_PROMPT = """ -Ты — автономный AI-агент, способный выполнять сложные задачи, используя доступные инструменты. -Твоя цель — успешно завершить поставленную задачу, шаг за шагом. - -# ПРАВИЛА ВЗАИМОДЕЙСТВИЯ: -1. **АНАЛИЗ**: Внимательно изучи предоставленный тебе контекст и задачу. -2. **ПЛАН**: Составь внутренний план действий. Какой инструмент использовать первым? -3. **ДЕЙСТВИЕ**: Вызови выбранный инструмент с правильными аргументами. -4. **РЕЗУЛЬТАТ**: Изучи результат вызова инструмента. -5. **ЦИКЛ**: Повторяй шаги 1-4, пока не достигнешь конечной цели. Когда задача будет полностью выполнена, сообщи об этом, предоставив финальный результат. - -# ВАЖНЫЕ ИНСТРУКЦИИ ПО РАБОТЕ: - -## 1. Работа с файловой системой: -- **Абсолютные импорты**: При написании Python кода всегда используй абсолютные импорты от корня проекта. Корень проекта - это директория, где лежит `main.py`. Например: `from app.agents.tools import read_file_tool`. **НИКОГДА** не используй относительные импорты (`from .. import ...`) или хаки с `sys.path`. -- **Редактирование файлов**: Инструмент `edit_file_tool` имеет два режима: - - `mode='overwrite'` (по умолчанию): **Полностью перезаписывает** файл. Используй с осторожностью. - - `mode='append'`: **Добавляет контент в конец файла**. Используй этот режим, когда тебе нужно добавить новую функцию, класс или текст в уже существующий файл, не удаляя его содержимое. - -## 2. Мышление и логика: -- **Код прежде тестов**: Если твоя цель — написать функцию и тесты к ней, всегда сначала полностью реализуй **корректную и финальную** логику самой функции. Только после этого приступай к написанию тестов. Не тестируй заготовки или неполный код. -- Будь методичен. Не торопись. -- Если результат не соответствует ожиданиям, попробуй другой подход. -- Если ты застрял, сделай шаг назад и пересмотри свой план. +class ToolExecutionError(Exception): + """Custom exception for errors during tool execution.""" + pass + +# Новый системный промпт, основанный на ReAct (Reason + Act) +REACT_SYSTEM_PROMPT = """ +You are a smart, autonomous AI agent. Your name is {agent_name}, and your role is {agent_role}. +Your ultimate goal is: {agent_goal}. + +You operate in a loop of Thought, Action, and Observation. +At each step, you MUST respond in a specific JSON format. Your entire response must be a single JSON object. + +1. **Thought**: First, think about your plan. Analyze the user's request, your goal, and the previous steps. Describe your reasoning for the current action. This is a mandatory field. + +2. **Action** or **Answer**: Based on your thought, you must choose ONE of the following: + a. `action`: An object representing the tool to use. It must contain: + - `name`: The name of the tool to execute. + - `input`: An object with the parameters for the tool. + b. `answer`: A final, comprehensive answer to the user's request. Use this ONLY when the task is fully complete. + +# AVAILABLE TOOLS: +You have access to the following tools. Use them to gather information and perform actions. + +{tools_description} + +# LONG-TERM MEMORY: +Before you begin, here are some facts you have previously saved to your long-term memory. +Use them to inform your decisions, but do not state them unless relevant. + +{memory_context} + +# EXAMPLE OF A SINGLE STEP: + +```json +{{ + "thought": "I need to understand the project structure first. I'll list the files in the current directory.", + "action": {{ + "name": "list_files_tool", + "input": {{ + "path": "." + }} + }} +}} +``` + +# IMPORTANT RULES: +- Your ENTIRE output MUST be a single, valid JSON object. Do not add any text before or after the JSON. +- You must choose either `action` or `answer`, not both. +- The `thought` field is always required. +- Think step by step. Your goal is to complete the task, not just use tools. """ # Настройка логирования @@ -74,11 +103,18 @@ def __init__( self.tool_definitions: List[Dict[str, Any]] = [] self.conversation_history: List[Dict[str, Any]] = [] self.max_iterations = max_iterations - self.system_prompt = "Ты — универсальный AI-ассистент." + self.system_prompt_template = REACT_SYSTEM_PROMPT + + # Initialize Memory Manager + self.memory_manager = MemoryManager() + # Initialize Guardrails + self.guardrail_manager = CustomGuardrailManager(api_key=api_key) # Add default tools self.add_tool(read_file_tool, read_file_tool_def) self.add_tool(list_files_tool, list_files_tool_def) + self.add_tool(web_search_tool, web_search_tool_def) + self.add_tool(save_memory_tool, save_memory_tool_def) # RAG specific attributes self.use_rag = use_rag @@ -159,6 +195,15 @@ def _enrich_with_knowledge(self, query: str) -> str: logging.info(f"Knowledge context added for agent '{self.name}'.") return knowledge_context + def _get_memory_context(self) -> str: + """Retrieves relevant facts from long-term memory.""" + recent_facts = self.memory_manager.get_recent_facts(limit=10) + if not recent_facts: + return "No relevant facts found in memory." + + formatted_facts = "\n".join([f"- {fact}" for fact in recent_facts]) + return f"--- RELEVANT FACTS ---\n{formatted_facts}\n----------------------" + def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: """Выполняет указанный инструмент с аргументами.""" if tool_name in self.tools: @@ -168,10 +213,29 @@ def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: return result except Exception as e: logging.error("Ошибка при выполнении инструмента '%s': %s", tool_name, e, exc_info=True) - return f"Ошибка: Не удалось выполнить инструмент '{tool_name}'. Причина: {e}" + # Выбрасываем кастомное исключение, чтобы его можно было поймать выше + raise ToolExecutionError(f"Error: Failed to execute tool '{tool_name}'. Reason: {e}") else: logging.warning("Попытка вызова неизвестного инструмента: '%s'", tool_name) - return f"Ошибка: Инструмент '{tool_name}' не найден." + raise ToolExecutionError(f"Error: Tool '{tool_name}' not found.") + + def _get_tools_description(self) -> str: + """Generates a description of available tools for the system prompt.""" + if not self.tool_definitions: + return "No tools available." + + desc = [] + for tool in self.tool_definitions: + func_desc = tool.get('function', {}) + params_desc = func_desc.get('parameters', {}).get('properties', {}) + + # Описание параметров + params_str = ", ".join([f"{name}: {info.get('type')}" for name, info in params_desc.items()]) + + desc.append( + f"- `{func_desc.get('name', 'N/A')}({params_str})`: {func_desc.get('description', 'No description.')}" + ) + return "\n".join(desc) def _get_user_input(self) -> Optional[str]: """Обрабатывает получение ввода от пользователя.""" @@ -182,128 +246,143 @@ def _get_user_input(self) -> Optional[str]: except EOFError: return None - def _get_model_response(self) -> ChatCompletionMessage: + def _get_model_response(self) -> str: """ - Отправляет текущую историю беседы в OpenAI и возвращает ответ модели. + Отправляет текущую историю беседы в OpenAI и возвращает текстовый ответ модели. """ try: - logging.info(f"Отправка запроса в OpenAI с {len(self.conversation_history)} сообщениями.") + logging.info(f"Sending request to OpenAI with {len(self.conversation_history)} messages.") response = self.client.chat.completions.create( model=self.model, messages=self.conversation_history, - tools=self.get_openai_tools(), - tool_choice="auto", + # Убираем tools и tool_choice, так как теперь мы парсим JSON ) - return response.choices[0].message + return response.choices[0].message.content or "" except Exception as e: - logging.error(f"Ошибка при вызове OpenAI API: {e}", exc_info=True) - # Возвращаем "пустое" сообщение с контентом об ошибке, чтобы цикл мог его обработать - return ChatCompletionMessage(role="assistant", content=f"Произошла ошибка API: {e}") + logging.error(f"Error calling OpenAI API: {e}", exc_info=True) + return json.dumps({ + "thought": "An API error occurred. I cannot proceed.", + "answer": f"API Error: {e}" + }) def execute_task(self, briefing: str) -> str: """ - Выполняет одну задачу на основе предоставленного брифинга. + Выполняет одну задачу на основе предоставленного брифинга, используя ReAct цикл. """ - logging.info(f"Агент {self.name} получил задачу.") - - knowledge_context = "" - if self.use_rag: - focused_query = self._create_rag_query(briefing) - knowledge_context = self._enrich_with_knowledge(query=focused_query) + logging.info(f"Agent {self.name} received task.") + + # 1. Prepare Prompts + knowledge_context = self._enrich_with_knowledge(self._create_rag_query(briefing)) if self.use_rag else "" + tools_description = self._get_tools_description() + memory_context = self._get_memory_context() - # Системный промпт определяет "личность" агента, а брифинг - контекст задачи. - # Контекст из базы знаний добавляется в начало системного промпта. - final_system_prompt = f"{knowledge_context}\n{self.system_prompt}" + system_prompt = self.system_prompt_template.format( + agent_name=self.name, + agent_role=self.role, + agent_goal=self.goal, + tools_description=tools_description, + memory_context=memory_context + ) + + final_system_prompt = f"{knowledge_context}\n{system_prompt}" self.conversation_history = [ {"role": "system", "content": final_system_prompt}, {"role": "user", "content": briefing} ] - for _ in range(self.max_iterations): - response_message = self._get_model_response() + # 2. Start ReAct Loop + for i in range(self.max_iterations): + logging.info(f"--- Iteration {i+1}/{self.max_iterations} ---") + + model_response_str = self._get_model_response() + self.conversation_history.append({"role": "assistant", "content": model_response_str}) - if response_message.content and "ошибка API" in response_message.content.lower(): - return response_message.content + try: + parsed_response = json.loads(model_response_str) + thought = parsed_response.get("thought") + if thought: + logging.info(f"🤖 Thought: {thought}") + else: + raise ValueError("Missing 'thought' in response.") - if not response_message.tool_calls: - final_answer = response_message.content or "Задача выполнена." - logging.info(f"Агент {self.name} завершил задачу с ответом: {final_answer}") - return final_answer - - self.conversation_history.append(response_message) - - for tool_call in response_message.tool_calls: - tool_name = tool_call.function.name - tool_args_str = tool_call.function.arguments - logging.info(f"Вызов инструмента: {tool_name} с аргументами: {tool_args_str}") - - try: - tool_args = json.loads(tool_args_str) - tool_result = self._execute_tool(tool_name, tool_args) - except json.JSONDecodeError: - error_msg = f"Ошибка: неверный JSON в аргументах инструмента: {tool_args_str}" - logging.error(error_msg) - tool_result = error_msg - - self.conversation_history.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": str(tool_result), - } + if parsed_response.get("answer"): + final_answer = parsed_response["answer"] + # Validate final answer with guardrails before returning + validated_answer = self.guardrail_manager.validate_and_format_response(final_answer) + logging.info(f"✅ Agent {self.name} finished task with answer: {validated_answer}") + return validated_answer + + elif parsed_response.get("tool"): + tool_name = parsed_response["tool"]["name"] + tool_input = parsed_response["tool"]["input"] + + if not tool_name: + raise ValueError("Missing 'name' in action.") + + logging.info(f"🛠️ Action: Calling tool '{tool_name}' with input: {tool_input}") + tool_result = self._execute_tool(tool_name, tool_input) + + observation = f"Tool '{tool_name}' returned:\n```\n{tool_result}\n```" + logging.info(f"👀 Observation: {observation}") + self.conversation_history.append({"role": "user", "content": observation}) + else: + raise ValueError("Response must contain 'action' or 'answer'.") + + except ToolExecutionError as e: + error_message = f"A tool failed to execute: {e}" + logging.error(error_message) + # Запускаем цикл рефлексии + reflection_prompt = ( + f"CRITICAL_ERROR: Your last action failed with the following error: '{e}'.\n" + "You MUST analyze this error and the execution history to understand what went wrong.\n" + "Then, devise a new plan. Either try a different approach, use a different tool, or modify the input to the tool.\n" + "Your next 'thought' MUST explain how you are correcting your course of action." ) + self.conversation_history.append({"role": "user", "content": reflection_prompt}) + continue # Продолжаем цикл, чтобы агент мог ответить на сообщение об ошибке + + except (json.JSONDecodeError, ValueError) as e: + error_message = f"Error parsing model response: {e}. Response was: '{model_response_str}'" + logging.error(error_message) + # Даем агенту шанс исправиться + error_feedback = ( + f"Error: Your last response was not a valid JSON object. " + f"Please correct your output to strictly follow the required format. " + f"The `thought` field is mandatory, and you must include either an `action` or an `answer`. " + f"Error details: {e}" + ) + self.conversation_history.append({"role": "user", "content": error_feedback}) + continue - warning_message = f"Агент {self.name} достиг лимита итераций ({self.max_iterations}), не завершив задачу." + warning_message = f"Agent {self.name} reached max iterations ({self.max_iterations}) without a final answer." logging.warning(warning_message) + # Validate the warning message as well, in case it contains sensitive info (less likely but good practice) return warning_message def run(self) -> None: """Запускает основной цикл общения с агентом в интерактивном режиме.""" - print("Общение с агентом (используйте Ctrl+D или Ctrl+C для выхода)") - conversation: List[Dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}] - - try: - while True: - user_input = self._get_user_input() - if user_input is None: + print(f"Запуск агента '{self.name}' в интерактивном режиме. Введите 'exit' для завершения.") + + while True: + try: + user_input = input("\033[94mВы > \033[0m") + if user_input.lower() == 'exit': + print("Завершение сеанса.") break - # Здесь мы используем ту же логику, что и в execute_task - # для выполнения пользовательского запроса. - # Это можно будет в будущем вынести в отдельный метод. - conversation.append({"role": "user", "content": user_input}) - - max_tool_calls = 5 - for _ in range(max_tool_calls): - response: ChatCompletionMessage = self.client.chat.completions.create( - model=self.model, - messages=conversation, - tools=self.get_openai_tools(), - ).choices[0].message - - if not response.tool_calls: - assistant_response = response.content or "" - print(f"\033[93m{self.name}:\033[0m {assistant_response}") - conversation.append({"role": "assistant", "content": assistant_response}) - break - - conversation.append(response.model_dump()) - - for tool_call in response.tool_calls: - tool_name = tool_call.function.name - tool_args = json.loads(tool_call.function.arguments) - tool_result = self._execute_tool(tool_name, tool_args) - conversation.append({ - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": tool_result, - }) - else: - print(f"\033[91m{self.name}: Достигнут лимит вызовов инструментов.\033[0m") - - except (KeyboardInterrupt, EOFError): - print("\nВыход из чата.") + if not user_input.strip(): + continue + + # Используем execute_task для обработки ввода пользователя + print(f"\033[93m{self.name} >\033[0m", end="", flush=True) + final_response = self.execute_task(user_input) + + # Печатаем финальный ответ, который execute_task вернул + # execute_task уже логирует промежуточные шаги + print(final_response) + + except (KeyboardInterrupt, EOFError): + print("\nЗавершение сеанса.") + break diff --git a/app/agents/memory_tool.py b/app/agents/memory_tool.py new file mode 100644 index 0000000..733c70c --- /dev/null +++ b/app/agents/memory_tool.py @@ -0,0 +1,47 @@ +""" +This module defines the memory tool for the agent. +""" +from typing import Dict, Any +from app.memory.memory_manager import MemoryManager + +# We will instantiate the manager when the agent loads the tool. +# This avoids creating a new connection for every call. +memory_manager = MemoryManager() + +# 1. Tool Definition (JSON Schema) +save_memory_tool_def = { + "type": "function", + "function": { + "name": "save_memory_tool", + "description": "Saves a specific piece of information to your long-term memory for future use. " + "Use this when you learn a new, important fact that needs to be remembered across sessions.", + "parameters": { + "type": "object", + "properties": { + "fact": { + "type": "string", + "description": "The specific fact or piece of information to save." + } + }, + "required": ["fact"], + }, + }, +} + +# 2. Tool Function +def save_memory_tool(args: Dict[str, Any]) -> str: + """ + Saves a fact to the long-term memory. + + Args: + args: A dictionary containing the 'fact' to be saved. + + Returns: + A confirmation string of the save operation. + """ + fact = args.get("fact") + if not fact: + return "Error: The 'fact' parameter is required to save to memory." + + result = memory_manager.add_fact(fact) + return result.get("message", "An unknown error occurred.") \ No newline at end of file diff --git a/app/agents/roles/coding_agent.py b/app/agents/roles/coding_agent.py index 423f205..f1a69ef 100644 --- a/app/agents/roles/coding_agent.py +++ b/app/agents/roles/coding_agent.py @@ -12,10 +12,18 @@ class CodingAgent(Agent): """An agent specializing in writing and refactoring code.""" def __init__(self, name: str, role: str, goal: str, **kwargs): + # Добавляем специфику роли в цель, чтобы она попала в основной промпт + refined_goal = ( + f"{goal}\n\n" + "IMPORTANT: Your primary goal is to make precise, targeted changes (surgical edits). " + "Do not rewrite entire files. Instead, identify the specific function, method, " + "or block of code that needs changing and use the 'edit_file_tool' to replace only that part." + ) + super().__init__( name=name, role=role, - goal=goal, + goal=refined_goal, # Передаем уточненную цель **kwargs, ) # Self-register tools @@ -23,43 +31,5 @@ def __init__(self, name: str, role: str, goal: str, **kwargs): self.add_tool(edit_file_tool, edit_file_tool_def) self.add_tool(list_files_tool, list_files_tool_def) - self.system_prompt = ( - "You are a CodingAgent, an elite AI developer. Your task is to write, modify, and fix code." - "You will be provided with the full plan, the history of previous steps, and your current task." - "\n\n" - "## Operating Principles:\n" - "1. **Think Before You Code:** Carefully study the task and context. Plan your actions." - "2. **Surgical Edits:** Your primary goal is to make precise, targeted changes. Do not rewrite entire files. Instead, identify the specific function, method, or block of code that needs changing and replace only that part.\n" - "3. **Use Tools Wisely:**\n" - " - Prefer `edit_file_tool` with `mode='replace'`. This is the safest and most professional way to work.\n" - " - Use `mode='append'` for adding new functions or tests to the end of a file.\n" - " - Use `mode='overwrite'` ONLY when creating a brand new file.\n" - "4. **Code Quality:** Write clean, efficient, and well-documented code that adheres to PEP8." - "\n\n" - "## Example of a Surgical Edit:\n" - "Your task is to fix a bug in the `add` function in `math_utils.py`.\n\n" - "1. **First, read the file:** `read_file_tool(path='app/utils/math_utils.py')`\n" - "2. **Identify the flawed function:**\n" - " ```python\n" - " # This is the old, incorrect code block you will replace\n" - " def add(a, b):\n" - " return a - b # Bug is here\n" - " ```\n" - "3. **Call the edit tool to replace ONLY that function:**\n" - " ```python\n" - " edit_file_tool(\n" - " path='app/utils/math_utils.py',\n" - " mode='replace',\n" - " old_content='''def add(a, b):\\n return a - b # Bug is here''',\n" - " new_content='''def add(a, b):\\n # Fix: Correctly perform addition\\n return a + b'''\n" - " )\n" - " ```\n\n" - "Your goal is to successfully complete your part of the plan, preparing the way for the next agent." - ) - - def _create_initial_messages(self, task_briefing: str) -> List[Dict[str, str]]: - """Creates the initial list of messages for the model dialogue.""" - return [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": task_briefing}, - ] \ No newline at end of file + # Старый system_prompt и _create_initial_messages больше не нужны, + # так как вся логика теперь в базовом классе Agent. \ No newline at end of file diff --git a/app/agents/roles/evaluator_agent.py b/app/agents/roles/evaluator_agent.py index 6725bf8..43f9b85 100644 --- a/app/agents/roles/evaluator_agent.py +++ b/app/agents/roles/evaluator_agent.py @@ -16,35 +16,5 @@ def __init__(self, name: str, role: str, goal: str, **kwargs): # Self-register tools self.add_tool(read_file_tool, read_file_tool_def) - self.system_prompt = """ -You are EvaluatorAgent, an experienced QA engineer and systems analyst. -Your primary task is to analyze the log from failed pytest runs and formulate a clear, concise, and single task for the CodingAgent to fix the code. - -# WORKFLOW: -1. **Study the Log**: Carefully read the pytest output provided to you. Pay close attention to the `FAILURES`, `ERRORS`, and traceback sections. -2. **Identify the Root Cause**: Determine the core reason for the failure. Is it a bug in the function's logic? An error in the test itself? Incorrect data? -3. **Formulate the Task**: Your response must be a SINGLE sentence that is a direct and understandable instruction for the developer. You do not need to suggest code; just describe WHAT needs to be fixed. - -# EXAMPLES: - -**Example 1 (Log with a logic error):** -``` -... -FAILED tests/test_tools.py::test_subtract_negative - assert subtract(5, 10) == 5 -AssertionError: assert -5 == 5 -... -``` -**Your response:** -"Fix the `subtract` function in `app/agents/tools.py`, as it incorrectly calculates the difference when the second argument is larger than the first." - -**Example 2 (Log with an import error):** -``` -... -ImportError: cannot import name 'substract' from 'app.agents.tools' (did you mean 'subtract'?) -... -``` -**Your response:** -"Fix the import error in `tests/test_tools.py`; the function name `subtract` is likely misspelled." - -Your output is the task for the other agent. Be as precise as possible. -""" \ No newline at end of file + # The system_prompt is now handled by the base Agent class. + # Specific instructions can be added to the 'goal' in the config file. \ No newline at end of file diff --git a/app/agents/roles/reviewer_agent.py b/app/agents/roles/reviewer_agent.py index 9c41480..0bb95f3 100644 --- a/app/agents/roles/reviewer_agent.py +++ b/app/agents/roles/reviewer_agent.py @@ -1,9 +1,15 @@ """Reviewer Agent.""" -from app.agents.agent import Agent +from app.agents.agent import Agent, REACT_SYSTEM_PROMPT from app.agents.tools import ( read_file_tool, read_file_tool_def, ) +# Custom prompt for Reviewer +REVIEWER_SYSTEM_PROMPT = REACT_SYSTEM_PROMPT.replace( + "You are a smart, autonomous AI agent.", + "You are a Senior Software Engineer, a meticulous code reviewer." +) + class ReviewerAgent(Agent): """An agent specializing in strict Code Review.""" def __init__(self, name: str, role: str, goal: str, **kwargs): @@ -13,13 +19,9 @@ def __init__(self, name: str, role: str, goal: str, **kwargs): goal=goal, **kwargs, ) + self.system_prompt_template = REVIEWER_SYSTEM_PROMPT # Self-register tools self.add_tool(read_file_tool, read_file_tool_def) - self.system_prompt = ( - "You are a Senior Software Engineer acting as a code reviewer. " - "Your task is to provide a thorough review of the code based on the " - "provided file path. Use your available tools to read the file content.\n\n" - "If the code meets all standards, respond with only the word 'LGTM'.\n" - "Otherwise, provide clear, constructive feedback on what needs to be improved." - ) \ No newline at end of file + # The system_prompt is now handled by the base Agent class. + # Specific instructions can be added to the 'goal' in the config file. \ No newline at end of file diff --git a/app/agents/roles/standard_roles.py b/app/agents/roles/standard_roles.py new file mode 100644 index 0000000..723e462 --- /dev/null +++ b/app/agents/roles/standard_roles.py @@ -0,0 +1,43 @@ +""" +This module defines the standard roles for specialized agents in the system. +""" + +# Configuration for an agent specialized in file system operations +FILESYSTEM_EXPERT = { + "name": "FileSystemExpert", + "role": "An expert in browsing and reading files on a local file system.", + "goal": "To help users understand the project structure by listing and reading files.", + "tools": ["list_files", "read_file", "save_memory"] +} + +# Configuration for an agent specialized in web searching +WEB_SEARCH_EXPERT = { + "name": "WebSearchExpert", + "role": "An expert in searching the web for real-time information.", + "goal": "To find the most relevant and up-to-date information online in response to a user's query.", + "tools": ["web_search", "save_memory"] +} + +# Add other specialized agent configurations here as needed. +# For example, a CodeWriterAgent, a DatabaseExpert, etc. + +ALL_ROLES = { + "FileSystemExpert": FILESYSTEM_EXPERT, + "WebSearchExpert": WEB_SEARCH_EXPERT, +} + +# Configuration for the planner agent +PLANNER_AGENT = { + "name": "PlannerAgent", + "role": "A master planner who specializes in breaking down complex goals into a sequence of actionable steps for a team of specialized agents.", + "goal": "To create a clear, step-by-step JSON plan that efficiently leads to the user's desired outcome.", + "tools": [] # The planner does not use tools, it only thinks. +} + +# Configuration for the evaluator agent +EVALUATOR_AGENT = { + "name": "EvaluatorAgent", + "role": "A meticulous evaluator who analyzes multiple execution plans and selects the most optimal one.", + "goal": "To choose the most efficient, logical, and safe plan from a given set of options.", + "tools": [] # The evaluator only thinks and chooses. +} \ No newline at end of file diff --git a/app/agents/roles/task_decomposer.py b/app/agents/roles/task_decomposer.py index 3c2626c..867aa19 100644 --- a/app/agents/roles/task_decomposer.py +++ b/app/agents/roles/task_decomposer.py @@ -28,36 +28,40 @@ def __init__(self, name: str, role: str, goal: str, **kwargs): - EvaluatorAgent: Analyzes test failures and creates bug reports. # OUTPUT FORMAT: -Your output must be a list of steps in JSON format. Do not include any other text or explanation. +Your output MUST be a single JSON object with a single key "plan", which contains a list of steps. Do not include any other text, explanation, or markdown code fences. # EXAMPLE: Goal: "Create a function to add two numbers and test it." Your output: -```json -[ - { - "step": 1, - "assignee": "CodingAgent", - "task": "Create a new function 'add(a, b)' in 'app/utils/math.py'." - }, - { - "step": 2, - "assignee": "ReviewerAgent", - "task": "Review the 'add' function in 'app/utils/math.py'." - }, - { - "step": 3, - "assignee": "CodingAgent", - "task": "Create a new test file 'tests/test_math.py' with tests for the 'add' function." - }, - { - "step": 4, - "assignee": "TestingAgent", - "task": "Run the tests in 'tests/test_math.py'." - } -] -``` +{ + "plan": [ + { + "step": 1, + "assignee": "CodingAgent", + "task": "Create a new function 'add(a, b)' in 'app/utils/math.py'", + "description": "Implement the core logic for the addition function." + }, + { + "step": 2, + "assignee": "ReviewerAgent", + "task": "Review the 'add' function in 'app/utils/math.py'", + "description": "Ensure the code quality and correctness of the new function." + }, + { + "step": 3, + "assignee": "CodingAgent", + "task": "Create a new test file 'tests/test_math.py' with tests for the 'add' function", + "description": "Write unit tests to verify the behavior of the 'add' function." + }, + { + "step": 4, + "assignee": "TestingAgent", + "task": "Run the tests in 'tests/test_math.py'", + "description": "Execute the newly created tests to confirm the function works as expected." + } + ] +} """ def get_plan(self, goal: str) -> list: @@ -75,23 +79,25 @@ def get_plan(self, goal: str) -> list: {"role": "system", "content": self.system_prompt}, {"role": "user", "content": task_briefing}, ], - response_format={"type": "json_object"}, ) response_content = response.choices[0].message.content logging.info("Received raw plan: %s", response_content) # The response is a JSON string, so we need to parse it. - plan = json.loads(response_content) - # Sometimes the model returns a dictionary with a "plan" key - if isinstance(plan, dict) and "plan" in plan: - return plan["plan"] - return plan + parsed_json = json.loads(response_content) + + # The model should return a dict with a "plan" key + if isinstance(parsed_json, dict) and "plan" in parsed_json: + return parsed_json["plan"] + else: + logging.error("Invalid plan format. 'plan' key not found in response.") + raise ValueError("Invalid plan format.") except json.JSONDecodeError as e: logging.error(f"Failed to decode JSON from OpenAI response: {e}") logging.error(f"Raw response was: {response_content}") - return [{"step": 1, "assignee": "DefaultAgent", "task": "Failed to create a valid plan due to JSON error."}] + return [{"step": 1, "assignee": "DefaultAgent", "task": "Failed to create a valid plan due to JSON error.", "description": "The model returned malformed JSON."}] except Exception as e: logging.error(f"An unexpected error occurred while getting the plan: {e}", exc_info=True) - return [{"step": 1, "assignee": "DefaultAgent", "task": "Failed to create a plan due to an unexpected error."}] \ No newline at end of file + return [{"step": 1, "assignee": "DefaultAgent", "task": "Failed to create a plan due to an unexpected error.", "description": f"An unexpected error occurred: {e}"}] \ No newline at end of file diff --git a/app/agents/roles/testing_agent.py b/app/agents/roles/testing_agent.py index 6a01ba5..7a301ac 100644 --- a/app/agents/roles/testing_agent.py +++ b/app/agents/roles/testing_agent.py @@ -18,8 +18,4 @@ def __init__(self, name: str, role: str, goal: str, **kwargs): self.add_tool(read_file_tool, read_file_tool_def) self.add_tool(run_tests_tool, run_tests_tool_def) - self.system_prompt = """ -You are TestingAgent, an automated test-running robot. -Your sole task is to call the `run_tests_tool` with the correct path to the tests. -After receiving the result, you must briefly and accurately report on the success or failure. -""" \ No newline at end of file + # The system_prompt is now handled by the base Agent class. \ No newline at end of file diff --git a/app/agents/web_search_tool.py b/app/agents/web_search_tool.py new file mode 100644 index 0000000..b8fc67b --- /dev/null +++ b/app/agents/web_search_tool.py @@ -0,0 +1,60 @@ +""" +This module defines the web search tool for the agent. +""" +from typing import Dict, Any, List +from duckduckgo_search import DDGS + +# 1. Tool Definition (JSON Schema) +web_search_tool_def = { + "type": "function", + "function": { + "name": "web_search_tool", + "description": "Searches the web for up-to-date information on a given topic. " + "Use this to find current events, facts, or information not present in the knowledge base.", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to send to the search engine." + } + }, + "required": ["query"], + }, + }, +} + +# 2. Tool Function +def web_search_tool(args: Dict[str, Any]) -> str: + """ + Performs a web search using the DuckDuckGo search engine. + + Args: + args: A dictionary containing the 'query' for the search. + + Returns: + A formatted string containing the top search results, or an error message. + """ + query = args.get("query") + if not query: + return "Error: The 'query' parameter is required for web_search_tool." + + try: + with DDGS() as ddgs: + # max_results=5 to keep the context concise + results: List[Dict[str, str]] = list(ddgs.text(query, max_results=5)) + if not results: + return f"No results found for '{query}'." + + # Format the results for the LLM + formatted_results = [] + for i, result in enumerate(results): + formatted_results.append( + f"Result {i+1}:\n" + f" Title: {result.get('title')}\n" + f" Snippet: {result.get('body')}\n" + f" URL: {result.get('href')}\n" + ) + return "\n---\n".join(formatted_results) + except Exception as e: + return f"Error during web search for '{query}': {e}" \ No newline at end of file diff --git a/app/evaluation/evaluator.py b/app/evaluation/evaluator.py new file mode 100644 index 0000000..1956cf0 --- /dev/null +++ b/app/evaluation/evaluator.py @@ -0,0 +1,95 @@ +# app/evaluation/evaluator.py +import os +import json +from typing import Dict, Any + +from openai import OpenAI + +EVALUATION_SYSTEM_PROMPT = """ +You are an impartial AI evaluator. Your task is to assess the quality of an AI agent's response based on a given criterion and a reference answer. + +You must follow these instructions strictly: +1. Compare the "Actual Answer" to the "Reference Answer". +2. Evaluate the answer ONLY on the provided "Evaluation Criterion". +3. Provide a score from 1 to 5, where 1 is the worst and 5 is the best. +4. Provide a brief, single-sentence justification for your score. +5. Your final output MUST be a JSON object with two keys: "score" (int) and "justification" (str). + +Example: +Actual Answer: "The capital of France is Paris." +Reference Answer: "Paris is the capital of France." +Evaluation Criterion: "Correctness" + +Your output: +{ + "score": 5, + "justification": "The answer is factually correct and directly answers the question." +} +""" + +class CustomEvaluator: + """ + A simple, custom evaluator that uses an LLM to score an agent's response + against a reference answer based on a specific criterion. + """ + + def __init__(self, api_key: str = None): + api_key = api_key or os.getenv("OPENAI_API_KEY") + if not api_key: + raise ValueError("OpenAI API key is required.") + self.client = OpenAI(api_key=api_key) + + def evaluate( + self, + actual_answer: str, + reference_answer: str, + criterion: str, + model: str = "gpt-4o-mini", + ) -> Dict[str, Any]: + """ + Evaluates the agent's response. + + Args: + actual_answer: The answer generated by the agent. + reference_answer: The "golden" or expected answer. + criterion: The specific aspect to evaluate (e.g., "Correctness", "Clarity"). + model: The LLM to use for evaluation. + + Returns: + A dictionary containing the score and justification. + """ + if not all([actual_answer, reference_answer, criterion]): + raise ValueError("actual_answer, reference_answer, and criterion cannot be empty.") + + user_prompt = f""" + Actual Answer: "{actual_answer}" + Reference Answer: "{reference_answer}" + Evaluation Criterion: "{criterion}" + """ + + try: + response = self.client.chat.completions.create( + model=model, + messages=[ + {"role": "system", "content": EVALUATION_SYSTEM_PROMPT}, + {"role": "user", "content": user_prompt}, + ], + response_format={"type": "json_object"}, + temperature=0.0, + ) + evaluation_result = json.loads(response.choices[0].message.content) + + if "score" not in evaluation_result or "justification" not in evaluation_result: + raise KeyError("Evaluation response is missing 'score' or 'justification'.") + + return evaluation_result + + except KeyError as e: + print(f"Error processing evaluation keys: {e}") + return {"score": 0, "justification": f"Evaluation response is missing 'score' or 'justification'."} + except json.JSONDecodeError as e: + print(f"Error parsing evaluation response: {e}") + return {"score": 0, "justification": "Failed to parse the evaluation response from the LLM."} + except Exception as e: + print(f"An unexpected error occurred during evaluation: {e}") + return {"score": 0, "justification": "An unexpected error occurred."} \ No newline at end of file diff --git a/app/factory/agent_factory.py b/app/factory/agent_factory.py index d3e956c..36a148c 100644 --- a/app/factory/agent_factory.py +++ b/app/factory/agent_factory.py @@ -5,10 +5,21 @@ import yaml import os from importlib import import_module -from typing import Dict, Tuple +from typing import Dict, Tuple, Any, List from app.agents.agent import Agent from app.agents.roles.task_decomposer import TaskDecomposer +from app.agents.tools import read_file_tool, read_file_tool_def, list_files_tool, list_files_tool_def +from app.agents.web_search_tool import web_search_tool, web_search_tool_def +from app.agents.memory_tool import save_memory_tool, save_memory_tool_def + +# A mapping of tool names to their functions and definitions +AVAILABLE_TOOLS = { + "read_file": (read_file_tool, read_file_tool_def), + "list_files": (list_files_tool, list_files_tool_def), + "web_search": (web_search_tool, web_search_tool_def), + "save_memory": (save_memory_tool, save_memory_tool_def), +} def load_config(path: str) -> dict: """Loads a YAML configuration file.""" @@ -57,4 +68,50 @@ def create_agent_team(main_config_path: str) -> Tuple[TaskDecomposer, Dict[str, workers[agent_name] = agent_instance task_decomposer = workers.pop("TaskDecomposer") - return task_decomposer, workers \ No newline at end of file + return task_decomposer, workers + +class AgentFactory: + """ + A factory class for creating different types of specialized agents. + """ + @staticmethod + def create_agent( + agent_config: Dict[str, Any], + api_key: str, + model: str + ) -> Agent: + """ + Creates an agent based on a configuration dictionary. + + Args: + agent_config: A dictionary containing the agent's name, role, goal, + and a list of tool names it should have. + api_key: The OpenAI API key. + model: The name of the model to use. + + Returns: + An instance of the Agent class, configured as specified. + """ + name = agent_config.get("name", "SpecializedAgent") + role = agent_config.get("role", "An assistant") + goal = agent_config.get("goal", "To complete tasks efficiently.") + tool_names = agent_config.get("tools", []) + + agent = Agent( + name=name, + role=role, + goal=goal, + api_key=api_key, + model=model + ) + + # Clear default tools and add only the specified ones + agent.tools = {} + agent.tool_definitions = [] + + for tool_name in tool_names: + if tool_name in AVAILABLE_TOOLS: + tool_func, tool_def = AVAILABLE_TOOLS[tool_name] + agent.add_tool(tool_func, tool_def) + + return agent \ No newline at end of file diff --git a/app/memory/memory_manager.py b/app/memory/memory_manager.py new file mode 100644 index 0000000..7626d7e --- /dev/null +++ b/app/memory/memory_manager.py @@ -0,0 +1,79 @@ +""" +This module defines the MemoryManager for the agent, responsible for +handling long-term memory storage and retrieval using SQLite. +""" +import sqlite3 +import os +from typing import List, Dict, Any, Optional + +class MemoryManager: + """ + Manages the agent's long-term memory using a SQLite database. + """ + def __init__(self, db_path: str = "db/memory.db"): + """ + Initializes the MemoryManager. + + Args: + db_path: The path to the SQLite database file. + """ + os.makedirs(os.path.dirname(db_path), exist_ok=True) + self.conn = sqlite3.connect(db_path) + self.create_table() + + def create_table(self): + """Creates the memory table if it doesn't exist.""" + with self.conn: + self.conn.execute(""" + CREATE TABLE IF NOT EXISTS long_term_memory ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + fact TEXT NOT NULL, + timestamp DATETIME DEFAULT CURRENT_TIMESTAMP + ) + """) + + def add_fact(self, fact: str) -> Dict[str, Any]: + """ + Adds a new fact to the long-term memory. + + Args: + fact: The piece of information to remember. + + Returns: + A confirmation message. + """ + if not fact or not isinstance(fact, str): + return {"status": "error", "message": "Fact must be a non-empty string."} + + try: + with self.conn: + self.conn.execute( + "INSERT INTO long_term_memory (fact) VALUES (?)", + (fact,) + ) + return {"status": "success", "message": f"Fact '{fact}' was successfully saved."} + except sqlite3.Error as e: + return {"status": "error", "message": f"Database error: {e}"} + + def get_recent_facts(self, limit: int = 10) -> List[str]: + """ + Retrieves the most recent facts from memory. + + Args: + limit: The maximum number of facts to retrieve. + + Returns: + A list of recent facts. + """ + try: + cursor = self.conn.cursor() + cursor.execute("SELECT fact FROM long_term_memory ORDER BY timestamp DESC LIMIT ?", (limit,)) + facts = [row[0] for row in cursor.fetchall()] + return facts + except sqlite3.Error: + return [] + + def close(self): + """Closes the database connection.""" + if self.conn: + self.conn.close() \ No newline at end of file diff --git a/app/orchestration/orchestrator.py b/app/orchestration/orchestrator.py index 3fc5f01..f50a655 100644 --- a/app/orchestration/orchestrator.py +++ b/app/orchestration/orchestrator.py @@ -1,196 +1,221 @@ """ -Orchestrator module that manages the execution of a task plan. +This module defines the Orchestrator, the central brain of the multi-agent team. +It creates a plan and manages its execution by delegating tasks to specialized agents. """ -import logging import json -from typing import Dict, Any, List, Tuple, Optional -from app.agents.agent import Agent -from app.agents.roles.task_decomposer import TaskDecomposer +import logging +from typing import Dict, Any, List + +from openai import OpenAI +from app.factory.agent_factory import AgentFactory +from app.agents.roles.standard_roles import ALL_ROLES, PLANNER_AGENT, EVALUATOR_AGENT + +# Updated prompt to ask for multiple plans (Tree-of-Thoughts) +PLANNER_PROMPT_TEMPLATE = """ +You are a master planner for a team of AI agents. Your job is to create THREE DISTINCT step-by-step plans to accomplish the user's goal. + +**User's Goal:** +"{user_goal}" + +**Available Team of Specialists:** +{agents_description} + +Based on the goal, create a JSON object with a key "plans", containing a list of THREE different plans. Each plan is a JSON array of tasks. Each task must have: +- `step`: An integer for the step number (e.g., 1, 2, 3). +- `agent`: The name of the single most appropriate agent from the available team. +- `task`: A clear and specific instruction for the agent. + +**Important Rules:** +- The three plans should represent different strategies to achieve the goal. +- Your entire response MUST be a single, valid JSON object like: `{{"plans": [[...plan1...], [...plan2...], [...plan3...]]}}` + +**Example:** +{{ + "plans": [ + [ + {{"step": 1, "agent": "WebSearchExpert", "task": "Find official pytest docs."}} + ], + [ + {{"step": 1, "agent": "FileSystemExpert", "task": "Check existing test files for pytest usage examples."}} + ], + [ + {{"step": 1, "agent": "WebSearchExpert", "task": "Search for tutorials on how to use pytest with FastAPI."}} + ] + ] +}} +""" + +EVALUATOR_PROMPT_TEMPLATE = """ +You are a meticulous and rational Evaluator. Your task is to analyze a list of proposed plans and select the single best one to achieve the user's goal. + +**User's Goal:** +"{user_goal}" + +**Proposed Plans:** +{plans_json_string} + +**Evaluation Criteria:** +1. **Efficiency:** Which plan is likely to achieve the goal in the fewest steps? +2. **Robustness:** Which plan is least likely to fail or run into errors? +3. **Clarity:** Which plan is the most logical and straightforward? + +Based on your analysis, respond with a JSON object containing the index (starting from 0) of the best plan. + +**Example:** +{{ + "best_plan_index": 1 +}} +""" class Orchestrator: """ - The Orchestrator manages the entire process: from goal decomposition to task execution by agents. + Creates multiple plans, evaluates them, and manages the execution of the best one. """ - def __init__(self, task_decomposer: TaskDecomposer, worker_agents: Dict[str, Agent]): - self.task_decomposer = task_decomposer - self.worker_agents = worker_agents + def __init__(self, api_key: str, model: str = "gpt-4o-mini"): + self.client = OpenAI(api_key=api_key) + self.model = model + self.agent_factory = AgentFactory() + self.api_key = api_key self.execution_history: List[Dict[str, Any]] = [] - def _get_briefing(self, current_task: Dict[str, Any], plan: List[Dict[str, Any]], goal: str) -> str: - """Creates a full context (briefing) for an agent before task execution.""" + def _get_agents_description(self) -> str: + """Creates a formatted string describing the available worker agents.""" + descriptions = [] + for name, config in ALL_ROLES.items(): + descriptions.append(f"- Agent: {name}\n - Role: {config['role']}\n - Best for: {config['goal']}") + return "\n".join(descriptions) + + def _create_plan(self, user_goal: str) -> List[List[Dict[str, Any]]]: + """ + Uses the PlannerAgent's logic to create multiple distinct plans. + """ + logging.info(f"Orchestrator is creating multiple plans for the goal: '{user_goal}'") + agents_description = self._get_agents_description() + prompt = PLANNER_PROMPT_TEMPLATE.format( + user_goal=user_goal, + agents_description=agents_description + ) + try: + response = self.client.chat.completions.create( + model=self.model, + messages=[{"role": "user", "content": prompt}], + response_format={"type": "json_object"} + ) + # We expect a dict with a "plans" key, which is a list of lists + plan_variants = json.loads(response.choices[0].message.content or "{}").get("plans", []) + logging.info(f"Planner proposed {len(plan_variants)} plans.") + return plan_variants + except Exception as e: + logging.error(f"Failed to create plans: {e}", exc_info=True) + return [] + + def _evaluate_and_select_plan(self, user_goal: str, plans: List[List[Dict[str, Any]]]) -> List[Dict[str, Any]]: + """Uses the Evaluator's logic to select the best plan.""" + if not plans: + return [] + if len(plans) == 1: + logging.info("Only one plan was generated, selecting it by default.") + return plans[0] + + logging.info("Evaluating plans to select the best one...") + prompt = EVALUATOR_PROMPT_TEMPLATE.format( + user_goal=user_goal, + plans_json_string=json.dumps(plans, indent=2) + ) + + try: + response = self.client.chat.completions.create( + model=self.model, + messages=[{"role": "user", "content": prompt}], + response_format={"type": "json_object"} + ) + choice = json.loads(response.choices[0].message.content or "{}") + best_plan_index = choice.get("best_plan_index", 0) + + if 0 <= best_plan_index < len(plans): + logging.info(f"Evaluator selected plan #{best_plan_index + 1}.") + return plans[best_plan_index] + else: + logging.warning("Evaluator returned an invalid index, defaulting to the first plan.") + return plans[0] + except Exception as e: + logging.error(f"Failed to evaluate plans: {e}. Defaulting to the first plan.", exc_info=True) + return plans[0] + + def _create_briefing(self, goal: str, plan: List[Dict[str, Any]], current_task: Dict[str, Any]) -> str: + """Creates a detailed context (briefing) for an agent.""" history_log = "" if not self.execution_history: - history_log = "This is the first step, there are no previous results.\n" + history_log = "This is the first step. No execution history yet." else: - history_log += "**Execution History:**\n" + history_log = "Here are the results from the previous steps:\n" for record in self.execution_history: - history_log += f"- **Step {record['step']} ({record['assignee']})**: {record['task']}\n" - history_log += f" - **Result**: {record['result']}\n" - - plan_log = "" - for step in plan: - marker = "-->" if step['step'] == current_task['step'] else " " - plan_log += f"{marker} Step {step['step']}: {step['task']} (Assignee: {step['assignee']})\n" - - briefing = f""" -TASK CONTEXT ---------------------------------- -**Global Goal:** {goal} - -**ENTIRE TASK PLAN:** -{plan_log} -**EXECUTION HISTORY:** -{history_log} ---------------------------------- -**YOUR CURRENT TASK (Step {current_task['step']}):** - -**Task:** {current_task['task']} -**Description:** {current_task['description']} - -Analyze all the provided information, especially the results of previous steps, and execute your task. -""" - return briefing - - def _perform_code_review(self, task: Dict[str, Any], plan: List[Dict[str, Any]], goal: str) -> Tuple[bool, Optional[str]]: - """Performs the Code Review step.""" - logging.info("--- Starting Code Review ---") - reviewer_agent = self.worker_agents.get("ReviewerAgent") - if not reviewer_agent: - logging.error("ReviewerAgent not found!") - return False, "ReviewerAgent was not initialized." - - briefing = self._get_briefing(task, plan, goal) - review_result = reviewer_agent.execute_task(briefing) - - logging.info(f"Code Review Result: {review_result}") - - if "LGTM" in review_result.upper(): - logging.info("--- Code Review Passed Successfully ---") - return True, None - else: - logging.warning("--- Code Review identified issues ---") - return False, review_result + history_log += ( + f"Step {record['step']} ({record['agent']}) completed.\n" + f"Task: {record['task']}\n" + f"Result: {record['result']}\n\n" + ) - def run(self, goal: str): + return ( + f"**Overall Goal:** {goal}\n\n" + f"**Full Plan:** {json.dumps(plan, indent=2)}\n\n" + f"**Execution History:**\n{history_log}\n" + f"-----------------------------------\n" + f"**Your Current Task (Step {current_task['step']}):**\n" + f"Your task is: \"{current_task['task']}\".\n" + f"Analyze the goal, plan, and history, then execute your task to produce the required result." + ) + + def run(self, user_goal: str) -> str: + """ + Runs the full orchestration process: generate multiple plans, evaluate them, and execute the best one. + """ self.execution_history = [] - plan = self._get_plan(goal) - current_step_index = 0 - - while current_step_index < len(plan): - task = plan[current_step_index] - logging.info(f"--- Executing Step {task['step']}: {task['task']} ---") - - if task.get("assignee") == "ReviewerAgent": - review_passed, suggestions = self._perform_code_review(task, plan, goal) - - history_record = { - "step": task['step'], - "task": task['task'], - "assignee": task['assignee'], - } - - if review_passed: - history_record["result"] = "Success (LGTM)." - self.execution_history.append(history_record) - current_step_index += 1 - continue - else: - history_record["result"] = f"Failed. Feedback: {suggestions}" - self.execution_history.append(history_record) - - logging.warning("Review failed. Creating a task for correction...") - new_task_description = ( - "The reviewer agent found issues in the code you wrote. " - f"Here are the comments: '{suggestions}'. " - "Your task is to fix the code according to these recommendations." - ) - - new_task = { - "step": float(task["step"]) + 0.1, - "assignee": "CodingAgent", - "task": "Fix code based on Code Review feedback.", - "description": new_task_description, - } - - plan.insert(current_step_index + 1, new_task) - logging.info(f"A new task has been added to the plan: {new_task}") - - current_step_index += 1 - continue - - assignee_name = task.get("assignee", "DefaultAgent") - agent = self.worker_agents.get(assignee_name) + + # 1. Create multiple plan variants + plan_variants = self._create_plan(user_goal) + if not plan_variants: + return "I'm sorry, I couldn't create any plans to address your request. Please try rephrasing it." + + # 2. Evaluate and select the best plan + best_plan = self._evaluate_and_select_plan(user_goal, plan_variants) + if not best_plan: + return "I'm sorry, I couldn't select a valid plan to execute. Please try again." + + print(f"\033[95mOrchestrator's Selected Plan:\033[0m\n{json.dumps(best_plan, indent=2)}") + + # 3. Execute the selected plan step-by-step + for task in best_plan: + agent_name = task.get("agent") + task_description = task.get("task") - if not agent: - logging.error(f"Agent {assignee_name} not found! Skipping step.") - result = f"Skipped (agent '{assignee_name}' not found)." - else: - briefing = self._get_briefing(task, plan, goal) - result = agent.execute_task(briefing) + if not agent_name or not task_description or agent_name not in ALL_ROLES: + logging.warning(f"Skipping invalid task in plan: {task}") + continue - logging.info(f"Result of step {task['step']}: {result}") + logging.info(f"--- Executing Step {task['step']}: {task_description} (Agent: {agent_name}) ---") + + # Create the specialist agent + agent_config = ALL_ROLES[agent_name] + specialist_agent = self.agent_factory.create_agent(agent_config, self.api_key, self.model) - if "reached iteration limit" in result: - logging.critical(f"Agent {assignee_name} failed to complete task {task['step']} and reached the iteration limit. Execution aborted.") - print(f"CRITICAL ERROR: Agent {assignee_name} failed the task. Check agent_activity.log for details.") - return + # Create the briefing for the agent + briefing = self._create_briefing(user_goal, best_plan, task) - history_record = { + # Execute the task + result = specialist_agent.execute_task(briefing) + + # Save the result to history + self.execution_history.append({ "step": task['step'], - "task": task['task'], - "assignee": task.get("assignee", "DefaultAgent"), + "agent": agent_name, + "task": task_description, "result": result - } - self.execution_history.append(history_record) - - if agent and agent.name == "TestingAgent" and "FAIL" in result.upper(): - logging.error("Tests failed! Initiating fix process.") - evaluator_agent = self.worker_agents.get("EvaluatorAgent") - if not evaluator_agent: - logging.error("EvaluatorAgent not found, cannot analyze the error.") - current_step_index += 1 - continue - - evaluator_briefing = ( - "The tests failed. Analyze the following log and formulate a task to fix it.\n\n" - "ERROR LOG:\n" - f"{result}" - ) - - fix_task_description = evaluator_agent.execute_task(evaluator_briefing) - logging.info(f"EvaluatorAgent suggested the following task: {fix_task_description}") - - new_task = { - "step": float(task["step"]) + 0.1, - "assignee": "CodingAgent", - "task": "Fix code based on failed tests.", - "description": fix_task_description, - } - - plan.insert(current_step_index + 1, new_task) - logging.info(f"A new fix task has been added to the plan: {new_task}") - - current_step_index += 1 - - logging.info("All tasks completed. Goal achieved.") - - def _get_plan(self, goal: str) -> List[Dict[str, Any]]: - """Gets the plan from the TaskDecomposer.""" - logging.info("Getting plan from TaskDecomposer...") - plan = self.task_decomposer.get_plan(goal) - if not plan: - logging.error("Failed to create a plan. Orchestrator is stopping.") - return [] - - logging.info("Plan successfully received.") - print("\nThe following plan has been created:") - for step in plan: - print(f"- Step {step['step']}: {step['task']} (Assignee: {step['assignee']})") - return plan - - # def _create_evaluator_briefing(self, failed_test_result: str, last_briefing: str) -> str: - # """Создает системный промпт для EvaluatorAgent.""" - # return f""" - # ... (здесь был промпт) - # """ \ No newline at end of file + }) + + logging.info(f"--- Step {task['step']} Result: {result} ---") + + # Return the result of the final step + final_result = self.execution_history[-1]["result"] if self.execution_history else "The plan was executed, but there is no final result." + return final_result \ No newline at end of file diff --git a/app/rag/retriever.py b/app/rag/retriever.py index 5f857cf..604fc19 100644 --- a/app/rag/retriever.py +++ b/app/rag/retriever.py @@ -5,13 +5,14 @@ import os import pickle from pathlib import Path -from typing import List, Dict, Any, Tuple +from typing import List, Dict, Any, Tuple, Optional import numpy as np from dotenv import load_dotenv from openai import OpenAI from rank_bm25 import BM25Okapi from sklearn.metrics.pairwise import cosine_similarity +from sentence_transformers import CrossEncoder # --- Configuration --- load_dotenv() @@ -25,119 +26,103 @@ class KnowledgeRetriever: """ - A class to retrieve relevant knowledge chunks from a database - using a hybrid search approach (vector search + keyword search). + Handles retrieving information from the knowledge base using a hybrid approach + (BM25 for sparse keyword search and dense vector search) followed by a + re-ranking step with a Cross-Encoder model. """ def __init__(self): - """Initializes the retriever, loading all necessary data from disk.""" - self.embeddings: np.ndarray = None + self.embeddings: Optional[np.ndarray] = None self.chunks: List[Dict[str, Any]] = [] - self.bm25: BM25Okapi = None + self.bm25_index: Optional[BM25Okapi] = None + self.cross_encoder: Optional[CrossEncoder] = None + self.openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) self._load_knowledge_base() - - api_key = os.getenv("OPENAI_API_KEY") - if not api_key: - raise ValueError("OPENAI_API_KEY is not set for the retriever.") - self.client = OpenAI(api_key=api_key) def _load_knowledge_base(self): - """Loads embeddings, text chunks, and BM25 index from disk.""" - logging.info("Loading knowledge base...") + """Loads all necessary files for the retriever from the db directory.""" if not all([EMBEDDINGS_FILE.exists(), CHUNKS_FILE.exists(), BM25_INDEX_FILE.exists()]): - msg = ( - f"Knowledge base file not found. Please run " - f"'scripts/build_knowledge_base.py' first." - ) - logging.error(msg) - raise FileNotFoundError(msg) + logging.warning("Knowledge base files not found. Please run build_knowledge_base.py.") + return + logging.info("Loading knowledge base...") + self.embeddings = np.load(EMBEDDINGS_FILE) + with open(CHUNKS_FILE, "r", encoding="utf-8") as f: + self.chunks = json.load(f) + with open(BM25_INDEX_FILE, "rb") as f: + self.bm25_index = pickle.load(f) + + # Load the Cross-Encoder model try: - self.embeddings = np.load(EMBEDDINGS_FILE) - with open(CHUNKS_FILE, "r", encoding="utf-8") as f: - self.chunks = json.load(f) - with open(BM25_INDEX_FILE, "rb") as f: - self.bm25 = pickle.load(f) - logging.info( - f"Knowledge base loaded successfully. " - f"({len(self.chunks)} chunks, BM25 index, Embeddings)" - ) + self.cross_encoder = CrossEncoder('cross-encoder/ms-marco-minilm-l-6-v2', max_length=512) + logging.info("Cross-Encoder model loaded successfully.") except Exception as e: - logging.error(f"Failed to load knowledge base: {e}") - raise - - def _vector_search(self, query: str, k: int) -> List[Tuple[int, float]]: - """Performs a pure vector search.""" - response = self.client.embeddings.create(input=[query], model=EMBEDDING_MODEL) - query_embedding = np.array([response.data[0].embedding]) - similarities = cosine_similarity(query_embedding, self.embeddings).flatten() - top_k_indices = np.argsort(similarities)[-k:][::-1] - return [(idx, similarities[idx]) for idx in top_k_indices] + logging.error(f"Failed to load Cross-Encoder model: {e}", exc_info=True) + self.cross_encoder = None + + logging.info(f"Knowledge base loaded successfully ({len(self.chunks)} chunks).") - def _keyword_search(self, query: str, k: int) -> List[Tuple[int, float]]: - """Performs a pure keyword search using BM25.""" - tokenized_query = query.lower().split(" ") - doc_scores = self.bm25.get_scores(tokenized_query) - top_k_indices = np.argsort(doc_scores)[-k:][::-1] - return [(idx, doc_scores[idx]) for idx in top_k_indices] + def _get_embedding(self, text: str) -> np.ndarray: + """Helper to get embedding for a single text query.""" + response = self.openai_client.embeddings.create(input=[text], model="text-embedding-3-small") + return np.array(response.data[0].embedding) - def _reciprocal_rank_fusion(self, search_results: List[List[Tuple[int, float]]], k: int = 60) -> Dict[int, float]: - """Merges search results using Reciprocal Rank Fusion.""" - fused_scores = {} - for result_list in search_results: - for rank, (doc_id, _) in enumerate(result_list): - if doc_id not in fused_scores: - fused_scores[doc_id] = 0 - fused_scores[doc_id] += 1 / (rank + k) - return fused_scores + def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray: + """Calculates cosine similarity between a vector and a matrix of vectors.""" + vec1 = vec1.reshape(1, -1) + return (vec1 @ vec2.T) / (np.linalg.norm(vec1) * np.linalg.norm(vec2, axis=1)) - def retrieve( - self, query: str, top_k: int = 5, filters: Dict[str, Any] = None - ) -> List[Dict[str, Any]]: - """ - Retrieves the top_k most relevant chunks using hybrid search. - """ - if self.embeddings is None or not self.chunks or not self.bm25: - logging.warning("Knowledge base is not fully loaded. Cannot retrieve.") + def retrieve(self, query: str, top_k: int = 5, filters: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: + if not self.chunks or self.bm25_index is None or self.embeddings is None: + logging.warning("Knowledge base is not loaded. Cannot retrieve.") return [] - # 1. Get results from both search methods - candidate_k = top_k * 10 - vector_results = self._vector_search(query, candidate_k) - keyword_results = self._keyword_search(query, candidate_k) - - # 2. Fuse the results - fused_scores = self._reciprocal_rank_fusion([vector_results, keyword_results]) + # 1. Sparse Search (BM25) + tokenized_query = query.lower().split() + bm25_scores = self.bm25_index.get_scores(tokenized_query) - # 3. Sort by fused score - sorted_doc_ids = sorted(fused_scores.keys(), key=lambda id: fused_scores[id], reverse=True) - - # 4. Filter candidates based on metadata + # 2. Dense Search (Vector) + query_embedding = self._get_embedding(query) + embedding_scores = self._cosine_similarity(query_embedding, self.embeddings).flatten() + + # 3. Hybrid Scoring (Combine and rank initial candidates) + # Normalize scores to be in a similar range before combining + norm_bm25_scores = bm25_scores / (np.max(bm25_scores) + 1e-9) + norm_embedding_scores = embedding_scores / (np.max(embedding_scores) + 1e-9) + combined_scores = 0.4 * norm_bm25_scores + 0.6 * norm_embedding_scores + + # Get a larger pool of candidates for re-ranking + candidate_pool_size = min(len(self.chunks), 25) + top_candidate_indices = np.argsort(combined_scores)[-candidate_pool_size:][::-1] + + # 4. Re-ranking with Cross-Encoder + if self.cross_encoder: + cross_encoder_pairs = [[query, self.chunks[i]["text"]] for i in top_candidate_indices] + rerank_scores = self.cross_encoder.predict(cross_encoder_pairs, show_progress_bar=False) + + # Sort candidate indices based on the new re-ranking scores + reranked_indices = [idx for _, idx in sorted(zip(rerank_scores, top_candidate_indices), reverse=True)] + final_indices_unfiltered = reranked_indices + logging.info(f"Re-ranked {len(top_candidate_indices)} candidates.") + else: + logging.warning("Cross-encoder not available. Falling back to simple hybrid search.") + final_indices_unfiltered = top_candidate_indices + + # 5. Filter results based on metadata final_indices = [] required_tags = set(filters.get("tags", [])) if filters else set() - - for doc_id in sorted_doc_ids: + for idx in final_indices_unfiltered: if len(final_indices) >= top_k: break - if required_tags: - chunk_tags = set(self.chunks[doc_id].get("metadata", {}).get("tags", [])) - if required_tags.issubset(chunk_tags): - final_indices.append(doc_id) - else: - final_indices.append(doc_id) - - # 5. Format results - results = [] - for idx in final_indices: - result = { - "text": self.chunks[idx]["text"], - "source": self.chunks[idx]["source"], - "metadata": self.chunks[idx].get("metadata", {}), - "score": fused_scores[idx], # Use the fused score - } - results.append(result) - - if not results: - logging.warning(f"No documents found for query '{query[:50]}...' with filters {filters}") - + chunk_tags = set(self.chunks[idx].get("metadata", {}).get("tags", [])) + if not required_tags or required_tags.issubset(chunk_tags): + final_indices.append(idx) + + # 6. Format and return final results + results = [{ + "text": self.chunks[i]["text"], + "source": self.chunks[i]["source"], + "metadata": self.chunks[i].get("metadata", {}), + } for i in final_indices] + return results \ No newline at end of file diff --git a/app/safety/custom_guardrails.py b/app/safety/custom_guardrails.py new file mode 100644 index 0000000..8a782ac --- /dev/null +++ b/app/safety/custom_guardrails.py @@ -0,0 +1,112 @@ +# app/safety/custom_guardrails.py +import os +import json +from typing import List, Dict, Any, Optional + +from openai import OpenAI + +GUARDRAIL_SYSTEM_PROMPT = """ +You are a content safety moderator. Your task is to determine if a given text violates any of the provided safety policies. + +You must follow these instructions strictly: +1. Read the "Agent's Answer" and the "Forbidden Topics". +2. Determine if the answer discusses any of the forbidden topics. +3. Your final output MUST be a JSON object with two keys: + - "is_safe" (boolean): `true` if the answer is safe, `false` if it violates a policy. + - "reason" (string): A brief, single-sentence justification for your decision. If the answer is safe, the reason should be "The answer is compliant with all policies." + +Example 1: +Agent's Answer: "To build a simple web server, you can use Python's Flask library." +Forbidden Topics: ["Financial advice", "Medical advice", "Illegal activities"] + +Your output: +{ + "is_safe": true, + "reason": "The answer is compliant with all policies." +} + +Example 2: +Agent's Answer: "You should consider investing all your money in the latest cryptocurrency." +Forbidden Topics: ["Financial advice", "Medical advice", "Illegal activities"] + +Your output: +{ + "is_safe": false, + "reason": "The answer provides financial advice, which is a forbidden topic." +} +""" + +DEFAULT_FORBIDDEN_TOPICS = [ + "Financial advice", + "Medical advice", + "Illegal activities or promotion of dangerous acts", + "Hate speech or harassment", + "Explicit or adult content", +] + +class CustomGuardrailManager: + """ + A simple, custom guardrail that uses an LLM to check if an agent's + response violates a set of predefined policies. + """ + + def __init__(self, api_key: str = None, forbidden_topics: Optional[List[str]] = None): + api_key = api_key or os.getenv("OPENAI_API_KEY") + if not api_key: + raise ValueError("OpenAI API key is required.") + self.client = OpenAI(api_key=api_key) + self.forbidden_topics = forbidden_topics or DEFAULT_FORBIDDEN_TOPICS + + def moderate(self, answer: str, model: str = "gpt-4o-mini") -> Dict[str, Any]: + """ + Moderates the agent's final answer. + + Args: + answer: The final answer generated by the agent. + model: The LLM to use for moderation. + + Returns: + A dictionary containing the safety check result. + """ + user_prompt = f""" + Agent's Answer: "{answer}" + Forbidden Topics: {json.dumps(self.forbidden_topics)} + """ + + try: + response = self.client.chat.completions.create( + model=model, + messages=[ + {"role": "system", "content": GUARDRAIL_SYSTEM_PROMPT}, + {"role": "user", "content": user_prompt}, + ], + response_format={"type": "json_object"}, + temperature=0.0, + ) + moderation_result = json.loads(response.choices[0].message.content) + + if "is_safe" not in moderation_result or "reason" not in moderation_result: + raise KeyError("Moderation response is missing 'is_safe' or 'reason'.") + + return moderation_result + + except (json.JSONDecodeError, KeyError) as e: + print(f"Error parsing moderation response: {e}") + # Fail safe: if we can't parse the response, assume it's not safe. + return {"is_safe": False, "reason": "Failed to parse the moderation response from the LLM."} + except Exception as e: + print(f"An unexpected error occurred during moderation: {e}") + return {"is_safe": False, "reason": "An unexpected error occurred."} + + def validate_and_format_response(self, final_answer: str) -> str: + """ + Checks the final answer against guardrails and returns it if safe, + or a fallback message if not. + """ + moderation_result = self.moderate(final_answer) + if moderation_result.get("is_safe", False): + return final_answer + else: + reason = moderation_result.get("reason", "an unknown policy.") + print(f"Blocked response due to: {reason}") + return "I am sorry, but I cannot provide a response on this topic. It violates my safety guidelines." \ No newline at end of file diff --git a/db/memory.db b/db/memory.db new file mode 100644 index 0000000..551fff0 Binary files /dev/null and b/db/memory.db differ diff --git a/docs/01_multi_agent_architecture.md b/docs/01_multi_agent_architecture.md new file mode 100644 index 0000000..ffd888e --- /dev/null +++ b/docs/01_multi_agent_architecture.md @@ -0,0 +1,30 @@ +# Архитектура многоагентной системы (Multi-Agent Architecture) + +В основе **PyAgentX** лежит отказ от идеи единого "агента-швейцарского ножа" в пользу **команды узкоспециализированных экспертов**. Такой подход повышает надежность, общую производительность системы и значительно упрощает её дальнейшее расширение и поддержку. + +Центральным элементом этой архитектуры является **Оркестратор**, который управляет всей командой. + +### Компоненты архитектуры + +1. **Оркестратор (`Orchestrator`)**: "Мозг" и project manager всей команды. Он не выполняет задачи сам, а получает цель от пользователя и управляет процессом ее достижения. Его ключевая функция — реализация продвинутого цикла планирования, вдохновленного подходом **"Дерево мыслей" (Tree-of-Thoughts, ToT)**. + +2. **Планировщик (`PlannerAgent`)**: "Стратег" команды. Получив высокоуровневую цель от Оркестратора, его задача — декомпозировать её на **несколько различных, многошаговых планов**. Это позволяет системе исследовать несколько потенциальных путей решения, а не придерживаться первого же сгенерированного варианта. + +3. **Оценщик (`EvaluatorAgent`)**: "Критик" и Quality Assurance в команде. Он получает несколько планов от Планировщика и анализирует их по заданным критериям (например, эффективность, надежность, простота, соответствие задаче). Его задача — выбрать **самый оптимальный план** и вернуть его Оркестратору. Этот шаг гарантирует, что к выполнению будет принята наилучшая из доступных стратегий. + +4. **Агенты-специалисты (`Specialized Agents`)**: "Исполнители" с узкой зоной ответственности. Каждый такой агент обладает ограниченным, но мощным набором инструментов для своей роли. Примеры: + * `FileSystemExpert`: отвечает за чтение, запись и анализ файлов в проекте. + * `WebSearchExpert`: отвечает за поиск актуальной информации в интернете. + Такая специализация делает их действия предсказуемыми, качественными и легко тестируемыми. + +### Жизненный цикл запроса + +1. **Пользователь** ставит высокоуровневую задачу. +2. **Оркестратор** принимает задачу и передает её **Планировщику**. +3. **Планировщик** генерирует 3-5 различных планов решения и отдает их **Оценщику**. +4. **Оценщик** выбирает лучший план и возвращает его **Оркестратору**. +5. **Оркестратор** последовательно выполняет шаги из выбранного плана, на каждом шаге делегируя задачу наиболее подходящему **Агенту-специалисту**. +6. **Агент-специалист** выполняет свою часть работы и возвращает результат Оркестратору. +7. Процесс повторяется, пока все шаги плана не будут выполнены. + +[Вернуться к README](../README.md) \ No newline at end of file diff --git a/docs/02_agent_core_loop.md b/docs/02_agent_core_loop.md new file mode 100644 index 0000000..2635cc8 --- /dev/null +++ b/docs/02_agent_core_loop.md @@ -0,0 +1,27 @@ +# Ядро агента: Цикл ReAct и Рефлексия + +Каждый агент-исполнитель в системе **PyAgentX** работает на основе мощного и гибкого внутреннего цикла, который делает его действия прозрачными, логичными и устойчивыми к ошибкам. + +### 1. ReAct (Reason + Act) + +Мы используем подход **ReAct**, который заставляет агента действовать не импульсивно, а обдуманно. Вместо простого выполнения команды, агент следует циклу **"Мысль → Действие → Наблюдение"**: + +1. **Мысль (Reasoning)**: Получив задачу, агент сначала генерирует рассуждение. Он анализирует цель, определяет, какой инструмент из своего арсенала лучше всего подходит для текущего шага, и в каком формате нужно подготовить для него входные данные. Эта мысль выводится в лог, что делает процесс принятия решений полностью прозрачным. + +2. **Действие (Action)**: Агент выполняет вызов выбранного инструмента с подготовленными параметрами. Это может быть чтение файла, поиск в интернете или сохранение информации в память. + +3. **Наблюдение (Observation)**: Агент получает результат выполнения инструмента. Это может быть как успешный результат (содержимое файла, поисковая выдача), так и сообщение об ошибке. + +Этот цикл повторяется до тех пор, пока агент не выполнит поставленную перед ним подзадачу. + +### 2. Самокоррекция (Self-Correction / Reflection) + +Одной из самых сильных сторон наших агентов является их способность к **самокоррекции**. Если на шаге "Наблюдение" инструмент возвращает ошибку, агент не прекращает работу, а входит в специальный цикл **рефлексии**: + +1. **Анализ ошибки**: Агент получает полный текст ошибки. +2. **Рефлексия**: Он "задумывается" над причиной неудачи. ("Я пытался прочитать файл, которого не существует. Возможно, я ошибся в названии или пути. Мне следует сначала проверить список файлов в директории.") +3. **Коррекция плана**: На основе этого анализа агент корректирует свой следующий шаг. Вместо того чтобы снова пытаться выполнить провальное действие, он генерирует новую "Мысль", нацеленную на решение проблемы (например, вызвать инструмент для листинга файлов). + +Этот механизм делает систему значительно более надежной и автономной, позволяя ей справляться с непредвиденными ситуациями без вмешательства человека. + +[Вернуться к README](../README.md) \ No newline at end of file diff --git a/docs/03_advanced_rag.md b/docs/03_advanced_rag.md new file mode 100644 index 0000000..bcd7449 --- /dev/null +++ b/docs/03_advanced_rag.md @@ -0,0 +1,33 @@ +# Продвинутый RAG-пайплайн (Advanced RAG) + +Для того чтобы агенты **PyAgentX** могли принимать решения, основанные на актуальных данных и контексте проекта, мы реализовали многоступенчатый RAG-пайплайн (Retrieval-Augmented Generation). Он нацелен на извлечение максимально релевантной информации из базы знаний. + +Пайплайн состоит из трех ключевых этапов: + +### 1. Семантическое чанкирование (Semantic Chunking) + +На этапе построения базы знаний (`scripts/build_knowledge_base.py`) мы отказались от наивного разделения текста на фрагменты (чанки) фиксированной длины. Вместо этого мы используем **семантический подход** с помощью библиотеки `semchunk`. + +Алгоритм анализирует семантическую близость между соседними предложениями и "разрезает" текст в тех местах, где происходит смена темы. Это позволяет нам получать на выходе осмысленные, логически завершенные чанки, которые содержат целостный контекст. Такой подход значительно повышает качество информации, которую мы сможем извлечь на последующих этапах. + +### 2. Двухэтапный гибридный поиск (Two-stage Hybrid Retrieval) + +Когда агенту требуется информация, наш ретривер (`app/rag/retriever.py`) запускает двухэтапный поиск для нахождения наиболее релевантных чанков: + +1. **Этап 1: Быстрый отбор кандидатов**. Мы одновременно используем два метода поиска, чтобы получить широкий пул потенциально релевантных документов: + * **Dense (Векторный) поиск**: Находит чанки, близкие к запросу по **смыслу**, даже если в них нет точных ключевых слов. Идеально подходит для поиска концепций и идей. + * **Sparse (Ключевой) поиск (BM25)**: Находит чанки, содержащие точные **ключевые слова** из запроса. Идеально подходит для поиска конкретных терминов, функций или имен переменных. + +Результаты обоих поисков объединяются для формирования списка из ~25 лучших кандидатов. + +### 3. Переранжирование с помощью Cross-Encoder (Cross-Encoder Re-ranking) + +2. **Этап 2: Точное переранжирование**. Получив список кандидатов, мы передаем его на вход более мощной, но и более медленной модели — **Cross-Encoder** (`sentence-transformers`). + +В отличие от стандартных моделей, которые создают векторы для запроса и чанка по отдельности, Cross-Encoder получает на вход **пару (запрос, чанк)** и оценивает степень их релевантности напрямую. Это позволяет ему улавливать гораздо более тонкие смысловые связи. + +Модель детально анализирует каждого кандидата в контексте исходного запроса и пересортировывает список, помещая на самые верхние позиции наиболее релевантные чанки. + +В итоге агент получает не просто набор разрозненных фактов, а короткий, но максимально концентрированный и релевантный контекст для принятия решения. + +[Вернуться к README](../README.md) \ No newline at end of file diff --git a/docs/04_evaluation_and_safety.md b/docs/04_evaluation_and_safety.md new file mode 100644 index 0000000..56219c6 --- /dev/null +++ b/docs/04_evaluation_and_safety.md @@ -0,0 +1,26 @@ +# Оценка качества и Безопасность +Для создания надежной и предсказуемой AI-системы мы реализовали собственные, легковесные модули для оценки качества и модерации контента. Такой подход избавляет нас от тяжелых внешних зависимостей и дает полный контроль над процессом. + +### 1. Собственный фреймворк Оценки (`app/evaluation`) +Вместо интеграции громоздких библиотек типа `deepeval`, мы создали `CustomEvaluator`. Этот модуль использует LLM для выполнения объективной оценки ответов агента. + +**Как это работает:** +- Мы подаем на вход `CustomEvaluator` три вещи: фактический ответ агента, "идеальный" (эталонный) ответ и критерий оценки (например, "Корректность" или "Следование инструкциям"). +- Evaluator формирует специальный промпт для LLM, прося его выступить в роли беспристрастного судьи и выставить оценку по шкале от 1 до 5. +- LLM возвращает JSON с оценкой и кратким обоснованием. + +Это позволяет нам создавать гибкие тесты для проверки производительности агентов на конкретных задачах, не привязываясь к сложным внешним инструментам. + +### 2. Собственные "Ограждения" (Guardrails) (`app/safety`) +Аналогично, вместо `guardrails-ai` мы реализовали `CustomGuardrailManager` для контроля за безопасностью ответов. + +**Как это работает:** +- В менеджере определен список запрещенных тем (например, финансовые, медицинские советы, разжигание ненависти). +- Перед тем, как отдать финальный ответ пользователю, агент передает его в `CustomGuardrailManager`. +- Менеджер с помощью LLM определяет, не нарушает ли ответ заданные правила. +- Если ответ безопасен, он передается пользователю. Если нет — ответ блокируется, и пользователь получает стандартное сообщение о том, что тема не может быть обсуждена. + +Такой подход обеспечивает базовый, но надежный уровень безопасности и контроля над контентом, генерируемым системой. + +--- +[<-- Назад к README](../README.md) \ No newline at end of file diff --git a/docs/04_long_term_memory.md b/docs/04_long_term_memory.md new file mode 100644 index 0000000..2dddc5d --- /dev/null +++ b/docs/04_long_term_memory.md @@ -0,0 +1,23 @@ +# Долгосрочная память (Long-Term Memory) + +В отличие от стандартной краткосрочной памяти, которая хранит только историю текущего диалога, система **PyAgentX** наделена **долгосрочной памятью**. Это позволяет агентам накапливать знания между сессиями, "учиться" на прошлом опыте и использовать его в будущих задачах. + +### Реализация + +Механизм долгосрочной памяти реализован с помощью двух ключевых компонентов: + +1. **Менеджер памяти (`app/memory/memory_manager.py`)**: Это специализированный класс, который инкапсулирует всю логику работы с базой данных памяти. Он использует простую и надежную базу данных **SQLite** для хранения фактов. Каждая запись в базе представляет собой отдельный "факт" или "воспоминание". Менеджер предоставляет простые методы для сохранения новых фактов и поиска релевантных старых фактов. + +2. **Инструмент для сохранения памяти (`save_memory_tool`)**: Мы не заставляем агента сохранять все подряд. Вместо этого мы предоставляем ему специальный инструмент. Агент в ходе своей работы (`ReAct-цикла`) может сам принять решение, что какая-то часть информации (например, вывод успешной команды, важный факт из поиска или ключевое решение в задаче) является достаточно ценной для сохранения на будущее. В этом случае он вызывает `save_memory_tool` и передает ему текст, который необходимо запомнить. + +### Как это работает? + +1. **Сохранение**: Во время выполнения задачи агент сталкивается с важной информацией. В своей "мысли" он решает: "Этот факт важен, его нужно сохранить". Затем он вызывает инструмент `save_memory_tool` с этой информацией. Менеджер памяти записывает этот факт в базу данных SQLite. + +2. **Извлечение**: В начале каждой новой задачи Оркестратор автоматически делает запрос к долгосрочной памяти. Он ищет в базе факты, которые семантически близки к новой поставленной цели. + +3. **Использование**: Найденные "воспоминания" добавляются в системный промпт агента, который будет выполнять задачу. Таким образом, еще до начала работы агент получает релевантный контекст из прошлого опыта, что помогает ему принимать более информированные и эффективные решения. + +Этот подход позволяет системе не решать одну и ту же проблему дважды и со временем становиться все более "опытной". + +[Вернуться к README](../README.md) \ No newline at end of file diff --git a/main.py b/main.py index 958f25f..0335b92 100644 --- a/main.py +++ b/main.py @@ -1,69 +1,55 @@ -import logging +""" +Main entry point for the multi-agent system. +""" import os -import typer +import logging from dotenv import load_dotenv -from typing_extensions import Annotated - -from app.factory.agent_factory import create_agent_team from app.orchestration.orchestrator import Orchestrator -# Create a typer app -app = typer.Typer() - -def setup_logging(): - """Configures the logging for the application.""" - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s", - handlers=[ - logging.FileHandler("agent_activity.log", mode='w'), - logging.StreamHandler() - ] - ) - -@app.command() -def run( - goal: Annotated[str, typer.Argument(help="The high-level goal for the agent team to accomplish.")], - config: Annotated[str, typer.Option(help="Path to the main configuration file.")] = "configs/config.yaml" -): +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + handlers=[ + logging.StreamHandler(), + logging.FileHandler("agent_activity.log", mode='w') + ] +) + +def main(): """ - Runs the multi-agent system to accomplish a given goal. + Initializes and runs the agent orchestration system. """ load_dotenv() - setup_logging() - api_key = os.getenv("OPENAI_API_KEY") if not api_key: - logging.error("OPENAI_API_KEY not found in .env file.") - typer.echo("Error: Please make sure your .env file contains OPENAI_API_KEY.") - raise typer.Exit(code=1) + logging.error("FATAL: OPENAI_API_KEY environment variable not set.") + return - try: - # Create the agent team using the factory - task_decomposer, worker_agents = create_agent_team(config) + model = "gpt-4o-mini" + + print("--- Multi-Agent System Initialized ---") + print("Enter your request. Type 'exit' to close.") - # Initialize and run the orchestrator - orchestrator = Orchestrator( - task_decomposer=task_decomposer, - worker_agents=worker_agents - ) - - typer.echo(f"🚀 Starting agent team to accomplish goal: {goal}") - orchestrator.run(goal) - typer.echo("✅ Goal accomplished successfully!") + orchestrator = Orchestrator(api_key=api_key, model=model) - except FileNotFoundError as e: - logging.error(f"Configuration file not found: {e}") - typer.echo(f"Error: Configuration file not found at '{config}'. Please check the path.") - raise typer.Exit(code=1) - except KeyboardInterrupt: - typer.echo("\nOperation cancelled by user. Exiting...") - logging.info("User cancelled the operation.") - raise typer.Exit() - except Exception as e: - logging.critical("A critical error occurred: %s", e, exc_info=True) - typer.echo(f"\n🚨 Critical Error: {e}") - raise typer.Exit(code=1) + try: + while True: + user_query = input("\033[94mYou > \033[0m") + if user_query.lower() in ['exit', 'quit']: + print("Shutting down...") + break + + if not user_query.strip(): + continue + + # The orchestrator handles the entire process + final_result = orchestrator.run(user_query) + + print(f"\033[92mFinal Answer >\033[0m {final_result}") + + except (KeyboardInterrupt, EOFError): + print("\nShutting down...") if __name__ == "__main__": - app() \ No newline at end of file + main() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 3bf24f0..4c76ac6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,18 +1,22 @@ fastapi==0.111.1 uvicorn==0.30.1 python-dotenv==1.0.1 -openai==1.35.13 +openai~=1.10.0 tiktoken unstructured unstructured-client markdown==3.6 numpy scikit-learn -sentence-transformers torch pytest==8.3.2 pytest-mock==3.14.0 httpx==0.27.0 rank-bm25 pyyaml -typer[all] \ No newline at end of file +typer[all] +pydantic +python-multipart +duckduckgo-search +semchunk +sentence-transformers \ No newline at end of file diff --git a/scripts/build_knowledge_base.py b/scripts/build_knowledge_base.py index 820b384..acf4692 100644 --- a/scripts/build_knowledge_base.py +++ b/scripts/build_knowledge_base.py @@ -14,6 +14,7 @@ from openai import OpenAI from unstructured.partition.md import partition_md from rank_bm25 import BM25Okapi +from semchunk.chunker import SemanticChunker # --- Configuration --- load_dotenv() @@ -38,6 +39,17 @@ client = OpenAI(api_key=api_key) tokenizer = tiktoken.get_encoding("cl100k_base") +# Initialize the semantic chunker +semantic_chunker = SemanticChunker( + embed_model=client, + model_name=EMBEDDING_MODEL, + max_chunk_size=TEXT_CHUNK_MAX_TOKENS, + # The 'breakpoint_percentile_threshold' is a key parameter to tune. + # It determines how different two sentences must be to create a split. + # Lower value = more splits, higher value = fewer splits. + # Let's start with a value that often works well. + breakpoint_percentile_threshold=90 +) def load_and_partition_documents(directory: Path) -> List[Dict[str, Any]]: """ @@ -71,48 +83,30 @@ def load_and_partition_documents(directory: Path) -> List[Dict[str, Any]]: logging.error(f"Failed to process {file_path}: {e}") return documents -def chunk_text( - doc: Dict[str, Any], - max_tokens: int = TEXT_CHUNK_MAX_TOKENS, -) -> List[Dict[str, Any]]: - """Splits a document into smaller, semantically meaningful chunks.""" - chunks = [] - doc_text = doc["text"] - source_name = doc["source"] - metadata = doc["metadata"] - - header_splits = re.split(r'(^## .+$|^### .+$)', doc_text, flags=re.MULTILINE) - texts_to_process = [header_splits[0]] - if len(header_splits) > 1: - texts_to_process.extend([h + c for h, c in zip(header_splits[1::2], header_splits[2::2])]) - - chunk_id_counter = 0 - for text_block in texts_to_process: - if not text_block.strip(): - continue - paragraphs = [p.strip() for p in text_block.split('\n\n') if p.strip()] - for paragraph in paragraphs: - sentences = re.split(r'(?<=[.!?])\s+', paragraph) - current_chunk = "" - for sentence in sentences: - if not sentence: - continue - if len(tokenizer.encode(current_chunk + " " + sentence)) <= max_tokens: - current_chunk += " " + sentence - else: - if current_chunk: - # Add metadata to each chunk - chunk_metadata = metadata.copy() - chunk_metadata["chunk_id"] = f"{source_name}_{chunk_id_counter}" - chunks.append({"text": current_chunk.strip(), "source": source_name, "metadata": chunk_metadata}) - chunk_id_counter += 1 - current_chunk = sentence - if current_chunk: - chunk_metadata = metadata.copy() - chunk_metadata["chunk_id"] = f"{source_name}_{chunk_id_counter}" - chunks.append({"text": current_chunk.strip(), "source": source_name, "metadata": chunk_metadata}) - chunk_id_counter += 1 - return chunks +def semantic_chunk_document(doc: Dict[str, Any]) -> List[Dict[str, Any]]: + """ + Splits a document into semantically coherent chunks using semchunk. + """ + logging.info(f"Semantically chunking '{doc['source']}'...") + try: + # The chunker returns a list of text strings + chunk_texts = semantic_chunker.chunk(doc["text"]) + + chunks_with_metadata = [] + for i, chunk_text in enumerate(chunk_texts): + chunk_metadata = doc["metadata"].copy() + chunk_metadata["chunk_id"] = f"{doc['source']}_{i}" + chunks_with_metadata.append({ + "text": chunk_text, + "source": doc["source"], + "metadata": chunk_metadata + }) + + logging.info(f"Successfully created {len(chunks_with_metadata)} semantic chunks for '{doc['source']}'.") + return chunks_with_metadata + except Exception as e: + logging.error(f"Failed to chunk document {doc['source']}: {e}", exc_info=True) + return [] def create_embeddings(texts: List[str]) -> np.ndarray: """Creates embeddings for a list of texts using OpenAI API.""" @@ -132,7 +126,8 @@ def main(): return all_chunks_with_metadata = [] for doc in documents: - doc_chunks = chunk_text(doc) + # Use the new semantic chunking function + doc_chunks = semantic_chunk_document(doc) all_chunks_with_metadata.extend(doc_chunks) if not all_chunks_with_metadata: logging.warning("Could not create any chunks from the documents. Exiting.") diff --git a/temp_test_run_data.json b/temp_test_run_data.json new file mode 100644 index 0000000..c19e7ec --- /dev/null +++ b/temp_test_run_data.json @@ -0,0 +1 @@ +{"dict_test_cases": {}, "testCases": [], "metricScores": [], "configurations": {}} \ No newline at end of file diff --git a/tests/agents/test_reviewer_agent.py b/tests/agents/test_reviewer_agent.py index 3347f02..3dba303 100644 --- a/tests/agents/test_reviewer_agent.py +++ b/tests/agents/test_reviewer_agent.py @@ -2,8 +2,10 @@ import pytest from unittest.mock import MagicMock, patch +import json from app.agents.roles.reviewer_agent import ReviewerAgent +from app.rag.retriever import KnowledgeRetriever @pytest.fixture def mock_retriever_fixture(): @@ -39,7 +41,14 @@ def test_reviewer_agent_uses_rag_context(mock_retriever_fixture, mocker): code_to_review = "my_variable = 1" # Mock the model's response to stop the execution loop after one turn - mocker.patch.object(agent.client.chat.completions, 'create', return_value=MagicMock()) + # Создаем корректную структуру ответа, которую ожидает json.loads + mock_response_content = json.dumps({ + "thought": "The user wants me to review code. I will provide a final answer.", + "answer": "The code `my_variable = 1` looks simple and correct." + }) + mock_completion = MagicMock() + mock_completion.choices[0].message.content = mock_response_content + mocker.patch.object(agent.client.chat.completions, 'create', return_value=mock_completion) # Act # Запускаем execute_task, который теперь содержит логику обогащения промпта @@ -68,7 +77,8 @@ def test_reviewer_agent_uses_rag_context(mock_retriever_fixture, mocker): use_rag=True, api_key="fake_api_key", ) - mocker.patch.object(agent_no_knowledge.client.chat.completions, 'create', return_value=MagicMock()) + # Применяем тот же самый мок и ко второму агенту + mocker.patch.object(agent_no_knowledge.client.chat.completions, 'create', return_value=mock_completion) agent_no_knowledge.execute_task(briefing=code_to_review) system_prompt_no_knowledge = agent_no_knowledge.conversation_history[0]['content'] diff --git a/tests/rag/test_retriever.py b/tests/rag/test_retriever.py index a51c7ba..0675e6d 100644 --- a/tests/rag/test_retriever.py +++ b/tests/rag/test_retriever.py @@ -30,6 +30,7 @@ def mock_retriever(mocker) -> KnowledgeRetriever: ] retriever.embeddings = np.random.rand(4, 1536) # 4 чанка, 1536 измерений retriever.client = MagicMock() # Нам не нужен клиент, так как мы мокаем cosine_similarity + retriever.openai_client = MagicMock() # 4. Инициализируем BM25, как это делается в методе load() tokenized_corpus = [chunk["text"].split(" ") for chunk in retriever.chunks] @@ -54,25 +55,28 @@ def test_retriever_ranking_logic( """ # Arrange num_chunks = len(mock_retriever.chunks) - mock_similarities = np.linspace(0, 1, num_chunks) # [0.0, 0.25, 0.5, 1.0] - np.random.shuffle(mock_similarities) # Перемешиваем - + # Ищем индекс нужного чанка и ставим ему максимальную схожесть target_index = [i for i, c in enumerate(mock_retriever.chunks) if c["source"] == expected_source][0] - mock_similarities[target_index] = 1.1 # Гарантированно максимальное значение + + # Мокаем _get_embedding, чтобы он возвращал предопределенный вектор для запроса + # и разные векторы для чанков, где у целевого чанка будет наибольшая схожесть + mock_query_embedding = np.random.rand(1536) + + # Создаем моки для эмбеддингов чанков + mock_chunk_embeddings = np.random.rand(num_chunks, 1536) + # Устанавливаем эмбеддинг целевого чанка так, чтобы он был идентичен запросу + mock_chunk_embeddings[target_index] = mock_query_embedding + mock_retriever.embeddings = mock_chunk_embeddings - mocker.patch( - 'app.rag.retriever.cosine_similarity', - return_value=np.array([mock_similarities]) - ) - # Мокаем BM25 так, чтобы он не возвращал результатов и не влиял на фьюжн - mocker.patch.object(mock_retriever, '_keyword_search', return_value=[]) + mocker.patch('app.rag.retriever.KnowledgeRetriever._get_embedding', return_value=mock_query_embedding) - # Так как `retrieve` вызывает `_get_embedding`, который вызывает `client`, - # нам достаточно замокать вызов клиента. - mock_retriever.client.embeddings.create.return_value = MagicMock( - data=[MagicMock(embedding=np.random.rand(1536).tolist())] - ) + # Мокаем BM25 так, чтобы он не возвращал результатов и не влиял на фьюжн + mock_retriever.bm25_index = MagicMock() + mocker.patch.object(mock_retriever.bm25_index, 'get_scores', return_value=np.zeros(num_chunks)) + + # Мокаем CrossEncoder, чтобы он не пытался скачаться и просто возвращал "как есть" + mock_retriever.cross_encoder = None # Act retrieved_chunks = mock_retriever.retrieve(query, top_k=1) diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py new file mode 100644 index 0000000..d0b7a2e --- /dev/null +++ b/tests/test_evaluation.py @@ -0,0 +1,80 @@ +# tests/test_evaluation.py +import pytest +from unittest.mock import MagicMock, patch +import json + +from app.evaluation.evaluator import CustomEvaluator + +@pytest.fixture +def evaluator(mocker): + """Fixture to create a CustomEvaluator with a mocked OpenAI client.""" + mocker.patch('os.getenv', return_value="fake_api_key") + evaluator_instance = CustomEvaluator() + # Mock the client within the instance + evaluator_instance.client = MagicMock() + return evaluator_instance + +def test_evaluate_success(evaluator): + """Tests a successful evaluation call.""" + # Arrange + mock_response_content = json.dumps({ + "score": 5, + "justification": "The answer is perfect." + }) + mock_completion = MagicMock() + mock_completion.choices[0].message.content = mock_response_content + evaluator.client.chat.completions.create.return_value = mock_completion + + # Act + result = evaluator.evaluate( + actual_answer="The capital of France is Paris.", + reference_answer="Paris is the capital of France.", + criterion="Correctness" + ) + + # Assert + assert result["score"] == 5 + assert result["justification"] == "The answer is perfect." + evaluator.client.chat.completions.create.assert_called_once() + +def test_evaluate_json_decode_error(evaluator): + """Tests handling of a JSON decoding error from the LLM.""" + # Arrange + mock_completion = MagicMock() + mock_completion.choices[0].message.content = "This is not valid JSON." + evaluator.client.chat.completions.create.return_value = mock_completion + + # Act + result = evaluator.evaluate("a", "b", "c") + + # Assert + assert result["score"] == 0 + assert "Failed to parse" in result["justification"] + +def test_evaluate_missing_keys(evaluator): + """Tests handling of a response with missing keys.""" + # Arrange + mock_response_content = json.dumps({"score": 4}) # Missing 'justification' + mock_completion = MagicMock() + mock_completion.choices[0].message.content = mock_response_content + evaluator.client.chat.completions.create.return_value = mock_completion + + # Act + result = evaluator.evaluate("a", "b", "c") + + # Assert + assert result["score"] == 0 + assert "missing 'score' or 'justification'" in result["justification"] + +def test_evaluate_empty_inputs(): + """Tests that the evaluator raises an error on empty inputs.""" + # No need to mock the client here, as it should fail before the API call. + evaluator_instance = CustomEvaluator(api_key="fake_key") + with pytest.raises(ValueError, match="cannot be empty"): + evaluator_instance.evaluate(actual_answer="", reference_answer="b", criterion="c") + + with pytest.raises(ValueError, match="cannot be empty"): + evaluator_instance.evaluate(actual_answer="a", reference_answer="", criterion="c") + + with pytest.raises(ValueError, match="cannot be empty"): + evaluator_instance.evaluate(actual_answer="a", reference_answer="b", criterion="") \ No newline at end of file