diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py index f803b59..4b4e841 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py @@ -1,3 +1,4 @@ +# rag_agent.py import os from pathlib import Path from typing import List @@ -6,161 +7,197 @@ 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.agents import create_react_agent, AgentExecutor +from langchain.prompts import PromptTemplate from langchain.schema import Document -from rag_tools import search_knowledge_base, add_to_knowledge_base +from qdrant_client import QdrantClient as _QdrantClient +from qdrant_client.http.models import Distance, VectorParams -# --------------------------------------------------------------------------- # -# 1. Векторное хранилище на базе Qdrant -# --------------------------------------------------------------------------- # +# --------------------------------------------------------------------------- +# Конфигурация +# --------------------------------------------------------------------------- -class KnowledgeBase: - """ - Обёртка над QdrantVectorStore с автоматическим эмбеддингом через Ollama. - """ +QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost") +QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333")) +COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection") +EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-embed-text") +LLM_MODEL = os.getenv("LLM_MODEL", "llama3") + +# --------------------------------------------------------------------------- +# Векторное хранилище (Qdrant + Ollama-эмбеддинги) +# --------------------------------------------------------------------------- + +class RagVectorStore: + """Обёртка над QdrantVectorStore с эмбеддингами через Ollama.""" + + def __init__(self, collection_name: str = COLLECTION_NAME): + self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) + + # Низкоуровневый клиент для управления коллекцией + self._client = _QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT) + + # Создаём коллекцию, если её ещё нет + if not self._collection_exists(collection_name): + self._client.create_collection( + collection_name=collection_name, + vectors_config=VectorParams(size=768, distance=Distance.COSINE), + ) - 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}", + client=self._client, collection_name=collection_name, + embedding=self.embeddings, ) + def _collection_exists(self, name: str) -> bool: + try: + self._client.get_collection(name) + return True + except Exception: + return False + 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) + def search(self, query: str, k: int = 5) -> List[tuple]: + results = self.store.similarity_search_with_score(query, k=k) + return [(doc.page_content, score) for doc, score in results] -# --------------------------------------------------------------------------- # -# 2. Чанкинг документов -# --------------------------------------------------------------------------- # +# --------------------------------------------------------------------------- +# Чанкинг +# --------------------------------------------------------------------------- -def chunk_text(text: str, title: str = "") -> List[Document]: - """ - Делит текст на чанки и возвращает список Document с метаданными. - """ +def chunk_document(text: str, title: str) -> List[Document]: + """Разбивает текст на чанки с метаданными.""" splitter = RecursiveCharacterTextSplitter( - chunk_size=500, - chunk_overlap=50, - length_function=len + chunk_size=1000, + chunk_overlap=200, + separators=["\n\n", "\n", " ", ""], ) - chunks = splitter.split_text(text) - docs = [ + return [ Document(page_content=chunk, metadata={"title": title}) - for chunk in chunks + for chunk in splitter.split_text(text) ] - return docs -# --------------------------------------------------------------------------- # -# 3. Инструменты для агента -# --------------------------------------------------------------------------- # +# --------------------------------------------------------------------------- +# Инициализация хранилища (глобальный синглтон для инструментов) +# --------------------------------------------------------------------------- -kb = KnowledgeBase() +rag_store = RagVectorStore() -@tool("search_knowledge_base") + +# --------------------------------------------------------------------------- +# LangChain-инструменты агента (@tool) +# --------------------------------------------------------------------------- + +@tool def search_knowledge_base(query: str, max_results: int = 5) -> List[str]: """ Семантический поиск в базе знаний. - Возвращает список строк с содержимым найденных чанков. + + Args: + query: поисковый запрос. + max_results: максимальное количество результатов. + + Returns: + Список найденных фрагментов с оценкой релевантности. """ - results = kb.search(query=query, limit=max_results) - return [doc.page_content for doc in results] + results = rag_store.search(query, k=max_results) + return [f"[score={score:.4f}] {text}" for text, score in results] -@tool("add_to_knowledge_base") +@tool def add_to_knowledge_base(content: str, title: str) -> str: """ - Добавляет новый документ в базу знаний после чанкинга. - Возвращает подтверждение количества добавленных чанков. + Добавляет новый документ в базу знаний. + + Args: + content: полный текст документа. + title: заголовок документа. + + Returns: + Сообщение об успехе. """ - docs = chunk_text(content, title=title) - kb.add_documents(docs) - return f"Добавлено {len(docs)} чанков из документа '{title}'." + docs = chunk_document(content, title) + rag_store.add_documents(docs) + return f"Документ '{title}' успешно добавлен в базу знаний." -# --------------------------------------------------------------------------- # -# 4. Создание агента -# --------------------------------------------------------------------------- # -LLM_MODEL = "llama3" +# --------------------------------------------------------------------------- +# Агент на базе LangChain (create_react_agent) +# --------------------------------------------------------------------------- + llm = OllamaLLM(model=LLM_MODEL) tools = [search_knowledge_base, add_to_knowledge_base] prompt = PromptTemplate.from_template( - "Вы — интеллектуальный агент с доступом к базе знаний.\n" - "Используйте инструменты 'search_knowledge_base' и 'add_to_knowledge_base' при необходимости.\n\n" + "Вы — интеллектуальный агент с доступом к локальной базе знаний.\n" + "При необходимости используйте инструменты search_knowledge_base и add_to_knowledge_base.\n\n" "Доступные инструменты:\n{tools}\n\n" - "Формат ответа:\n" + "Формат ответа (строго соблюдайте):\n" + "Question: вопрос пользователя\n" + "Thought: что нужно сделать\n" + "Action: инструмент из [{tool_names}]\n" + "Action Input: входные данные для инструмента\n" + "Observation: результат инструмента\n" + "... (можно повторять Thought/Action/Observation)\n" + "Thought: теперь я знаю ответ\n" + "Final Answer: итоговый ответ пользователю\n\n" + "Begin!\n\n" "Question: {input}\n" - "Thought: {agent_scratchpad}\n" - "Action: действие из [{tool_names}]\n" - "Action Input: входные данные\n" - "Observation: результат\n" - "Final Answer: итоговый ответ" + "Thought: {agent_scratchpad}" ) agent = create_react_agent(llm=llm, tools=tools, prompt=prompt) agent_executor = AgentExecutor(agent=agent, tools=tools, 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) +# --------------------------------------------------------------------------- +# Загрузка документов из директории (для init_loader.py) +# --------------------------------------------------------------------------- -# --------------------------------------------------------------------------- # -# 6. Интерактивный клиент -# --------------------------------------------------------------------------- # +def load_documents_from_directory(directory: str) -> None: + """Загружает все .txt и .md файлы из папки в базу знаний.""" + 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") + add_to_knowledge_base.invoke({"content": text, "title": file_path.stem}) + print(f"Загружен: {file_path.name}") -def interactive_cli() -> None: - """ - Простая CLI с командами /add, /search и /quit. - """ - init_agent() - print("=== RAG-агент готов к работе ===") + +# --------------------------------------------------------------------------- +# Запуск напрямую — интерактивный CLI +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + print("=== RAG Агент ===") + print("Команды: /add , /search , /quit") while True: try: - user_input = input("\n> ").strip() + user_input = input("> ").strip() except (EOFError, KeyboardInterrupt): break if not user_input: continue - if user_input.lower() == "/quit": - print("Завершение работы.") + 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}") + elif user_input.startswith("/add "): + file_path = Path(user_input[5:].strip()) + if file_path.is_file(): + text = file_path.read_text(encoding="utf-8") + print(add_to_knowledge_base.invoke({"content": text, "title": file_path.stem})) + else: + print(f"Файл не найден: {file_path}") + elif user_input.startswith("/search "): + query = user_input[8:].strip() + for r in search_knowledge_base.invoke({"query": query, "max_results": 3}): + print(r) else: - # Любой другой ввод считается запросом к агенту response = agent_executor.invoke({"input": user_input}) - print(response["output"]) - -if __name__ == "__main__": - interactive_cli() \ No newline at end of file + print(response["output"]) \ No newline at end of file