Агент с RAG‑памятью: rag_agent.py
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
import os
|
||||
<|channel|>final code<|message|>import os
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
@@ -6,169 +6,135 @@ from langchain_ollama import OllamaEmbeddings, OllamaLLM
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from langchain.tools import tool
|
||||
from langchain.agents import create_agent, AgentExecutor, Tool
|
||||
from langchain.agents import create_agent, AgentExecutor
|
||||
from langchain.schema import Document
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Конфигурация подключения к Qdrant и Ollama
|
||||
# 1. Векторное хранилище на базе Qdrant
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
|
||||
QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
|
||||
COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")
|
||||
class KnowledgeBase:
|
||||
"""
|
||||
Обёртка над QdrantVectorStore с автоматическим эмбеддингом через Ollama.
|
||||
"""
|
||||
|
||||
EMBEDDING_MODEL = "nomic-embed-text"
|
||||
LLM_MODEL = "llama3"
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Векторное хранилище
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class RagVectorStore:
|
||||
"""Обёртка над QdrantVectorStore с автоматическим эмбеддингом Ollama."""
|
||||
|
||||
def __init__(self, collection_name: str = COLLECTION_NAME):
|
||||
self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||||
def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333):
|
||||
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
self.store = QdrantVectorStore(
|
||||
url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
|
||||
collection_name=collection_name,
|
||||
embeddings=self.embeddings,
|
||||
url=f"http://{host}:{port}",
|
||||
collection_name=collection_name,
|
||||
)
|
||||
# Создаём коллекцию, если её ещё нет
|
||||
if not self.store.client.get_collection(collection_name).exists:
|
||||
self.store.create_collection()
|
||||
|
||||
def add_documents(self, docs: List[Document]) -> None:
|
||||
"""Добавляет список документов в хранилище."""
|
||||
"""Добавляет список документов в коллекцию."""
|
||||
self.store.add_documents(docs)
|
||||
|
||||
def search(
|
||||
self,
|
||||
query: str,
|
||||
k: int = 5,
|
||||
) -> List[tuple[str, float]]:
|
||||
"""Семантический поиск по запросу. Возвращает кортежи (текст, релевантность)."""
|
||||
results = self.store.similarity_search_with_score(query, k=k)
|
||||
return [(doc.page_content, score) for doc, score in results]
|
||||
def search(self, query: str, limit: int = 5) -> List[Document]:
|
||||
"""Поиск по запросу с ограничением на количество результатов."""
|
||||
return self.store.similarity_search(query=query, k=limit)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Чанкинг документов
|
||||
# 2. Чанкинг документов
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def chunk_document(text: str, title: str) -> List[Document]:
|
||||
"""Разбивает текст на чанки и возвращает список Document."""
|
||||
def chunk_text(text: str, title: str = "") -> List[Document]:
|
||||
"""
|
||||
Делит текст на чанки и возвращает список Document с метаданными.
|
||||
"""
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000,
|
||||
chunk_overlap=200,
|
||||
chunk_size=500,
|
||||
chunk_overlap=50,
|
||||
length_function=len
|
||||
)
|
||||
chunks = splitter.split_text(text)
|
||||
return [
|
||||
docs = [
|
||||
Document(page_content=chunk, metadata={"title": title})
|
||||
for chunk in chunks
|
||||
]
|
||||
return docs
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Инструменты для агента
|
||||
# 3. Инструменты для агента
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
kb = KnowledgeBase()
|
||||
|
||||
@tool("search_knowledge_base")
|
||||
def search_knowledge_base(
|
||||
query: str,
|
||||
max_results: int = 5,
|
||||
) -> List[str]:
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> List[str]:
|
||||
"""
|
||||
Семантический поиск в базе знаний.
|
||||
|
||||
Parameters:
|
||||
query (str): поисковый запрос.
|
||||
max_results (int): максимальное количество результатов.
|
||||
|
||||
Returns:
|
||||
list of strings: найденные тексты с оценкой релевантности.
|
||||
Возвращает список строк с содержимым найденных чанков.
|
||||
"""
|
||||
results = rag_store.search(query, k=max_results)
|
||||
return [f"{text} (score={score:.4f})" for text, score in results]
|
||||
results = kb.search(query=query, limit=max_results)
|
||||
return [doc.page_content for doc in results]
|
||||
|
||||
|
||||
@tool("add_to_knowledge_base")
|
||||
def add_to_knowledge_base(
|
||||
content: str,
|
||||
title: str,
|
||||
) -> str:
|
||||
def add_to_knowledge_base(content: str, title: str) -> str:
|
||||
"""
|
||||
Добавляет новый документ в базу знаний.
|
||||
|
||||
Parameters:
|
||||
content (str): полный текст документа.
|
||||
title (str): заголовок/название документа.
|
||||
|
||||
Returns:
|
||||
str: сообщение об успехе.
|
||||
Добавляет новый документ в базу знаний после чанкинга.
|
||||
Возвращает подтверждение количества добавленных чанков.
|
||||
"""
|
||||
docs = chunk_document(content, title)
|
||||
rag_store.add_documents(docs)
|
||||
return f"Документ '{title}' успешно добавлен в базу знаний."
|
||||
docs = chunk_text(content, title=title)
|
||||
kb.add_documents(docs)
|
||||
return f"Добавлено {len(docs)} чанков из документа '{title}'."
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Инициализация хранилища
|
||||
# 4. Создание агента
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
rag_store = RagVectorStore()
|
||||
llm = OllamaLLM(model="llama3")
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Создание агента с RAG‑интеграцией
|
||||
# --------------------------------------------------------------------------- #
|
||||
agent_executor: AgentExecutor
|
||||
|
||||
llm = OllamaLLM(model=LLM_MODEL)
|
||||
|
||||
tools: List[Tool] = [
|
||||
Tool.from_function(
|
||||
func=search_knowledge_base,
|
||||
name="search_knowledge_base",
|
||||
description="Используйте для поиска информации в базе знаний.",
|
||||
),
|
||||
Tool.from_function(
|
||||
func=add_to_knowledge_base,
|
||||
name="add_to_knowledge_base",
|
||||
description="Добавьте новый документ в базу знаний.",
|
||||
),
|
||||
]
|
||||
|
||||
agent = create_agent(
|
||||
def init_agent() -> None:
|
||||
"""
|
||||
Инициализирует агент с системным промптом, указывающим использовать RAG.
|
||||
"""
|
||||
global agent_executor
|
||||
system_prompt = (
|
||||
"Ты AI-ассистент. При необходимости ищи информацию в базе знаний "
|
||||
"с помощью инструмента search_knowledge_base и добавляй новые данные "
|
||||
"через add_to_knowledge_base."
|
||||
)
|
||||
tools = [search_knowledge_base, add_to_knowledge_base]
|
||||
agent_executor = create_agent(
|
||||
llm=llm,
|
||||
tools=tools,
|
||||
system_message=(
|
||||
"Вы — интеллектуальный агент, способный искать и добавлять информацию "
|
||||
"в локальную базу знаний. При необходимости используйте инструменты "
|
||||
"'search_knowledge_base' и 'add_to_knowledge_base'."
|
||||
),
|
||||
)
|
||||
|
||||
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
|
||||
system_prompt=system_prompt,
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Клиент для загрузки документов из директории
|
||||
# 5. Инициализация из директории (для загрузки файлов)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def load_documents_from_directory(directory: str) -> None:
|
||||
"""Загружает все текстовые файлы из указанной папки в базу знаний."""
|
||||
path = Path(directory)
|
||||
for file_path in path.rglob("*"):
|
||||
if file_path.suffix.lower() in {".txt", ".md"}:
|
||||
text = file_path.read_text(encoding="utf-8")
|
||||
def load_documents_from_dir(directory: str) -> None:
|
||||
"""
|
||||
Загружает все .txt файлы из указанной папки в базу знаний.
|
||||
"""
|
||||
for file_path in Path(directory).glob("*.txt"):
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
title = file_path.stem
|
||||
add_to_knowledge_base(content=text, title=title)
|
||||
add_to_knowledge_base(content=content, title=title)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Интерактивный клиент
|
||||
# 6. Интерактивный клиент
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def interactive_cli():
|
||||
"""Простой CLI для взаимодействия с агентом."""
|
||||
print("=== RAG Агент ===")
|
||||
print("Команды: /add <path>, /search <query>, /quit")
|
||||
def interactive_cli() -> None:
|
||||
"""
|
||||
Простая CLI с командами /add, /search и /quit.
|
||||
"""
|
||||
init_agent()
|
||||
print("=== RAG-агент готов к работе ===")
|
||||
while True:
|
||||
try:
|
||||
user_input = input("> ").strip()
|
||||
user_input = input("\n> ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
break
|
||||
|
||||
@@ -176,31 +142,23 @@ def interactive_cli():
|
||||
continue
|
||||
|
||||
if user_input.lower() == "/quit":
|
||||
print("Выход.")
|
||||
print("Завершение работы.")
|
||||
break
|
||||
elif user_input.startswith("/add"):
|
||||
_, *path_parts = user_input.split(maxsplit=1)
|
||||
if path_parts:
|
||||
file_path = Path(path_parts[0])
|
||||
if file_path.is_file():
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
title = file_path.stem
|
||||
# /add путь_к_файлу
|
||||
_, path = user_input.split(maxsplit=1)
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
title = Path(path).stem
|
||||
print(add_to_knowledge_base(content, title))
|
||||
else:
|
||||
print(f"Файл не найден: {file_path}")
|
||||
else:
|
||||
print("Укажите путь к файлу.")
|
||||
elif user_input.startswith("/search"):
|
||||
_, *query_parts = user_input.split(maxsplit=1)
|
||||
if query_parts:
|
||||
query = query_parts[0]
|
||||
# /search запрос
|
||||
_, query = user_input.split(maxsplit=1)
|
||||
results = search_knowledge_base(query, max_results=3)
|
||||
for r in results:
|
||||
print(r)
|
||||
for i, res in enumerate(results, 1):
|
||||
print(f"\n--- Результат {i} ---\n{res}")
|
||||
else:
|
||||
print("Укажите поисковый запрос.")
|
||||
else:
|
||||
# Любой другой ввод передаём агенту
|
||||
# Любой другой ввод считается запросом к агенту
|
||||
response = agent_executor.invoke({"input": user_input})
|
||||
print(response["output"])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user