203 lines
8.2 KiB
Python
203 lines
8.2 KiB
Python
# rag_agent.py
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
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_react_agent, AgentExecutor
|
|
from langchain.prompts import PromptTemplate
|
|
from langchain.schema import Document
|
|
from qdrant_client import QdrantClient as _QdrantClient
|
|
from qdrant_client.http.models import Distance, VectorParams
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Конфигурация
|
|
# ---------------------------------------------------------------------------
|
|
|
|
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
|
|
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
|
|
COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")
|
|
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
|
|
LLM_MODEL = os.getenv("LLM_MODEL", "llama3")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Векторное хранилище (Qdrant + Ollama-эмбеддинги)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class RagVectorStore:
|
|
"""Обёртка над QdrantVectorStore с эмбеддингами через Ollama."""
|
|
|
|
def __init__(self, collection_name: str = COLLECTION_NAME):
|
|
self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
|
|
|
# Низкоуровневый клиент для управления коллекцией
|
|
self._client = _QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
|
|
|
|
# Создаём коллекцию, если её ещё нет
|
|
if not self._collection_exists(collection_name):
|
|
self._client.create_collection(
|
|
collection_name=collection_name,
|
|
vectors_config=VectorParams(size=768, distance=Distance.COSINE),
|
|
)
|
|
|
|
self.store = QdrantVectorStore(
|
|
client=self._client,
|
|
collection_name=collection_name,
|
|
embedding=self.embeddings,
|
|
)
|
|
|
|
def _collection_exists(self, name: str) -> bool:
|
|
try:
|
|
self._client.get_collection(name)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def add_documents(self, docs: List[Document]) -> None:
|
|
self.store.add_documents(docs)
|
|
|
|
def search(self, query: str, k: int = 5) -> List[tuple]:
|
|
results = self.store.similarity_search_with_score(query, k=k)
|
|
return [(doc.page_content, score) for doc, score in results]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Чанкинг
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def chunk_document(text: str, title: str) -> List[Document]:
|
|
"""Разбивает текст на чанки с метаданными."""
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=1000,
|
|
chunk_overlap=200,
|
|
separators=["\n\n", "\n", " ", ""],
|
|
)
|
|
return [
|
|
Document(page_content=chunk, metadata={"title": title})
|
|
for chunk in splitter.split_text(text)
|
|
]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Инициализация хранилища (глобальный синглтон для инструментов)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
rag_store = RagVectorStore()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LangChain-инструменты агента (@tool)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@tool
|
|
def search_knowledge_base(query: str, max_results: int = 5) -> List[str]:
|
|
"""
|
|
Семантический поиск в базе знаний.
|
|
|
|
Args:
|
|
query: поисковый запрос.
|
|
max_results: максимальное количество результатов.
|
|
|
|
Returns:
|
|
Список найденных фрагментов с оценкой релевантности.
|
|
"""
|
|
results = rag_store.search(query, k=max_results)
|
|
return [f"[score={score:.4f}] {text}" for text, score in results]
|
|
|
|
|
|
@tool
|
|
def add_to_knowledge_base(content: str, title: str) -> str:
|
|
"""
|
|
Добавляет новый документ в базу знаний.
|
|
|
|
Args:
|
|
content: полный текст документа.
|
|
title: заголовок документа.
|
|
|
|
Returns:
|
|
Сообщение об успехе.
|
|
"""
|
|
docs = chunk_document(content, title)
|
|
rag_store.add_documents(docs)
|
|
return f"Документ '{title}' успешно добавлен в базу знаний."
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Агент на базе LangChain (create_react_agent)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
llm = OllamaLLM(model=LLM_MODEL)
|
|
|
|
tools = [search_knowledge_base, add_to_knowledge_base]
|
|
|
|
prompt = PromptTemplate.from_template(
|
|
"Вы — интеллектуальный агент с доступом к локальной базе знаний.\n"
|
|
"При необходимости используйте инструменты search_knowledge_base и add_to_knowledge_base.\n\n"
|
|
"Доступные инструменты:\n{tools}\n\n"
|
|
"Формат ответа (строго соблюдайте):\n"
|
|
"Question: вопрос пользователя\n"
|
|
"Thought: что нужно сделать\n"
|
|
"Action: инструмент из [{tool_names}]\n"
|
|
"Action Input: входные данные для инструмента\n"
|
|
"Observation: результат инструмента\n"
|
|
"... (можно повторять Thought/Action/Observation)\n"
|
|
"Thought: теперь я знаю ответ\n"
|
|
"Final Answer: итоговый ответ пользователю\n\n"
|
|
"Begin!\n\n"
|
|
"Question: {input}\n"
|
|
"Thought: {agent_scratchpad}"
|
|
)
|
|
|
|
agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
|
|
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Загрузка документов из директории (для init_loader.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_documents_from_directory(directory: str) -> None:
|
|
"""Загружает все .txt и .md файлы из папки в базу знаний."""
|
|
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")
|
|
add_to_knowledge_base.invoke({"content": text, "title": file_path.stem})
|
|
print(f"Загружен: {file_path.name}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Запуск напрямую — интерактивный CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
if __name__ == "__main__":
|
|
print("=== RAG Агент ===")
|
|
print("Команды: /add <path>, /search <query>, /quit")
|
|
while True:
|
|
try:
|
|
user_input = input("> ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
break
|
|
|
|
if not user_input:
|
|
continue
|
|
if user_input.lower() == "/quit":
|
|
print("Выход.")
|
|
break
|
|
elif user_input.startswith("/add "):
|
|
file_path = Path(user_input[5:].strip())
|
|
if file_path.is_file():
|
|
text = file_path.read_text(encoding="utf-8")
|
|
print(add_to_knowledge_base.invoke({"content": text, "title": file_path.stem}))
|
|
else:
|
|
print(f"Файл не найден: {file_path}")
|
|
elif user_input.startswith("/search "):
|
|
query = user_input[8:].strip()
|
|
for r in search_knowledge_base.invoke({"query": query, "max_results": 3}):
|
|
print(r)
|
|
else:
|
|
response = agent_executor.invoke({"input": user_input})
|
|
print(response["output"]) |