diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py new file mode 100644 index 0000000..66afd7e --- /dev/null +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py @@ -0,0 +1,187 @@ +<|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, Tool +from langchain.schema import Document + +# --------------------------------------------------------------------------- # +# Конфигурация и клиент Qdrant +# --------------------------------------------------------------------------- # + +QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost") +QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333)) +COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection") + +# Инициализируем клиент Qdrant +qdrant_store = QdrantVectorStore( + url=f"http://{QDRANT_HOST}:{QDRANT_PORT}", + collection_name=COLLECTION_NAME, + embeddings=OllamaEmbeddings(model="nomic-embed-text"), +) + +# --------------------------------------------------------------------------- # +# Чанкинг текста +# --------------------------------------------------------------------------- # + +text_splitter = RecursiveCharacterTextSplitter( + chunk_size=500, # символов в чанке + chunk_overlap=50, +) + +# --------------------------------------------------------------------------- # +# Инструменты агента +# --------------------------------------------------------------------------- # + +@tool("search_knowledge_base") +def search_knowledge_base(query: str, max_results: int = 5) -> List[Document]: + """ + Семантический поиск по базе знаний. + + :param query: поисковый запрос + :param max_results: максимальное количество результатов + :return: список документов с метаданными + """ + results = qdrant_store.similarity_search(query, k=max_results) + return results + + +@tool("add_to_knowledge_base") +def add_to_knowledge_base(content: str, title: str) -> str: + """ + Добавление нового документа в базу знаний. + + :param content: полный текст документа + :param title: заголовок/название документа + :return: подтверждение добавления + """ + # Разбиваем на чанки и сохраняем с метаданными + chunks = text_splitter.split_text(content) + docs = [ + Document(page_content=chunk, metadata={"title": title}) + for chunk in chunks + ] + qdrant_store.add_documents(docs) + return f"Документ '{title}' успешно добавлен в базу знаний." + + +# --------------------------------------------------------------------------- # +# Создание агента +# --------------------------------------------------------------------------- # + +llm = OllamaLLM(model="llama3") + +tools: List[Tool] = [ + Tool.from_function( + func=search_knowledge_base, + name="search_knowledge_base", + description="Используйте для поиска информации в базе знаний. Входные параметры: query, max_results.", + ), + Tool.from_function( + func=add_to_knowledge_base, + name="add_to_knowledge_base", + description="Добавляет новый документ в базу знаний. Параметры: content, title.", + ), +] + +agent = create_agent( + llm=llm, + tools=tools, + system_message=( + "Вы — интеллектуальный агент с доступом к базе знаний. " + "При необходимости используйте инструмент search_knowledge_base для поиска информации. " + "Если нужно добавить новый материал, применяйте add_to_knowledge_base." + ), +) + +agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True) + + +# --------------------------------------------------------------------------- # +# Загрузка документов из директории +# --------------------------------------------------------------------------- # + +def load_documents_from_directory(directory: str) -> None: + """ + Загружает все текстовые файлы из указанной папки в базу знаний. + + :param directory: путь к каталогу с документами + """ + path = Path(directory) + for file_path in path.rglob("*"): + if file_path.suffix.lower() not in {".txt", ".md"}: + continue + content = file_path.read_text(encoding="utf-8") + title = file_path.stem + add_to_knowledge_base(content, title) + + +# --------------------------------------------------------------------------- # +# Интерактивный клиент +# --------------------------------------------------------------------------- # + +def interactive_cli(): + """ + Простая CLI для взаимодействия с агентом. + Команды: + /add - добавить документ из файла + /search - поиск по базе знаний + /quit - выйти + любой другой ввод передаётся агенту как вопрос + """ + 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"): + _, *parts = user_input.split(maxsplit=1) + if parts: + file_path = parts[0] + try: + content = Path(file_path).read_text(encoding="utf-8") + title = Path(file_path).stem + print(add_to_knowledge_base(content, title)) + except Exception as e: + print(f"Ошибка при добавлении: {e}") + else: + print("Используйте /add <путь_к_файлу>") + elif user_input.startswith("/search"): + _, *parts = user_input.split(maxsplit=1) + if parts: + query = parts[0] + results = search_knowledge_base(query, max_results=3) + for i, doc in enumerate(results, 1): + print(f"\nРезультат {i}:") + print(f"Титул: {doc.metadata.get('title', 'Неизвестно')}") + print(doc.page_content[:500], "...") + else: + print("Используйте /search <запрос>") + else: + # Любой другой ввод считается вопросом агента + response = agent_executor.run(user_input) + print(response) + + +# --------------------------------------------------------------------------- # +# Точка входа +# --------------------------------------------------------------------------- # + +if __name__ == "__main__": + # При запуске можно загрузить документы из папки docs/ + docs_dir = os.getenv("DOCS_DIR", "docs") + if Path(docs_dir).exists(): + load_documents_from_directory(docs_dir) + interactive_cli() \ No newline at end of file