199 lines
8.6 KiB
Python
199 lines
8.6 KiB
Python
# main.py
|
||
# Полностью рабочий пример агента с RAG‑памятью на базе Qdrant и Ollama.
|
||
# Весь код написан в одном файле для простоты демонстрации.
|
||
#
|
||
# Требования:
|
||
# - Python 3.10+
|
||
# - Qdrant (работает по умолчанию на localhost:6333)
|
||
# - Ollama (llama3 + nomic-embed-text)
|
||
# - LangChain 1.2.10+ и связанные пакеты
|
||
#
|
||
# Запуск:
|
||
# 1. Убедитесь, что Qdrant и Ollama запущены.
|
||
# 2. pip install -r requirements.txt
|
||
# 3. python main.py
|
||
#
|
||
# После запуска появится интерактивный клиент с командами:
|
||
# /add <path_to_file> – добавить документ в базу знаний.
|
||
# /search <query> – выполнить поиск.
|
||
# /quit – выйти.
|
||
# Любой другой ввод – агент обработает как обычный запрос.
|
||
|
||
import os
|
||
import sys
|
||
import textwrap
|
||
from pathlib import Path
|
||
from typing import List, Dict, Any
|
||
|
||
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
||
from langchain_qdrant import QdrantVectorStore
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain.tools import tool
|
||
from langchain.agents import create_agent, AgentExecutor, AgentType
|
||
from langchain_core.messages import HumanMessage
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Конфигурация
|
||
# ---------------------------------------------------------------------------
|
||
QDRANT_URL = "http://localhost:6333"
|
||
COLLECTION_NAME = "knowledge"
|
||
EMBEDDING_MODEL = "nomic-embed-text"
|
||
LLM_MODEL = "llama3"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Векторное хранилище
|
||
# ---------------------------------------------------------------------------
|
||
# Инициализируем эмбеддер и клиент Qdrant
|
||
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||
vector_store = QdrantVectorStore(
|
||
url=QDRANT_URL,
|
||
collection_name=COLLECTION_NAME,
|
||
embedding=embeddings,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Чанкинг
|
||
# ---------------------------------------------------------------------------
|
||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Инструменты
|
||
# ---------------------------------------------------------------------------
|
||
@tool("search_knowledge_base", "Semantic search in the knowledge base")
|
||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||
"""Возвращает топ‑N релевантных фрагментов из Qdrant."""
|
||
results = vector_store.similarity_search_with_score(query, k=max_results)
|
||
if not results:
|
||
return "No relevant documents found."
|
||
# Формируем читаемый ответ
|
||
formatted = []
|
||
for i, (doc, score) in enumerate(results, 1):
|
||
formatted.append(f"{i}. (score={score:.3f})\n{doc.page_content[:500]}\n---")
|
||
return "\n".join(formatted)
|
||
|
||
@tool("add_to_knowledge_base", "Add a document to the knowledge base")
|
||
def add_to_knowledge_base(content: str, title: str = "") -> str:
|
||
"""Разбивает документ на чанки, эмбеддит и сохраняет в Qdrant."""
|
||
# Разбиваем на чанки
|
||
chunks = text_splitter.split_text(content)
|
||
# Создаём Document объекты с метаданными
|
||
from langchain.schema import Document
|
||
docs = [Document(page_content=chunk, metadata={"title": title, "source": title}) for chunk in chunks]
|
||
# Добавляем в хранилище
|
||
vector_store.add_documents(docs)
|
||
return f"Added {len(docs)} chunks from '{title}'."
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Агент
|
||
# ---------------------------------------------------------------------------
|
||
# Системный промпт, который подсказывает агенту использовать инструменты
|
||
SYSTEM_PROMPT = textwrap.dedent("""
|
||
You are an intelligent assistant with access to a knowledge base.
|
||
Use the following tools when you need to retrieve or store information:
|
||
- search_knowledge_base: Perform a semantic search in the knowledge base.
|
||
- add_to_knowledge_base: Add new content to the knowledge base.
|
||
When you answer a user query, first decide if you need to search the base.
|
||
If you do, call search_knowledge_base with a relevant query.
|
||
If you need to add new information, call add_to_knowledge_base.
|
||
Otherwise, answer directly using your internal knowledge.
|
||
""")
|
||
|
||
# Создаём LLM
|
||
llm = ChatOllama(model=LLM_MODEL, temperature=0.2)
|
||
|
||
# Создаём агент
|
||
agent = create_agent(
|
||
llm=llm,
|
||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||
system_message=SYSTEM_PROMPT,
|
||
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
||
)
|
||
|
||
# Обёртка для удобного вызова
|
||
agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Загрузка документов из директории (инициализация)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def load_documents_from_dir(directory: str) -> None:
|
||
"""Загружает все .txt и .md файлы из указанной папки в базу."""
|
||
p = Path(directory)
|
||
if not p.is_dir():
|
||
print(f"Directory {directory} does not exist.")
|
||
return
|
||
for file_path in p.rglob("*.txt") | p.rglob("*.md"):
|
||
try:
|
||
content = file_path.read_text(encoding="utf-8")
|
||
title = file_path.stem
|
||
add_to_knowledge_base(content, title)
|
||
print(f"Loaded {file_path}")
|
||
except Exception as e:
|
||
print(f"Failed to load {file_path}: {e}")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Интерактивный клиент
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def interactive_client():
|
||
print("Welcome to the RAG agent. Type /help for commands.")
|
||
while True:
|
||
try:
|
||
user_input = input("\n> ")
|
||
except (EOFError, KeyboardInterrupt):
|
||
print("\nExiting.")
|
||
break
|
||
if not user_input:
|
||
continue
|
||
if user_input.startswith("/"):
|
||
# Команды
|
||
if user_input.startswith("/add "):
|
||
path = user_input[5:].strip()
|
||
if not path:
|
||
print("Usage: /add <path_to_file>")
|
||
continue
|
||
try:
|
||
content = Path(path).read_text(encoding="utf-8")
|
||
title = Path(path).stem
|
||
result = add_to_knowledge_base(content, title)
|
||
print(result)
|
||
except Exception as e:
|
||
print(f"Error reading file: {e}")
|
||
elif user_input.startswith("/search "):
|
||
query = user_input[8:].strip()
|
||
if not query:
|
||
print("Usage: /search <query>")
|
||
continue
|
||
result = search_knowledge_base(query)
|
||
print(result)
|
||
elif user_input in {"/quit", "/exit"}:
|
||
print("Goodbye!")
|
||
break
|
||
elif user_input == "/help":
|
||
print(textwrap.dedent("""
|
||
Commands:
|
||
/add <path> – add a document to the knowledge base
|
||
/search <q> – search the knowledge base
|
||
/quit /exit – exit the program
|
||
/help – show this help
|
||
Any other input is treated as a normal user query for the agent.
|
||
"""))
|
||
else:
|
||
print("Unknown command. Type /help for list of commands.")
|
||
else:
|
||
# Передаём запрос агенту
|
||
try:
|
||
response = agent_executor.invoke({"input": user_input})
|
||
print(response["output"])
|
||
except Exception as e:
|
||
print(f"Agent error: {e}")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Точка входа
|
||
# ---------------------------------------------------------------------------
|
||
if __name__ == "__main__":
|
||
# Если передан аргумент – загрузить документы из указанной папки
|
||
if len(sys.argv) > 1:
|
||
load_documents_from_dir(sys.argv[1])
|
||
interactive_client()
|