""" 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, ChatOllama from langchain_qdrant import QdrantVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_core.tools import tool from langchain_core.documents import Document from langchain.agents import create_agent from qdrant_client import QdrantClient as _QdrantClient 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 + 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), ) self.store = QdrantVectorStore( 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, k: int = 5) -> List[tuple]: # 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] # --------------------------------------------------------------------------- # Чанкинг # --------------------------------------------------------------------------- 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, "source": title, "chunk_index": i}, ) for i, chunk in enumerate(splitter.split_text(text)) ] # --------------------------------------------------------------------------- # Синглтон хранилища (используется инструментами агента) # --------------------------------------------------------------------------- rag_store = RagVectorStore() # --------------------------------------------------------------------------- # RAG-инструменты агента (@tool) # --------------------------------------------------------------------------- @tool def search_knowledge_base(query: str, max_results: int = 5) -> str: """ Выполняет семантический поиск в базе знаний. Args: query: поисковый запрос на естественном языке. max_results: максимальное количество результатов (по умолчанию 5). Returns: Строка с найденными фрагментами и оценками релевантности. """ results = rag_store.search(query, k=max_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: """ Добавляет новый документ в базу знаний. Args: content: текстовое содержимое документа. title: название / источник документа (необязательно). Returns: Сообщение об успехе с количеством созданных чанков. """ docs = chunk_document(content, title=title) rag_store.add_documents(docs) return ( f"Документ «{title or 'без названия'}» успешно добавлен. " f"Создано чанков: {len(docs)}." ) # --------------------------------------------------------------------------- # Агент (create_agent из langchain.agents) # --------------------------------------------------------------------------- llm = ChatOllama(model=LLM_MODEL, temperature=0.3) SYSTEM_PROMPT = """Ты — интеллектуальный AI-ассистент с доступом к локальной базе знаний. Правила работы: 1. ВСЕГДА начинай с поиска через search_knowledge_base перед ответом на вопрос. 2. Если найдена релевантная информация — используй её и ссылайся на источник. 3. Если пользователь просит сохранить информацию — используй add_to_knowledge_base. 4. Если база знаний пуста по теме — честно сообщи об этом и помоги по своим знаниям. 5. Отвечай на том же языке, на котором задан вопрос. """ agent = create_agent( model=llm, tools=[search_knowledge_base, add_to_knowledge_base], system_prompt=SYSTEM_PROMPT, ) 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 "Агент не вернул ответ." # --------------------------------------------------------------------------- # Загрузка документов из директории # --------------------------------------------------------------------------- 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", 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-клиент # --------------------------------------------------------------------------- if __name__ == "__main__": 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("\n> ").strip() except (EOFError, KeyboardInterrupt): console.print("\n[yellow]Выход.[/yellow]") break if not user_input: continue # /quit if user_input.lower() in ("/quit", "/exit"): console.print("[yellow]Выход.[/yellow]") break # /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", errors="ignore") result = add_to_knowledge_base.invoke( {"content": text, "title": file_path.name} ) console.print(f"[green]✓ {result}[/green]") else: 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() 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: 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]")