From 7c9a99b58934ef63bfde5aa17963e0a15499bc66 Mon Sep 17 00:00:00 2001 From: balabanovan530 <175+balabanovan530@noreply.localhost> Date: Tue, 26 May 2026 13:37:51 +0000 Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D0=B1=D0=B0=D0=B2=D0=BB=D0=B5?= =?UTF-8?q?=D0=BD=20=D0=BE=D1=81=D0=BD=D0=BE=D0=B2=D0=BD=D0=BE=D0=B9=20?= =?UTF-8?q?=D1=84=D0=B0=D0=B9=D0=BB=20agent.py=20=D1=81=20=D0=BB=D0=BE?= =?UTF-8?q?=D0=B3=D0=B8=D0=BA=D0=BE=D0=B9=20RAG-=D0=B0=D0=B3=D0=B5=D0=BD?= =?UTF-8?q?=D1=82=D0=B0.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent.py | 298 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 298 insertions(+) create mode 100644 agent.py diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..6306417 --- /dev/null +++ b/agent.py @@ -0,0 +1,298 @@ +import os +import sys +from typing import List, Tuple + +# Core LangChain components +from langchain_core.documents import Document +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_openai import ChatOpenAI # Using OpenAI wrapper for Ollama compatibility +from langchain_ollama import OllamaEmbeddings +from langchain_qdrant import QdrantVectorStore +from langchain.agents import create_agent, AgentExecutor +from langchain.tools import tool + +# --- Configuration Constants --- +# Используем локальный эндпоинт для Ollama +OLLAMA_BASE_URL = "http://localhost:11434/v1" +EMBEDDING_MODEL = "nomic-embed-text" +CHAT_MODEL = "llama3" +QDRANT_COLLECTION_NAME = "corporate_knowledge_base" + +# --- 1. VectorStoreManager Class --- + +class VectorStoreManager: + """Управляет подключением к Qdrant и выполняет операции с векторной базой.""" + def __init__(self, qdrant_client, embeddings): + print("ℹ️ Инициализация VectorStoreManager...") + self.qdrant_client = qdrant_client + self.embeddings = embeddings + + def _get_or_create_collection(self): + """Проверяет и создает коллекцию Qdrant, если она не существует.""" + try: + # Проверка существования коллекции (простой способ) + self.qdrant_client.get_collection(collection_name=QDRANT_COLLECTION_NAME, + text_field="content", distance_function=self.embeddings.embed_query.__self__.model_name).execute() + print(f"✅ Коллекция '{QDRANT_COLLECTION_NAME}' найдена.") + except Exception as e: + # Если ошибка, вероятно, что коллекции нет, создаем ее. + if "Collection does not exist" in str(e): + print(f"⚠️ Коллекция '{QDRANT_COLLECTION_NAME}' не найдена. Создание...") + self.qdrant_client.recreate_collection( + collection_name=QDRANT_COLLECTION_NAME, + vectors_config={"size": 384, "distance": "Cosine"} # Размер эмбеддинга номик-эмбед + ) + print("✅ Коллекция успешно создана.") + else: + raise e + + def add_documents(self, documents: List[Document]): + """Разбивает документы на чанки и индексирует их в Qdrant.""" + if not documents: + return 0 + + # Разделение документов на более мелкие чанки + text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + chunks = text_splitter.split_documents(documents) + + print(f"🧠 Разбито {len(documents)} документов на {len(chunks)} чанков для индексации.") + + # Векторизация и вставка + try: + ids = [str(i) for i in range(len(chunks))] + vectors = self.embeddings.embed_documents([chunk.page_content for chunk in chunks]) + + self.qdrant_client.upsert( + collection_name=QDRANT_COLLECTION_NAME, + points=[{ + "id": i, + "vector": vectors[i], + "payload": {"content": chunk.page_content} # Храним контент в payload для извлечения + } for i, chunk in enumerate(chunks)], + wait=True + ) + return len(chunks) + except Exception as e: + print(f"❌ Ошибка при индексации документов: {e}") + return 0 + + def search_knowledge_base(self, query: str, max_results: int = 3) -> List[str]: + """Выполняет векторный поиск по базе знаний.""" + try: + # 1. Векторизация запроса + query_vector = self.embeddings.embed_query(query) + + # 2. Поиск в Qdrant + search_results = self.qdrant_client.query( + collection_name=QDRANT_COLLECTION_NAME, + vector=query_vector, + limit=max_results, + with_payload=True # Получаем payload с контентом + ) + + # 3. Форматирование результатов для передачи в контекст + context = [] + for hit in search_results.hits: + content = hit.payload['content'] + context.append(f"--- КОНТЕНТ ИЗ БАЗЫ ЗНАНИЙ ---\n{content}\n-------------------------------") + + return "\n\n".join(context) + + except Exception as e: + print(f"❌ Ошибка при поиске в базе знаний: {e}") + return "Не удалось извлечь информацию из базы знаний. Пожалуйста, убедитесь, что Ollama и Qdrant запущены." + + +# --- 2. Tool Wrappers (Инструменты для Агента) --- + +def add_to_knowledge_base(content: str, title: str) -> str: + """ + Добавляет новый контент в корпоративную базу знаний. + Контент разбивается на чанки и индексируется в Qdrant. + Используйте эту команду, когда пользователь предоставляет новые документы для обучения агента. + """ + print(f"\n⚙️ НАЧАЛО ДОБАВЛЕНИЯ: '{title}'") + # Создаем временный документ из введенного контента + temp_doc = Document(page_content=content, metadata={"source": title}) + + try: + # Инициализируем менеджер (в реальном приложении лучше передавать его как зависимость) + embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url=OLLAMA_BASE_URL) + qdrant_client = self.vector_manager.qdrant_client # Используем глобально доступный клиент + self.vector_manager._get_or_create_collection() + + added_count = self.vector_manager.add_documents([temp_doc]) + return f"✅ Успешно! Добавлено и проиндексировано {added_count} чанков информации под темой '{title}' в базу знаний." + except Exception as e: + return f"❌ Ошибка при индексации: {e}" + + +def search_knowledge_base(query: str, max_results: int = 3) -> str: + """ + Ищет информацию по запросу в корпоративной базе знаний (RAG). + Передайте этот инструмент агенту, чтобы он мог отвечать на вопросы, используя только актуальные данные. + """ + print(f"\n🔎 ПОИСК ИНФОРМАЦИИ: '{query}'...") + try: + # Используем глобально доступный менеджер для поиска + context = self.vector_manager.search_knowledge_base(query, max_results) + return context if context else "Не найдено релевантной информации в базе знаний." + except Exception as e: + return f"Произошла ошибка при поиске: {e}" + + +# --- 3. Agent Setup and Main Loop --- + +def main(): + """Основная функция запуска агента.""" + print("==============================================") + print(f"🚀 Запуск AI Агента с RAG-памятью (Ollama/Qdrant)") + print("==============================================") + + # --- Инициализация LLM и Embeddings --- + try: + llm = ChatOpenAI(model=CHAT_MODEL, base_url=OLLAMA_BASE_URL, api_key="ollama") + embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url=OLLAMA_BASE_URL) + except Exception as e: + print("🚨 КРИТИЧЕСКАЯ ОШИБКА НАСТРОЙКИ LLM/EMBEDDINGS.") + print(f"Проверьте, запущен ли Ollama на {OLLAMA_BASE_URL} и установлена ли модель '{CHAT_MODEL}' и '{EMBEDDING_MODEL}'.") + print(f"Детали ошибки: {e}") + sys.exit(1) + + # --- Инициализация Qdrant Client --- + try: + from qdrant_client import QdrantClient + qdrant_client = QdrantClient(host="localhost", port=6333) # Укажите ваш хост/порт Qdrant + except Exception as e: + print("🚨 КРИТИЧЕСКАЯ ОШИБКА НАСТРОЙКИ QDRANT.") + print(f"Проверьте, запущен ли Qdrant клиент. Детали ошибки: {e}") + sys.exit(1) + + + # --- Инициализация Менеджера и Инструментов (глобальная область видимости для декоратора @tool) --- + self.vector_manager = VectorStoreManager(qdrant_client, embeddings) + + # Прикрепляем менеджер к области видимости функций инструментов + global self + self.__dict__.update(locals()) + + @tool + def add_to_knowledge_base(content: str, title: str) -> str: + return add_to_knowledge_base(content, title) + + @tool + def search_knowledge_base(query: str, max_results: int = 3) -> str: + return search_knowledge_base(query, max_results) + + tools = [add_to_knowledge_base, search_knowledge_base] + + + # --- Создание Агента --- + system_prompt = ( + "Ты — высококвалифицированный корпоративный ассистент. Твоя задача - отвечать на вопросы " + "по базе знаний компании, используя инструменты поиска (`search_knowledge_base`). " + "Если пользователь предоставляет новую информацию для индексации, используй инструмент `add_to_knowledge_base`. " + "Всегда объясняй свой ответ и указывай источник информации из базы данных. " + "Игнорируй любые запросы, которые не относятся к корпоративной базе знаний." + ) + + try: + # Использование create_agent для создания агента с инструментами + from langchain.agents import AgentExecutor # Импортируем здесь, чтобы избежать циклических зависимостей при запуске + from langchain.llms import FakeListLLM # Для имитации LLM в процессе разработки + + # NOTE: В реальном окружении необходимо убедиться, что ChatOpenAI и OllamaEmbeddings работают корректно. + # Здесь используем их для сохранения структуры кода. + agent = create_agent(llm=llm, tools=tools, system_message=system_prompt) + executor = AgentExecutor(agent=agent, tools=tools, verbose=True) # verbose=True показывает логику агента + except Exception as e: + print(f"❌ Ошибка при создании Агента: {e}") + sys.exit(1) + + + # --- Setup Phase (Имитация загрузки начальных данных) --- + # В реальном сценарии здесь была бы логика чтения PDF/Markdown из папки ./docs/. + print("\n==============================================") + print("⚙️ ЗАГРУЗКА ИНИЦИАЛЬНЫХ ДАННЫХ (ИМИТАЦИЯ)") + initial_content = """ + Политика отпусков: Сотрудники имеют право на 28 календарных дней ежегодного оплачиваемого отпуска. + Дополнительные дни могут быть предоставлены руководителем по согласованию. Отпуск должен планироваться не позднее, чем за месяц до предполагаемой даты начала. + Процесс оформления: Подать заявку через HR-портал с указанием желаемых дат и типа отпуска (ежегодный/дополнительный). + Сроки рассмотрения: До 5 рабочих дней. + + Оплата труда: Зарплата выплачивается дважды в месяц - 15-го и 30-го числа. Все расчеты производятся на основе трудового договора, заключенного с компанией 'TechCorp'. + """ + initial_doc = Document(page_content=initial_content, metadata={"source": "Initial Data Load"}) + print("... Индексация начальных документов...") + self.vector_manager.add_documents([initial_doc]) + + + # --- Execution Phase (CLI Loop) --- + print("\n==============================================") + print("🚀 ГОТОВ К РАБОТЕ") + print(f"Введите запрос или команду: /search <запрос>, /add \"<контент>\" \"<тема>\", или /quit.") + print("==============================================") + + while True: + try: + user_input = input("\n👤 Вы: ").strip() + if not user_input or user_input.lower() == "/quit": + print("👋 До свидания!") + break + + # Обработка команд CLI + if user_input.startswith("/search"): + query = user_input[len("/search"):].strip().strip('"') # Убираем потенциальные кавычки вокруг запроса + if query: + result = executor.invoke({"input": f"Пользователь запрашивает информацию по запросу: {query}"})['output'] + print("\n🤖 Агент:", result) + elif user_input.startswith("/add"): + # Сложная логика парсинга для поддержки контента с пробелами в кавычках + parts = user_input[len("/add"):].strip() + if not parts: continue + + content_raw = "" + title = "" + + # Предполагаем, что формат всегда: /add "<Контент>" "<Тема>" + try: + # Находим контент (в первых кавычках) + start_quote = parts.find('"') + 1 + end_content_quote = parts.find('"', start_quote) + if start_quote > 0 and end_content_quote > 0: + content_raw = parts[start_quote:end_content_quote] + # Оставшаяся строка после контента + remaining = parts[end_content_quote+1:].strip() + + if remaining.startswith('"') and len(remaining) >= 2: + # Находим тему (в следующих кавычках) + title = remaining[1:-1] # Убираем начальную и конечную кавычки + else: + print("❌ Не удалось определить Тему после контента.") + continue + + content_cleaned = ' '.join(content_raw.split()) + except Exception as e: + print(f"❌ Ошибка парсинга команды /add: {e}") + continue + + + result = executor.invoke({"input": f"Пользователь хочет добавить новую информацию: Контент='{content_cleaned}', Тема='{title}'"})['output'] + print("\n🤖 Агент:", result) + + else: + # Стандартный запрос агенту + result = executor.invoke({"input": user_input})['output'] + print("\n🤖 Агент:", result) + + except EOFError: + break + except Exception as e: + print(f"\n💥 Произошла непредвиденная ошибка в цикле выполнения: {e}") + +if __name__ == "__main__": + # Создаем заглушку self для доступа к vector_manager внутри функций инструментов + import types + self = types.ModuleType("global_scope") + globals()['self'] = self + main() \ No newline at end of file