Update agent.py
This commit is contained in:
@@ -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()
|
||||
print("Agent initialized.")
|
||||
Reference in New Issue
Block a user