diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py index 69e38fb..29334bb 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py @@ -2,19 +2,20 @@ # -*- coding: utf-8 -*- """ -solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama. +solution.py – основной скрипт с реализацией RAG‑агента и инструментов. """ -import argparse +import os +import sys from pathlib import Path from typing import List, Dict -from langchain_ollama import OllamaEmbeddings +from langchain_ollama import OllamaEmbeddings, OllamaLLM from langchain_qdrant import QdrantVectorStore +from langchain.schema import Document from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.tools import tool -from langchain.agents import create_agent -from langchain.schema import Document +from langchain.agents import create_agent, AgentExecutor, Tool # --------------------------------------------------------------------------- # # Конфигурация @@ -24,16 +25,20 @@ 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) +EMBEDDING_MODEL = "nomic-embed-text" +LLM_MODEL = "llama3" + +MAX_CHUNK_SIZE = 1000 # символов +CHUNK_OVERLAP = 200 # символов + # --------------------------------------------------------------------------- # -# Векторное хранилище и вспомогательные функции +# Векторное хранилище (обёртка над QdrantVectorStore) # --------------------------------------------------------------------------- # class KnowledgeBase: """ - Обёртка над QdrantVectorStore с поддержкой чанкинга и инструментов. + Класс, инкапсулирующий работу с Qdrant и Ollama. """ def __init__(self, @@ -41,108 +46,121 @@ class KnowledgeBase: 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, + embedding=self.embeddings ) - # Чанкер для разбивки документов - self.splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + # Создаём коллекцию, если её ещё нет + if not self.store.collection_exists(): + self.store.create_collection() - def add_document(self, content: str, title: str) -> None: + def add_documents(self, documents: List[Document]) -> 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(documents) + + def similarity_search( + self, + query: str, + k: int = 5 + ) -> List[Dict]: + """ + Семантический поиск по запросу. Возвращает списки словарей с полями: + 'content', 'metadata'. + """ + results = self.store.similarity_search(query, k=k) + return [ + {"content": doc.page_content, + "metadata": doc.metadata} + for doc in results ] - 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() +def split_text_to_documents(text: str, title: str) -> List[Document]: + """ + Делит текст на чанки и возвращает список Document. + """ + splitter = RecursiveCharacterTextSplitter( + chunk_size=MAX_CHUNK_SIZE, + chunk_overlap=CHUNK_OVERLAP + ) + chunks = splitter.split_text(text) + return [ + Document(page_content=chunk, metadata={"title": title}) + for chunk in chunks + ] + # --------------------------------------------------------------------------- # # Инструменты для агента # --------------------------------------------------------------------------- # -@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) +kb = KnowledgeBase() -@tool +@tool("search_knowledge_base") +def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict]: + """ + Семантический поиск в базе знаний. + Возвращает список найденных документов с их содержимым и метаданными. + """ + return kb.similarity_search(query, k=max_results) + + +@tool("add_to_knowledge_base") def add_to_knowledge_base(content: str, title: str) -> str: """ - Добавление нового документа в базу знаний. + Добавляет новый документ в базу знаний после разбиения на чанки. + Возвращает сообщение об успешном добавлении. """ - kb.add_document(content, title) - return f"Документ «{title}» успешно добавлен." + docs = split_text_to_documents(content, title) + kb.add_documents(docs) + return f"Документ '{title}' успешно добавлен в базу." # --------------------------------------------------------------------------- # # Создание агента # --------------------------------------------------------------------------- # -def build_agent(): - """ - Возвращает готовый агент с инструментами RAG‑памяти. - Используется актуальный API: create_agent(model, tools, system_prompt). - """ - system_prompt = ( - "Ты – интеллектуальный ассистент, который использует локальную базу знаний. " - "При ответе на запросы сначала ищи релевантную информацию в базе через инструмент " - "`search_knowledge_base`. Если нужно сохранить новую информацию, добавляй её через " - "`add_to_knowledge_base`." - ) +llm = OllamaLLM(model=LLM_MODEL) - agent = create_agent( - LLM_MODEL, - tools=[search_knowledge_base, add_to_knowledge_base], - system_prompt=system_prompt, +system_prompt = """ +Ты – интеллектуальный агент с доступом к локальной базе знаний. +При ответе на запросы используй только информацию из базы, если она есть. +Если необходима дополнительная информация, сначала выполни поиск в базе, +затем сформируй ответ. Не выдавай вымышленных фактов. +""" + +tools = [ + Tool.from_function( + func=search_knowledge_base, + name="search_knowledge_base", + description="Используется для поиска информации в базе знаний." + ), + Tool.from_function( + func=add_to_knowledge_base, + name="add_to_knowledge_base", + description="Добавляет новый документ в базу знаний." ) - return agent +] + +agent = create_agent( + llm=llm, + tools=tools, + system_prompt=system_prompt +) + +executor = AgentExecutor(agent=agent, tools=tools, verbose=True) # --------------------------------------------------------------------------- # -# Загрузка документов из директории +# Инициализация: загрузка документов из директории # --------------------------------------------------------------------------- # def load_documents_from_dir(directory: Path) -> None: @@ -150,95 +168,86 @@ 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}") + if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}: + text = file_path.read_text(encoding="utf-8") + title = file_path.stem + docs = split_text_to_documents(text, title) + kb.add_documents(docs) + print(f"Загружен документ: {title}") # --------------------------------------------------------------------------- # -# CLI клиент +# Интерактивный клиент # --------------------------------------------------------------------------- # -def interactive_cli(agent) -> None: +def interactive_cli() -> None: """ - Простая REPL‑интерфейс с командами /add, /search и /quit. - Любой другой ввод считается запросом к агенту. + Простая CLI с командами /add, /search и /quit. """ - print("=== Агент RAG ===") + print("=== RAG‑Агент ===") print("Команды:") - print("/add – добавить новый документ (будет запрошен контент)") + print("/add <title> – добавить новый документ (будет запрошен ввод текста)") print("/search <query> – поиск в базе знаний") print("/quit – выйти") while True: try: - inp = input("\n> ").strip() + user_input = input("\n> ").strip() except EOFError: break - if not inp: + + if not user_input: continue - if inp.lower() == "/quit": + if user_input.lower() == "/quit": + print("До свидания!") 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}) + if user_input.startswith("/add"): + parts = user_input.split(maxsplit=1) + title = parts[1] if len(parts) > 1 else "Untitled" + print(f"Введите текст для документа '{title}'. Завершите ввод пустой строкой.") + lines = [] + while True: + line = input() + if line == "": + break + lines.append(line) + content = "\n".join(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) + elif user_input.startswith("/search"): + query = user_input[len("/search"):].strip() + if not query: + print("Введите запрос после /search") + continue + results = search_knowledge_base(query, max_results=3) if not results: print("Ничего не найдено.") - else: - for r in results: - print(f"\n**{r['title']}** ({r['source']}) – {r['score']:.2f}") - print(r["content"][:400] + "...") + continue + for i, res in enumerate(results, 1): + print(f"\nРезультат {i}:") + print(f"Титул: {res['metadata'].get('title', 'Неизвестно')}") + print(res["content"][:500] + ("..." if len(res["content"]) > 500 else "")) else: - # Любой другой ввод – запрос к агенту - # Актуальный API: invoke принимает {"messages": [...]} - result = agent.invoke( - {"messages": [{"role": "user", "content": inp}]} - ) - # Извлекаем текст последнего сообщения агента - messages = result.get("messages", []) - if messages: - print(messages[-1].content) + # Любой другой ввод считается запросом к агенту + response = executor.invoke({"input": user_input}) + 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 = build_agent() - interactive_cli(agent) - - if __name__ == "__main__": - main() \ No newline at end of file + if len(sys.argv) > 1 and sys.argv[1] == "--init": + # Инициализация: загрузка документов из папки docs/ + docs_dir = Path("docs") + if not docs_dir.exists(): + print(f"Папка {docs_dir} не найдена.") + sys.exit(1) + load_documents_from_dir(docs_dir) + else: + interactive_cli() \ No newline at end of file