166 lines
6.5 KiB
Python
166 lines
6.5 KiB
Python
<|channel|>final code<|message|>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_agent, AgentExecutor
|
|
from langchain.schema import Document
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 1. Векторное хранилище на базе Qdrant
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
class KnowledgeBase:
|
|
"""
|
|
Обёртка над QdrantVectorStore с автоматическим эмбеддингом через Ollama.
|
|
"""
|
|
|
|
def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333):
|
|
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
self.store = QdrantVectorStore(
|
|
embeddings=self.embeddings,
|
|
url=f"http://{host}:{port}",
|
|
collection_name=collection_name,
|
|
)
|
|
|
|
def add_documents(self, docs: List[Document]) -> None:
|
|
"""Добавляет список документов в коллекцию."""
|
|
self.store.add_documents(docs)
|
|
|
|
def search(self, query: str, limit: int = 5) -> List[Document]:
|
|
"""Поиск по запросу с ограничением на количество результатов."""
|
|
return self.store.similarity_search(query=query, k=limit)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 2. Чанкинг документов
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def chunk_text(text: str, title: str = "") -> List[Document]:
|
|
"""
|
|
Делит текст на чанки и возвращает список Document с метаданными.
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=500,
|
|
chunk_overlap=50,
|
|
length_function=len
|
|
)
|
|
chunks = splitter.split_text(text)
|
|
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]:
|
|
"""
|
|
Семантический поиск в базе знаний.
|
|
Возвращает список строк с содержимым найденных чанков.
|
|
"""
|
|
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:
|
|
"""
|
|
Добавляет новый документ в базу знаний после чанкинга.
|
|
Возвращает подтверждение количества добавленных чанков.
|
|
"""
|
|
docs = chunk_text(content, title=title)
|
|
kb.add_documents(docs)
|
|
return f"Добавлено {len(docs)} чанков из документа '{title}'."
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 4. Создание агента
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
llm = OllamaLLM(model="llama3")
|
|
|
|
agent_executor: AgentExecutor
|
|
|
|
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_prompt=system_prompt,
|
|
verbose=True
|
|
)
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 5. Инициализация из директории (для загрузки файлов)
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
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=content, title=title)
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 6. Интерактивный клиент
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def interactive_cli() -> None:
|
|
"""
|
|
Простая CLI с командами /add, /search и /quit.
|
|
"""
|
|
init_agent()
|
|
print("=== RAG-агент готов к работе ===")
|
|
while True:
|
|
try:
|
|
user_input = input("\n> ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
break
|
|
|
|
if not user_input:
|
|
continue
|
|
|
|
if user_input.lower() == "/quit":
|
|
print("Завершение работы.")
|
|
break
|
|
elif user_input.startswith("/add"):
|
|
# /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))
|
|
elif user_input.startswith("/search"):
|
|
# /search запрос
|
|
_, query = user_input.split(maxsplit=1)
|
|
results = search_knowledge_base(query, max_results=3)
|
|
for i, res in enumerate(results, 1):
|
|
print(f"\n--- Результат {i} ---\n{res}")
|
|
else:
|
|
# Любой другой ввод считается запросом к агенту
|
|
response = agent_executor.invoke({"input": user_input})
|
|
print(response["output"])
|
|
|
|
if __name__ == "__main__":
|
|
interactive_cli() |