Обновить solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py
This commit is contained in:
@@ -2,43 +2,46 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
solution.py – основной скрипт с реализацией RAG‑агента и инструментов.
|
||||
solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List, Dict
|
||||
|
||||
from langchain_ollama import OllamaEmbeddings, OllamaLLM
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from langchain.schema import Document
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain.tools import tool
|
||||
from langchain.agents import create_agent, AgentExecutor, Tool
|
||||
from langchain.agents import create_agent
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Конфигурация
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
QDRANT_HOST = "localhost"
|
||||
QDRANT_PORT = 6333
|
||||
COLLECTION_NAME = "knowledge_base"
|
||||
QDRANT_HOST = "localhost"
|
||||
QDRANT_PORT = 6333
|
||||
COLLECTION_NAME = "knowledge_base"
|
||||
|
||||
EMBEDDING_MODEL = "nomic-embed-text"
|
||||
LLM_MODEL = "llama3"
|
||||
EMBEDDING_MODEL = "nomic-embed-text"
|
||||
LLM_MODEL = "ollama:llama3" # формат model-string для create_agent
|
||||
|
||||
MAX_CHUNK_SIZE = 1000 # символов
|
||||
CHUNK_OVERLAP = 200 # символов
|
||||
MAX_CHUNK_SIZE = 1000 # символов
|
||||
CHUNK_OVERLAP = 200 # символов
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Векторное хранилище (обёртка над QdrantVectorStore)
|
||||
# Векторное хранилище
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class KnowledgeBase:
|
||||
"""
|
||||
Класс, инкапсулирующий работу с Qdrant и Ollama.
|
||||
Инкапсулирует работу с Qdrant и эмбеддингами Ollama.
|
||||
|
||||
Инициализация через QdrantVectorStore.from_existing_collection —
|
||||
подключается к уже существующей коллекции (или создаёт при первом
|
||||
add_documents, когда Qdrant сам управляет коллекцией через URL).
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
@@ -46,35 +49,31 @@ class KnowledgeBase:
|
||||
port: int = QDRANT_PORT,
|
||||
collection_name: str = COLLECTION_NAME):
|
||||
self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||||
self.store = QdrantVectorStore(
|
||||
url=f"http://{host}:{port}",
|
||||
# Подключаемся к существующей коллекции (или используем URL-режим,
|
||||
# при котором QdrantVectorStore создаёт коллекцию автоматически).
|
||||
self.store = QdrantVectorStore.from_existing_collection(
|
||||
embedding=self.embeddings,
|
||||
collection_name=collection_name,
|
||||
embedding=self.embeddings
|
||||
url=f"http://{host}:{port}",
|
||||
)
|
||||
# Создаём коллекцию, если её ещё нет
|
||||
if not self.store.collection_exists():
|
||||
self.store.create_collection()
|
||||
|
||||
def add_documents(self, documents: List[Document]) -> None:
|
||||
"""
|
||||
Добавляет список документов в хранилище.
|
||||
"""
|
||||
"""Добавляет список Document в хранилище."""
|
||||
self.store.add_documents(documents)
|
||||
|
||||
def similarity_search(
|
||||
self,
|
||||
query: str,
|
||||
k: int = 5
|
||||
) -> List[Dict]:
|
||||
def similarity_search(self, query: str, k: int = 5) -> List[Dict]:
|
||||
"""
|
||||
Семантический поиск по запросу. Возвращает списки словарей с полями:
|
||||
'content', 'metadata'.
|
||||
Семантический поиск с метрикой релевантности.
|
||||
Возвращает список словарей с полями 'content', 'metadata', 'score'.
|
||||
"""
|
||||
results = self.store.similarity_search(query, k=k)
|
||||
results = self.store.similarity_search_with_score(query, k=k)
|
||||
return [
|
||||
{"content": doc.page_content,
|
||||
"metadata": doc.metadata}
|
||||
for doc in results
|
||||
{
|
||||
"content": doc.page_content,
|
||||
"metadata": doc.metadata,
|
||||
"score": score,
|
||||
}
|
||||
for doc, score in results
|
||||
]
|
||||
|
||||
|
||||
@@ -84,157 +83,164 @@ class KnowledgeBase:
|
||||
|
||||
def split_text_to_documents(text: str, title: str) -> List[Document]:
|
||||
"""
|
||||
Делит текст на чанки и возвращает список Document.
|
||||
Разбивает текст на чанки через RecursiveCharacterTextSplitter.
|
||||
Каждый чанк получает метаданные: title и порядковый номер чанка.
|
||||
"""
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=MAX_CHUNK_SIZE,
|
||||
chunk_overlap=CHUNK_OVERLAP
|
||||
chunk_overlap=CHUNK_OVERLAP,
|
||||
)
|
||||
chunks = splitter.split_text(text)
|
||||
return [
|
||||
Document(page_content=chunk, metadata={"title": title})
|
||||
for chunk in chunks
|
||||
Document(
|
||||
page_content=chunk,
|
||||
metadata={"title": title, "chunk_index": i},
|
||||
)
|
||||
for i, chunk in enumerate(chunks)
|
||||
]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Инструменты для агента
|
||||
# Глобальный экземпляр базы знаний + инструменты агента (@tool)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
kb = KnowledgeBase()
|
||||
|
||||
|
||||
@tool("search_knowledge_base")
|
||||
@tool
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict]:
|
||||
"""
|
||||
Семантический поиск в базе знаний.
|
||||
Возвращает список найденных документов с их содержимым и метаданными.
|
||||
Возвращает список найденных документов с содержимым, метаданными и оценкой релевантности.
|
||||
"""
|
||||
return kb.similarity_search(query, k=max_results)
|
||||
|
||||
|
||||
@tool("add_to_knowledge_base")
|
||||
@tool
|
||||
def add_to_knowledge_base(content: str, title: str) -> str:
|
||||
"""
|
||||
Добавляет новый документ в базу знаний после разбиения на чанки.
|
||||
Возвращает сообщение об успешном добавлении.
|
||||
Возвращает подтверждение успешного добавления.
|
||||
"""
|
||||
docs = split_text_to_documents(content, title)
|
||||
kb.add_documents(docs)
|
||||
return f"Документ '{title}' успешно добавлен в базу."
|
||||
return f"Документ '{title}' успешно добавлен в базу ({len(docs)} чанков)."
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Создание агента
|
||||
# Агент с RAG‑интеграцией
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
llm = OllamaLLM(model=LLM_MODEL)
|
||||
|
||||
system_prompt = """
|
||||
Ты – интеллектуальный агент с доступом к локальной базе знаний.
|
||||
При ответе на запросы используй только информацию из базы, если она есть.
|
||||
Если необходима дополнительная информация, сначала выполни поиск в базе,
|
||||
затем сформируй ответ. Не выдавай вымышленных фактов.
|
||||
"""
|
||||
|
||||
tools = [
|
||||
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(
|
||||
llm=llm,
|
||||
tools=tools,
|
||||
system_prompt=system_prompt
|
||||
system_prompt = (
|
||||
"Ты – интеллектуальный агент с доступом к локальной базе знаний.\n"
|
||||
"При ответе на запросы сначала ищи релевантную информацию через инструмент "
|
||||
"`search_knowledge_base`. Если нужно сохранить новую информацию — используй "
|
||||
"`add_to_knowledge_base`. Не выдумывай факты, которых нет в базе."
|
||||
)
|
||||
|
||||
executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
|
||||
# create_agent(model, tools, system_prompt) — актуальный API LangChain.
|
||||
# Первый аргумент — строка модели в формате "provider:model".
|
||||
# Возвращает граф-агент с .invoke() / .stream() — AgentExecutor не нужен.
|
||||
agent = create_agent(
|
||||
LLM_MODEL,
|
||||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||||
system_prompt=system_prompt,
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Инициализация: загрузка документов из директории
|
||||
# Клиент инициализации: загрузка документов из директории
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def load_documents_from_dir(directory: Path) -> None:
|
||||
"""
|
||||
Загружает все текстовые файлы из указанной папки в базу знаний.
|
||||
Загружает все .txt и .md файлы из указанной папки в базу знаний.
|
||||
"""
|
||||
loaded = 0
|
||||
for file_path in directory.rglob("*"):
|
||||
if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}:
|
||||
text = file_path.read_text(encoding="utf-8")
|
||||
text = file_path.read_text(encoding="utf-8")
|
||||
title = file_path.stem
|
||||
docs = split_text_to_documents(text, title)
|
||||
docs = split_text_to_documents(text, title)
|
||||
kb.add_documents(docs)
|
||||
print(f"Загружен документ: {title}")
|
||||
print(f"Загружен: {file_path} ({len(docs)} чанков)")
|
||||
loaded += 1
|
||||
print(f"\nИтого загружено файлов: {loaded}")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Интерактивный клиент
|
||||
# Интерактивный тестовый клиент
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def interactive_cli() -> None:
|
||||
"""
|
||||
Простая CLI с командами /add, /search и /quit.
|
||||
REPL с командами /add, /search, /quit.
|
||||
Любой другой ввод передаётся агенту.
|
||||
"""
|
||||
print("=== RAG‑Агент ===")
|
||||
print("Команды:")
|
||||
print("/add <title> – добавить новый документ (будет запрошен ввод текста)")
|
||||
print("/search <query> – поиск в базе знаний")
|
||||
print("/quit – выйти")
|
||||
print(" /add <title> – добавить документ (ввод текста, конец — пустая строка)")
|
||||
print(" /search <query> – поиск в базе знаний")
|
||||
print(" /quit – выйти")
|
||||
print(" <любой текст> – запрос к агенту\n")
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input("\n> ").strip()
|
||||
user_input = input("> ").strip()
|
||||
except EOFError:
|
||||
break
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
# /quit
|
||||
if user_input.lower() == "/quit":
|
||||
print("До свидания!")
|
||||
break
|
||||
|
||||
if user_input.startswith("/add"):
|
||||
# /add <title>
|
||||
elif user_input.startswith("/add"):
|
||||
parts = user_input.split(maxsplit=1)
|
||||
title = parts[1] if len(parts) > 1 else "Untitled"
|
||||
print(f"Введите текст для документа '{title}'. Завершите ввод пустой строкой.")
|
||||
lines = []
|
||||
title = parts[1].strip() if len(parts) > 1 else "Untitled"
|
||||
print(f"Вводите текст для «{title}». Пустая строка — конец ввода.")
|
||||
lines: List[str] = []
|
||||
while True:
|
||||
line = input()
|
||||
if line == "":
|
||||
break
|
||||
lines.append(line)
|
||||
content = "\n".join(lines)
|
||||
result = add_to_knowledge_base(content, title)
|
||||
# Вызов @tool через .invoke() — правильный способ вызова инструмента
|
||||
result = add_to_knowledge_base.invoke({"content": content, "title": title})
|
||||
print(result)
|
||||
|
||||
# /search <query>
|
||||
elif user_input.startswith("/search"):
|
||||
query = user_input[len("/search"):].strip()
|
||||
if not query:
|
||||
print("Введите запрос после /search")
|
||||
print("Укажите запрос: /search <query>")
|
||||
continue
|
||||
results = search_knowledge_base(query, max_results=3)
|
||||
results = search_knowledge_base.invoke({"query": query, "max_results": 3})
|
||||
if not results:
|
||||
print("Ничего не найдено.")
|
||||
continue
|
||||
for i, res in enumerate(results, 1):
|
||||
print(f"\nРезультат {i}:")
|
||||
print(f"Титул: {res['metadata'].get('title', 'Неизвестно')}")
|
||||
print(res["content"][:500] + ("..." if len(res["content"]) > 500 else ""))
|
||||
else:
|
||||
for i, res in enumerate(results, 1):
|
||||
title = res["metadata"].get("title", "—")
|
||||
score = res.get("score", 0)
|
||||
print(f"\nРезультат {i} | {title} | score={score:.4f}")
|
||||
snippet = res["content"]
|
||||
print(snippet[:500] + ("..." if len(snippet) > 500 else ""))
|
||||
|
||||
# Запрос к агенту
|
||||
else:
|
||||
# Любой другой ввод считается запросом к агенту
|
||||
response = executor.invoke({"input": user_input})
|
||||
print(response.get("output", "Нет ответа."))
|
||||
# Актуальный API: invoke принимает {"messages": [...]},
|
||||
# ответ — в result["messages"][-1].content
|
||||
result = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
|
||||
messages = result.get("messages", [])
|
||||
if messages:
|
||||
print(messages[-1].content)
|
||||
else:
|
||||
print("Нет ответа.")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
@@ -243,10 +249,10 @@ def interactive_cli() -> None:
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "--init":
|
||||
# Инициализация: загрузка документов из папки docs/
|
||||
docs_dir = Path("docs")
|
||||
# python solution.py --init → загрузить документы из ./docs/
|
||||
docs_dir = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("docs")
|
||||
if not docs_dir.exists():
|
||||
print(f"Папка {docs_dir} не найдена.")
|
||||
print(f"Папка '{docs_dir}' не найдена.")
|
||||
sys.exit(1)
|
||||
load_documents_from_dir(docs_dir)
|
||||
else:
|
||||
|
||||
Reference in New Issue
Block a user