244 lines
9.1 KiB
Python
244 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama.
|
||
"""
|
||
|
||
import argparse
|
||
from pathlib import Path
|
||
from typing import List, Dict
|
||
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_qdrant import QdrantVectorStore
|
||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||
from langchain.tools import tool
|
||
from langchain.agents import create_agent
|
||
from langchain.schema import Document
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Конфигурация
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
QDRANT_HOST = "localhost"
|
||
QDRANT_PORT = 6333
|
||
COLLECTION_NAME = "knowledge_base"
|
||
|
||
EMBEDDING_MODEL = "nomic-embed-text" # Ollama embedding model
|
||
LLM_MODEL = "ollama:llama3" # Ollama LLM (формат для create_agent)
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Векторное хранилище и вспомогательные функции
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
class KnowledgeBase:
|
||
"""
|
||
Обёртка над QdrantVectorStore с поддержкой чанкинга и инструментов.
|
||
"""
|
||
|
||
def __init__(self,
|
||
host: str = QDRANT_HOST,
|
||
port: int = QDRANT_PORT,
|
||
collection_name: str = COLLECTION_NAME):
|
||
self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||
# Создаём/подключаем коллекцию
|
||
self.store = QdrantVectorStore(
|
||
embedding=self.embeddings,
|
||
url=f"http://{host}:{port}",
|
||
collection_name=collection_name,
|
||
)
|
||
# Чанкер для разбивки документов
|
||
self.splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||
|
||
def add_document(self, content: str, title: str) -> None:
|
||
"""
|
||
Добавляет документ в базу после разбиения на чанки.
|
||
"""
|
||
chunks = self.splitter.split_text(content)
|
||
docs = [
|
||
Document(
|
||
page_content=chunk,
|
||
metadata={"title": title, "source": f"{title} (chunk {i})"}
|
||
)
|
||
for i, chunk in enumerate(chunks)
|
||
]
|
||
self.store.add_documents(docs)
|
||
|
||
def search(self, query: str, max_results: int = 5) -> List[Dict]:
|
||
"""
|
||
Семантический поиск по базе знаний.
|
||
Возвращает список словарей с полями 'title', 'source' и 'content'.
|
||
"""
|
||
results = self.store.similarity_search_with_score(query, k=max_results)
|
||
output = []
|
||
for doc, score in results:
|
||
output.append({
|
||
"score": score,
|
||
"title": doc.metadata.get("title", ""),
|
||
"source": doc.metadata.get("source", ""),
|
||
"content": doc.page_content,
|
||
})
|
||
return output
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Глобальный экземпляр базы знаний
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
kb = KnowledgeBase()
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Инструменты для агента
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
@tool
|
||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||
"""
|
||
Поиск в базе знаний.
|
||
Возвращает строку с найденными результатами в формате Markdown.
|
||
"""
|
||
results = kb.search(query, max_results)
|
||
if not results:
|
||
return "Ничего не найдено."
|
||
lines = []
|
||
for r in results:
|
||
lines.append(
|
||
f"**{r['title']}** ({r['source']}) – {r['score']:.2f}\n> {r['content'][:200]}..."
|
||
)
|
||
return "\n\n".join(lines)
|
||
|
||
|
||
@tool
|
||
def add_to_knowledge_base(content: str, title: str) -> str:
|
||
"""
|
||
Добавление нового документа в базу знаний.
|
||
"""
|
||
kb.add_document(content, title)
|
||
return f"Документ «{title}» успешно добавлен."
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Создание агента
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def build_agent():
|
||
"""
|
||
Возвращает готовый агент с инструментами RAG‑памяти.
|
||
Используется актуальный API: create_agent(model, tools, system_prompt).
|
||
"""
|
||
system_prompt = (
|
||
"Ты – интеллектуальный ассистент, который использует локальную базу знаний. "
|
||
"При ответе на запросы сначала ищи релевантную информацию в базе через инструмент "
|
||
"`search_knowledge_base`. Если нужно сохранить новую информацию, добавляй её через "
|
||
"`add_to_knowledge_base`."
|
||
)
|
||
|
||
agent = create_agent(
|
||
LLM_MODEL,
|
||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||
system_prompt=system_prompt,
|
||
)
|
||
return agent
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Загрузка документов из директории
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def load_documents_from_dir(directory: Path) -> None:
|
||
"""
|
||
Загружает все текстовые файлы из указанной папки в базу знаний.
|
||
"""
|
||
for file_path in directory.rglob("*"):
|
||
if file_path.suffix.lower() not in {".txt", ".md"}:
|
||
continue
|
||
content = file_path.read_text(encoding="utf-8")
|
||
title = file_path.stem
|
||
kb.add_document(content, title)
|
||
print(f"Загружено: {file_path}")
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# CLI клиент
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def interactive_cli(agent) -> None:
|
||
"""
|
||
Простая REPL‑интерфейс с командами /add, /search и /quit.
|
||
Любой другой ввод считается запросом к агенту.
|
||
"""
|
||
print("=== Агент RAG ===")
|
||
print("Команды:")
|
||
print("/add <title> – добавить новый документ (будет запрошен контент)")
|
||
print("/search <query> – поиск в базе знаний")
|
||
print("/quit – выйти")
|
||
|
||
while True:
|
||
try:
|
||
inp = input("\n> ").strip()
|
||
except EOFError:
|
||
break
|
||
if not inp:
|
||
continue
|
||
|
||
if inp.lower() == "/quit":
|
||
break
|
||
|
||
elif inp.startswith("/add "):
|
||
title = inp[5:].strip()
|
||
print("Введите содержимое документа (завершите ввод Ctrl+D / Ctrl+Z):")
|
||
try:
|
||
content_lines = []
|
||
while True:
|
||
line = input()
|
||
content_lines.append(line)
|
||
except EOFError:
|
||
pass
|
||
content = "\n".join(content_lines)
|
||
result = add_to_knowledge_base.invoke({"content": content, "title": title})
|
||
print(result)
|
||
|
||
elif inp.startswith("/search "):
|
||
query = inp[8:].strip()
|
||
results = kb.search(query, max_results=5)
|
||
if not results:
|
||
print("Ничего не найдено.")
|
||
else:
|
||
for r in results:
|
||
print(f"\n**{r['title']}** ({r['source']}) – {r['score']:.2f}")
|
||
print(r["content"][:400] + "...")
|
||
|
||
else:
|
||
# Любой другой ввод – запрос к агенту
|
||
# Актуальный API: invoke принимает {"messages": [...]}
|
||
result = agent.invoke(
|
||
{"messages": [{"role": "user", "content": inp}]}
|
||
)
|
||
# Извлекаем текст последнего сообщения агента
|
||
messages = result.get("messages", [])
|
||
if messages:
|
||
print(messages[-1].content)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Точка входа
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="Агент с RAG‑памятью")
|
||
parser.add_argument(
|
||
"--load-dir",
|
||
type=str,
|
||
help="Путь к директории с документами для предварительной загрузки"
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
if args.load_dir:
|
||
load_documents_from_dir(Path(args.load_dir))
|
||
|
||
agent = build_agent()
|
||
interactive_cli(agent)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |