diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py new file mode 100644 index 0000000..106d4e6 --- /dev/null +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama. +""" + +import os +import sys +import argparse +from pathlib import Path +from typing import List, Dict + +from langchain_ollama import ChatOllama, OllamaEmbeddings +from langchain_qdrant import QdrantVectorStore +from langchain.text_splitter import RecursiveCharacterTextSplitter +from langchain.tools import tool +from langchain.agents import create_agent, AgentExecutor, Tool +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 = "llama3" # Ollama LLM + +# --------------------------------------------------------------------------- # +# Векторное хранилище и вспомогательные функции +# --------------------------------------------------------------------------- # + +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) + ] + # Добавляем в Qdrant + 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: + out = { + "score": score, + "title": doc.metadata.get("title", ""), + "source": doc.metadata.get("source", ""), + "content": doc.page_content + } + output.append(out) + return output + + +# --------------------------------------------------------------------------- # +# Инструменты для агента +# --------------------------------------------------------------------------- # + +kb = KnowledgeBase() + +@tool +def search_knowledge_base(query: str, max_results: int) -> 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() -> AgentExecutor: + """ + Возвращает готовый агент с инструментами RAG‑памяти. + """ + system_prompt = ( + "Ты – интеллектуальный ассистент, который использует локальную базу знаний. " + "При ответе на запросы сначала ищи релевантную информацию в базе через tool " + "`search_knowledge_base`. Если нужна новая информация, добавляй её через " + "`add_to_knowledge_base`." + ) + tools = [ + Tool.from_function( + func=search_knowledge_base, + name="search_knowledge_base", + description="Семантический поиск в базе знаний. Вход: query, max_results. Выход: Markdown с результатами." + ), + Tool.from_function( + func=add_to_knowledge_base, + name="add_to_knowledge_base", + description="Добавить новый документ в базу знаний. Вход: content, title. Выход: подтверждение." + ) + ] + llm = ChatOllama(model=LLM_MODEL) + agent = create_agent( + llm=llm, + tools=tools, + system_prompt=system_prompt, + verbose=True + ) + return AgentExecutor(agent=agent, tools=tools, verbose=True) + + +# --------------------------------------------------------------------------- # +# Загрузка документов из директории +# --------------------------------------------------------------------------- # + +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_executor: AgentExecutor) -> None: + """ + Простая REPL‑интерфейс с командами /add, /search и /quit. + Любой другой ввод считается запросом к агенту. + """ + print("=== Агент RAG ===") + print("Команды:") + print("/add – добавить новый документ (будет запрошен контент)") + print("/search <query> – поиск в базе знаний") + print("/quit – выйти") + + while True: + try: + inp = input("\n> ").strip() + except EOFError: + break + if not inp: + continue + + if inp.lower() == "/quit": + break + + if inp.startswith("/add "): + title = inp[5:].strip() + print("Введите содержимое документа (конец строки с EOF Ctrl+D):") + try: + content_lines = [] + while True: + line = input() + content_lines.append(line) + except EOFError: + pass + content = "\n".join(content_lines) + result = add_to_knowledge_base(content, 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: + # Любой другой ввод – запрос к агенту + response = agent_executor.invoke({"input": inp}) + print(response.get("output", "")) + + +# --------------------------------------------------------------------------- # +# Точка входа +# --------------------------------------------------------------------------- # + +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_executor = build_agent() + interactive_cli(agent_executor) + + +if __name__ == "__main__": + main() \ No newline at end of file