diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py index 4b4e841..d780008 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py @@ -1,15 +1,21 @@ -# rag_agent.py +""" +rag_agent.py — AI-агент с RAG-памятью на базе Ollama + Qdrant. + +Стек: langchain-ollama, langchain-qdrant, langgraph, qdrant-client, rich +Установка: pip install langchain langchain-core langchain-ollama langchain-qdrant + langchain-text-splitters langgraph qdrant-client rich +""" + import os from pathlib import Path from typing import List -from langchain_ollama import OllamaEmbeddings, OllamaLLM +from langchain_ollama import OllamaEmbeddings, ChatOllama from langchain_qdrant import QdrantVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain.tools import tool -from langchain.agents import create_react_agent, AgentExecutor -from langchain.prompts import PromptTemplate -from langchain.schema import Document +from langchain_core.tools import tool +from langchain_core.documents import Document +from langgraph.prebuilt import create_react_agent from qdrant_client import QdrantClient as _QdrantClient from qdrant_client.http.models import Distance, VectorParams @@ -17,11 +23,11 @@ from qdrant_client.http.models import Distance, VectorParams # Конфигурация # --------------------------------------------------------------------------- -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_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-эмбеддинги) @@ -33,10 +39,8 @@ class RagVectorStore: 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, @@ -60,7 +64,9 @@ class RagVectorStore: self.store.add_documents(docs) def search(self, query: str, k: int = 5) -> List[tuple]: - results = self.store.similarity_search_with_score(query, k=k) + # similarity_search_with_relevance_scores возвращает score в [0..1], + # где 1.0 = максимальная релевантность (исправлено с with_score) + results = self.store.similarity_search_with_relevance_scores(query, k=k) return [(doc.page_content, score) for doc, score in results] @@ -69,135 +75,222 @@ class RagVectorStore: # --------------------------------------------------------------------------- def chunk_document(text: str, title: str) -> List[Document]: - """Разбивает текст на чанки с метаданными.""" + """Разбивает текст на чанки с метаданными через RecursiveCharacterTextSplitter.""" splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200, separators=["\n\n", "\n", " ", ""], ) return [ - Document(page_content=chunk, metadata={"title": title}) - for chunk in splitter.split_text(text) + Document( + page_content=chunk, + metadata={"title": title, "source": title, "chunk_index": i}, + ) + for i, chunk in enumerate(splitter.split_text(text)) ] # --------------------------------------------------------------------------- -# Инициализация хранилища (глобальный синглтон для инструментов) +# Синглтон хранилища (используется инструментами агента) # --------------------------------------------------------------------------- rag_store = RagVectorStore() # --------------------------------------------------------------------------- -# LangChain-инструменты агента (@tool) +# RAG-инструменты агента (@tool) # --------------------------------------------------------------------------- @tool -def search_knowledge_base(query: str, max_results: int = 5) -> List[str]: +def search_knowledge_base(query: str, max_results: int = 5) -> str: """ - Семантический поиск в базе знаний. + Выполняет семантический поиск в базе знаний. Args: - query: поисковый запрос. - max_results: максимальное количество результатов. + query: поисковый запрос на естественном языке. + max_results: максимальное количество результатов (по умолчанию 5). Returns: - Список найденных фрагментов с оценкой релевантности. + Строка с найденными фрагментами и оценками релевантности. """ results = rag_store.search(query, k=max_results) - return [f"[score={score:.4f}] {text}" for text, score in results] + if not results: + return "В базе знаний ничего не найдено по данному запросу." + lines = [ + f"[{i}] (score={score:.3f}) {text}" + for i, (text, score) in enumerate(results, start=1) + ] + return "\n\n".join(lines) @tool -def add_to_knowledge_base(content: str, title: str) -> str: +def add_to_knowledge_base(content: str, title: str = "") -> str: """ Добавляет новый документ в базу знаний. Args: - content: полный текст документа. - title: заголовок документа. + content: текстовое содержимое документа. + title: название / источник документа (необязательно). Returns: - Сообщение об успехе. + Сообщение об успехе с количеством созданных чанков. """ - docs = chunk_document(content, title) + docs = chunk_document(content, title=title) rag_store.add_documents(docs) - return f"Документ '{title}' успешно добавлен в базу знаний." + return ( + f"Документ «{title or 'без названия'}» успешно добавлен. " + f"Создано чанков: {len(docs)}." + ) # --------------------------------------------------------------------------- -# Агент на базе LangChain (create_react_agent) +# Агент (create_react_agent из langgraph — актуальный API v0.3+) # --------------------------------------------------------------------------- -llm = OllamaLLM(model=LLM_MODEL) +llm = ChatOllama(model=LLM_MODEL, temperature=0.3) -tools = [search_knowledge_base, add_to_knowledge_base] +SYSTEM_PROMPT = """Ты — интеллектуальный AI-ассистент с доступом к локальной базе знаний. -prompt = PromptTemplate.from_template( - "Вы — интеллектуальный агент с доступом к локальной базе знаний.\n" - "При необходимости используйте инструменты search_knowledge_base и add_to_knowledge_base.\n\n" - "Доступные инструменты:\n{tools}\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}" +Правила работы: +1. ВСЕГДА начинай с поиска через search_knowledge_base перед ответом на вопрос. +2. Если найдена релевантная информация — используй её и ссылайся на источник. +3. Если пользователь просит сохранить информацию — используй add_to_knowledge_base. +4. Если база знаний пуста по теме — честно сообщи об этом и помоги по своим знаниям. +5. Отвечай на том же языке, на котором задан вопрос. +""" + +agent = create_react_agent( + model=llm, + tools=[search_knowledge_base, add_to_knowledge_base], + prompt=SYSTEM_PROMPT, ) -agent = create_react_agent(llm=llm, tools=tools, prompt=prompt) -agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True) + +def run_agent(user_input: str) -> str: + """Запускает агента и возвращает финальный ответ.""" + result = agent.invoke({"messages": [("user", user_input)]}) + for msg in reversed(result["messages"]): + if msg.type == "ai" and msg.content: + return msg.content + return "Агент не вернул ответ." # --------------------------------------------------------------------------- -# Загрузка документов из директории (для init_loader.py) +# Загрузка документов из директории # --------------------------------------------------------------------------- -def load_documents_from_directory(directory: str) -> None: - """Загружает все .txt и .md файлы из папки в базу знаний.""" +def load_documents_from_directory(directory: str) -> int: + """ + Загружает все .txt и .md файлы из папки в базу знаний. + Возвращает количество загруженных файлов. + """ path = Path(directory) + count = 0 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}") + text = file_path.read_text(encoding="utf-8", errors="ignore") + if text.strip(): + docs = chunk_document(text, title=file_path.name) + rag_store.add_documents(docs) + print(f" Загружен: {file_path.name} ({len(docs)} чанков)") + count += 1 + return count # --------------------------------------------------------------------------- -# Запуск напрямую — интерактивный CLI +# Интерактивный CLI-клиент # --------------------------------------------------------------------------- if __name__ == "__main__": - print("=== RAG Агент ===") - print("Команды: /add , /search , /quit") + from rich.console import Console + from rich.panel import Panel + from rich.markdown import Markdown + from rich.rule import Rule + + console = Console() + + console.print(Panel( + "[bold]RAG-агент[/bold] · Qdrant + Ollama\n\n" + "[cyan]/add <текст>[/cyan] — добавить текст в базу знаний\n" + "[cyan]/add-file <путь>[/cyan] — загрузить файл в базу знаний\n" + "[cyan]/search <запрос>[/cyan] — прямой поиск (без агента)\n" + "[cyan]/quit[/cyan] — выйти", + border_style="dim", + )) + while True: try: - user_input = input("> ").strip() + user_input = input("\n> ").strip() except (EOFError, KeyboardInterrupt): + console.print("\n[yellow]Выход.[/yellow]") break if not user_input: continue - if user_input.lower() == "/quit": - print("Выход.") + + # /quit + if user_input.lower() in ("/quit", "/exit"): + console.print("[yellow]Выход.[/yellow]") break - elif user_input.startswith("/add "): - file_path = Path(user_input[5:].strip()) + + # /add-file <путь> + elif user_input.lower().startswith("/add-file "): + file_path = Path(user_input[10:].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})) + text = file_path.read_text(encoding="utf-8", errors="ignore") + result = add_to_knowledge_base.invoke( + {"content": text, "title": file_path.name} + ) + console.print(f"[green]✓ {result}[/green]") else: - print(f"Файл не найден: {file_path}") - elif user_input.startswith("/search "): + console.print(f"[red]Файл не найден: {file_path}[/red]") + + # /add <текст> ← исправлено: теперь принимает текст напрямую, + # а не только путь к файлу + elif user_input.lower().startswith("/add "): + raw = user_input[5:].strip() + path = Path(raw) + if path.is_file(): + # если передан путь к существующему файлу — загружаем файл + text = path.read_text(encoding="utf-8", errors="ignore") + result = add_to_knowledge_base.invoke( + {"content": text, "title": path.name} + ) + else: + # иначе сохраняем как текст напрямую + result = add_to_knowledge_base.invoke( + {"content": raw, "title": "manual"} + ) + console.print(f"[green]✓ {result}[/green]") + + # /search <запрос> + elif user_input.lower().startswith("/search "): query = user_input[8:].strip() - for r in search_knowledge_base.invoke({"query": query, "max_results": 3}): - print(r) + console.print(Rule("Результаты поиска")) + results = rag_store.search(query, k=5) + if not results: + console.print("[yellow]Ничего не найдено.[/yellow]") + else: + for i, (text, score) in enumerate(results, 1): + console.print(Panel( + text, + title=f"[cyan]#{i} score={score:.3f}[/cyan]", + border_style="dim", + )) + + # /search без аргумента + elif user_input.lower() == "/search": + console.print("[yellow]Укажите запрос: /search <текст>[/yellow]") + + # запрос к агенту else: - response = agent_executor.invoke({"input": user_input}) - print(response["output"]) \ No newline at end of file + console.print("[dim]Агент думает...[/dim]") + try: + answer = run_agent(user_input) + console.print(Panel( + Markdown(answer), + title="[green]Агент[/green]", + border_style="green", + )) + except Exception as e: + console.print(f"[red]Ошибка агента: {e}[/red]") \ No newline at end of file