From 7781965c65ee66a0ef1c048d8c8a9e817df0dd34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9C=D0=B0=D1=80=D0=B8=D1=8F=20=D0=91=D0=B5=D1=80=D0=B4?= =?UTF-8?q?=D0=BD=D0=B8=D0=BA=D0=BE=D0=B2=D0=B0?= Date: Wed, 27 May 2026 06:52:07 +0000 Subject: [PATCH] =?UTF-8?q?=D0=90=D0=B3=D0=B5=D0=BD=D1=82=20=D1=81=20RAG?= =?UTF-8?q?=E2=80=91=D0=BF=D0=B0=D0=BC=D1=8F=D1=82=D1=8C=D1=8E:=20rag=5Fag?= =?UTF-8?q?ent.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../rag_agent.py | 230 +++++++----------- 1 file changed, 94 insertions(+), 136 deletions(-) diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py index 72d90ad..f6dab55 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py @@ -1,4 +1,4 @@ -import os +<|channel|>final code<|message|>import os from pathlib import Path from typing import List @@ -6,169 +6,135 @@ 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.agents import create_agent, AgentExecutor from langchain.schema import Document # --------------------------------------------------------------------------- # -# Конфигурация подключения к Qdrant и Ollama +# 1. Векторное хранилище на базе Qdrant # --------------------------------------------------------------------------- # -QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost") -QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333)) -COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection") +class KnowledgeBase: + """ + Обёртка над QdrantVectorStore с автоматическим эмбеддингом через Ollama. + """ -EMBEDDING_MODEL = "nomic-embed-text" -LLM_MODEL = "llama3" - -# --------------------------------------------------------------------------- # -# Векторное хранилище -# --------------------------------------------------------------------------- # - -class RagVectorStore: - """Обёртка над QdrantVectorStore с автоматическим эмбеддингом Ollama.""" - - def __init__(self, collection_name: str = COLLECTION_NAME): - self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) + def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333): + self.embeddings = OllamaEmbeddings(model="nomic-embed-text") self.store = QdrantVectorStore( - url=f"http://{QDRANT_HOST}:{QDRANT_PORT}", - collection_name=collection_name, embeddings=self.embeddings, + url=f"http://{host}:{port}", + collection_name=collection_name, ) - # Создаём коллекцию, если её ещё нет - if not self.store.client.get_collection(collection_name).exists: - self.store.create_collection() def add_documents(self, docs: List[Document]) -> None: - """Добавляет список документов в хранилище.""" + """Добавляет список документов в коллекцию.""" self.store.add_documents(docs) - def search( - self, - query: str, - k: int = 5, - ) -> List[tuple[str, float]]: - """Семантический поиск по запросу. Возвращает кортежи (текст, релевантность).""" - results = self.store.similarity_search_with_score(query, k=k) - return [(doc.page_content, score) for doc, score in results] + def search(self, query: str, limit: int = 5) -> List[Document]: + """Поиск по запросу с ограничением на количество результатов.""" + return self.store.similarity_search(query=query, k=limit) + # --------------------------------------------------------------------------- # -# Чанкинг документов +# 2. Чанкинг документов # --------------------------------------------------------------------------- # -def chunk_document(text: str, title: str) -> List[Document]: - """Разбивает текст на чанки и возвращает список Document.""" +def chunk_text(text: str, title: str = "") -> List[Document]: + """ + Делит текст на чанки и возвращает список Document с метаданными. + """ splitter = RecursiveCharacterTextSplitter( - chunk_size=1000, - chunk_overlap=200, + chunk_size=500, + chunk_overlap=50, + length_function=len ) chunks = splitter.split_text(text) - return [ + 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]: +def search_knowledge_base(query: str, max_results: int = 5) -> List[str]: """ Семантический поиск в базе знаний. - - Parameters: - query (str): поисковый запрос. - max_results (int): максимальное количество результатов. - - Returns: - list of strings: найденные тексты с оценкой релевантности. + Возвращает список строк с содержимым найденных чанков. """ - results = rag_store.search(query, k=max_results) - return [f"{text} (score={score:.4f})" for text, score in results] + 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: +def add_to_knowledge_base(content: str, title: str) -> str: """ - Добавляет новый документ в базу знаний. - - Parameters: - content (str): полный текст документа. - title (str): заголовок/название документа. - - Returns: - str: сообщение об успехе. + Добавляет новый документ в базу знаний после чанкинга. + Возвращает подтверждение количества добавленных чанков. """ - docs = chunk_document(content, title) - rag_store.add_documents(docs) - return f"Документ '{title}' успешно добавлен в базу знаний." + docs = chunk_text(content, title=title) + kb.add_documents(docs) + return f"Добавлено {len(docs)} чанков из документа '{title}'." # --------------------------------------------------------------------------- # -# Инициализация хранилища +# 4. Создание агента # --------------------------------------------------------------------------- # -rag_store = RagVectorStore() +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 + ) # --------------------------------------------------------------------------- # -# Создание агента с RAG‑интеграцией +# 5. Инициализация из директории (для загрузки файлов) # --------------------------------------------------------------------------- # -llm = OllamaLLM(model=LLM_MODEL) - -tools: List[Tool] = [ - 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="Добавьте новый документ в базу знаний.", - ), -] - -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_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 load_documents_from_directory(directory: str) -> None: - """Загружает все текстовые файлы из указанной папки в базу знаний.""" - 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") - title = file_path.stem - add_to_knowledge_base(content=text, title=title) - -# --------------------------------------------------------------------------- # -# Интерактивный клиент -# --------------------------------------------------------------------------- # - -def interactive_cli(): - """Простой CLI для взаимодействия с агентом.""" - print("=== RAG Агент ===") - print("Команды: /add , /search , /quit") +def interactive_cli() -> None: + """ + Простая CLI с командами /add, /search и /quit. + """ + init_agent() + print("=== RAG-агент готов к работе ===") while True: try: - user_input = input("> ").strip() + user_input = input("\n> ").strip() except (EOFError, KeyboardInterrupt): break @@ -176,31 +142,23 @@ def interactive_cli(): continue if user_input.lower() == "/quit": - print("Выход.") + print("Завершение работы.") break elif user_input.startswith("/add"): - _, *path_parts = user_input.split(maxsplit=1) - if path_parts: - file_path = Path(path_parts[0]) - if file_path.is_file(): - content = file_path.read_text(encoding="utf-8") - title = file_path.stem - print(add_to_knowledge_base(content, title)) - else: - print(f"Файл не найден: {file_path}") - else: - print("Укажите путь к файлу.") + # /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"): - _, *query_parts = user_input.split(maxsplit=1) - if query_parts: - query = query_parts[0] - results = search_knowledge_base(query, max_results=3) - for r in results: - print(r) - else: - print("Укажите поисковый запрос.") + # /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"])