From 07a0c4f98150aa59e2fc9c94f4d9c5594b851bd3 Mon Sep 17 00:00:00 2001 From: balabanovan530 <175+balabanovan530@noreply.localhost> Date: Thu, 28 May 2026 13:58:52 +0000 Subject: [PATCH] Update agent.py --- agent.py | 308 +++---------------------------------------------------- 1 file changed, 17 insertions(+), 291 deletions(-) diff --git a/agent.py b/agent.py index 6306417..d51ed54 100644 --- a/agent.py +++ b/agent.py @@ -1,298 +1,24 @@ -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 import ChatPromptTemplate +from langchain.agents import create_agent from langchain.tools import tool +from .tools import add_content, search_content +import os -# --- Configuration Constants --- -# Используем локальный эндпоинт для Ollama -OLLAMA_BASE_URL = "http://localhost:11434/v1" -EMBEDDING_MODEL = "nomic-embed-text" -CHAT_MODEL = "llama3" -QDRANT_COLLECTION_NAME = "corporate_knowledge_base" +# LLM configuration – use Ollama via langchain-ollama +MODEL = os.getenv("OLLAMA_MODEL", "llama3.1") +BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1") -# --- 1. VectorStoreManager Class --- +from langchain_ollama import ChatOllama +llm = ChatOllama(model=MODEL, base_url=BASE_URL) -class VectorStoreManager: - """Управляет подключением к Qdrant и выполняет операции с векторной базой.""" - def __init__(self, qdrant_client, embeddings): - print("ℹ️ Инициализация VectorStoreManager...") - self.qdrant_client = qdrant_client - self.embeddings = embeddings +# System prompt for the agent +SYSTEM_PROMPT = """You are an assistant that answers user queries using a knowledge base. Use the provided tools to search and add content.""" +prompt = ChatPromptTemplate.from_messages([ + ("system", SYSTEM_PROMPT), +]) - 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}") +tools = [add_content, search_content] +agent = create_agent(llm=llm, prompt=prompt, tools=tools) if __name__ == "__main__": - # Создаем заглушку self для доступа к vector_manager внутри функций инструментов - import types - self = types.ModuleType("global_scope") - globals()['self'] = self - main() \ No newline at end of file + print("Agent initialized.") \ No newline at end of file