Update agent.py
This commit is contained in:
@@ -1,298 +1,24 @@
|
|||||||
import os
|
from langchain import ChatPromptTemplate
|
||||||
import sys
|
from langchain.agents import create_agent
|
||||||
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
|
from langchain.tools import tool
|
||||||
|
from .tools import add_content, search_content
|
||||||
|
import os
|
||||||
|
|
||||||
# --- Configuration Constants ---
|
# LLM configuration – use Ollama via langchain-ollama
|
||||||
# Используем локальный эндпоинт для Ollama
|
MODEL = os.getenv("OLLAMA_MODEL", "llama3.1")
|
||||||
OLLAMA_BASE_URL = "http://localhost:11434/v1"
|
BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")
|
||||||
EMBEDDING_MODEL = "nomic-embed-text"
|
|
||||||
CHAT_MODEL = "llama3"
|
|
||||||
QDRANT_COLLECTION_NAME = "corporate_knowledge_base"
|
|
||||||
|
|
||||||
# --- 1. VectorStoreManager Class ---
|
from langchain_ollama import ChatOllama
|
||||||
|
llm = ChatOllama(model=MODEL, base_url=BASE_URL)
|
||||||
|
|
||||||
class VectorStoreManager:
|
# System prompt for the agent
|
||||||
"""Управляет подключением к Qdrant и выполняет операции с векторной базой."""
|
SYSTEM_PROMPT = """You are an assistant that answers user queries using a knowledge base. Use the provided tools to search and add content."""
|
||||||
def __init__(self, qdrant_client, embeddings):
|
prompt = ChatPromptTemplate.from_messages([
|
||||||
print("ℹ️ Инициализация VectorStoreManager...")
|
("system", SYSTEM_PROMPT),
|
||||||
self.qdrant_client = qdrant_client
|
])
|
||||||
self.embeddings = embeddings
|
|
||||||
|
|
||||||
def _get_or_create_collection(self):
|
tools = [add_content, search_content]
|
||||||
"""Проверяет и создает коллекцию Qdrant, если она не существует."""
|
agent = create_agent(llm=llm, prompt=prompt, tools=tools)
|
||||||
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__":
|
if __name__ == "__main__":
|
||||||
# Создаем заглушку self для доступа к vector_manager внутри функций инструментов
|
print("Agent initialized.")
|
||||||
import types
|
|
||||||
self = types.ModuleType("global_scope")
|
|
||||||
globals()['self'] = self
|
|
||||||
main()
|
|
||||||
Reference in New Issue
Block a user